Add files using upload-large-folder tool

.gitattributes

@@ -33,3 +33,17 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+image/blobs/sha256/3de1f5eb93e54b4561afa733cf844d1914a9ba260f49e86fe600f344c5dd025c filter=lfs diff=lfs merge=lfs -text
+image/blobs/sha256/ca0b072b65f8c21199f96e7498f2ccbc252490ce39060800647332751f857287 filter=lfs diff=lfs merge=lfs -text
+image/blobs/sha256/cbb77a738c7df827819a8b8bf87682eeca8bdb434d41411a9682c38717f2f187 filter=lfs diff=lfs merge=lfs -text
+image/blobs/sha256/530b0e35f44c6f963e06fdaacdbfecb2021d4c55d49fe4c0019e161d94c18de3 filter=lfs diff=lfs merge=lfs -text
+image/blobs/sha256/540cf00275e913a9bccc49fe7beba58037696661bde6ae083a3ac84c5b160e67 filter=lfs diff=lfs merge=lfs -text
+image/blobs/sha256/b6df468b82a4b2f9ee3ca3a79a6bbe99b5cda02ddc63b7e3c89e7bd08ef41706 filter=lfs diff=lfs merge=lfs -text
+image/blobs/sha256/0926a8eb0e608a5c6888d1cd5594184bdf3ed3aa311dba5b42a547caefdc6f2e filter=lfs diff=lfs merge=lfs -text
+image/blobs/sha256/c3bc7b373b4523cabcdd9c64ab10d32510ef61bace60873a279ccc4902738989 filter=lfs diff=lfs merge=lfs -text
+image/blobs/sha256/fdc1ed79ffd24d66f8be3754ec4dc80b1ab0fcc8c8165a6007748c464c94f897 filter=lfs diff=lfs merge=lfs -text
+image/blobs/sha256/c18d0f3c8022bcd5a8059f67fc0e5cfd53ef36a44997335d6cd0aa6b19db140d filter=lfs diff=lfs merge=lfs -text
+image/blobs/sha256/8753e0cfbd424e962ccaf50aaaf02fd06ff2efb8677219657e751a53922efa9f filter=lfs diff=lfs merge=lfs -text
+image/blobs/sha256/24cba7375920bef8d4cc4f0ce4294f8f70c65b9a37c4ed6f6c2d63405c76ba3c filter=lfs diff=lfs merge=lfs -text
+image/blobs/sha256/e19b6f1fb65dd2888d9003ef9513a21d129041a328bb8a9a4164d29ef0382b16 filter=lfs diff=lfs merge=lfs -text
+image/blobs/sha256/c00314a02c644cf2aea5a8ae4e3ee5ee383f4072d3e35248fe20dea57a90c46c filter=lfs diff=lfs merge=lfs -text

code/models/demos/blackhole/qwen3_coder_30b_a3b/README.md

@@ -0,0 +1,143 @@
+# Qwen3-Coder-30B-A3B-Instruct on Blackhole
+
+This directory implements Tenstorrent Blackhole inference for
+**`Qwen/Qwen3-Coder-30B-A3B-Instruct`** — a 48-layer `Qwen3MoeForCausalLM`
+sparse mixture-of-experts model (30B total / ~3B active parameters) with an
+advertised context of 262144 tokens.
+
+| Model | `HF_MODEL` | Mesh / `--mesh-device` | Parallelism |
+| ----- | ---------- | ---------------------- | ----------- |
+| Qwen3-Coder-30B-A3B-Instruct | `Qwen/Qwen3-Coder-30B-A3B-Instruct` | 4 Blackhole dies — `P300x2` (a `(1, 4)` mesh) | 4-way tensor parallel |
+
+The 4-die path needs `FABRIC_1D_RING` for the cross-device collectives and a
+trace region for the captured decode and chunked-prefill traces
+(`DEFAULT_TRACE_REGION_SIZE = 300_000_000`, see [tt/model.py](tt/model.py)).
+
+## Architecture
+
+Assembly: `tok_embeddings → 48 × decoder layer → RMSNorm → LM head → on-device
+sampling`.
+
+Each decoder layer is GQA attention (rotary, per-head QK-norm) followed by a
+sparse MoE block — a router plus top-k expert MLPs — replacing the dense MLP.
+Everything shape-related (layer count, expert count, top-k, head dims, vocab,
+rope base) is read from the parsed HF config, so the code follows the
+checkpoint rather than hard-coding it.
+
+| File | Role |
+| ---- | ---- |
+| [tt/model.py](tt/model.py) | Full 48-layer model, weight load, KV cache, trace capture, LM head + sampling |
+| [tt/functional_decoder.py](tt/functional_decoder.py) | Reference-shaped single-layer decoder (attention + MoE) — the correctness baseline |
+| [tt/optimized_decoder.py](tt/optimized_decoder.py) | Single-device optimized layer (fused QKV, sharded matmuls, program configs) |
+| [tt/multichip_decoder.py](tt/multichip_decoder.py) | Tensor-parallel layer and the CCL schedule for the 4-die mesh |
+| [tt/precision.py](tt/precision.py) | `PrecisionConfig` — per-tensor dtypes and math fidelity, overridable via `QWEN3_PRECISION_CONFIG` |
+| [config/](config/) | Runtime policy the serving path reads: the selected precision config and the served-context contract |
+| [tt/weight_mapping.py](tt/weight_mapping.py) | HF checkpoint → device tensor layout (QKV permutes, expert stacking) |
+| [tt/generator.py](tt/generator.py) | `build_generator()` + the high-level `generate()` loop (owns KV cache and page table) |
+| [tt/generator_vllm.py](tt/generator_vllm.py) | vLLM adapter — `prefill_forward` / `decode_forward` against a caller-owned cache |
+| [vllm_bundle/](vllm_bundle/) | `EXTRA_MODELS_DIR` bundle that registers the adapter with the TT vLLM plugin |
+
+The generator implements the `Generator` ABC in
+[models/common/readiness_check/contract.py](../../../common/readiness_check/contract.py),
+so the same object serves the host-side demo/readiness path and vLLM.
+
+## Precision
+
+A 29-row datatype sweep over the attention, MoE, KV, CCL, norm and LM-head
+tensors selected [config/selected_precision_config.json](config/selected_precision_config.json),
+which the vLLM path loads on every serve so that serving and readiness cannot
+run different numerics. `DEFAULT_PRECISION` in [tt/precision.py](tt/precision.py)
+is the equivalent in-code default. Override for experiments with
+`QWEN3_PRECISION_CONFIG=<path-to-json>`.
+
+Served context is capped by
+[config/context_contract.json](config/context_contract.json) rather than by the
+`--max-model-len` you pass, so a request for more context than has been
+validated fails loudly instead of serving a quietly-clipped model.
+
+## Running the tests
+
+All device tests target the 4-die mesh. From the repository root:
+
+```bash
+source python_env/bin/activate
+export HF_MODEL=Qwen/Qwen3-Coder-30B-A3B-Instruct
+D=models/demos/blackhole/qwen3_coder_30b_a3b
+
+# module + model correctness (excludes the perf-only tests)
+pytest $D/tests/ -m "not models_performance_bare_metal" -q
+
+# perf tests (decode/prefill timings; writes CSVs under doc/)
+pytest $D/tests/test_perf.py -q
+```
+
+`test_full_model.py` runs a 2-layer model by default so it stays cheap; set
+`QWEN3_FULL_MODEL_LAYERS=48` for the complete model.
+
+Host-only tests (no device): `tests/test_reference.py`,
+`tests/test_precision_config.py`.
+
+## Readiness checks and vLLM serving
+
+The shared harness in [models/common/readiness_check/](../../../common/readiness_check/)
+drives this model through `tt/generator.py`:
+
+```bash
+# teacher-forced accuracy against a reference completion
+python -m models.common.readiness_check.run_prefill_check \
+  --model-dir models/demos/blackhole/qwen3_coder_30b_a3b \
+  --reference <reference.refpt> \
+  --mesh-device P300X2 --fabric-config FABRIC_1D_RING --trace-region-size 300000000
+```
+
+vLLM serving registers through the plugin's `EXTRA_MODELS_DIR` hook — no edit
+to the vLLM checkout is required:
+
+```bash
+export EXTRA_MODELS_DIR=$PWD/models/demos/blackhole/qwen3_coder_30b_a3b/vllm_bundle
+
+python -m models.common.readiness_check.run_vllm_server \
+  --model-dir <output-dir> --hf-model Qwen/Qwen3-Coder-30B-A3B-Instruct \
+  --mesh-device P300x2 --max-num-seqs 1 --max-model-len 262144 \
+  --block-size 32 --port 8100 --stages serve \
+  --tt-config '{"trace_region_size": 50331648, "fabric_config": "FABRIC_1D_RING"}' \
+  --additional-server-args "--generation-config vllm"
+```
+
+`--generation-config vllm` matters: this checkpoint's `generation_config.json`
+injects `repetition_penalty=1.05` into every request that does not override it,
+and a penalised request costs ~14% TPOT because the penalty operands are staged
+per step on the host.
+
+## Measured performance
+
+4 Blackhole dies, 48 layers, traced decode, on-device sampling, greedy,
+128-token input / 128-token output, batch 1.
+
+| Path | TTFT | Decode |
+| ---- | ---- | ------ |
+| Standalone traced (`generate()`) | 129.9 ms | 19.213 ms — **52.05 t/s/u** |
+| Through vLLM (`max_num_seqs=1`, `--max-concurrency 1`) | 307–312 ms | 19.78 ms — **50.3–50.6 t/s/u** |
+
+vLLM adds ~0.57 ms per decoded token (2.9%) over the standalone traced path.
+The TTFT gap is request handling, tokenisation and detokenisation, not decode.
+
+Evaluated end to end by tt-inference-server 0.20.0 against a live server:
+mbpp 77.2%, humaneval 92.7%, ifeval 81.1%/87.1%, gpqa_diamond_cot 56.1%.
+
+## Known limitations
+
+- **Long-prefill scaling is bad.** A 131072-token prefill completes and returns
+  valid output but takes 94.4 minutes — 3.05× worse per token than 65536, and
+  well off what the single-layer sweep predicts. The suspected cause is
+  per-chunk tensor accumulation in the MoE prefill path; this is an unproven
+  hypothesis, not a measured root cause.
+- **No full-model 262144 prefill has been verified.** 262144 is allocated,
+  page-tabled and served, and prefills through a single layer, but the largest
+  48-layer prefill measured end to end is 131072. The advertised context is
+  left at 262144.
+- **Not yet registered in the tiered models CI.** No entry exists in
+  `models/model_ci_tiers.md`, `tests/pipeline_reorg/models_*_tests.yaml`,
+  `models/model_targets.yaml` or the vLLM test registry, and there is no
+  `demo/` entry point in the shape those pipelines invoke. See
+  [models/MIGRATING_TO_TIERED_CI.md](../../../MIGRATING_TO_TIERED_CI.md).

code/models/demos/blackhole/qwen3_coder_30b_a3b/config/context_contract.json

@@ -0,0 +1,360 @@
+{
+  "hf_model": "Qwen/Qwen3-Coder-30B-A3B-Instruct",
+  "hf_advertised_context": 262144,
+  "current_supported_context": 256000,
+  "capability_reduction": true,
+  "limiting_reason": "Serving prefill. TWO independent limits, both above this bound. (1) An UNRESOLVED prefill cliff: 259,000 tokens completes in 1130.8 s, tracking the O(S^2) curve to within 0.3 %, but 262,136 tokens does not complete in 3600 s -- one continuous blocking prefill, no watchdog throw, engine healthy. The cliff is confined to the last ~3,100 tokens of the advertised context and its mechanism was deliberately not chased (see stage11_serving_context). (2) LATENCY: prefill at this length is minutes, not seconds -- 253,000 tokens costs 1049.3 s (17.5 min) of TTFT on the path that WORKS. The usable interactive range is bounded far below this figure by latency, not by the cliff.",
+  "stage": "optimized_full_model",
+  "scope": "the complete 48-layer causal LM on the full 4-die P300_X2 mesh (1x4), batch 1, AFTER the stage-06 optimizations: embeddings, 48 stage-04 multichip decoder layers, final norm, column-parallel lm_head and on-device traced sampling with a distributed argmax. The single-layer stage-02/03/04 numbers and the stage-05 full-model numbers below are retained unchanged; the stage06_* fields are the shipped ones.",
+  "device": {
+    "arch": "blackhole",
+    "board": "p300 x2 (ClusterType.P300_X2), 4 dies",
+    "mesh": "1x4",
+    "dram_per_die_gb": 34.18,
+    "dram_per_die_note": "34.18 GB is what the TTNN allocator reports as total DRAM per die on this host (doc/multichip_decoder/footprint_probe.log, first line). Stages 01 and 02 recorded the nominal board figure of 32 GB; nothing they concluded changes.",
+    "topology": "ring, 2 ethernet links per hop, FABRIC_1D_RING"
+  },
+  "measured": {
+    "decode_context_tokens": 262144,
+    "decode_probe": "paged KV cache allocated to depth 262144 (block_size=32); decode step executed at position 262143; output finite; 2.42 s",
+    "prefill_context_tokens": 262144,
+    "prefill_probe": "single-shot prefill of a 262144-token sequence through the full decoder layer; output finite; 192.23 s",
+    "prefill_sweep_seconds": {
+      "512": 0.51,
+      "1024": 4.7,
+      "2048": 5.35,
+      "4096": 7.92,
+      "8192": 10.24,
+      "16384": 14.87,
+      "32768": 23.79,
+      "65536": 42.05,
+      "131072": 85.85,
+      "262144": 192.23
+    },
+    "decode_sweep_seconds": {
+      "4096": 1.19,
+      "16384": 1.21,
+      "65536": 1.45,
+      "131072": 1.77,
+      "262144": 2.42
+    },
+    "prefill_sweep_seconds_PROVENANCE": "SINGLE-LAYER, PRE-OPTIMISATION. Not comparable to the shipped 48-layer model and not to be divided into any full-model figure. This file's own `scope` says 'The single-layer stage-02/03/04 numbers ... are retained unchanged', and `prefill_probe` describes 'a single-shot prefill ... through the full decoder LAYER' -- one complete layer, not the stack. The evidence list points at doc/functional_decoder/ and doc/multichip_decoder/, i.e. stages 02-04, before the stage-04/06 optimisations. Measured on the current tree, the whole 48-layer model prefills 32,768 tokens in 37.638 s, i.e. 0.78 s per layer, against 23.79 s for the single layer recorded here -- ~30x apart, so the two are measuring different things on different code. In particular the 262144 -> 192.23 s entry does NOT describe the shipped model, and no current-tree measurement of that length exists.",
+    "prefill_same_tree_seconds": {
+      "note": "Full 48-layer model on the current tree. Standalone and served at a MATCHED length, which is the comparison that retires the suspected serving penalty.",
+      "32768": {
+        "standalone": 37.638,
+        "served": 37.778,
+        "served_over_standalone": 1.004
+      },
+      "65536": {
+        "served": 98.557,
+        "standalone": null,
+        "warning": "Do NOT divide 98.557 by the 42.05 s single-layer entry above; that repeats the exact error this block exists to correct."
+      },
+      "evidence": [
+        "doc/batch_scaling/probes/expert_chunk_sweep_v3.json (standalone 32,768)",
+        "doc/batch_scaling/probes/long_prompt_gap_32768.json",
+        "doc/batch_scaling/probes/long_prompt_gap_65536.json"
+      ]
+    }
+  },
+  "kv_cache": {
+    "dtype": "bfloat16",
+    "bytes_per_token_per_layer": 2048,
+    "formula": "num_key_value_heads(4) * head_dim(128) * 2 bytes * 2 tensors (K and V)",
+    "bytes_at_full_context_per_layer": 536870912,
+    "paged": true,
+    "default_block_size": 32
+  },
+  "pcc_validated_context_tokens": 512,
+  "pcc_validation_note": "PCC against the HuggingFace reference is verified up to 512 tokens; longer lengths are capacity/liveness probes only, because a full-length torch reference for 262144 tokens does not fit in host RAM. This is a limit of the reference, not of the TTNN decoder.",
+  "forward_looking_note": "DISCHARGED by stage 03. The stage-01 note read: 'the full 48-layer model would need ~24 GiB on a 32 GiB die, so the full-model stage will have to weigh KV dtype, paging across dice, or a served-context cap.' None of those trade-offs is needed. TP=4 across the four dies puts the whole 48-layer model plus a full 262144-token paged KV cache in 11.759 GB per die with 22.119 GB free -- measured in stage 05 by allocating the real 48-layer model, and the figure this note now carries (see full_model_measured and comparison_to_stage03_prediction). Stage 03 predicted 11.829 / 22.350 from a synthetic allocation of the same shapes; the prediction and the measurement agree to within 0.6%. KV dtype is unchanged (bfloat16), paging is unchanged (block_size 32), and the served context is unchanged. The note is kept rather than deleted because it is the question this stage answers.",
+  "evidence": [
+    "doc/functional_decoder/work_log.md (context probe section)",
+    "doc/multichip_decoder/README.md and work_log.md",
+    "doc/multichip_decoder/footprint_probe.log (the allocation this file quotes)",
+    "doc/multichip_decoder/probes/footprint_probe.py",
+    "doc/multichip_decoder/pcc_log.txt",
+    "doc/optimized_multichip_decoder/README.md and work_log.md",
+    "doc/optimized_multichip_decoder/pcc_log.txt",
+    "doc/full_model/README.md and work_log.md",
+    "doc/full_model/probes/footprint_262144.json (the allocation this file quotes)",
+    "doc/full_model/probes/perf_full_model.csv and .json",
+    "doc/full_model/run_prefill_check.log, run_teacher_forcing.log, run_autoregressive.log",
+    "doc/optimized_full_model/README.md and work_log.md",
+    "doc/optimized_full_model/profile_48layer_work_log.md (the lever analysis; its op-level figures are the pre-adoption ones)",
+    "doc/optimized_full_model/ops_perf_full_model_48layer_decode.csv.gz (the verified one-iteration decode window)",
+    "doc/optimized_full_model/ops_perf_full_model_48layer_prefill_s128.csv.gz (the 48-layer prefill window)",
+    "doc/optimized_full_model/probes/profile_summary_decode.json and profile_summary_prefill.json",
+    "doc/optimized_full_model/probes/perf_full_model_p{128,1024,4096}_{before,after,argmaxrows}.json",
+    "doc/optimized_full_model/probes/footprint_262144.json (the allocation this file quotes)",
+    "doc/optimized_full_model/probes/runtime_fallback_audit.json",
+    "doc/optimized_full_model/probes/sdpa_sweep_confirm_bf16.json and sdpa_hf_pcc_at_depth.json",
+    "doc/optimized_full_model/logs/run_prefill_check_argmaxrows.log, run_teacher_forcing_argmaxrows.log, run_autoregressive_argmaxrows.log, check_degenerate_argmaxrows.log"
+  ],
+  "expert_weight_dtype": "bfloat4_b",
+  "optimized_note": "Stage 02 changed expert weight dtype to bfloat4_b, attention projection dtype to bfloat8_b, the expert matmul block widths, expert math fidelity to LoFi, and the decode attention projections to a DRAM-sharded program config. None of these touches the KV cache layout, dtype, or paging, so the measured context limits carry over unchanged. Per-layer weight DRAM: expert weights shrink ~3.6x (bf16 to bfloat4_b, which is 0.5625 B/elem once each 16-element block's exponent byte is counted). Decode additionally needs a second, DRAM-width-sharded copy of wqkv/wo; at bfloat8_b (1.0625 B/elem) each copy is 20.05 MB, so the pair is 40.11 MB against the single bf16 copy stage 01 held at 37.75 MB -- a 2.36 MB increase, not a wash, and an earlier revision of this file called it one by rounding bfloat8_b to 1 B/elem. Against 24 GB of usable DRAM it changes nothing: no capability reduction.",
+  "attention_weight_dtype": "bfloat8_b",
+  "parallelism": {
+    "attention": "TP=4 (8 Q heads, 1 K head, 1 V head per die)",
+    "experts": "EP=4 (32 of 128 experts per die)",
+    "router": "replicated",
+    "norms": "replicated",
+    "residual": "replicated [1,1,B,2048]",
+    "collectives_per_layer": "2 all-reduces, each reduce-scatter then all-gather on dim 3",
+    "padding": "none -- 2048/4, 32/4, 4/4, 128/4 and 151936/4 are all exact"
+  },
+  "multichip_measured": {
+    "method": "doc/multichip_decoder/probes/footprint_probe.py -- allocates the real per-die shapes rather than computing their size. Output archived verbatim at doc/multichip_decoder/footprint_probe.log.",
+    "decoder_weights_48_layers_gb_per_die": 4.596,
+    "plus_embed_and_lm_head_gb_per_die": 5.374,
+    "embed_lm_head_note": "embed_tokens replicated (bf16), lm_head column-parallel, per-die N = 37984 = 151936/4.",
+    "kv_bytes_per_token_per_layer_per_die": 512,
+    "kv_formula": "1 local kv head * head_dim(128) * 2 bytes * 2 tensors (K and V); one quarter of the single-die 2048",
+    "total_at_full_context_batch1_gb_per_die": 11.829,
+    "total_note": "MEASURED: 48 layers of sharded weights + embed + lm_head + 48 paged KV caches at 262144 tokens, all allocated simultaneously on the mesh. 22.350 GB/die free afterwards. The arithmetic prediction in mesh_plan.md section 8 was 11.80 GB; the allocator says 11.829.",
+    "headroom_gb_per_die": 22.35,
+    "single_die_equivalent_gb": 45.28,
+    "single_die_note": "One die would need 19.52 GB of weights plus 25.77 GB of KV at the advertised context. That does not fit in 34.18 GB, so the mesh is a capability requirement here and not only a speed one. This is what discharges the stage-01 forward_looking_note.",
+    "single_die_convention": "Recomputed from the probe's own per-die tensor list (footprint_probe.log, GB = 1e9 as the probe reports it). Convention: tensors that are sharded across the mesh are counted 4x -- the two expert tensors, both wqkv copies, both wo copies, and lm_head -- and tensors that are replicated are counted ONCE, since a single die would hold exactly one of each: the router (2048x128 bf16), the two RMSNorm vectors, and embed_tokens. That gives 48 * (4 * 94.962 MB sharded + 0.786 MB replicated) = 18.270 GB of decoder weights, + 0.622 GB embed + 4 * 0.156 GB lm_head = 19.515 GB, i.e. 19.52. Counting the replicated router/norms 4x as well would read 19.63; the difference is 0.11 GB and changes nothing. The 19.50 this field previously carried came from mesh_plan.md section 8, whose per-layer table omits the two tile-padded RMSNorm vectors (0.26 MB/layer, 0.013 GB over 48 layers) and rounds each row; recomputing from the probe's tensor list closes that. KV is 4 kv heads * 128 * 2 B * 2 tensors = 2048 B/token/layer * 48 layers * 262144 tokens = 25.770 GB. Total 45.28 GB against 34.18 GB of DRAM -- the conclusion is unaffected by which convention is used."
+  },
+  "multichip_largest_feasible": {
+    "context_at_batch1": 262144,
+    "context_at_batch1_basis": "the HF-advertised context, allocated and held on the mesh; not a limit reached here",
+    "batch_at_full_context": 4,
+    "batch_at_full_context_basis": "floor((34.18 - 5.374 weights - ~2 GB trace/activation) / 6.442 GB of KV per user) = floor(26.806 / 6.442) = floor(4.16) = 4. KV per user is 48 layers * 512 B/token/layer/die * 262144 tokens = 6.442 GB, from the measured per-die per-token figure above. Arithmetic on the measured per-user KV, not allocated; an earlier revision of this file recorded 3, which its own formula does not give.",
+    "batch_ceiling_from_ops": 32,
+    "batch_ceiling_basis": "nlp_create_qkv_heads_decode_device_operation.cpp:51 asserts num_users <= 32. A TTNN op limit, unchanged by TP, and reached long before DRAM is.",
+    "context_at_batch32": 28788,
+    "context_at_batch32_basis": "~22.6 GB usable / (32 users * 48 layers * 512 B/token); arithmetic"
+  },
+  "multichip_pcc_note": "Multichip PCC is validated against the single-chip TTNN optimized decoder run replicated on the same mesh, 0.99962-0.99997 across prefill S = 32/33/100/128/257/512 and decode in both cache modes, and against HF at 0.9990 (prefill) and 0.9988-0.9993 (decode). See doc/multichip_decoder/pcc_log.txt. The 512-token PCC validation ceiling of the stage-01/02 entry above is unchanged and has the same cause: a full-length torch reference does not fit in host RAM. Stage 04 re-ran every one of those gates on the optimized path: 112 passed, 0 failed (stage 03's 111 plus the host-only test_meta_rope_weights_match_hf), with prefill 0.999620-0.999868 and decode 0.999935-0.999941 against the single-chip TTNN baseline. The one movement is decode vs single-chip, 0.99997 -> 0.99994, and it is the right sign: the reference is optimized_decoder.py, whose numerics stage 04 left unchanged (its only edit is an optional rope= seam that defaults to the shipped op) and which still accumulates its RMSNorm sum of squares in bf16, while the multichip norm now measures 1.686e-02 against a torch fp64 reference where the old one measured 6.711e-02. Against HF the four decode steps move by at most 8e-05 and in both directions (0.9992853->0.9992450, 0.9987817->0.9987388, 0.9993126->0.9993928, 0.9989592->0.9989421), i.e. three are marginally lower and one is higher; that is bf16 rounding under a changed norm, not a regression, and the earlier wording 'nothing regressed' was literally false. See doc/optimized_multichip_decoder/pcc_log.txt.",
+  "stage04_note": "Stage 04 optimized the multichip decoder in place and changed NOTHING this file measures. No capability reduction. What it changed: both residual RMSNorms now run width-sharded over 8 L1 cores instead of one core, the router projection reads that L1 shard instead of DRAM-interleaved, and the two collectives use caller-owned persistent buffers. KV cache dtype, layout, paging and bytes/token are untouched, and so is the activation sharding of the cache. Two things grew: each layer carries a second, ROW_MAJOR copy of each residual RMSNorm vector (the layout the sharded rms_norm program factory reads) at 4 KB each, i.e. 8 KB per layer and 0.384 MB per die over all 48 layers, against 4.596 GB of sharded decoder weights; and the mesh context owns two sets of persistent reduce-scatter/all-gather buffers of about 0.5 MB each, once per mesh rather than once per layer. Both are far below the 22.350 GB/die of measured headroom, so the 262144-token contract, the batch-4-at-full-context figure and the single-die comparison all stand unchanged. Evidence: doc/optimized_multichip_decoder/README.md and work_log.md section 9.",
+  "full_model_measured": {
+    "method": "doc/full_model/probes/footprint_probe.py --context 262144. Unlike the stage-03 probe this builds the REAL model -- real weights, real embed_tokens, real lm_head, the real paged KV cache at the advertised context, the real RoPE tables -- captures both decode traces and runs a token through it. Output archived at doc/full_model/probes/footprint_262144.json and .log. GB = 1e9, as the allocator reports it.",
+    "weights_embed_lm_head_rope_gb_per_die": 5.311,
+    "kv_cache_at_262144_batch1_gb_per_die": 6.443,
+    "traces_and_persistent_buffers_gb_per_die": 0.006,
+    "total_gb_per_die": 11.759,
+    "headroom_gb_per_die": 22.119,
+    "dram_per_die_gb_reported": 33.879,
+    "dram_per_die_note": "33.879 rather than the 34.18 recorded for stage 03/04 because this probe opens the mesh with trace_region_size=300 MB, which the allocator carves out of DRAM before reporting. It is the same hardware; the difference is 0.30 GB and is the trace region the decode traces live in.",
+    "sum_note": "the three stage rows in doc/full_model/probes/footprint_262144.json sum to 11.759415296, BIT-IDENTICAL to total_gb_per_die -- there is no residual and nothing is omitted. The 0.001 that appears if the 3dp values in this file or in the README's table are added up is introduced by that display rounding alone. An earlier revision of this field attributed it to 'the allocator's rounding across the three separate reads', which the raw JSON contradicts; doc/full_model/probes/check_published_figures.py now asserts the raw rows equal the raw total exactly.",
+    "comparison_to_stage03_prediction": "stage 03's footprint probe ALLOCATED 11.829 GB/die for the same contents and never ran through them. Executing it measures 11.759, 0.070 GB lower: the running model carries one shared pair of RoPE tables and one shared set of persistent collective buffers where the allocation probe's per-layer tensor list double-counted nothing, and the two paged caches are allocated once per layer at exactly ceil(262144/32) blocks. The stage-03 figure remains the honest allocation-time answer and is not restated here as wrong.",
+    "weight_load_seconds": 182.66,
+    "weight_load_note": "48 layers streamed one at a time out of the sharded safetensors checkpoint; the 61 GB checkpoint is never fully materialised on the host.",
+    "rope_table_note": "the probe builds the model with the default rope_cache_len of 8192, so the RoPE tables inside weights_embed_lm_head_rope_gb_per_die are the 8192-row pair (0.004 GB/die). Growing them to the full 262144 context costs 0.134 GB/die (262144 * 128 * 2 B * 2 tables), taking the total to 11.889 GB/die and the headroom to 21.985. No capability reduction; the tables grow on demand via Qwen3CoderModel.ensure_rope_capacity, and tt/generator.py's decode_forward takes decode_horizon= so a low-level caller grows them once before any trace is captured."
+  },
+  "full_model_context_note": "NO REDUCTION. The advertised 262144 context is allocated and held by the running 48-layer model with 22.119 GB/die free. What has NOT been done at the full stack is a 262144-token prefill: prefill is single-shot (no internal chunking), and the longest prompt actually pushed through all 48 layers in this stage is 1000 tokens. The 262144-token single-layer prefill probe in measured.prefill_probe above still stands, and the KV cache, page tables and positions are sized and exercised for the full context. This is a coverage gap in the evidence, not a capability reduction, and it is named as limitation 1 of doc/full_model/README.md. STAGE 06: still no reduction, and the 262144 context is now usable rather than only allocatable -- decode is nearly flat in context (1.04x from 128 to 4096, against 1.96x before). See stage06_context_flatness for exactly how deep that was measured and what is still unmeasured.",
+  "full_model_batch": {
+    "primary": 1,
+    "primary_note": "every performance figure in doc/full_model/README.md is batch 1.",
+    "largest_tested": 4,
+    "largest_tested_note": "tests/test_full_model.py::test_mixed_length_batch_prefill_and_decode runs four users at prompt lengths 7 / 33 / 64 / 129 through one prefill and one decode with disjoint physical cache pages. Batch was not pushed to 32 at the full stack in this stage.",
+    "hard_ceiling": 32,
+    "hard_ceiling_basis": "nlp_create_qkv_heads_decode_device_operation.cpp:51 asserts num_users <= 32. A TTNN op limit, unchanged by TP, and reached long before DRAM is."
+  },
+  "full_model_performance": {
+    "workload": "prompt 128 / generate 128 / batch 1, 48 layers, 1x4 P300_X2, FABRIC_1D_RING",
+    "source": "doc/full_model/probes/perf_full_model.csv and .json",
+    "ttft_ms_warmed": 126.695,
+    "ttft_ms_cold": 1115.157,
+    "decode_logits_only_ms": 20.211,
+    "decode_logits_only_tps_user": 49.479,
+    "decode_token_out_ms": 22.079,
+    "decode_token_out_tps_user": 45.292,
+    "decode_token_out_with_readback_ms": 22.748,
+    "decode_token_out_with_readback_tps_user": 43.959,
+    "layer_stack_lower_bound_ms": 20.573,
+    "layer_stack_lower_bound_basis": "48 x the stage-04 traced decode layer at ctx128, 0.4286 ms",
+    "ttft_cold_note": "first pass through a freshly opened mesh, so dominated by JIT kernel compilation and by the state of ~/.cache/ttnn. An earlier revision recorded 188.866 from a run with a warmer cache. Warmed TTFT is the served figure.",
+    "sampler_note": "greedy routes to Sampling1D's force-argmax strategy, measured at 1.125 ms against the split path's 6.155 ms on the same logits, both returning token 16. The 1.125 figure is after tt/model.py's _WatcherCleanSampling1D stopped pinning num_workers_per_link=1 on the argmax gather; it was 1.859 with the upstream spelling, which also tripped a device ASSERT under the watcher. See doc/full_model/README.md and watcher_ab.log.",
+    "watcher": "clean. TT_METAL_WATCHER=10 TT_METAL_WATCHER_DISABLE_ETH=1 pytest tests/ -m 'not models_performance_bare_metal' -q is 145 passed with zero tripped asserts (doc/full_model/pytest_watcher_clean.log.gz)."
+  },
+  "full_model_accuracy": {
+    "reference": "readiness_aime24_chat.refpt, AIME24 prompt 0, HF chat template, 158 prompt tokens, gen_len 100, top_k 100, generated fresh by stage 05",
+    "prefill": {
+      "top1": 0.98,
+      "top5": 1.0,
+      "top100": 1.0,
+      "log": "doc/full_model/run_prefill_check.log"
+    },
+    "decode_teacher_forced": {
+      "top1": 0.99,
+      "top5": 1.0,
+      "top100": 1.0,
+      "log": "doc/full_model/run_teacher_forcing.log"
+    },
+    "bar": "top5 >= 0.98 and top100 == 1.00; both met at 1.000"
+  },
+  "full_model_policy_note": "Stage 05 changed NOTHING about the decoder layer's dtype, fidelity, KV, activation or CCL policy, the paged-cache dtype split, the rejection ledger, or the inter-layer residual layout. The one edit to tt/multichip_decoder.py is an optional rope= parameter defaulting to the op it already called, and the stage-04 suite re-runs at 112 passed / 0 failed on this tree. What the full model adds is a replicated bf16 embedding (0.622 GB/die, no collective), a replicated final norm, a column-parallel bfloat8_b lm_head with zero vocabulary padding (151936 = 4 x 37984), growable replicated RoPE tables, and two captured traces plus their persistent inputs -- 0.006 GB/die of traces and buffers in total. Decode rotary moved from ttnn.experimental.rotary_embedding (Python-int position, unreplayable in a trace) to ttnn.experimental.rotary_embedding_hf over a device-gathered cos/sin pair; both are HF rotate_half and the swap is bit-identical at max|diff| 0.000e+00 and PCC 1.0 at every position tested, so the KV cache channel convention is untouched.",
+  "stage06_note": "Stage 06 changed NOTHING this file measures about capability. No capability reduction. What it changed, all three measured: (1) the greedy sampler stopped all-gathering the 151936-wide logit row and now reduces per die and all-gathers four candidate values and indices (tt/model.py, _WatcherCleanSampling1D._sample_argmax); (2) the PAGED SDPA-decode call got the program config it never had -- q_chunk_size 32, k_chunk_size min(256, per-user cache depth), max_cores_per_head_batch 16, memoised on the compute grid (tt/multichip_decoder.py, _sdpa_program_config); (3) that same reduction now runs over the live user rows instead of the 32 fixed sampler slots. KV-cache dtype, layout, paging and bytes/token are untouched; so are expert and attention weight dtypes, math fidelity, activation memory configs, the CCL policy (Topology.Ring, no num_workers_per_link pinned, 2 links prefill / 1 decode) and the inter-layer residual layout. tt/functional_decoder.py gained an sdpa_program_config= seam on attention_prefill that defaults to None and is passed None -- the prefill lever is built, measured and NOT adopted. Evidence: doc/optimized_full_model/README.md and work_log.md.",
+  "stage06_performance": {
+    "workload": "prompt 128 / generate 128 / batch 1, 48 layers, 1x4 P300_X2, FABRIC_1D_RING",
+    "source": "doc/optimized_full_model/probes/perf_full_model_p128_argmaxrows.json (128 timed reps, median)",
+    "ttft_ms_warmed": 125.431,
+    "ttft_ms_cold": 221.494,
+    "decode_logits_only_ms": 19.567,
+    "decode_logits_only_tps_user": 51.107,
+    "decode_token_out_ms": 19.693,
+    "decode_token_out_tps_user": 50.781,
+    "decode_token_out_with_readback_ms": 19.71,
+    "decode_token_out_with_readback_tps_user": 50.735,
+    "teacher_forcing_decode_tps_user": 42.25,
+    "teacher_forcing_note": "run_teacher_forcing uploads a forced token and reads the prediction back every step, so it is a correctness gate that prints a rate and is NOT the same measurement as token-out. It moved 38.50 -> 42.25 t/s/u (doc/full_model/run_teacher_forcing.log, doc/optimized_full_model/logs/run_teacher_forcing_argmaxrows.log).",
+    "against_stage05": {
+      "source": "doc/full_model/probes/perf_full_model.json",
+      "decode_token_out_ms": 22.079,
+      "decode_token_out_tps_user": 45.292,
+      "ttft_ms_warmed": 126.695,
+      "token_out_speedup": 1.1212,
+      "note": "the stage-05 run allocated a 4096-position KV cache and the stage-06 runs allocate 8192. SDPA-decode cost is independent of ALLOCATED depth and linear in cur_pos, measured directly at doc/optimized_full_model/probes/sdpa_depth_probe.json, so the comparison holds; the stage-06 before/after legs are all at 8192 and are like-for-like throughout."
+    },
+    "layer_stack_lower_bound_ms": 18.47,
+    "layer_stack_lower_bound_basis": "48 x the OPTIMIZED IN-MODEL per-layer device-kernel time, 384.791 us, from the verified one-iteration window at doc/optimized_full_model/ops_perf_full_model_48layer_decode.csv.gz (summary: probes/profile_summary_decode.json, regions_us.layer_stack). This REPLACES the stage-05 basis of 48 x 0.4286 ms = 20.573 ms, which multiplied a WALL figure for a one-layer traced model and so charged 48 layers for one iteration's dispatch overhead -- which is why stage 05 appeared to be under its own lower bound. The stage-04 layer's isolated device-kernel content is 362.83 us (doc/optimized_multichip_decoder/window_decode.txt); the in-model layer is 6.1% dearer than that.",
+    "bound_plus_terminal_ms": 18.889,
+    "gap_to_token_out_ms": 0.803,
+    "gap_to_token_out_percent": 4.08,
+    "gap_note": "dispatch and op-to-op gap across 3512 device ops -- 0.23 us each. The stage goal flags >10-15% as needing action; this is 4.08%.",
+    "prefill_profile": {
+      "source": "doc/optimized_full_model/probes/profile_summary_prefill.json",
+      "note": "stage 05 shipped with prefill unprofiled and disclosed it as a gap; stage 06 closes it. One verified 48-layer prefill of a 128-token prompt, boundary-checked by requiring the preceding pass to be the identical sequence of op codes row for row on all four devices, plus 56 per-device tallies.",
+      "device_kernel_ms": 122.921,
+      "share_of_ttft_percent": 98.0,
+      "expert_sparse_matmul_percent": 61.44,
+      "collectives_percent": 3.13,
+      "sdpa_percent": 0.58
+    },
+    "watcher": "clean. TT_METAL_WATCHER=10 TT_METAL_WATCHER_DISABLE_ETH=1 pytest tests/ -m 'not models_performance_bare_metal' -q is 145 passed with zero tripped asserts (doc/optimized_full_model/logs/watcher_argmaxrows.log.gz)."
+  },
+  "stage06_context_flatness": {
+    "claim": "The advertised 262144-token context is now USABLE, not merely allocatable. Before stage 06 decode cost grew 1.96x between a 128-token context and a 4096-token one; after it, 1.04x. That is the difference between a context you can hold and a context you can serve.",
+    "token_out_ms": {
+      "128": 19.6925,
+      "1024": 19.9787,
+      "4096": 20.505
+    },
+    "token_out_tps_user": {
+      "128": 50.781,
+      "1024": 50.053,
+      "4096": 48.768
+    },
+    "token_out_ms_before_stage06": {
+      "128": 21.4776,
+      "1024": 26.1432,
+      "4096": 42.0623
+    },
+    "ratio_4096_over_128": 1.0413,
+    "ratio_4096_over_128_before_stage06": 1.9584,
+    "cause": "the PAGED SDPA-decode call ran at the op default, whose cost is linear in cur_pos. With k_chunk_size=256 / max_cores_per_head_batch=16 it is nearly flat. Sources: doc/optimized_full_model/probes/sdpa_sweep_confirm_bf16.json (the op) and the three perf_full_model_p*_{before,argmaxrows}.json pairs (the model).",
+    "measured_to_context_tokens": 4096,
+    "measured_to_context_basis": "END TO END, this is how deep it was actually measured: prompt 4096, generate 128, batch 1, through the real 48-layer model, 128 timed reps (doc/optimized_full_model/probes/perf_full_model_p4096_argmaxrows.json). Decode beyond a 4096-token context has NOT been run end to end at 48 layers. The 262144 claim beyond that point rests on the three narrower measurements below, and this field exists so nobody reads the flatness claim as an end-to-end measurement at the advertised context.",
+    "op_level_evidence_to_cur_pos": 32767,
+    "op_level_evidence_basis": "the shipped SDPA-decode configuration measured at the real per-die decode shapes and the real bfloat16 cache dtype out to cur_pos 32767: 74.13 us against the op default's 3545.05 us, a 47.8x, with the configured leg's PCC holding at 0.9997 where the default's has decayed to 0.9897. doc/optimized_full_model/probes/sdpa_sweep_confirm_bf16.json. This is the op, not the model.",
+    "in_model_pcc_to_cur_pos": 16383,
+    "in_model_pcc_basis": "the real multichip layer against a HuggingFace reference with a prefill-primed paged cache: PCC 0.999363 at cur_pos 16383 on the adopted path (0.999294 on the default). doc/optimized_full_model/logs/sdpa_hf_pcc_at_depth_deep.log; the 128/1024/4096 legs are in probes/sdpa_hf_pcc_at_depth.json. 32768 was attempted and the HOST was OOM-killed building the torch reference, not the device -- the same reference-side ceiling pcc_validation_note records.",
+    "capacity_evidence": "the full 262144-token paged KV cache is allocated and held by the running 48-layer model, with both decode traces captured and a token run through it: doc/optimized_full_model/probes/footprint_262144.json.",
+    "still_not_measured": "a 262144-token PREFILL through all 48 layers (prefill is single-shot, and the longest 48-layer prompt run is 4096); end-to-end decode at contexts above 4096; and any of this at batch > 1. These are coverage gaps in the evidence, not capability reductions, and they are named as limitations in doc/optimized_full_model/README.md."
+  },
+  "stage06_measured": {
+    "method": "doc/optimized_full_model/probes/footprint_probe.py --context 262144, re-run on the SHIPPED tree. Same probe as stage 05's, copied rather than imported so it writes beside this stage's evidence. It builds the real model -- real weights, real embed_tokens, real lm_head, the real paged KV cache at the advertised context, the real RoPE tables -- captures both decode traces and runs a token through it. GB = 1e9, as the allocator reports it.",
+    "weights_embed_lm_head_rope_gb_per_die": 5.311,
+    "kv_cache_at_262144_batch1_gb_per_die": 6.443,
+    "traces_and_persistent_buffers_gb_per_die": 0.006,
+    "total_gb_per_die": 11.76,
+    "headroom_gb_per_die": 22.119,
+    "dram_per_die_gb_reported": 33.879,
+    "sum_note": "the three stage rows in doc/optimized_full_model/probes/footprint_262144.json sum to 11.759906816, bit-identical to total_gb_per_die. Any residual visible when the 3dp values above are added by hand is this file's display rounding and nothing else; doc/optimized_full_model/probes/check_published_figures.py asserts the raw rows equal the raw total exactly.",
+    "kv_formula_check": "1 local kv head * head_dim(128) * 2 bytes * 2 tensors (K and V) = 512 B/token/layer/die, x 48 layers x 262144 tokens = 6.442450944 GB/die of KV at batch 1 -- which is what the allocator reports, so the paged cache carries no per-block overhead worth a digit. Unchanged by stage 06: the KV dtype, the block size and the head split are all untouched.",
+    "comparison_to_stage05": "stage 05 measured 11.759 GB/die total and 22.119 free for the same contents. Stage 06 adds one small constant to the model -- the distributed argmax's per-die vocabulary offset, a [1,1,rows,4] int32 sharded to one tile-padded column per die -- and no new activation, cache or collective buffer. Any difference between the two totals is that constant plus allocator placement, and it is far below the headroom either way.",
+    "rope_table_note": "the probe builds the model with the default rope_cache_len of 8192, so the RoPE tables inside weights_embed_lm_head_rope_gb_per_die are the 8192-row pair (0.004 GB/die). Growing them to the full 262144 context costs 0.134 GB/die (262144 * 128 * 2 B * 2 tables). No capability reduction; the tables grow on demand via Qwen3CoderModel.ensure_rope_capacity, and tt/generator.py's decode_forward takes decode_horizon= so a low-level caller grows them once before any trace is captured."
+  },
+  "stage06_runtime_audit": {
+    "source": "doc/optimized_full_model/probes/runtime_fallback_audit.json",
+    "host_logit_readback_on_token_out_path": false,
+    "host_argmax_on_token_out_path": false,
+    "sampling_greedy": "Sampling1D force-argmax, distributed: per-die untilize/argmax/gather -> all-gather 4 candidates -> masked-min, traced, writes tt_out_tok",
+    "sdpa_decode_program_config": "SDPAProgramConfig(compute_with_storage_grid_size=(11, 10), q_chunk_size=32, k_chunk_size=256, max_cores_per_head_batch=16)",
+    "sdpa_decode_k_chunk_clamped_at_shipped_context": false,
+    "sdpa_prefill_program_config_passed": "None",
+    "steady_state_only_replays_moved": true,
+    "note": "runtime_fallback_audit() itself does not yet carry the two properties stage 06 introduced -- the paged SDPA program config and the sampler's live-row count -- because adding fields would change a dict tests/test_full_model.py::test_runtime_fallback_audit_is_clean pins field by field. They are recorded here and in the JSON above instead, read off the modules that own them."
+  },
+  "stage07_note": "Stage 07 (the datatype sweep) changed NOTHING this file measures about capability. No capability reduction; advertised and supported context both remain 262144. The selected precision config moves only the two expert matmul inner block widths (experts_gate_up_in0_block_w 16 -> 64, experts_down_in0_block_w 12 -> 24). A block width is a program-config field, not a tensor dtype or shape: device_expert_bytes_per_die is identical (84 934 656 B) before and after, no allocation moves, and the KV cache dtype, layout, paging and bytes/token are untouched. Evidence: doc/datatype_sweep/README.md and work_log.md, doc/datatype_sweep/sweep_results.json.",
+  "stage07_kv_bfp8_candidate": {
+    "status": "measured, not selected -- the shipped kv_cache_dtype stays bfloat16",
+    "why_recorded": "This is the only candidate in the sweep that could move capacity, so the contract records what it would buy even though it was not taken. It moves capacity UPWARD (fewer bytes per token), so it could never have forced a reduction. Its first measurement was invalid -- a bfloat8_b paged cache filled from this model's bfloat16 K/V reads back as NaN, so the row scored at chance -- and the prefill writer now casts K/V to the cache dtype (tt/functional_decoder.match_cache_dtype). The numbers below are from the post-fix runs.",
+    "kv_cache_dtype": "bfloat8_b",
+    "bytes_per_token_per_layer": 1088,
+    "bytes_per_token_per_layer_per_die": 272,
+    "bytes_per_token_per_layer_per_die_bfloat16": 512,
+    "formula": "num_key_value_heads_per_die(1) * head_dim(128) * 1.0625 B/elem * 2 tensors (K and V); 1.0625 is bfloat8_b once each 16-element block's shared exponent byte is counted",
+    "device_kv_bytes_at_full_context_per_die": 3422552064,
+    "device_kv_bytes_at_full_context_per_die_bfloat16": 6442450944,
+    "device_kv_bytes_saved_per_die": 3019898880,
+    "device_kv_gb_at_full_context_per_die": 3.423,
+    "device_kv_gb_at_full_context_per_die_bfloat16": 6.442,
+    "share_of_the_11_759_gb_per_die_footprint_bfloat16": 0.548,
+    "measured_rows": {
+      "R19_kv_bfp8": "bfp8 KV as a delta from the stage-06 baseline: top-1 0.980 / top-5 1.000 / top-100 1.000, decode 42.29 t/s/u (-0.12% vs 42.34, inside the 0.368% band), TTFT 7842.57 ms against 3250.90",
+      "R28_kv_bfp8_bw64_24": "bfp8 KV on top of the selected block widths: top-1 0.980 / top-5 1.000 / top-100 1.000, decode 43.45 t/s/u (-0.21% vs the selected 43.54, inside the band), TTFT 7792.06 ms against 3287.61"
+    },
+    "verdict": "Not selected. bfp8 KV buys 3.020 GB/die of headroom at 262144 tokens -- real capacity -- but it buys NO decode throughput (both rows land inside the run-to-run band), costs one top-1 point, and costs 2.4x TTFT (7.8 s against 3.3 s on the 158-token gate prompt). The stage ranks on decode and the advertised context is already met at bfloat16 with 22.119 GB/die free, so there is no capacity problem for this to solve. It is recorded here so that a future stage needing more KV headroom -- a larger batch, or a longer served context -- has the price already measured. The TTFT cost would need explaining first; it is not the cast (it was present before the cast existed) and is unattributed.",
+    "evidence": [
+      "doc/datatype_sweep/sweep_results.json (rows R19_kv_bfp8, R28_kv_bfp8_bw64_24)",
+      "doc/datatype_sweep/probes/kv_bfp8_diagnosis.json (both cache writers, six cache/input combinations)",
+      "doc/datatype_sweep/logs/rows/R19_kv_bfp8.log, R28_kv_bfp8_bw64_24.log"
+    ]
+  },
+  "stage11_note": "Stage 11 (batch scaling, doc/batch_scaling/) added the variable-width decode ladder and re-validated the context bound under the SHIPPED SERVING configuration -- vLLM KV pool 263168, max_num_seqs 32, QWEN3_DECODE_WIDTHS 1,2,4,8,16,32 -- which is a different setup from the standalone batch-1 probe recorded in `measured` above. current_supported_context is unchanged at 262144 and there is still no capability reduction: `get_max_tokens_all_users` serves min(max_model_len, 262144). The fields below record exactly how far that bound has been exercised in serving, so that this file -- which tt/generator_vllm.py reads as the source of truth -- and doc/batch_scaling/README.md's limitations state the same thing rather than each carrying half of it.",
+  "stage11_serving_context": {
+    "bound_tokens": 256000,
+    "bound_blocks": 8000,
+    "bound_rationale": "Block-aligned (8000 blocks x 32) and 3,000 tokens below the largest length actually validated in serving (259,000). 259,000 is not block-aligned (8093.75 blocks); the largest aligned value at or below it is 258,976, but that leaves NO margin under a cliff whose onset is not located and whose mechanism is unresolved. 256,000 is a clean figure for a catalog, sits 6,136 tokens below the known-bad 262,136, and costs 1.2 % of the validated range.",
+    "validated_by": "a real vLLM server, prompt sent as exact token ids",
+    "evidence": [
+      "doc/batch_scaling/probes/long_prompt_t120_253k.json",
+      "doc/batch_scaling/probes/long_prompt_t120_259k.json",
+      "doc/batch_scaling/probes/long_prompt_t120_full.json",
+      "doc/batch_scaling/logs/bd_engine_len259.log"
+    ],
+    "serving_measurements": {
+      "253000": {
+        "elapsed_s": 1049.3,
+        "ok": true
+      },
+      "259000": {
+        "elapsed_s": 1130.8,
+        "ok": true,
+        "note": "O(S^2) predicted 1127 s; 0.3 % error"
+      },
+      "262136": {
+        "elapsed_s": null,
+        "ok": false,
+        "note": "did not complete in 3600 s; one blocking prefill; no watchdog throw"
+      }
+    },
+    "conditional_on_env": {
+      "TT_METAL_OPERATION_TIMEOUT_SECONDS": "120.0",
+      "why": "This bound is only reachable with the tt-metal dispatch STALL watchdog set generously. It is not an elapsed-time limit: loop_and_wait_with_timeout resets its clock whenever the device dispatch progress counter moves, and fires only after that counter has been still for the whole duration. tt-metal's own default is 0.0 (off). tt-inference-server/vllm-tt-metal/src/run_vllm_api_server.py:521 sets 5.0 unconditionally, and at 5.0 a 253,000-token prompt KILLED a deployment and required tt-smi -r. Measured: 259,000 passes at 120.0 and at 30.0, and FAILS at 10.0 (fired at 17.7 s). Recommended 120.0 -- see timeout_recommendation."
+    },
+    "timeout_recommendation": {
+      "value": 120.0,
+      "largest_legitimate_gap_bracketed_between_s": [
+        10.0,
+        30.0
+      ],
+      "rationale": "A genuine hang NEVER advances dispatch progress, so ANY finite timeout detects it -- a larger value costs detection LATENCY, not detection ability. The costs are therefore asymmetric: too low misdiagnoses a legitimate long op as a hang, kills the server and requires tt-smi -r (observed); too high only delays detection of a real hang. 120.0 is 4x the largest known-passing value (30.0) and still catches a genuine hang within two minutes. DO NOT disable the watchdog: this port has a documented NoC hang (PRESERVE_DECODE_TRACES) that this watchdog is the thing to catch."
+    },
+    "standalone_result_is_separate": "The 262144-token result in `measured` is STANDALONE (batch 1, direct generator calls, no vLLM) and remains correct as such. It is not a serving result and must not be quoted as one.",
+    "correction": "A previous version of this block published serving_validated_prefill_tokens 262080 with a 64-token gap, citing kv_pool_ceiling_sixwidth_production.json as evidence and describing a vLLM serving_config. That probe NEVER went through vLLM -- its own `what` field says 'measured through configure_paging rather than vLLM' and kv_pool_ceiling.py calls gen.prefill_forward() directly. A standalone single-user measurement was labelled serving-validated. The 4x timing discrepancy that should have exposed it was sitting in this same file unexamined: 192 s standalone at 262144 against 387 s SERVED at 168,901 tokens, i.e. serving prefill is ~4x slower per token. The real serving bound is 256000, not 262080.",
+    "why_our_own_runs_never_saw_this": "Every serving measurement in doc/batch_scaling/ was taken with TT_METAL_OPERATION_TIMEOUT_SECONDS UNSET, i.e. tt-metal's default 0.0 = watchdog DISABLED, because models/common/readiness_check does not set it. The TTI deployment sets 5.0. Our serving numbers were therefore measured under a different safety configuration from the one users actually deploy.",
+    "unresolved": {
+      "what": "the prefill cliff between 259,000 and 262,136 tokens",
+      "leading_hypothesis": "KV-pool edge. The pool is 263,168 tokens; a 262,144-token request leaves 1,024 tokens (32 blocks) while 259,000 leaves 4,168. The sharpness -- quadratic to within 0.3 % at 259,000, then >=3.5x the prediction 3.6 % later -- fits an allocation-constrained step.",
+      "second_hypothesis": "An L1/memory-config step function in the prefill path, analogous to the measured width-19 decode cliff where _decode_expert_memory_config flips L1->DRAM.",
+      "what_is_already_ruled_out": "Preemption/recompute THRASH. Engine telemetry shows ZERO loggers.py lines across the 59-minute attempt, and the scheduler runs (and logs) between steps, so a preempt/recompute loop would have emitted a line per interval. This does NOT rule out a SINGLE allocation-constrained step -- do not re-run this refutation believing it closed the question.",
+      "why_not_chased": "At 253,000 tokens TTFT is already 17.5 minutes, so no interactive deployment operates near the cliff; the usable range is bounded by latency far below it. Locating the mechanism would have cost 2-3 further 20-60 minute board runs for near-zero operational value."
+    }
+  }
+}

code/models/demos/blackhole/qwen3_coder_30b_a3b/config/selected_precision_config.json

@@ -0,0 +1,22 @@
+{
+  "activation_dtype": "bfloat16",
+  "attention_fidelity": null,
+  "attention_qkv_dtype": "bfloat8_b",
+  "attention_wo_dtype": "bfloat8_b",
+  "ccl_dtype": null,
+  "embedding_dtype": "bfloat16",
+  "experts_down_dtype": "bfloat4_b",
+  "experts_down_in0_block_w": 24,
+  "experts_fidelity": "LoFi",
+  "experts_gate_up_dtype": "bfloat4_b",
+  "experts_gate_up_in0_block_w": 64,
+  "kv_cache_dtype": "bfloat16",
+  "lm_head_dtype": "bfloat8_b",
+  "lm_head_fidelity": "HiFi2",
+  "logits_dtype": "bfloat16",
+  "norm_fidelity": "HiFi4",
+  "norm_weight_dtype": "bfloat16",
+  "router_dtype": "bfloat16",
+  "router_window_fidelity": "HiFi4",
+  "sampling_dtype": "bfloat16"
+}

code/models/demos/blackhole/qwen3_coder_30b_a3b/tests/reference.py

@@ -0,0 +1,145 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+"""Layer-only HuggingFace reference for Qwen3-Coder-30B-A3B-Instruct.
+
+Loading the full 30.5B causal LM to test one decoder layer wastes ~57GB of host
+RAM and several minutes, so this reads just the tensors for a single layer
+straight out of the safetensors shards and populates one ``Qwen3MoeDecoderLayer``.
+
+Checkpoint-vs-module weight layout
+----------------------------------
+The checkpoint stores experts as 3 separate tensors per expert::
+
+    model.layers.L.mlp.experts.E.gate_proj.weight   [moe_inter, hidden]
+    model.layers.L.mlp.experts.E.up_proj.weight     [moe_inter, hidden]
+    model.layers.L.mlp.experts.E.down_proj.weight   [hidden, moe_inter]
+
+``Qwen3MoeExperts`` instead holds them batched and gate/up fused::
+
+    gate_up_proj  [num_experts, 2 * moe_inter, hidden]   # [gate ; up] along dim 0
+    down_proj     [num_experts, hidden, moe_inter]
+
+The fusion order matters: ``Qwen3MoeExperts.forward`` does
+``linear(x, gate_up_proj[e]).chunk(2, dim=-1)``, so the FIRST half is gate and
+the second is up. Concatenating in the other order silently swaps them and the
+layer still runs -- it just produces wrong numbers.
+"""
+
+from __future__ import annotations
+
+import json
+from collections import defaultdict
+
+import torch
+from huggingface_hub import hf_hub_download
+from safetensors import safe_open
+from transformers import AutoConfig
+from transformers.models.qwen3_moe.modeling_qwen3_moe import Qwen3MoeDecoderLayer, Qwen3MoeRotaryEmbedding
+
+HF_MODEL = "Qwen/Qwen3-Coder-30B-A3B-Instruct"
+
+
+def load_config(hf_model: str = HF_MODEL):
+    return AutoConfig.from_pretrained(hf_model)
+
+
+def layer_state_dict(layer_idx: int = 0, hf_model: str = HF_MODEL) -> dict[str, torch.Tensor]:
+    """Read only ``model.layers.<layer_idx>.*`` from the shards that hold them."""
+    index = json.load(open(hf_hub_download(hf_model, "model.safetensors.index.json")))["weight_map"]
+    prefix = f"model.layers.{layer_idx}."
+
+    per_shard: dict[str, list[str]] = defaultdict(list)
+    for name, shard in index.items():
+        if name.startswith(prefix):
+            per_shard[shard].append(name)
+    if not per_shard:
+        raise KeyError(f"no tensors found for layer {layer_idx}")
+
+    out: dict[str, torch.Tensor] = {}
+    for shard, names in per_shard.items():
+        path = hf_hub_download(hf_model, shard)
+        with safe_open(path, framework="pt") as f:
+            for name in names:
+                out[name[len(prefix) :]] = f.get_tensor(name)
+    return out
+
+
+def build_reference_layer(layer_idx: int = 0, hf_model: str = HF_MODEL):
+    """Return ``(layer, config)`` with real checkpoint weights, in eval mode."""
+    config = load_config(hf_model)
+    sd = layer_state_dict(layer_idx, hf_model)
+
+    with torch.device("meta"):
+        layer = Qwen3MoeDecoderLayer(config, layer_idx)
+    layer.to_empty(device="cpu")
+
+    direct = [
+        "input_layernorm.weight",
+        "post_attention_layernorm.weight",
+        "self_attn.q_proj.weight",
+        "self_attn.k_proj.weight",
+        "self_attn.v_proj.weight",
+        "self_attn.o_proj.weight",
+        "self_attn.q_norm.weight",
+        "self_attn.k_norm.weight",
+        "mlp.gate.weight",
+    ]
+    params = dict(layer.named_parameters())
+    for key in direct:
+        params[key].data.copy_(sd[key])
+
+    # Fuse + stack the experts. gate first, then up -- see module docstring.
+    n_experts = config.num_experts
+    gate_up = torch.stack(
+        [
+            torch.cat([sd[f"mlp.experts.{e}.gate_proj.weight"], sd[f"mlp.experts.{e}.up_proj.weight"]], dim=0)
+            for e in range(n_experts)
+        ]
+    )
+    down = torch.stack([sd[f"mlp.experts.{e}.down_proj.weight"] for e in range(n_experts)])
+    params["mlp.experts.gate_up_proj"].data.copy_(gate_up)
+    params["mlp.experts.down_proj"].data.copy_(down)
+
+    return layer.eval(), config
+
+
+def rotary_embeddings(config, seq_len: int, device="cpu"):
+    """Return the ``(cos, sin)`` pair the decoder layer expects."""
+    rope = Qwen3MoeRotaryEmbedding(config=config, device=device)
+    position_ids = torch.arange(seq_len, device=device).unsqueeze(0)
+    dummy = torch.zeros(1, seq_len, config.hidden_size, dtype=torch.float32, device=device)
+    return rope(dummy, position_ids)
+
+
+def weight_stats(sd: dict[str, torch.Tensor]) -> dict[str, dict]:
+    """Per-tensor name/shape/dtype/mean/std, for deterministic synthetic weights."""
+    stats = {}
+    for name, t in sd.items():
+        f = t.float()
+        stats[name] = {
+            "shape": list(t.shape),
+            "dtype": str(t.dtype),
+            "mean": f.mean().item(),
+            "std": f.std().item(),
+        }
+    return stats
+
+
+if __name__ == "__main__":
+    torch.manual_seed(0)
+
+    layer, config = build_reference_layer(0)
+    n_params = sum(p.numel() for p in layer.parameters())
+    print(f"layer 0 built: {n_params/1e9:.2f}B params, dtype={next(layer.parameters()).dtype}")
+
+    seq_len = 32
+    hidden = torch.randn(1, seq_len, config.hidden_size, dtype=torch.float32) * 0.02
+    cos, sin = rotary_embeddings(config, seq_len)
+
+    with torch.no_grad():
+        out = layer(hidden, position_embeddings=(cos, sin), attention_mask=None)
+    out = out[0] if isinstance(out, tuple) else out
+
+    print(f"forward OK: {tuple(hidden.shape)} -> {tuple(out.shape)}")
+    print(f"  out mean={out.mean():.6f} std={out.std():.6f} finite={torch.isfinite(out).all().item()}")

code/models/demos/blackhole/qwen3_coder_30b_a3b/tests/test_attention.py

@@ -0,0 +1,117 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+"""TTNN attention vs the HuggingFace reference layer's ``self_attn``.
+
+The PCC test is the real check, but two cheap host-side tests run first because
+they isolate failures that PCC alone reports as one undifferentiated low number:
+
+  * ``test_qk_norm_is_applied`` -- proves the per-head norm actually changed the
+    tensor. A QK-norm quietly skipped (wrong key, missing weight) still yields
+    a plausible ~0.9 PCC that looks like ordinary bf16 loss.
+  * ``test_causality`` -- perturbing a late token must not move an early one.
+    A non-causal run scores high PCC on short sequences, so ``is_causal`` being
+    dropped is otherwise invisible here and only surfaces as garbage generation
+    much later.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+from loguru import logger
+
+import ttnn
+from models.common.utility_functions import comp_allclose, comp_pcc
+
+from ..tt.functional_decoder import AttentionConfig, attention_prefill, build_rope_cache, upload_attention_weights
+from ..tt.weight_mapping import convert_attention_weights
+from .reference import build_reference_layer, layer_state_dict, rotary_embeddings
+
+LAYER_IDX = 0
+PCC_REQUIRED = 0.99  # attention_bias=False, so no large-bias PCC degradation applies
+
+
+@pytest.fixture(scope="module")
+def reference():
+    return build_reference_layer(LAYER_IDX)
+
+
+@pytest.fixture(scope="module")
+def torch_weights():
+    return convert_attention_weights(
+        {k: v for k, v in layer_state_dict(LAYER_IDX).items()},
+        n_heads=32,
+        n_kv_heads=4,
+        head_dim=128,
+    )
+
+
+def _hidden(config, seq_len, seed=0):
+    torch.manual_seed(seed)
+    return torch.randn(1, seq_len, config.hidden_size, dtype=torch.float32) * 0.02
+
+
+def _causal_mask(seq_len):
+    return torch.full((seq_len, seq_len), float("-inf")).triu(1).reshape(1, 1, seq_len, seq_len)
+
+
+def _reference_attention(layer, config, hidden):
+    seq_len = hidden.shape[1]
+    cos, sin = rotary_embeddings(config, seq_len)
+    with torch.no_grad():
+        out = layer.self_attn(
+            hidden_states=hidden,
+            position_embeddings=(cos, sin),
+            attention_mask=_causal_mask(seq_len),
+        )
+    return out[0] if isinstance(out, tuple) else out
+
+
+def test_qk_norm_weights_are_non_trivial(torch_weights):
+    """A QK-norm weight of all ones would make the norm undetectable in PCC."""
+    for name in ("q_norm", "k_norm"):
+        w = torch_weights[name]
+        assert w.shape == (128,), f"{name} has shape {tuple(w.shape)}, expected (head_dim,)"
+        assert not torch.allclose(w, torch.ones_like(w)), f"{name} is all ones -- cannot detect a skipped norm"
+
+
+def test_causality(reference):
+    """Changing the last token must leave earlier outputs untouched."""
+    layer, config = reference
+    hidden = _hidden(config, 32)
+    baseline = _reference_attention(layer, config, hidden)
+
+    perturbed = hidden.clone()
+    perturbed[:, -1, :] += 1.0
+    after = _reference_attention(layer, config, perturbed)
+
+    assert torch.allclose(baseline[:, :-1], after[:, :-1], atol=1e-5), "reference attention is not causal"
+
+
+@pytest.mark.parametrize("mesh_device", [(1, 1)], ids=["1x1"], indirect=True)
+@pytest.mark.parametrize("seq_len", [32, 128, 512], ids=["s32", "s128", "s512"])
+def test_attention_prefill_vs_reference(mesh_device, reference, torch_weights, seq_len):
+    layer, hf_config = reference
+    config = AttentionConfig.from_hf(hf_config)
+    hidden = _hidden(hf_config, seq_len)
+
+    ref_out = _reference_attention(layer, hf_config, hidden)
+
+    weights = upload_attention_weights(torch_weights, mesh_device)
+    cos_cache, sin_cache = build_rope_cache(hf_config, seq_len, mesh_device)
+
+    tt_in = ttnn.from_torch(
+        hidden.unsqueeze(0),  # [1, 1, S, hidden]
+        dtype=ttnn.bfloat16,
+        layout=ttnn.TILE_LAYOUT,
+        device=mesh_device,
+        memory_config=ttnn.DRAM_MEMORY_CONFIG,
+    )
+    tt_out = attention_prefill(tt_in, weights, config, cos_cache, sin_cache)
+    tt_out_torch = ttnn.to_torch(tt_out).squeeze(0)
+
+    passing, pcc_message = comp_pcc(ref_out, tt_out_torch, PCC_REQUIRED)
+    logger.info(comp_allclose(ref_out, tt_out_torch))
+    logger.info(f"attention prefill seq={seq_len}: {pcc_message}")
+    assert passing, f"attention prefill (seq={seq_len}) below {PCC_REQUIRED}: {pcc_message}"

code/models/demos/blackhole/qwen3_coder_30b_a3b/tests/test_attention_decode.py

@@ -0,0 +1,140 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+"""Single-token decode attention against the KV cache.
+
+Decode is validated against the *prefill reference*, not against a separate
+decode reference: prefill token S and decode token S must produce the same
+output, because causal attention at position S sees exactly the same context
+either way. That equivalence is the whole contract of a KV cache, so testing it
+directly catches the errors that matter -- an off-by-one write position, RoPE
+applied at the wrong index, or a cache that was never seeded by prefill.
+
+This is also the first exercise of the Blackhole ``nlp_create_qkv_heads_decode``
+DRAM bug (tt-metal #16667), which zeroes odd-indexed Q rows. The workaround
+lives in ``attention_decode``; ``test_decode_q_rows_are_all_live`` is the
+regression guard, because with half of Q zeroed the output is still finite,
+still plausible, and still scores a deceptively high PCC.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+from loguru import logger
+
+import ttnn
+from models.common.utility_functions import comp_allclose, comp_pcc
+
+from ..tt.functional_decoder import (
+    AttentionConfig,
+    attention_decode,
+    attention_prefill,
+    build_rope_cache,
+    create_kv_cache,
+    upload_attention_weights,
+)
+from ..tt.weight_mapping import convert_attention_weights
+from .reference import build_reference_layer, layer_state_dict, rotary_embeddings
+
+LAYER_IDX = 0
+PCC_REQUIRED = 0.99
+MAX_SEQ = 256
+
+
+@pytest.fixture(scope="module")
+def reference():
+    return build_reference_layer(LAYER_IDX)
+
+
+@pytest.fixture(scope="module")
+def torch_weights():
+    return convert_attention_weights(layer_state_dict(LAYER_IDX), n_heads=32, n_kv_heads=4, head_dim=128)
+
+
+def _hidden(hf_config, seq_len, seed=0):
+    torch.manual_seed(seed)
+    return torch.randn(1, seq_len, hf_config.hidden_size, dtype=torch.float32) * 0.02
+
+
+def _reference_attention(layer, hf_config, hidden):
+    seq_len = hidden.shape[1]
+    cos, sin = rotary_embeddings(hf_config, seq_len)
+    mask = torch.full((seq_len, seq_len), float("-inf")).triu(1).reshape(1, 1, seq_len, seq_len)
+    with torch.no_grad():
+        out = layer.self_attn(hidden_states=hidden, position_embeddings=(cos, sin), attention_mask=mask)
+    return out[0] if isinstance(out, tuple) else out
+
+
+def _to_device(t, mesh_device, dtype=ttnn.bfloat16):
+    return ttnn.from_torch(
+        t, dtype=dtype, layout=ttnn.TILE_LAYOUT, device=mesh_device, memory_config=ttnn.DRAM_MEMORY_CONFIG
+    )
+
+
+def _prefill_then_decode(mesh_device, hf_config, torch_weights, hidden_full, prompt_len, block_size=None):
+    """Prefill ``prompt_len`` tokens, then decode the token at ``prompt_len``."""
+    config = AttentionConfig.from_hf(hf_config)
+    weights = upload_attention_weights(torch_weights, mesh_device)
+    cos_cache, sin_cache = build_rope_cache(hf_config, MAX_SEQ, mesh_device)
+    kv_cache = create_kv_cache(mesh_device, config, max_batch=1, max_seq_len=MAX_SEQ, block_size=block_size)
+
+    prompt = hidden_full[:, :prompt_len, :]
+    attention_prefill(_to_device(prompt.unsqueeze(0), mesh_device), weights, config, cos_cache, sin_cache, kv_cache)
+
+    # [1, 1, batch=1, hidden]
+    next_tok = hidden_full[:, prompt_len, :].reshape(1, 1, 1, hf_config.hidden_size)
+    current_pos = ttnn.from_torch(torch.tensor([prompt_len], dtype=torch.int32), dtype=ttnn.int32, device=mesh_device)
+    out = attention_decode(
+        _to_device(next_tok, mesh_device),
+        weights,
+        config,
+        cos_cache,
+        sin_cache,
+        kv_cache,
+        current_pos,
+        token_index=prompt_len,
+    )
+    return ttnn.to_torch(out).reshape(1, hf_config.hidden_size)
+
+
+@pytest.mark.parametrize("mesh_device", [(1, 1)], ids=["1x1"], indirect=True)
+@pytest.mark.parametrize("prompt_len", [32, 128], ids=["p32", "p128"])
+@pytest.mark.parametrize("block_size", [None, 32, 64], ids=["contiguous", "paged32", "paged64"])
+def test_decode_matches_prefill_at_same_position(mesh_device, reference, torch_weights, prompt_len, block_size):
+    layer, hf_config = reference
+    hidden_full = _hidden(hf_config, prompt_len + 1)
+
+    # Reference: run the whole prompt+1 through prefill and take the last row.
+    ref_out = _reference_attention(layer, hf_config, hidden_full)[:, prompt_len, :]
+
+    tt_out = _prefill_then_decode(mesh_device, hf_config, torch_weights, hidden_full, prompt_len, block_size)
+
+    passing, pcc_message = comp_pcc(ref_out, tt_out, PCC_REQUIRED)
+    logger.info(comp_allclose(ref_out, tt_out))
+    kind = "contiguous" if block_size is None else f"paged(block={block_size})"
+    logger.info(f"decode at pos {prompt_len} [{kind}]: {pcc_message}")
+    assert passing, f"decode at position {prompt_len} [{kind}] below {PCC_REQUIRED}: {pcc_message}"
+
+
+@pytest.mark.parametrize("mesh_device", [(1, 1)], ids=["1x1"], indirect=True)
+def test_decode_q_rows_are_all_live(mesh_device, reference, torch_weights):
+    """Guard for tt-metal #16667: no systematic zeroing of alternating rows.
+
+    The Blackhole bug zeroes odd-indexed Q rows when the fused QKV is read from
+    DRAM. Rather than reach into the op, this checks the observable
+    consequence: the head-dim structure of the output must not show a
+    stripe of exactly-zero alternating entries.
+    """
+    layer, hf_config = reference
+    hidden_full = _hidden(hf_config, 33)
+    out = _prefill_then_decode(mesh_device, hf_config, torch_weights, hidden_full, 32).float()
+
+    assert torch.isfinite(out).all(), "decode produced non-finite values"
+    zero_fraction = (out == 0).float().mean().item()
+    logger.info(f"decode output zero fraction = {zero_fraction:.4f}")
+    assert zero_fraction < 0.1, (
+        f"{zero_fraction:.1%} of decode outputs are exactly zero -- looks like the "
+        "Blackhole nlp_create_qkv_heads_decode DRAM bug (#16667); the fused QKV "
+        "must be staged through L1 before the split"
+    )

code/models/demos/blackhole/qwen3_coder_30b_a3b/tests/test_decode_compaction_fifo.py

@@ -0,0 +1,596 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+"""The async boundary the compaction un-permutation has to survive.
+
+`Qwen3CoderForCausalLM` may decode a compacted batch in a graph narrower than
+`max_num_seqs`, which means the sampled tokens come back in **graph-row** order
+and have to be scattered back to vLLM's slots. The mapping that does that is
+chosen when a forward is *issued*; the scatter happens when the output is
+*read*. Under `--async-scheduling` those are different steps.
+
+So the hazard is not the permutation arithmetic -- that is covered by
+`doc/batch_scaling/probes/compaction_identity.py` on real weights -- it is the
+**pairing**: if a later step installs a new mapping before an earlier step's
+tokens are read, the earlier tokens get scattered with the wrong permutation and
+every one of them lands on the wrong request. Silently, and as a correctness
+bug rather than a slowdown.
+
+The vLLM plugin happens to order this safely today (it drains pending async
+decodes on a layout change, and only a layout change can move the mapping), but
+that invariant lives in a repository this one must not modify and cannot pin.
+`_pending_orders` removes the dependency by pairing each output with the mapping
+its own forward used, and these tests are what hold that property in place.
+
+Deliberately **device-free**: the whole hazard is adapter bookkeeping, so a fake
+generator exercises it exactly and the tests run in milliseconds. Three earlier
+pieces of evidence -- a real-weights identity probe, a 158-test suite and a
+serving A/B -- all passed while this bug was present, because every one of them
+either used the synchronous read path or never moved the mapping. That is the
+gap these tests close.
+"""
+
+from __future__ import annotations
+
+from unittest.mock import patch
+
+import pytest
+import torch
+
+from models.demos.blackhole.qwen3_coder_30b_a3b.tt import generator_vllm as gv
+from models.demos.blackhole.qwen3_coder_30b_a3b.tt.generator import Qwen3CoderGenerator
+from models.demos.blackhole.qwen3_coder_30b_a3b.tt.generator_vllm import Qwen3CoderForCausalLM
+
+SLOTS = 32
+
+
+class _Handle:
+    """Stands in for the device tensor a decode forward returns."""
+
+    def __init__(self, tag: int):
+        self.tag = tag
+
+
+class _FakeGenerator:
+    """Only the surface `decode_forward` / `process_decode_output_host` touch.
+
+    `read_sampled_tokens` returns a vector that encodes **which forward** it came
+    from and **which graph row** each entry is, so a mis-paired scatter is
+    visible in the values rather than having to be inferred.
+    """
+
+    def __init__(self):
+        self.model = object()
+        self.mesh_device = object()
+        self.pages_per_user = 8
+        self.page_block_size = 32
+        self.num_blocks = SLOTS * 8
+        self.issued: list[_Handle] = []
+        self.trace_stats: dict = {}
+
+    # -- the calls the adapter makes on the decode path --------------------
+    def set_sampling_params(self, **kwargs):
+        return None
+
+    def set_penalty_params(self, **kwargs):
+        return False, False
+
+    def decode_device_state(self):
+        return None  # first install every time: the adapter takes the host view
+
+    def prefill_forward(self, tokens, **kwargs):
+        # Host-sampled shape: the adapter reshapes this to [active, 1, vocab].
+        return torch.zeros((int(tokens.shape[0]), 8))
+
+    def decode_forward(self, *args, **kwargs):
+        if kwargs.get("sampling_mode") == "host":
+            # Host-sampled decode returns logits, not a device handle; the
+            # adapter reshapes them to [rows, 1, vocab].
+            return torch.zeros((SLOTS, 8))
+        handle = _Handle(len(self.issued))
+        self.issued.append(handle)
+        return handle
+
+    def read_sampled_tokens(self, tt_out, count):
+        # row r of forward `tag` -> 1000 * (tag + 1) + r
+        return torch.tensor([1000 * (tt_out.tag + 1) + r for r in range(count)], dtype=torch.long)
+
+
+class _Sampling:
+    def __init__(self, rows: int):
+        self.temperature = [0.0] * rows
+        self.top_k = [1] * rows
+        self.top_p = [1.0] * rows
+        self.seed = [None] * rows
+        self.repetition_penalty = [1.0] * rows
+        self.presence_penalty = [0.0] * rows
+        self.frequency_penalty = [0.0] * rows
+
+
+def _adapter(widths: str = "1,2,4,8,16,32") -> Qwen3CoderForCausalLM:
+    with patch.dict("os.environ", {"QWEN3_DECODE_WIDTHS": widths}):
+        adapter = Qwen3CoderForCausalLM(_FakeGenerator(), max_model_len=4096, max_num_seqs=SLOTS)
+    adapter.kv_cache = ["fake-cache"]
+    return adapter
+
+
+def _batch(live_rows):
+    """vLLM's padded decode batch: position -1 on every unoccupied slot."""
+    positions = torch.full((SLOTS,), -1, dtype=torch.int64)
+    tokens = torch.zeros((SLOTS, 1), dtype=torch.int64)
+    for row in live_rows:
+        positions[row] = 128
+        tokens[row, 0] = 5000 + row
+    return tokens, positions
+
+
+def _issue(adapter, live_rows, *, reset=True):
+    """One decode forward, output deliberately NOT read (the async path)."""
+    tokens, positions = _batch(live_rows)
+    return adapter.decode_forward(
+        tokens=tokens,
+        page_table=torch.zeros((SLOTS, 8), dtype=torch.int32),
+        kv_cache=adapter.kv_cache,
+        start_pos=positions,
+        sampling_params=_Sampling(SLOTS),
+        reset_batch=reset,
+        read_from_device=False,
+    )
+
+
+def _read(adapter, handle):
+    with patch.object(gv.ttnn, "is_tensor_storage_on_device", lambda _t: False):
+        return adapter.process_decode_output_host(handle, is_tokens=True).reshape(-1)
+
+
+def test_output_uses_the_mapping_its_own_forward_was_issued_with():
+    """Issue A, then issue B with a different mapping, then read A.
+
+    This is the exact interleaving `--async-scheduling` produces and the one
+    nothing else covers. With the mapping stored on the adapter rather than
+    queued, reading A after B scatters A's tokens through **B's** permutation.
+    """
+    adapter = _adapter()
+
+    handle_a = _issue(adapter, [3, 17, 29])  # width 4, order starts 3,17,29
+    order_a = adapter._compaction.clone()
+    handle_b = _issue(adapter, [0, 1, 2, 3, 4, 5])  # width 8, a different order
+    order_b = adapter._compaction.clone()
+
+    # The premise of the test: the adapter's live mapping is no longer A's.
+    assert not torch.equal(
+        order_a[: min(len(order_a), len(order_b))], order_b[: min(len(order_a), len(order_b))]
+    ), "the two steps must have different mappings or this test proves nothing"
+
+    tokens_a = _read(adapter, handle_a)
+    # Forward 0, graph rows 0,1,2 -> vLLM slots 3,17,29.
+    assert tokens_a[3] == 1000, tokens_a[[3, 17, 29]]
+    assert tokens_a[17] == 1001, tokens_a[[3, 17, 29]]
+    assert tokens_a[29] == 1002, tokens_a[[3, 17, 29]]
+
+    tokens_b = _read(adapter, handle_b)
+    # Forward 1, contiguous live rows -> identity over the first six slots.
+    for row in range(6):
+        assert tokens_b[row] == 2000 + row, tokens_b[:6]
+
+    assert adapter._audit["compaction_fifo_underflows"] == 0
+    assert len(adapter._pending_orders) == 0
+
+
+def test_three_forwards_in_flight_are_read_in_issue_order():
+    """FIFO depth > 2, so the pairing cannot be a lucky one-slot swap."""
+    adapter = _adapter()
+    handles = [_issue(adapter, rows) for rows in ([3, 17, 29], [0, 1], [7])]
+    assert adapter._audit["compaction_fifo_max_depth"] == 3
+
+    expected = ({3: 1000, 17: 1001, 29: 1002}, {0: 2000, 1: 2001}, {7: 3000})
+    for handle, wanted in zip(handles, expected):
+        got = _read(adapter, handle)
+        for slot, value in wanted.items():
+            assert got[slot] == value, (slot, value, got[slot])
+    assert adapter._audit["compaction_fifo_underflows"] == 0
+
+
+def test_full_width_steps_queue_a_null_mapping():
+    """At full occupancy there is no permutation, and that must still be paired.
+
+    A `None` entry is meaningful: it says "this step needs no un-permutation".
+    Skipping the push for full-width steps would misalign the queue for every
+    narrow step behind them.
+    """
+    adapter = _adapter()
+    handle_full = _issue(adapter, list(range(SLOTS)))
+    assert adapter._compaction is None
+    handle_narrow = _issue(adapter, [9])
+    assert adapter._compaction is not None
+
+    tokens_full = _read(adapter, handle_full)
+    for row in (0, 5, 31):
+        assert tokens_full[row] == 1000 + row, tokens_full[:3]
+
+    tokens_narrow = _read(adapter, handle_narrow)
+    assert tokens_narrow[9] == 2000, tokens_narrow[:3]
+    assert adapter._audit["compaction_fifo_underflows"] == 0
+
+
+def _unset_widths_adapter() -> Qwen3CoderForCausalLM:
+    """An adapter built with `QWEN3_DECODE_WIDTHS` genuinely absent."""
+    import os
+
+    with patch.dict("os.environ", {}, clear=False):
+        os.environ.pop("QWEN3_DECODE_WIDTHS", None)
+        adapter = Qwen3CoderForCausalLM(_FakeGenerator(), max_model_len=4096, max_num_seqs=SLOTS)
+    adapter.kv_cache = ["fake-cache"]
+    return adapter
+
+
+def test_the_ladder_is_on_when_the_variable_is_unset():
+    """Unset must mean the ladder, not the fixed-width graph.
+
+    The previous default was off, and it was *known wrong*: a `max_num_seqs=32`
+    server that simply does not set an environment variable decodes one user at
+    4.3464 t/s/u instead of 49.3636. This test is what stops that default coming
+    back by accident.
+    """
+    adapter = _unset_widths_adapter()
+
+    assert adapter._decode_widths == [1, 2, 4, 8, 16, SLOTS]
+    assert adapter._compaction_enabled is True
+    handle = _issue(adapter, [3, 17, 29])
+    assert adapter._compaction is not None, "unset must still compact"
+    assert len(adapter._pending_orders) == 1
+    tokens = _read(adapter, handle)
+    # Un-permuted back to vLLM slots, exactly as with the ladder set explicitly.
+    for row in (3, 17, 29):
+        assert tokens[row] == 1000 + [3, 17, 29].index(row), tokens[[3, 17, 29]]
+
+
+def test_a_single_width_restores_the_fixed_width_path_exactly():
+    """`QWEN3_DECODE_WIDTHS=32` is the escape hatch back to the old behaviour.
+
+    With one width there is nothing to compact, so the path takes none of the
+    bookkeeping and cannot raise any of the pairing errors below -- which is the
+    property the previous default provided and which must remain reachable.
+    """
+    adapter = _adapter(widths=str(SLOTS))
+
+    assert adapter._decode_widths == [SLOTS]
+    assert adapter._compaction_enabled is False
+    handle = _issue(adapter, [3, 17, 29])
+    assert adapter._compaction is None, "no compaction may happen with a single width"
+    assert len(adapter._pending_orders) == 0, "the single-width path must not touch the queue"
+    tokens = _read(adapter, handle)
+    # Straight through: graph row r is slot r.
+    for row in (0, 3, 17, 29):
+        assert tokens[row] == 1000 + row, tokens[[0, 3, 17, 29]]
+
+
+def test_widths_above_max_num_seqs_are_dropped_from_the_default():
+    """A smaller server must not try to capture a graph wider than its slots."""
+    import os
+
+    with patch.dict("os.environ", {}, clear=False):
+        os.environ.pop("QWEN3_DECODE_WIDTHS", None)
+        adapter = Qwen3CoderForCausalLM(_FakeGenerator(), max_model_len=4096, max_num_seqs=8)
+
+    assert adapter._decode_widths == [1, 2, 4, 8]
+
+
+@pytest.mark.parametrize("live_rows", [[0], [3, 17, 29], list(range(16)), list(range(SLOTS))])
+def test_queue_drains_exactly_once_per_forward(live_rows):
+    adapter = _adapter()
+    handle = _issue(adapter, live_rows)
+    assert len(adapter._pending_orders) == 1
+    _read(adapter, handle)
+    assert len(adapter._pending_orders) == 0
+    assert adapter._audit["compaction_fifo_underflows"] == 0
+
+
+# -- the guards -------------------------------------------------------------
+#
+# The pairing rests on decode forwards being finalized exactly once and in issue
+# order. Those are invariants of `vllm-tt-plugin`, which this repository must not
+# modify, so they are *checked* here rather than trusted. The route that makes
+# this concrete: `async_decode.py::ensure_finalized` sets `_finalized = True`
+# only after `_get_output_impl()` returns, and the pop happens inside that call
+# -- so a raise anywhere after the pop leaves the step un-finalized and a later
+# `wait_for_all_pending_async_steps` finalizes it a second time.
+
+
+def test_underflow_raises_instead_of_guessing(expect_error):
+    """No queued mapping means the pairing is broken; a wrong guess is worse.
+
+    Falling back to the adapter's current mapping here would apply exactly the
+    permutation the queue exists to prevent, in the one state where it is known
+    not to belong to these tokens -- every token to the wrong request, silently.
+    """
+    adapter = _adapter()
+    handle = _issue(adapter, [3, 17, 29])
+    _read(adapter, handle)
+    with expect_error(RuntimeError, "no queued row mapping"):
+        _read(adapter, handle)  # the second finalize of the same forward
+    assert adapter._audit["compaction_fifo_underflows"] == 1
+
+
+def test_double_finalize_is_caught_by_the_tag_before_it_can_mis_scatter(expect_error):
+    """A re-finalize consumes the *next* step's mapping; the tag says so.
+
+    This is the overflow-direction desync: the queue drains faster than it
+    fills. Depth alone cannot see it -- the tags can.
+    """
+    adapter = _adapter()
+    a = _issue(adapter, [3, 17, 29])
+    _issue(adapter, [0, 1])
+    _read(adapter, a)  # legitimate: pops tag 0
+    with expect_error(RuntimeError, "out of step"):
+        # A second finalize of forward A pops tag 1, which belongs to forward B.
+        # Without the tag it would silently scatter A's tokens through B's map.
+        adapter._pending_orders.appendleft((99, adapter._compaction))
+        _read(adapter, a)
+
+
+def test_queue_cap_raises_rather_than_growing_without_bound(expect_error):
+    """Outputs never read is a leak; fail at the cap instead of mis-pairing later."""
+    adapter = _adapter()
+    adapter._pending_orders_cap = 4
+    with expect_error(RuntimeError, "row-mapping queue reached"):
+        for _ in range(adapter._pending_orders_cap + 2):
+            _issue(adapter, [3, 17, 29])
+
+
+def test_reset_realigns_tags_so_later_pops_still_pair():
+    """A released trace makes queued outputs unreadable; the reset must not desync.
+
+    Clearing without realigning the tags would make the next legitimate pop look
+    like a skipped step and raise on a perfectly healthy server.
+    """
+    adapter = _adapter()
+    _issue(adapter, [3, 17, 29])
+    _issue(adapter, [0, 1])
+    assert len(adapter._pending_orders) == 2
+    adapter._reset_pending_orders()
+    assert len(adapter._pending_orders) == 0
+
+    handle = _issue(adapter, [7])
+    got = _read(adapter, handle)
+    assert got[7] == 3000, got[:8]  # third forward issued -> tag 1000*(2+1)
+    assert adapter._audit["compaction_fifo_underflows"] == 0
+
+
+def test_prefill_resets_the_queue():
+    """Prefill may release the decode traces, so queued outputs die with them."""
+    adapter = _adapter()
+    _issue(adapter, [3, 17, 29])
+    assert len(adapter._pending_orders) == 1
+    adapter.prefill_forward(
+        tokens=torch.zeros((1, 8), dtype=torch.int64),
+        page_table=torch.zeros((SLOTS, 8), dtype=torch.int32),
+        kv_cache=adapter.kv_cache,
+        prompt_lens=[8],
+        sampling_params=None,
+    )
+    assert len(adapter._pending_orders) == 0
+    handle = _issue(adapter, [5])
+    _read(adapter, handle)
+    assert adapter._audit["compaction_fifo_underflows"] == 0
+
+
+# ---------------------------------------------------------------------------
+# The penalty path under the ladder.
+#
+# `_apply_penalties` reorders the token history into graph-row order alongside
+# the per-row penalty scalars. The scalars are genuine python lists, so a list
+# comprehension is right for them; the histories are vLLM `[rows, L]` **torch
+# tensors**, and rebuilding one as a list of 1-D tensors makes
+# `Qwen3CoderGenerator._row_token_ids` raise `TypeError: only integer tensors of
+# a single element can be converted to an index` on its `torch.as_tensor` call.
+#
+# That crashed a real penalised request in production. Nothing caught it: the
+# ladder must be ON (with it off the history is passed through untouched) *and*
+# the request must carry a non-neutral penalty, and `_FakeGenerator` above
+# accepts `set_penalty_params(**kwargs)` without ever looking at the history.
+# These tests run the **real** `_row_token_ids` over whatever the adapter
+# actually passed, so a type that the generator cannot consume fails here.
+# ---------------------------------------------------------------------------
+
+
+class _PenaltyRecordingGenerator(_FakeGenerator):
+    """Captures the kwargs `_apply_penalties` hands to the generator."""
+
+    def __init__(self):
+        super().__init__()
+        self.penalty_calls: list[dict] = []
+
+    def set_penalty_params(self, **kwargs):
+        self.penalty_calls.append(kwargs)
+        return False, False
+
+
+def _penalty_adapter(widths: str = "1,2,4,8,16,32") -> Qwen3CoderForCausalLM:
+    with patch.dict("os.environ", {"QWEN3_DECODE_WIDTHS": widths}):
+        adapter = Qwen3CoderForCausalLM(_PenaltyRecordingGenerator(), max_model_len=4096, max_num_seqs=SLOTS)
+    adapter.kv_cache = ["fake-cache"]
+    return adapter
+
+
+def _history(rows: int, width: int = 6) -> torch.Tensor:
+    """A vLLM `[rows, L]` history: slot r holds tokens 100*r+1.., -1 padded.
+
+    The -1 padding and the batch padded to `max_num_reqs` are what
+    `_row_token_ids` documents, so this mirrors the real contract.
+    """
+    hist = torch.full((rows, width), -1, dtype=torch.int64)
+    for row in range(rows):
+        hist[row, :3] = torch.tensor([100 * row + 1, 100 * row + 2, 100 * row + 3])
+    return hist
+
+
+def _penalised(rows: int) -> _Sampling:
+    sampling = _Sampling(rows)
+    sampling.repetition_penalty = [1.2] * rows  # non-neutral: takes the staged path
+    return sampling
+
+
+def test_penalised_decode_under_the_ladder_does_not_crash_on_a_tensor_history():
+    """The production crash, reduced.
+
+    A `[rows, L]` tensor history reordered into a python list of 1-D tensors
+    reaches `_row_token_ids` as something `torch.as_tensor` cannot index.
+    """
+    adapter = _penalty_adapter()
+    live = [3, 17, 29]
+    order = adapter._compaction_order(_batch(live)[1], 4, SLOTS)
+
+    adapter._apply_penalties(_penalised(SLOTS), SLOTS, _history(SLOTS), _history(SLOTS), order=order, graph_rows=4)
+
+    call = adapter.generator.penalty_calls[-1]
+    for name in ("prompt_tokens", "output_tokens"):
+        # The real consumer, on the real object the adapter passed.
+        ids = Qwen3CoderGenerator._row_token_ids(call[name], 0)
+        assert ids.numel() == 3, f"{name} row 0 unreadable by the generator"
+
+
+def test_penalty_history_follows_its_own_slot_through_the_compaction():
+    """Graph row g must carry the history of the slot the mapping sent there.
+
+    A reorder that is merely type-correct but inverted would apply one user's
+    repetition penalty to another user's tokens -- wrong output, no crash.
+    """
+    adapter = _penalty_adapter()
+    live = [3, 17, 29]
+    order = adapter._compaction_order(_batch(live)[1], 4, SLOTS)
+
+    adapter._apply_penalties(_penalised(SLOTS), SLOTS, _history(SLOTS), _history(SLOTS), order=order, graph_rows=4)
+
+    call = adapter.generator.penalty_calls[-1]
+    for graph_row, slot in enumerate(int(v) for v in order.tolist()):
+        ids = Qwen3CoderGenerator._row_token_ids(call["prompt_tokens"], graph_row)
+        expected = torch.tensor([100 * slot + 1, 100 * slot + 2, 100 * slot + 3])
+        assert torch.equal(ids, expected), f"graph row {graph_row} carries slot {slot}'s history"
+
+
+def test_penalty_scalars_and_history_are_reordered_the_same_way():
+    """The scalar and the history for one slot must not come apart.
+
+    They are reordered by separate statements; if only one of them tracked the
+    mapping, a user would get another user's penalty strength.
+    """
+    adapter = _penalty_adapter()
+    live = [3, 17, 29]
+    order = adapter._compaction_order(_batch(live)[1], 4, SLOTS)
+
+    sampling = _Sampling(SLOTS)
+    # A distinct penalty per slot, so a mis-pairing is visible in the value.
+    sampling.repetition_penalty = [1.0 + 0.01 * r for r in range(SLOTS)]
+
+    adapter._apply_penalties(sampling, SLOTS, _history(SLOTS), None, order=order, graph_rows=4)
+
+    call = adapter.generator.penalty_calls[-1]
+    for graph_row, slot in enumerate(int(v) for v in order.tolist()):
+        assert call["repetition"][graph_row] == pytest.approx(1.0 + 0.01 * slot)
+        ids = Qwen3CoderGenerator._row_token_ids(call["prompt_tokens"], graph_row)
+        assert int(ids[0]) == 100 * slot + 1, "history and scalar disagree about the slot"
+
+
+@pytest.mark.parametrize("as_list", [False, True])
+def test_history_reorder_preserves_the_type_it_was_given(as_list):
+    """Tensors stay tensors; genuine python sequences keep working.
+
+    `_row_token_ids` accepts both, but only via `torch.as_tensor`, which is what
+    the list-of-tensors form breaks.
+    """
+    adapter = _penalty_adapter()
+    live = [3, 17, 29]
+    order = adapter._compaction_order(_batch(live)[1], 4, SLOTS)
+
+    tensor_history = _history(SLOTS)
+    history = tensor_history.tolist() if as_list else tensor_history
+
+    adapter._apply_penalties(_penalised(SLOTS), SLOTS, history, None, order=order, graph_rows=4)
+
+    passed = adapter.generator.penalty_calls[-1]["prompt_tokens"]
+    if as_list:
+        assert isinstance(passed, list)
+    else:
+        assert isinstance(passed, torch.Tensor), "a tensor history must stay a tensor"
+    ids = Qwen3CoderGenerator._row_token_ids(passed, 0)
+    assert int(ids[0]) == 100 * int(order[0]) + 1
+
+
+def test_ladder_off_leaves_the_history_exactly_as_vllm_sent_it():
+    """The shipped default path must not be touched by any of the above."""
+    adapter = _penalty_adapter(widths=str(SLOTS))
+    history = _history(SLOTS)
+
+    adapter._apply_penalties(_penalised(SLOTS), SLOTS, history, None, order=None, graph_rows=None)
+
+    passed = adapter.generator.penalty_calls[-1]["prompt_tokens"]
+    assert passed is history, "with the ladder off the history is passed through unchanged"
+
+
+# ---------------------------------------------------------------------------
+# The host-sampling demotion has to be audible.
+#
+# vLLM decides per request, in `check_perform_device_sampling`, whether a
+# request may sample on device. On this 4-die mesh any request carrying
+# `logprobs` -- including `logprobs: 0`, because the guard tests
+# `max_num_logprobs is not None` before it ever looks at the value -- is routed
+# to eager host sampling. That bypasses the captured trace and the width ladder
+# and costs ~14x (measured: 3.595 t/s/u against 49.345), and the plugin emits no
+# log line for it. The server-level `sample_on_device_mode: all` stays correct
+# and stays silent.
+#
+# A 14x cliff whose only symptom is "the model got slow" is precisely the
+# failure this port was first reported with. The adapter cannot prevent the
+# demotion -- the guard lives in a repository this one must not modify -- so the
+# least it must do is say so.
+# ---------------------------------------------------------------------------
+
+
+def _host_sampled_step(adapter):
+    """One decode step with `sampling_params=None`, i.e. vLLM's host-sampled route."""
+    tokens, positions = _batch([3, 17, 29])
+    return adapter.decode_forward(
+        tokens=tokens,
+        page_table=torch.zeros((SLOTS, 8), dtype=torch.int32),
+        kv_cache=adapter.kv_cache,
+        start_pos=positions,
+        sampling_params=None,
+        reset_batch=True,
+        read_from_device=False,
+    )
+
+
+def test_host_sampled_decode_warns_once_with_the_cause_and_the_cost():
+    """The demotion must name what happened, why, and what it costs."""
+    adapter = _adapter()
+    with patch.object(gv.logger, "warning") as warn:
+        _host_sampled_step(adapter)
+
+    assert warn.call_count == 1, "the demotion must be reported"
+    message = warn.call_args[0][0]
+    for needle in ("HOST sampling", "logprobs", "3.595", "49.345", "14x"):
+        assert needle in message, f"the warning must mention {needle!r}"
+
+    assert adapter._audit["host_sampled_decode_steps"] == 1
+
+
+def test_host_sampled_warning_does_not_repeat_every_step():
+    """One line, not one per token -- a per-step warning would be its own defect."""
+    adapter = _adapter()
+    with patch.object(gv.logger, "warning") as warn:
+        for _ in range(5):
+            _host_sampled_step(adapter)
+
+    assert warn.call_count == 1, "the warning must be once per server, not per step"
+    assert adapter._audit["host_sampled_decode_steps"] == 5
+
+
+def test_device_sampled_steps_do_not_warn():
+    """The traced path must stay silent; a false alarm here trains people to ignore it."""
+    adapter = _adapter()
+    with patch.object(gv.logger, "warning") as warn:
+        _issue(adapter, [3, 17, 29])
+
+    assert warn.call_count == 0
+    assert adapter._audit["host_sampled_decode_steps"] == 0

code/models/demos/blackhole/qwen3_coder_30b_a3b/tests/test_decoder_layer.py

@@ -0,0 +1,176 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+"""The composed decoder layer against the HuggingFace reference layer.
+
+This is the stage-01 deliverable: norm -> attention -> residual -> norm -> MoE
+-> residual, end to end, at PCC >= 0.995.
+
+The submodules already pass on their own (attention 0.9994+, MoE 0.9981+), so
+a shortfall here is a *composition* error -- a residual added in the wrong
+place, the router fed the un-normed tensor, the two norms swapped -- rather
+than accumulated precision. ``test_residual_path_is_present`` exists to
+separate those two explanations: it checks the layer's output actually depends
+on the residual stream, which is the composition mistake most likely to still
+score respectably on PCC.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+from loguru import logger
+
+import ttnn
+from models.common.utility_functions import comp_allclose, comp_pcc
+
+from ..tt.functional_decoder import (
+    DecoderLayerConfig,
+    build_expert_sparsity,
+    build_rope_cache,
+    decoder_layer_prefill,
+    upload_layer_weights,
+)
+from ..tt.weight_mapping import convert_layer_weights
+from .reference import build_reference_layer, layer_state_dict, rotary_embeddings
+
+LAYER_IDX = 0
+PCC_REQUIRED = 0.995  # the stage-01 functional-decoder bar
+
+
+@pytest.fixture(scope="module")
+def reference():
+    return build_reference_layer(LAYER_IDX)
+
+
+@pytest.fixture(scope="module")
+def torch_weights(reference):
+    _, hf_config = reference
+    return convert_layer_weights(layer_state_dict(LAYER_IDX), hf_config)
+
+
+def _hidden(hf_config, seq_len, seed=0):
+    torch.manual_seed(seed)
+    return torch.randn(1, seq_len, hf_config.hidden_size, dtype=torch.float32) * 0.02
+
+
+def _causal_mask(seq_len):
+    return torch.full((seq_len, seq_len), float("-inf")).triu(1).reshape(1, 1, seq_len, seq_len)
+
+
+def _reference_layer(layer, hf_config, hidden):
+    seq_len = hidden.shape[1]
+    cos, sin = rotary_embeddings(hf_config, seq_len)
+    with torch.no_grad():
+        out = layer(
+            hidden,
+            position_embeddings=(cos, sin),
+            attention_mask=_causal_mask(seq_len),
+        )
+    return out[0] if isinstance(out, tuple) else out
+
+
+def _run_layer(mesh_device, hf_config, torch_weights, hidden):
+    config = DecoderLayerConfig.from_hf(hf_config)
+    weights = upload_layer_weights(torch_weights, mesh_device, config)
+    cos_cache, sin_cache = build_rope_cache(hf_config, hidden.shape[1], mesh_device)
+    sparsity = build_expert_sparsity(mesh_device, config.moe.num_experts)
+
+    tt_in = ttnn.from_torch(
+        hidden.unsqueeze(0),
+        dtype=ttnn.bfloat16,
+        layout=ttnn.TILE_LAYOUT,
+        device=mesh_device,
+        memory_config=ttnn.DRAM_MEMORY_CONFIG,
+    )
+    tt_out = decoder_layer_prefill(tt_in, weights, config, cos_cache, sin_cache, sparsity)
+    return ttnn.to_torch(tt_out).squeeze(0)
+
+
+@pytest.mark.parametrize("mesh_device", [(1, 1)], ids=["1x1"], indirect=True)
+@pytest.mark.parametrize(
+    "seq_len",
+    [32, 128, 512, 33, 100, 257],
+    ids=["s32", "s128", "s512", "s33", "s100", "s257"],
+)
+def test_decoder_layer_vs_reference(mesh_device, reference, torch_weights, seq_len):
+    """Tile-aligned and deliberately non-aligned sequence lengths.
+
+    33, 100 and 257 exercise the zero-padding in ``moe_prefill``: one row past a
+    tile, a mid-tile length, and one past a large power of two. Real prompts are
+    almost never a multiple of 32, and padding bugs typically corrupt only the
+    tail tokens -- which a sequence-wide PCC can absorb, so these run as
+    separate cases rather than being folded into the aligned ones.
+    """
+    layer, hf_config = reference
+    hidden = _hidden(hf_config, seq_len)
+
+    ref_out = _reference_layer(layer, hf_config, hidden)
+    tt_out = _run_layer(mesh_device, hf_config, torch_weights, hidden)
+
+    passing, pcc_message = comp_pcc(ref_out, tt_out, PCC_REQUIRED)
+    logger.info(comp_allclose(ref_out, tt_out))
+    logger.info(f"decoder layer seq={seq_len}: {pcc_message}")
+    assert passing, f"decoder layer (seq={seq_len}) below {PCC_REQUIRED}: {pcc_message}"
+
+
+@pytest.mark.parametrize("mesh_device", [(1, 1)], ids=["1x1"], indirect=True)
+@pytest.mark.parametrize("seq_len", [33, 100], ids=["s33", "s100"])
+def test_non_aligned_tail_tokens(mesh_device, reference, torch_weights, seq_len):
+    """Tokens near the pad boundary must be no worse than the rest of the sequence.
+
+    Stated *relatively*, on purpose. An absolute per-token bar cannot tell
+    "zero-padding corrupted the tail" apart from "this token was always noisy":
+    a couple of tokens sit near 0.9946 regardless of length because their router
+    top-8 contains a near-tie, and they stay low at seq_len=32 where no padding
+    exists at all. Comparing the tail against the sequence's own distribution
+    isolates the padding question, which is the only thing this test is for.
+    """
+    layer, hf_config = reference
+    hidden = _hidden(hf_config, seq_len)
+
+    ref_out = _reference_layer(layer, hf_config, hidden)
+    tt_out = _run_layer(mesh_device, hf_config, torch_weights, hidden)
+
+    def token_pcc(pos):
+        pair = torch.stack([ref_out[:, pos, :].flatten().float(), tt_out[:, pos, :].flatten().float()])
+        return float(torch.corrcoef(pair)[0, 1])
+
+    per_token = [token_pcc(p) for p in range(seq_len)]
+    tail = per_token[-3:]
+    body_worst = min(per_token[:-3])
+
+    logger.info(
+        f"seq={seq_len}: tail={[round(v, 5) for v in tail]} "
+        f"body_worst={body_worst:.5f} median={sorted(per_token)[len(per_token) // 2]:.5f}"
+    )
+
+    assert min(tail) >= 0.99, f"tail tokens of seq_len {seq_len} are outright wrong: {tail}"
+    # Padding, if broken, would make the tail distinctly worse than the body.
+    assert min(tail) >= body_worst - 1e-3, (
+        f"tail tokens ({min(tail):.5f}) are worse than the worst body token "
+        f"({body_worst:.5f}) at seq_len {seq_len} -- suspect the moe_prefill zero-padding"
+    )
+
+
+@pytest.mark.parametrize("mesh_device", [(1, 1)], ids=["1x1"], indirect=True)
+def test_residual_path_is_present(mesh_device, reference, torch_weights):
+    """The output must track the input, and must not merely echo it.
+
+    A dropped residual still produces sane-looking activations; a layer that
+    returns its input unchanged does too. Both are composition bugs that PCC
+    against the reference would report as a single vague number.
+    """
+    layer, hf_config = reference
+    hidden = _hidden(hf_config, 32)
+    out = _run_layer(mesh_device, hf_config, torch_weights, hidden).float()
+    flat_in = hidden.squeeze(0).float()
+
+    assert torch.isfinite(out).all(), "layer produced non-finite values"
+    assert not torch.allclose(out, flat_in, atol=1e-3), "output equals input -- layer body is a no-op"
+
+    # With the residual intact the output stays correlated with the input;
+    # without it the sublayer outputs alone would decorrelate.
+    corr = torch.corrcoef(torch.stack([out.flatten(), flat_in.flatten()]))[0, 1]
+    logger.info(f"corr(output, input) = {corr:.4f}")
+    assert corr > 0.5, f"output barely tracks input (corr={corr:.4f}) -- residual likely dropped"

code/models/demos/blackhole/qwen3_coder_30b_a3b/tests/test_decoder_layer_decode.py

@@ -0,0 +1,155 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+"""The full decoder layer in decode mode: prefill a prompt, then step tokens.
+
+Validated against the prefill reference at the same absolute position, which is
+the KV cache's defining contract -- attending to a cached prompt must equal
+attending to it inline.
+
+``test_multi_step_decode`` matters more than the single-step case. One step can
+pass while the cache is subtly broken (a write that lands on the position being
+read this turn still looks right); errors in the write position only diverge
+once a later token has to read what an earlier step wrote. Three consecutive
+steps against a three-token-longer reference catches that.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+from loguru import logger
+
+import ttnn
+from models.common.utility_functions import comp_allclose, comp_pcc
+
+from ..tt.functional_decoder import (
+    DecoderLayerConfig,
+    build_expert_sparsity,
+    build_rope_cache,
+    create_kv_cache,
+    decoder_layer_decode,
+    decoder_layer_prefill,
+    upload_layer_weights,
+)
+from ..tt.weight_mapping import convert_layer_weights
+from .reference import build_reference_layer, layer_state_dict, rotary_embeddings
+
+LAYER_IDX = 0
+PCC_REQUIRED = 0.99
+MAX_SEQ = 256
+
+
+@pytest.fixture(scope="module")
+def reference():
+    return build_reference_layer(LAYER_IDX)
+
+
+@pytest.fixture(scope="module")
+def torch_weights(reference):
+    _, hf_config = reference
+    return convert_layer_weights(layer_state_dict(LAYER_IDX), hf_config)
+
+
+def _hidden(hf_config, seq_len, seed=0):
+    torch.manual_seed(seed)
+    return torch.randn(1, seq_len, hf_config.hidden_size, dtype=torch.float32) * 0.02
+
+
+def _reference_layer(layer, hf_config, hidden):
+    seq_len = hidden.shape[1]
+    cos, sin = rotary_embeddings(hf_config, seq_len)
+    mask = torch.full((seq_len, seq_len), float("-inf")).triu(1).reshape(1, 1, seq_len, seq_len)
+    with torch.no_grad():
+        out = layer(hidden, position_embeddings=(cos, sin), attention_mask=mask)
+    return out[0] if isinstance(out, tuple) else out
+
+
+def _to_device(t, mesh_device):
+    return ttnn.from_torch(
+        t, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=mesh_device, memory_config=ttnn.DRAM_MEMORY_CONFIG
+    )
+
+
+def _setup(mesh_device, hf_config, torch_weights, block_size=None):
+    config = DecoderLayerConfig.from_hf(hf_config)
+    weights = upload_layer_weights(torch_weights, mesh_device, config)
+    cos_cache, sin_cache = build_rope_cache(hf_config, MAX_SEQ, mesh_device)
+    sparsity = build_expert_sparsity(mesh_device, config.moe.num_experts)
+    kv_cache = create_kv_cache(mesh_device, config.attention, max_batch=1, max_seq_len=MAX_SEQ, block_size=block_size)
+    return config, weights, cos_cache, sin_cache, sparsity, kv_cache
+
+
+def _decode_step(mesh_device, hf_config, ctx, token_hidden, position):
+    config, weights, cos_cache, sin_cache, _, kv_cache = ctx
+    current_pos = ttnn.from_torch(torch.tensor([position], dtype=torch.int32), dtype=ttnn.int32, device=mesh_device)
+    tt_in = _to_device(token_hidden.reshape(1, 1, 1, hf_config.hidden_size), mesh_device)
+    out = decoder_layer_decode(
+        tt_in, weights, config, cos_cache, sin_cache, kv_cache, current_pos, token_index=position
+    )
+    return ttnn.to_torch(out).reshape(1, hf_config.hidden_size)
+
+
+@pytest.mark.parametrize("mesh_device", [(1, 1)], ids=["1x1"], indirect=True)
+@pytest.mark.parametrize("prompt_len", [32, 128], ids=["p32", "p128"])
+@pytest.mark.parametrize("block_size", [None, 32], ids=["contiguous", "paged32"])
+def test_decode_layer_matches_prefill(mesh_device, reference, torch_weights, prompt_len, block_size):
+    layer, hf_config = reference
+    hidden_full = _hidden(hf_config, prompt_len + 1)
+    ref_out = _reference_layer(layer, hf_config, hidden_full)[:, prompt_len, :]
+
+    ctx = _setup(mesh_device, hf_config, torch_weights, block_size)
+    config, weights, cos_cache, sin_cache, sparsity, kv_cache = ctx
+
+    # Prefill the prompt through the full layer so the cache is populated.
+    decoder_layer_prefill(
+        _to_device(hidden_full[:, :prompt_len, :].unsqueeze(0), mesh_device),
+        weights,
+        config,
+        cos_cache,
+        sin_cache,
+        sparsity,
+        kv_cache=kv_cache,
+    )
+
+    tt_out = _decode_step(mesh_device, hf_config, ctx, hidden_full[:, prompt_len, :], prompt_len)
+
+    passing, pcc_message = comp_pcc(ref_out, tt_out, PCC_REQUIRED)
+    logger.info(comp_allclose(ref_out, tt_out))
+    kind = "contiguous" if block_size is None else f"paged(block={block_size})"
+    logger.info(f"decode layer at pos {prompt_len} [{kind}]: {pcc_message}")
+    assert passing, f"decode layer at pos {prompt_len} [{kind}] below {PCC_REQUIRED}: {pcc_message}"
+
+
+@pytest.mark.parametrize("mesh_device", [(1, 1)], ids=["1x1"], indirect=True)
+def test_multi_step_decode(mesh_device, reference, torch_weights):
+    """Three sequential decode steps, each checked against the prefill reference.
+
+    A cache written one position off still passes a single step; it only shows
+    up when a later token reads what an earlier step stored.
+    """
+    prompt_len, steps = 32, 3
+    layer, hf_config = reference
+    hidden_full = _hidden(hf_config, prompt_len + steps)
+    ref_out = _reference_layer(layer, hf_config, hidden_full)
+
+    # Paged: multi-step is where a block-table mapping error would surface.
+    ctx = _setup(mesh_device, hf_config, torch_weights, block_size=32)
+    config, weights, cos_cache, sin_cache, sparsity, kv_cache = ctx
+
+    decoder_layer_prefill(
+        _to_device(hidden_full[:, :prompt_len, :].unsqueeze(0), mesh_device),
+        weights,
+        config,
+        cos_cache,
+        sin_cache,
+        sparsity,
+        kv_cache=kv_cache,
+    )
+
+    for step in range(steps):
+        pos = prompt_len + step
+        tt_out = _decode_step(mesh_device, hf_config, ctx, hidden_full[:, pos, :], pos)
+        passing, pcc_message = comp_pcc(ref_out[:, pos, :], tt_out, PCC_REQUIRED)
+        logger.info(f"decode step {step} (pos {pos}): {pcc_message}")
+        assert passing, f"decode step {step} at position {pos} below {PCC_REQUIRED}: {pcc_message}"

code/models/demos/blackhole/qwen3_coder_30b_a3b/tests/test_determinism.py

@@ -0,0 +1,167 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+"""Determinism and repeated-run stability of the TTNN decoder layer.
+
+Bit-exactness is asserted, not PCC. Identical inputs through identical kernels
+must give identical bits; anything less means the result depends on something
+not in the inputs -- uninitialised memory, a race between cores, or a reduction
+whose order varies run to run. Those defects are intermittent by nature, so a
+tolerance-based check would hide exactly the cases worth catching.
+
+``test_repeated_decode_steps_are_stable`` runs a long decode rollout instead.
+It is the counterpart test: nothing there is compared against a reference, it
+just has to keep producing finite, non-degenerate activations for 64 steps.
+Cache-indexing and accumulation faults tend to show up as slow drift rather
+than a hard failure, and a handful of steps will not surface them.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+from loguru import logger
+
+import ttnn
+
+from ..tt.functional_decoder import (
+    DecoderLayerConfig,
+    build_expert_sparsity,
+    build_rope_cache,
+    create_kv_cache,
+    decoder_layer_decode,
+    decoder_layer_prefill,
+    upload_layer_weights,
+)
+from ..tt.weight_mapping import convert_layer_weights
+from .reference import build_reference_layer, layer_state_dict
+
+LAYER_IDX = 0
+MAX_SEQ = 256
+BLOCK_SIZE = 32
+
+
+@pytest.fixture(scope="module")
+def reference():
+    return build_reference_layer(LAYER_IDX)
+
+
+@pytest.fixture(scope="module")
+def torch_weights(reference):
+    _, hf_config = reference
+    return convert_layer_weights(layer_state_dict(LAYER_IDX), hf_config)
+
+
+def _hidden(hf_config, seq_len, seed=0):
+    torch.manual_seed(seed)
+    return torch.randn(1, seq_len, hf_config.hidden_size, dtype=torch.float32) * 0.02
+
+
+def _to_device(t, mesh_device):
+    return ttnn.from_torch(
+        t, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=mesh_device, memory_config=ttnn.DRAM_MEMORY_CONFIG
+    )
+
+
+def _build(mesh_device, hf_config, torch_weights):
+    config = DecoderLayerConfig.from_hf(hf_config)
+    weights = upload_layer_weights(torch_weights, mesh_device, config)
+    cos_cache, sin_cache = build_rope_cache(hf_config, MAX_SEQ, mesh_device)
+    sparsity = build_expert_sparsity(mesh_device, config.moe.num_experts)
+    return config, weights, cos_cache, sin_cache, sparsity
+
+
+@pytest.mark.parametrize("mesh_device", [(1, 1)], ids=["1x1"], indirect=True)
+def test_prefill_is_deterministic(mesh_device, reference, torch_weights):
+    _, hf_config = reference
+    config, weights, cos_cache, sin_cache, sparsity = _build(mesh_device, hf_config, torch_weights)
+    hidden = _hidden(hf_config, 128)
+
+    outs = []
+    for _ in range(3):
+        tt_in = _to_device(hidden.unsqueeze(0), mesh_device)
+        out = decoder_layer_prefill(tt_in, weights, config, cos_cache, sin_cache, sparsity)
+        outs.append(ttnn.to_torch(out).clone())
+
+    assert torch.equal(outs[0], outs[1]), "prefill run 1 != run 2 (bitwise)"
+    assert torch.equal(outs[0], outs[2]), "prefill run 1 != run 3 (bitwise)"
+    logger.info("prefill: 3 runs bit-identical")
+
+
+@pytest.mark.parametrize("mesh_device", [(1, 1)], ids=["1x1"], indirect=True)
+def test_decode_is_deterministic(mesh_device, reference, torch_weights):
+    """Two independent prefill+decode sequences must agree bitwise.
+
+    Each repetition allocates a fresh paged cache, so this also checks that no
+    state leaks between runs through the cache or the page table.
+    """
+    _, hf_config = reference
+    config, weights, cos_cache, sin_cache, sparsity = _build(mesh_device, hf_config, torch_weights)
+    prompt_len = 32
+    hidden_full = _hidden(hf_config, prompt_len + 1)
+
+    outs = []
+    for _ in range(2):
+        kv_cache = create_kv_cache(
+            mesh_device, config.attention, max_batch=1, max_seq_len=MAX_SEQ, block_size=BLOCK_SIZE
+        )
+        decoder_layer_prefill(
+            _to_device(hidden_full[:, :prompt_len, :].unsqueeze(0), mesh_device),
+            weights,
+            config,
+            cos_cache,
+            sin_cache,
+            sparsity,
+            kv_cache=kv_cache,
+        )
+        current_pos = ttnn.from_torch(
+            torch.tensor([prompt_len], dtype=torch.int32), dtype=ttnn.int32, device=mesh_device
+        )
+        tt_in = _to_device(hidden_full[:, prompt_len, :].reshape(1, 1, 1, hf_config.hidden_size), mesh_device)
+        out = decoder_layer_decode(
+            tt_in, weights, config, cos_cache, sin_cache, kv_cache, current_pos, token_index=prompt_len
+        )
+        outs.append(ttnn.to_torch(out).clone())
+
+    assert torch.equal(outs[0], outs[1]), "decode from two fresh caches differs bitwise"
+    logger.info("decode: 2 independent prefill+decode sequences bit-identical")
+
+
+@pytest.mark.parametrize("mesh_device", [(1, 1)], ids=["1x1"], indirect=True)
+def test_repeated_decode_steps_are_stable(mesh_device, reference, torch_weights):
+    """A 64-step rollout must stay finite and non-degenerate.
+
+    Long rollouts are where cache-indexing and accumulation faults show up as
+    drift rather than an exception, so this watches the activation scale across
+    every step instead of only checking the last one.
+    """
+    _, hf_config = reference
+    config, weights, cos_cache, sin_cache, sparsity = _build(mesh_device, hf_config, torch_weights)
+    prompt_len, steps = 32, 64
+    hidden_full = _hidden(hf_config, prompt_len + steps)
+
+    kv_cache = create_kv_cache(mesh_device, config.attention, max_batch=1, max_seq_len=MAX_SEQ, block_size=BLOCK_SIZE)
+    decoder_layer_prefill(
+        _to_device(hidden_full[:, :prompt_len, :].unsqueeze(0), mesh_device),
+        weights,
+        config,
+        cos_cache,
+        sin_cache,
+        sparsity,
+        kv_cache=kv_cache,
+    )
+
+    stds = []
+    for step in range(steps):
+        pos = prompt_len + step
+        current_pos = ttnn.from_torch(torch.tensor([pos], dtype=torch.int32), dtype=ttnn.int32, device=mesh_device)
+        tt_in = _to_device(hidden_full[:, pos, :].reshape(1, 1, 1, hf_config.hidden_size), mesh_device)
+        out = decoder_layer_decode(tt_in, weights, config, cos_cache, sin_cache, kv_cache, current_pos, token_index=pos)
+        t = ttnn.to_torch(out).float()
+        assert torch.isfinite(t).all(), f"decode step {step} (pos {pos}) produced non-finite values"
+        stds.append(float(t.std()))
+
+    lo, hi = min(stds), max(stds)
+    logger.info(f"64-step rollout: activation std min={lo:.5f} max={hi:.5f} ratio={hi / lo:.3f}")
+    assert lo > 1e-6, f"activations collapsed to zero during the rollout (min std {lo:.3e})"
+    assert hi / lo < 5.0, f"activation scale drifted {hi / lo:.1f}x across 64 steps -- suspect cache indexing"

code/models/demos/blackhole/qwen3_coder_30b_a3b/tests/test_full_model.py

@@ -0,0 +1,976 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+"""Full-model and generator gates for Qwen3-Coder-30B-A3B on the 4-die mesh.
+
+Two tiers, selected by ``QWEN3_FULL_MODEL_LAYERS`` (default 2):
+
+* the **reduced** tier, one real layer of each kind (there is only one kind
+  here) with every other shape, memory config, cache/page-table layout, terminal
+  norm/LM head and sampler call identical to the shipped path. Two layers load
+  in ~10 s, which is what makes these runnable as a normal test suite;
+* the **all-layer** tier, ``QWEN3_FULL_MODEL_LAYERS=48``, which is the final
+  evidence and takes several minutes to load.
+
+Accuracy against HuggingFace is *not* asserted here. A 48-layer torch reference
+is a 61 GB CPU forward, so the accuracy gate is
+``models.common.readiness_check.run_prefill_check`` /
+``run_teacher_forcing`` against the AIME24 chat reference, reported in
+``doc/full_model/README.md``. What these tests own is everything that can be
+wrong *without* moving PCC: the trace/feedback contract, position coherence,
+page-table refresh policy, non-aligned prompt lengths, batch handling, cache
+ownership, reset semantics and the runtime fallback audit.
+"""
+
+from __future__ import annotations
+
+import math
+import os
+
+import pytest
+import torch
+
+import ttnn
+from models.common.utility_functions import comp_pcc
+
+from ..tt import functional_decoder as FD
+from ..tt import multichip_decoder as MC
+from ..tt.generator import Qwen3CoderGenerator, _first_device_to_torch, build_generator, prefill_bucket_ladder
+from ..tt.model import DEFAULT_TRACE_REGION_SIZE
+
+MODEL_DIR = "models/demos/blackhole/qwen3_coder_30b_a3b"
+CONTEXT = 8192
+
+
+@pytest.fixture(scope="module")
+def mesh_device():
+    """A **module-scoped** 4-die ring mesh, opened here rather than by conftest.
+
+    The repository's `mesh_device` fixture is function-scoped, so a module-scoped
+    generator cannot depend on it (`ScopeMismatch`). Reopening the mesh per test
+    would also mean reloading the model per test -- ten seconds at two layers and
+    over three minutes at forty-eight -- which would make the all-layer tier
+    unrunnable. `FABRIC_1D_RING` must be set before the open, exactly as the
+    stage-03/04 tests and every probe in `doc/full_model/probes/` do it.
+    """
+    ttnn.set_fabric_config(ttnn.FabricConfig.FABRIC_1D_RING)
+    mesh = ttnn.open_mesh_device(mesh_shape=ttnn.MeshShape(*MC.MESH_SHAPE), trace_region_size=DEFAULT_TRACE_REGION_SIZE)
+    yield mesh
+    ttnn.close_mesh_device(mesh)
+    ttnn.set_fabric_config(ttnn.FabricConfig.DISABLED)
+
+
+#: Layer count for this run. An environment variable rather than a pytest
+#: option so that the choice survives being run as part of the whole model
+#: suite from the repository root, where a subdirectory ``conftest.py`` would
+#: not have been loaded in time to register an option.
+NUM_LAYERS = int(os.environ.get("QWEN3_FULL_MODEL_LAYERS", "2"))
+
+
+@pytest.fixture(scope="module")
+def num_layers():
+    return NUM_LAYERS
+
+
+@pytest.fixture(scope="module")
+def generator(mesh_device, num_layers):
+    gen = build_generator(
+        MODEL_DIR,
+        mesh_device,
+        override_num_layers=num_layers,
+        max_context_len=CONTEXT,
+        max_batch_size=1,
+    )
+    yield gen
+    gen.teardown()
+
+
+@pytest.fixture(scope="module")
+def batch_generator(mesh_device, num_layers):
+    gen = build_generator(
+        MODEL_DIR,
+        mesh_device,
+        override_num_layers=num_layers,
+        max_context_len=1024,
+        max_batch_size=4,
+    )
+    yield gen
+    gen.teardown()
+
+
+@pytest.fixture(scope="module")
+def small_rope_generator(mesh_device, num_layers):
+    """A generator whose cos/sin tables are far shorter than its context.
+
+    ``rope_cache_len`` defaults to 8192 against a 262144-token contract, so the
+    gap between "the table is sized" and "the context is advertised" is real on
+    the shipped configuration; 64 makes it reachable in a handful of tokens.
+    """
+    gen = build_generator(
+        MODEL_DIR,
+        mesh_device,
+        override_num_layers=num_layers,
+        max_context_len=CONTEXT,
+        max_batch_size=1,
+        rope_cache_len=64,
+    )
+    yield gen
+    gen.teardown()
+
+
+def _prompt_ids(gen, text: str) -> list[int]:
+    rendered = gen.tokenizer.apply_chat_template(
+        [{"role": "user", "content": text}], add_generation_prompt=True, tokenize=False
+    )
+    return gen.tokenizer(rendered, add_special_tokens=False)["input_ids"]
+
+
+# --- the generator contract ---------------------------------------------------
+
+
+def test_generator_implements_the_contract(generator):
+    from models.common.readiness_check.contract import Generator
+
+    assert isinstance(generator, Generator)
+    assert isinstance(generator, Qwen3CoderGenerator)
+    assert generator.tokenizer is not None
+    import inspect
+
+    # The teacher-forcing runner requires an explicit keyword, not **kwargs.
+    assert "enable_trace" in inspect.signature(generator.generate).parameters
+
+
+# --- split sampling, token feedback, position coherence -----------------------
+
+
+def test_split_sampling_feeds_its_own_token_back_on_device(generator):
+    """Step N's sampled token *is* step N+1's token input, with no host copy."""
+    generator.reset()
+    prompt = _prompt_ids(generator, "List three prime numbers.")
+    kv_cache = generator._ensure_kv_cache()
+    page_table = generator.make_page_table([len(prompt) + 8])
+    sampled = generator.prefill_forward(
+        torch.tensor([prompt]),
+        page_table=page_table,
+        kv_cache=kv_cache,
+        prompt_lens=[len(prompt)],
+        sampling_mode="device",
+    )
+    first = int(generator._sampled_to_torch(sampled)[0].item())
+
+    def read(tensor):
+        return int(ttnn.to_torch(ttnn.get_device_tensors(tensor)[0]).reshape(-1)[0].item())
+
+    # Installing the trace also performs the first replay, so the prefill token
+    # has already been consumed by the time anything can be read back.
+    host_copies_before = generator.trace_stats["token_host_copies"]
+    generator.decode_forward(
+        None,
+        torch.tensor([len(prompt)]),
+        page_table=page_table,
+        kv_cache=kv_cache,
+        sampling_mode="device",
+        enable_trace=True,
+        active_batch=1,
+    )
+    token_in, current_pos, rotary_pos, _ = generator._trace_inputs
+
+    observed = [read(token_in)]
+    positions = [(read(current_pos), read(rotary_pos))]
+    # The sampler wrote through tt_out_tok, so the persistent decode token input
+    # already holds what the sampling trace produced -- same value, same tensor.
+    assert read(generator._trace_sampled) == observed[0]
+
+    for _ in range(3):
+        generator.decode_forward(
+            None, None, page_table=None, kv_cache=kv_cache, sampling_mode="device", enable_trace=True
+        )
+        observed.append(read(token_in))
+        positions.append((read(current_pos), read(rotary_pos)))
+        assert read(generator._trace_sampled) == observed[-1]
+
+    # Positions were advanced on device by the trace itself, one per replay,
+    # starting from the prompt length, and cache position and rotary position
+    # stayed in lockstep.
+    assert positions[0] == (len(prompt) + 1, len(prompt) + 1), positions
+    for step in range(len(positions) - 1):
+        assert positions[step + 1][0] == positions[step][0] + 1, positions
+        assert positions[step + 1][1] == positions[step][1] + 1, positions
+
+    # Nothing was written to the token input from the host across any of it.
+    assert generator.trace_stats["token_host_copies"] == host_copies_before
+
+    # And the tokens observed on device are exactly what the public generator
+    # returns for the same prompt: [prefill sample, then one per replay].
+    generator.reset()
+    produced = generator.generate(prompt, 1 + len(observed), enable_trace=True, sampling_mode="device")
+    assert produced == [first] + observed, (produced, first, observed)
+
+
+def test_steady_state_decode_does_no_host_work(generator):
+    """Only ``replays`` may move between two steady-state tokens."""
+    generator.reset()
+    prompt = _prompt_ids(generator, "Say hello.")
+    generator.generate(prompt, 4, enable_trace=True, sampling_mode="device")
+    before = dict(generator.trace_stats)
+    generator.decode_forward(
+        None, None, page_table=None, kv_cache=generator._kv_cache, sampling_mode="device", enable_trace=True
+    )
+    after = dict(generator.trace_stats)
+    moved = {k: (before[k], after[k]) for k in before if before[k] != after[k]}
+    assert moved == {"replays": (before["replays"], before["replays"] + 1)}, moved
+
+
+def test_unchanged_page_table_costs_no_host_copy(generator):
+    """A page table that has not changed must not be re-uploaded."""
+    generator.reset()
+    prompt = _prompt_ids(generator, "Say hello.")
+    generator.generate(prompt, 3, enable_trace=True, sampling_mode="device")
+    page_table = generator._trace_page_table_snapshot.clone()
+    before = generator.trace_stats["page_table_host_copies"]
+    generator._refresh_persistent_page_table(page_table, generator._trace_kv_cache, active_batch=1)
+    assert generator.trace_stats["page_table_host_copies"] == before
+
+    changed = page_table.clone()
+    changed[0, -1] = 0 if changed[0, -1] != 0 else 1
+    generator._refresh_persistent_page_table(changed, generator._trace_kv_cache, active_batch=1)
+    assert generator.trace_stats["page_table_host_copies"] == before + 1
+
+
+def test_device_split_sampling_matches_host_argmax(generator):
+    """Greedy split sampling is semantically greedy, not merely close."""
+    generator.reset()
+    prompt = _prompt_ids(generator, "The capital of France is")
+    device_tokens = generator.generate(prompt, 6, enable_trace=True, sampling_mode="device")
+    generator.reset()
+    host_tokens = generator.generate(prompt, 6, sampling_mode="host")
+    assert device_tokens == host_tokens, (device_tokens, host_tokens)
+
+
+def test_force_argmax_matches_split_sampling(generator):
+    """The rejected alternative gives the same token as the shipped one."""
+    generator.reset()
+    prompt = _prompt_ids(generator, "2 + 2 =")
+    kv_cache = generator._ensure_kv_cache()
+    page_table = generator.make_page_table([len(prompt) + 1])
+    split = generator.prefill_forward(
+        torch.tensor([prompt]),
+        page_table=page_table,
+        kv_cache=kv_cache,
+        prompt_lens=[len(prompt)],
+        sampling_mode="device",
+    )
+    split_token = int(generator._sampled_to_torch(split)[0].item())
+
+    generator.reset()
+    logits = generator.prefill_forward(
+        torch.tensor([prompt]),
+        page_table=page_table,
+        kv_cache=kv_cache,
+        prompt_lens=[len(prompt)],
+        sampling_mode="host",
+    )
+    assert split_token == int(logits[0, 0].argmax().item())
+
+
+# --- rotary capacity on the low-level API -------------------------------------
+
+
+def _low_level_decode(gen, prompt, steps, *, decode_horizon=None):
+    """Drive ``prefill_forward``/``decode_forward`` exactly as the docstring says.
+
+    This is the surface a serving adapter drives, and the only one that can
+    decode past the rotary table: ``generate`` sizes for its own horizon.
+    """
+    gen.reset()
+    kv_cache = gen._ensure_kv_cache()
+    page_table = gen.make_page_table([len(prompt) + steps + 1])
+    sampled = gen.prefill_forward(
+        torch.tensor([prompt]),
+        page_table=page_table,
+        kv_cache=kv_cache,
+        prompt_lens=[len(prompt)],
+        sampling_mode="device",
+    )
+    tokens = [int(gen._sampled_to_torch(sampled)[0].item())]
+    for step in range(steps):
+        initial = step == 0
+        sampled = gen.decode_forward(
+            None,
+            torch.tensor([len(prompt)]) if initial else None,
+            page_table=page_table if initial else None,
+            kv_cache=kv_cache,
+            sampling_mode="device",
+            enable_trace=True,
+            active_batch=1,
+            **({"decode_horizon": decode_horizon} if initial and decode_horizon is not None else {}),
+        )
+        tokens.append(int(gen._sampled_to_torch(sampled)[0].item()))
+    return tokens
+
+
+def test_decode_past_the_rope_cache_length_through_the_low_level_api(small_rope_generator, expect_error):
+    """Walking off the cos/sin table must raise, and must be preventable.
+
+    The traced loop advances ``rotary_position`` with ``ttnn.plus_one`` and
+    nothing on device clamps it, so an out-of-range ``ttnn.embedding`` gather
+    would rotate at a wrong position and return a plausible-looking token. The
+    contract advertises 262144; the tables default to 8192.
+    """
+    gen = small_rope_generator
+    rope_len = gen.model.rope_cache_len
+    prompt = _prompt_ids(gen, "Count upwards.")
+    assert len(prompt) < rope_len, (len(prompt), rope_len)
+    steps = rope_len - len(prompt) + 6  # comfortably past the table
+
+    # 1. Undeclared horizon: the run must stop rather than silently gather out
+    #    of range. The message has to name the fix.
+    with expect_error(RuntimeError, "rotary table"):
+        _low_level_decode(gen, prompt, steps)
+    assert gen.model.rope_cache_len == rope_len, "the failing path must not have grown the tables"
+
+    # 2. Declared horizon: the same run completes, and the tables grew.
+    horizon = len(prompt) + steps
+    declared = _low_level_decode(gen, prompt, steps, decode_horizon=horizon)
+    assert len(declared) == steps + 1
+    assert gen.model.rope_cache_len > rope_len
+    assert gen.model.rope_cache_len >= horizon
+
+    # 3. And it is the *right* answer: identical to what the high-level
+    #    ``generate``, which sizes its own rotary horizon, produces.
+    gen.reset()
+    reference = gen.generate(prompt, steps + 1, enable_trace=True, sampling_mode="device")
+    assert declared == reference, (declared, reference)
+
+
+def test_eager_decode_grows_the_rope_tables_for_its_position(small_rope_generator):
+    """The eager/host branch gathers cos/sin too, and holds no trace to protect."""
+    gen = small_rope_generator
+    gen.reset()
+    prompt_len = 8
+    position = gen.model.rope_cache_len + 5
+    kv_cache = gen._ensure_kv_cache()
+    page_table = gen.make_page_table([position + 1])
+    gen.prefill_forward(
+        torch.arange(1000, 1000 + prompt_len, dtype=torch.long).unsqueeze(0),
+        page_table=page_table,
+        kv_cache=kv_cache,
+        prompt_lens=[prompt_len],
+        sampling_mode="host",
+    )
+    logits = gen.decode_forward(
+        torch.tensor([[42]]),
+        torch.tensor([position]),
+        page_table=page_table,
+        kv_cache=kv_cache,
+        sampling_mode="host",
+        enable_trace=False,
+    )
+    assert torch.isfinite(logits).all()
+    assert gen.model.rope_cache_len > position
+
+
+def test_decode_beyond_the_advertised_context_is_refused(generator, expect_error):
+    with expect_error(ValueError, "exceeds the supported context"):
+        generator.decode_forward(
+            None,
+            torch.tensor([generator.model.max_cache_len + 1]),
+            page_table=generator.make_page_table([8]),
+            kv_cache=generator._ensure_kv_cache(),
+            sampling_mode="device",
+            enable_trace=True,
+            active_batch=1,
+        )
+
+
+# --- top-k / top-p sampling ---------------------------------------------------
+
+
+def test_top_k_top_p_sampling_runs_through_a_traced_generate(generator):
+    """The stochastic route is exercised, not merely reachable.
+
+    ``README.md`` calls the top-k/top-p route "a live code path, not a promise";
+    this is the assertion behind that sentence. It drives ``sample_split``
+    through a captured trace and checks the generator really switched strategy.
+    """
+    gen = generator
+    gen.reset()
+    prompt = _prompt_ids(gen, "Name a colour.")
+    calls = {"split": 0, "argmax": 0}
+    real_split = gen.model.sample_split
+    real_argmax = gen.model.sample_greedy_argmax
+
+    def counting_split(*a, **k):
+        calls["split"] += 1
+        return real_split(*a, **k)
+
+    def counting_argmax(*a, **k):
+        calls["argmax"] += 1
+        return real_argmax(*a, **k)
+
+    gen.model.sample_split = counting_split
+    gen.model.sample_greedy_argmax = counting_argmax
+    try:
+        tokens = gen.generate(prompt, 5, enable_trace=True, sampling_mode="device", top_k=8, top_p=0.9, temperature=0.8)
+    finally:
+        gen.model.sample_split = real_split
+        gen.model.sample_greedy_argmax = real_argmax
+
+    assert len(tokens) == 5
+    assert all(0 <= t < gen.model.vocab_size for t in tokens), tokens
+    assert gen._sampling_stochastic is True
+    # Every sampler dispatch on this run went to the split path, and none to
+    # force-argmax -- warm-up, capture and prefill included.
+    assert calls["split"] > 0 and calls["argmax"] == 0, calls
+    assert gen._trace_model_id is not None and gen._trace_sampling_id is not None
+
+
+def test_alternating_sampling_modes_recapture_the_traces(generator):
+    """The trace-id cache is keyed by sampling mode; prove the key is honoured.
+
+    A stale trace served across a greedy/stochastic flip would silently sample
+    with the wrong strategy, which no accuracy gate on this stage would catch.
+    """
+    gen = generator
+    gen.reset()
+    prompt = _prompt_ids(gen, "The capital of France is")
+
+    greedy_first = gen.generate(prompt, 4, enable_trace=True, sampling_mode="device", top_k=1)
+    assert gen._sampling_stochastic is False
+    releases_before = gen.trace_stats["releases"]
+    captures_before = gen.trace_stats["captures"]
+
+    gen.reset()
+    stochastic = gen.generate(
+        prompt, 4, enable_trace=True, sampling_mode="device", top_k=16, top_p=0.95, temperature=1.0
+    )
+    assert gen._sampling_stochastic is True
+    assert len(stochastic) == 4
+    assert gen.trace_stats["releases"] > releases_before, gen.trace_stats
+    assert gen.trace_stats["captures"] > captures_before, gen.trace_stats
+
+    gen.reset()
+    greedy_again = gen.generate(prompt, 4, enable_trace=True, sampling_mode="device", top_k=1)
+    assert gen._sampling_stochastic is False
+    # Flipping back must restore the greedy strategy exactly, not leave the
+    # stochastic trace installed.
+    assert greedy_again == greedy_first, (greedy_first, greedy_again)
+
+
+def test_temperature_zero_is_spelled_as_greedy(generator):
+    """A serving stack spells greedy ``temperature=0``; it must not go stochastic."""
+    gen = generator
+    gen.reset()
+    prompt = _prompt_ids(gen, "The capital of France is")
+    greedy = gen.generate(prompt, 4, enable_trace=True, sampling_mode="device", top_k=1)
+    gen.reset()
+    as_temp_zero = gen.generate(prompt, 4, enable_trace=True, sampling_mode="device", top_k=0, temperature=0.0)
+    assert gen._sampling_stochastic is False
+    assert as_temp_zero == greedy, (greedy, as_temp_zero)
+
+
+def test_set_sampling_params_releases_traces_only_on_a_mode_flip(generator):
+    """Changing k/p within one mode must not cost a recapture."""
+    gen = generator
+    gen.reset()
+    prompt = _prompt_ids(gen, "Say hello.")
+    gen.generate(prompt, 3, enable_trace=True, sampling_mode="device", top_k=8, top_p=0.9)
+    assert gen._sampling_stochastic is True
+    releases = gen.trace_stats["releases"]
+
+    gen.set_sampling_params(top_k=16, top_p=0.5, temperature=0.7, active_batch=1)
+    assert gen._sampling_stochastic is True
+    assert gen.trace_stats["releases"] == releases, "a k/p change inside one mode recaptured"
+    assert gen._trace_model_id is not None
+
+    gen.set_sampling_params(top_k=1, top_p=0.0, temperature=1.0, active_batch=1)
+    assert gen._sampling_stochastic is False
+    assert gen.trace_stats["releases"] == releases + 1, "the greedy flip did not recapture"
+
+
+# --- prompt lengths -----------------------------------------------------------
+
+
+@pytest.mark.parametrize("prompt_len", [1, 31, 33, 100, 127, 128, 129, 257, 1000])
+def test_non_aligned_prompt_lengths(generator, prompt_len):
+    """Every length up to the context, aligned or not, through the public API."""
+    generator.reset()
+    tokens = torch.arange(1000, 1000 + prompt_len, dtype=torch.long).unsqueeze(0)
+    kv_cache = generator._ensure_kv_cache()
+    page_table = generator.make_page_table([prompt_len + 1])
+    logits = generator.prefill_forward(
+        tokens,
+        page_table=page_table,
+        kv_cache=kv_cache,
+        prompt_lens=[prompt_len],
+        sampling_mode="host",
+    )
+    assert tuple(logits.shape) == (1, 1, generator.model.vocab_size)
+    assert torch.isfinite(logits).all()
+
+
+def test_return_all_logits_is_sliced_to_the_logical_length(generator):
+    generator.reset()
+    prompt_len = 37
+    tokens = torch.arange(500, 500 + prompt_len, dtype=torch.long).unsqueeze(0)
+    kv_cache = generator._ensure_kv_cache()
+    logits = generator.prefill_forward(
+        tokens,
+        page_table=generator.make_page_table([prompt_len]),
+        kv_cache=kv_cache,
+        prompt_lens=[prompt_len],
+        return_all_logits=True,
+        sampling_mode="host",
+    )
+    assert tuple(logits.shape) == (1, prompt_len, generator.model.vocab_size)
+    assert torch.isfinite(logits).all()
+
+
+# --- batch, fixed slots, inactive rows ---------------------------------------
+
+
+def test_mixed_length_batch_prefill_and_decode(batch_generator):
+    """Four users, four different prompt lengths, disjoint physical pages."""
+    gen = batch_generator
+    gen.reset()
+    lengths = [7, 33, 64, 129]
+    width = max(lengths)
+    tokens = torch.zeros(len(lengths), width, dtype=torch.long)
+    for user, length in enumerate(lengths):
+        tokens[user, :length] = torch.arange(100 + user * 50, 100 + user * 50 + length)
+    kv_cache = gen._ensure_kv_cache()
+    page_table = gen.make_page_table([length + 4 for length in lengths])
+    logits = gen.prefill_forward(
+        tokens, page_table=page_table, kv_cache=kv_cache, prompt_lens=lengths, sampling_mode="host"
+    )
+    assert tuple(logits.shape) == (len(lengths), 1, gen.model.vocab_size)
+    assert torch.isfinite(logits).all()
+
+    predicted = logits[:, 0].argmax(dim=-1)
+    decoded = gen.decode_forward(
+        predicted.reshape(-1, 1),
+        torch.tensor(lengths),
+        page_table=page_table,
+        kv_cache=kv_cache,
+        sampling_mode="host",
+        enable_trace=False,
+    )
+    assert tuple(decoded.shape) == (len(lengths), gen.model.vocab_size)
+    assert torch.isfinite(decoded).all()
+
+
+def test_inactive_rows_are_expressible(batch_generator):
+    """A negative position marks an inactive slot and must not be validated."""
+    gen = batch_generator
+    gen.reset()
+    lengths = [16, 16]
+    tokens = torch.arange(2 * 16, dtype=torch.long).reshape(2, 16)
+    kv_cache = gen._ensure_kv_cache()
+    page_table = gen.make_page_table([20, 20])
+    gen.prefill_forward(tokens, page_table=page_table, kv_cache=kv_cache, prompt_lens=lengths, sampling_mode="host")
+    positions = torch.tensor([16, -1])
+    gen._validate_page_coverage(gen._normalise_page_table(page_table, 2), positions, 2)
+
+
+def test_page_table_must_map_disjoint_pages(batch_generator, expect_error):
+    gen = batch_generator
+    table = gen.make_page_table([64, 64])
+    table[1, :] = table[0, :]
+    with expect_error(ValueError, "must map disjoint physical cache pages"):
+        gen._validate_page_coverage(gen._normalise_page_table(table, 2), torch.tensor([63, 63]), 2)
+
+
+def test_sdpa_rounded_page_count_covers_the_read_window(generator):
+    """The allocation must match the kernel's rounded read, not the token count."""
+    for tokens, expected in ((1, 1), (32, 1), (33, 2), (96, 4), (256, 8), (257, 16), (320, 16), (513, 24)):
+        assert generator._sdpa_rounded_page_count(tokens) == expected, tokens
+    for tokens in (1, 33, 100, 257, 1000, 4095):
+        assert generator._sdpa_rounded_page_count(tokens) >= math.ceil(tokens / generator.page_block_size)
+
+
+def test_distributed_argmax_is_exact_at_batch_above_one(batch_generator):
+    """The live-row slice must be right at every batch, not only at batch 1.
+
+    Stage 06 made ``_WatcherCleanSampling1D._sample_argmax`` **batch-dependent**:
+    it slices the 32-slot logit tile down to ``_dist_active_rows`` before the
+    per-die ``ttnn.argmax`` and pads the result back. Every other device-sampling
+    test in this file uses the ``max_batch_size=1`` fixture and the
+    ``max_batch_size=4`` fixture only ever samples on the host, so the branch
+    that the slice introduced was uncovered at batch > 1 -- which is exactly
+    where an off-by-one in the slice or the pad would live.
+
+    This drives the sampler directly with crafted logits, because the property
+    is about the reduction and not about what the model predicts:
+
+    * every **live** row returns the host argmax of the same bf16 logits;
+    * every **padding** row returns token 0, which is the value the shipped
+      32-row reduction produces for a zero-logit row and the value the pad
+      writes back, so the 32-slot buffer is unchanged slot for slot;
+    * the caller's ``tt_out_tok`` object survives -- the traced decode loop
+      feeds that exact tensor back, so a new tensor would break feedback
+      silently.
+
+    The all-negative leg matters on its own: it is the case where "the padding
+    rows are zero" stops being harmless, because a zero padding row would beat
+    every live row if the slice were not there.
+    """
+    gen = batch_generator
+    model = gen.model
+    sampler = model.sampler
+    mesh = gen.mesh_device
+    dies = mesh.get_num_devices()
+    local_vocab = model.vocab_size // dies
+
+    assert sampler._dist_active_rows == model.max_batch_size > 1, (
+        f"this test is only meaningful when the sampler is batched: "
+        f"_dist_active_rows={sampler._dist_active_rows}, max_batch_size={model.max_batch_size}"
+    )
+    sampler.load_device_buffers()
+    assert getattr(sampler, "_dist_die_offset", None) is not None, "the distributed path is not active"
+    assert sampler._dist_local_vocab == local_vocab
+
+    slots = 32
+    active = sampler._dist_active_rows
+    torch.manual_seed(0)
+    legs = {
+        # random logits: the ordinary case, and the winner lands on a different
+        # die for different rows
+        "random": torch.randn(1, 1, slots, model.vocab_size),
+        # every live logit strictly negative: a padding row's exact 0.0 would win
+        # every row if the live-row slice were not doing its job
+        "all_negative": -1.0 - torch.rand(1, 1, slots, model.vocab_size),
+    }
+    for name, logits in legs.items():
+        # bf16 on the way in, so the host reference sees the same values the
+        # device compares -- otherwise near-ties round differently.
+        logits = logits.to(torch.bfloat16).to(torch.float32)
+        expected = logits[0, 0, :active].argmax(dim=-1).tolist()
+
+        device_logits = ttnn.from_torch(
+            logits,
+            dtype=ttnn.bfloat16,
+            layout=ttnn.TILE_LAYOUT,
+            device=mesh,
+            memory_config=ttnn.DRAM_MEMORY_CONFIG,
+            mesh_mapper=ttnn.ShardTensorToMesh(mesh, dim=-1),
+        )
+        out_tok = ttnn.from_torch(
+            torch.full((1, 1, 1, slots), 12345, dtype=torch.int32),
+            dtype=ttnn.uint32,
+            layout=ttnn.ROW_MAJOR_LAYOUT,
+            device=mesh,
+            memory_config=ttnn.DRAM_MEMORY_CONFIG,
+            mesh_mapper=ttnn.ReplicateTensorToMesh(mesh),
+        )
+        returned, logprobs = sampler._sample_argmax(device_logits, out_tok)
+        assert logprobs is None
+        assert returned is out_tok, f"{name}: the sampler returned a new tensor, breaking token feedback"
+
+        tokens = [int(v) for v in ttnn.to_torch(ttnn.get_device_tensors(returned)[0]).reshape(-1)[:slots].tolist()]
+        assert tokens[:active] == expected, f"{name}: live rows {tokens[:active]} != host argmax {expected}"
+        assert tokens[active:] == [0] * (slots - active), f"{name}: padding rows are {tokens[active:]}, not 0"
+        ttnn.deallocate(device_logits)
+        ttnn.deallocate(out_tok)
+
+
+# --- cache ownership, reset, determinism -------------------------------------
+
+
+def test_caller_owned_cache_is_used_verbatim(generator):
+    """A caller-allocated cache must be honoured, not silently replaced."""
+    generator.reset()
+    caller_cache = generator.model.allocate_kv_cache(max_cache_len=1024, num_blocks=64)
+    prompt = _prompt_ids(generator, "Hi.")
+    generator.prefill_forward(
+        torch.tensor([prompt]),
+        page_table=generator.make_page_table([len(prompt)]),
+        kv_cache=caller_cache,
+        prompt_lens=[len(prompt)],
+        sampling_mode="host",
+    )
+    written = ttnn.to_torch(ttnn.get_device_tensors(caller_cache[0].k)[0])
+    assert written.abs().sum() > 0, "prefill did not write into the caller's cache"
+    for cache in caller_cache:
+        ttnn.deallocate(cache.k, True)
+        ttnn.deallocate(cache.v, True)
+
+
+def test_reset_makes_generation_reproducible(generator):
+    prompt = _prompt_ids(generator, "Count to five.")
+    generator.reset()
+    first = generator.generate(prompt, 6, enable_trace=True, sampling_mode="device")
+    generator.reset()
+    second = generator.generate(prompt, 6, enable_trace=True, sampling_mode="device")
+    assert first == second, (first, second)
+
+
+def test_reset_zeroes_the_cache(generator):
+    generator.reset()
+    prompt = _prompt_ids(generator, "Hello there.")
+    generator.generate(prompt, 2, enable_trace=True, sampling_mode="device")
+    generator.reset()
+    for cache in generator._kv_cache:
+        assert ttnn.to_torch(ttnn.get_device_tensors(cache.k)[0]).abs().sum() == 0
+        assert ttnn.to_torch(ttnn.get_device_tensors(cache.v)[0]).abs().sum() == 0
+
+
+def test_prefill_logits_are_deterministic_across_runs(generator):
+    generator.reset()
+    prompt = _prompt_ids(generator, "Deterministic?")
+    args = dict(
+        page_table=generator.make_page_table([len(prompt)]),
+        kv_cache=generator._ensure_kv_cache(),
+        prompt_lens=[len(prompt)],
+        sampling_mode="host",
+    )
+    first = generator.prefill_forward(torch.tensor([prompt]), **args)
+    generator.reset()
+    second = generator.prefill_forward(torch.tensor([prompt]), **args)
+    assert torch.equal(first, second)
+
+
+# --- the carried-forward decoder contract ------------------------------------
+
+
+def test_runtime_fallback_audit_is_clean(generator):
+    audit = generator.model.runtime_fallback_audit()
+    assert audit["dram_sharded_taken"] is True
+    # Stage 07 retuned both to their full-K ceilings (+2.83% decode, no
+    # accuracy change); see doc/datatype_sweep/README.md.
+    assert audit["gate_up_in0_block_w"] == 64
+    assert audit["down_in0_block_w"] == 24
+    assert audit["expert_intermediate_buffer"] == "L1"
+    assert audit["local_heads"] == (8, 1)
+    assert audit["local_experts"] == 32
+    assert audit["norm_shard_feeds_qkv_directly"] is True
+    assert audit["decode_ccl_buffers_persistent"] is True
+    assert audit["host_logit_readback_on_token_out_path"] is False
+    assert audit["host_argmax_on_token_out_path"] is False
+    assert audit["vocab_padding"] == 0
+    assert audit["kv_cache_dtype"] == "bfloat16"
+    assert audit["collective_topology"] == "Topology.Ring"
+    assert (audit["prefill_num_links"], audit["decode_num_links"]) == (2, 1)
+
+
+def test_inter_layer_residual_contract_is_preserved(generator):
+    """A layer's output must be indistinguishable from its input as a tensor.
+
+    This is the stage-04 contract restated at the full-model boundary: if it
+    holds, 48 layers stack with no conversion, which is what
+    ``decode_hidden``'s bare ``for`` loop assumes.
+    """
+    model = generator.model
+    generator.reset()
+    tokens = ttnn.from_torch(
+        torch.zeros((1, 1, 1, 32), dtype=torch.int32),
+        device=model.mesh_device,
+        dtype=ttnn.uint32,
+        layout=ttnn.ROW_MAJOR_LAYOUT,
+        mesh_mapper=ttnn.ReplicateTensorToMesh(model.mesh_device),
+    )
+    hidden = model.embed_decode(tokens)
+    assert tuple(hidden.shape) == (1, 1, model.max_batch_size, model.hidden_size)
+    assert hidden.dtype == ttnn.bfloat16
+    assert hidden.layout == ttnn.TILE_LAYOUT
+    assert hidden.memory_config() == ttnn.DRAM_MEMORY_CONFIG
+
+    current_pos = ttnn.from_torch(
+        torch.zeros(model.max_batch_size, dtype=torch.int32),
+        device=model.mesh_device,
+        dtype=ttnn.int32,
+        layout=ttnn.ROW_MAJOR_LAYOUT,
+        mesh_mapper=ttnn.ReplicateTensorToMesh(model.mesh_device),
+    )
+    rotary = ttnn.from_torch(
+        torch.zeros((1, model.max_batch_size), dtype=torch.int32),
+        device=model.mesh_device,
+        dtype=ttnn.uint32,
+        layout=ttnn.ROW_MAJOR_LAYOUT,
+        mesh_mapper=ttnn.ReplicateTensorToMesh(model.mesh_device),
+    )
+    caches = generator._ensure_kv_cache()
+    model.bind_page_table(caches, generator._prefill_page_table)
+    cos, sin = model.rope_decode_tables(rotary)
+    out = MC.decoder_layer_decode_multichip(
+        hidden,
+        model.layers[0],
+        model.config,
+        model.ctx,
+        cos,
+        sin,
+        caches[0],
+        current_pos,
+        0,
+        rope=model._rope_decode,
+    )
+    assert tuple(out.shape) == tuple(hidden.shape)
+    assert out.dtype == hidden.dtype
+    assert out.layout == hidden.layout
+    assert out.memory_config() == hidden.memory_config()
+
+
+# --- prefill bucketing -------------------------------------------------------
+#
+# A prefill program is compiled per sequence length, so serving at the exact
+# logical length makes the shape space 1..max_cache_len and leaves every new
+# prompt length paying a fresh compile. Bucketing collapses that space to a
+# ladder that ``warmup_model_prefill`` can enumerate. These tests hold the line
+# that it is a *pure* latency change: the padding must be invisible in the
+# result, which is what makes it safe to leave on by default.
+
+
+@pytest.mark.parametrize("name,expected_head", [("pow2", 1024), ("pow2_half", 1024), ("1k", 1024)])
+def test_prefill_bucket_ladders_are_finite_ascending_and_total(name, expected_head):
+    """Every admissible length must land on a rung, and the ladder must be small."""
+    ladder = prefill_bucket_ladder(name, 256000)
+    assert ladder == tuple(sorted(set(ladder))), "rungs must be ascending and unique"
+    assert ladder[0] == 128 and ladder[-1] == 256000, "the ladder must span the context"
+    assert 128 in ladder and expected_head in ladder
+    # Finiteness is the whole point -- an unbounded ladder cannot be warmed.
+    assert len(ladder) < 300
+
+
+def test_exact_ladder_is_empty_and_a_typo_is_rejected(expect_error):
+    """A misspelled ladder must fail loudly: silently ignoring it would cost the
+    very multi-second compiles the ladder exists to remove."""
+    assert prefill_bucket_ladder("exact", 256000) == ()
+    with expect_error(ValueError, "QWEN3_PREFILL_BUCKETS"):
+        prefill_bucket_ladder("pow_2", 256000)
+
+
+@pytest.mark.parametrize("prompt_len", [1, 33, 100, 127, 128, 129, 257, 700, 1000, 1537])
+def test_bucketed_prefill_gives_the_same_token_as_exact_prefill(generator, prompt_len):
+    """The padding rows are causally invisible, so the sampled token must not move.
+
+    Measured on 4 dies at two layers: token-identical at every length above,
+    aligned or not, and PCC >= 0.9996 against the exact-length logits -- the
+    residual is bf16 reduction order, the same order-of-magnitude difference two
+    different SDPA chunkings of the same prompt produce.
+    """
+    tokens = torch.arange(1000, 1000 + prompt_len, dtype=torch.long).unsqueeze(0)
+
+    def run(ladder):
+        generator._prefill_buckets = ladder
+        generator.reset()
+        return generator.prefill_forward(
+            tokens,
+            page_table=generator.make_page_table([prompt_len + 1]),
+            kv_cache=generator._ensure_kv_cache(),
+            prompt_lens=[prompt_len],
+            sampling_mode="host",
+        )
+
+    try:
+        exact = run(())
+        bucketed = run(prefill_bucket_ladder("pow2_half", generator.model.max_cache_len))
+    finally:
+        generator._prefill_buckets = None
+    assert int(bucketed.argmax(-1).item()) == int(exact.argmax(-1).item())
+    passing, message = comp_pcc(exact.float(), bucketed.float(), 0.999)
+    assert passing, message
+
+
+def test_bucketed_split_prefill_matches_exact_at_every_chunk_size(generator):
+    """The cached-suffix branch buckets too, at all of its reachable chunk sizes.
+
+    ``sdpa_chunk_size(start) == min(256, start & -start)`` over block-aligned
+    starts is a closed set, and this walks it. The assertion on
+    ``PREFILL_ATTENTION_BRANCHES`` is not decoration: a PCC of ~1.0 on a split
+    prefill has two explanations, and "the chunked branch never ran" is one.
+    """
+    block = generator.page_block_size
+    ladder = prefill_bucket_ladder("pow2_half", generator.model.max_cache_len)
+    for start in (block, 2 * block, 3 * block, 4 * block, 6 * block):
+        prompt_len = start + 17  # a suffix aligned to nothing
+        tokens = torch.arange(2000, 2000 + prompt_len, dtype=torch.long).unsqueeze(0)
+
+        def run(chosen):
+            generator._prefill_buckets = chosen
+            generator.reset()
+            FD.PREFILL_ATTENTION_BRANCHES["chunked"] = 0
+            out = generator.prefill_forward(
+                tokens,
+                page_table=generator.make_page_table([prompt_len + 1]),
+                kv_cache=generator._ensure_kv_cache(),
+                prompt_lens=[prompt_len],
+                sampling_mode="host",
+                start_pos=[start],
+            )
+            assert FD.PREFILL_ATTENTION_BRANCHES["chunked"] > 0, "the chunked branch never ran"
+            return out
+
+        try:
+            exact = run(())
+            bucketed = run(ladder)
+        finally:
+            generator._prefill_buckets = None
+        assert int(bucketed.argmax(-1).item()) == int(exact.argmax(-1).item()), f"start={start}"
+        passing, message = comp_pcc(exact.float(), bucketed.float(), 0.999)
+    assert passing, message
+
+
+def test_bucketing_keeps_the_cache_write_at_the_real_length(generator):
+    """Padding must not reach the KV cache -- it would run off the page table.
+
+    This is the one way bucketing could corrupt a *different* request rather
+    than merely slow this one down. vLLM allocates ``ceil(real/block)`` blocks
+    and not one more, so a write at the padded length would land in whatever
+    physical block comes next, which belongs to somebody else. ``fill_len`` is
+    what prevents that, and this asserts the block after the prompt's own is
+    still byte-for-byte what it was.
+    """
+    prompt_len = 129  # 5 blocks at a 32-token block; the 512 rung wants 16
+    block = generator.page_block_size
+    owned = math.ceil(prompt_len / block)
+    caches = generator._ensure_kv_cache()
+    assert generator.prefill_padded_len(prompt_len, start=0) > owned * block, "test needs real padding"
+
+    # Give the user exactly the blocks it is entitled to, and nothing after.
+    page_table = torch.full((generator.batch, generator.pages_per_user), -1, dtype=torch.int32)
+    page_table[0, :owned] = torch.arange(owned, dtype=torch.int32)
+    sentinel_block = owned
+
+    def read_sentinel():
+        host = _first_device_to_torch(caches[0].k)
+        return host[sentinel_block].clone()
+
+    generator.reset()
+    before = read_sentinel()
+    generator._prefill_buckets = prefill_bucket_ladder("pow2_half", generator.model.max_cache_len)
+    try:
+        generator.prefill_forward(
+            torch.arange(4000, 4000 + prompt_len, dtype=torch.long).unsqueeze(0),
+            page_table=page_table,
+            kv_cache=caches,
+            prompt_lens=[prompt_len],
+            sampling_mode="host",
+        )
+    finally:
+        generator._prefill_buckets = None
+    assert torch.equal(read_sentinel(), before), "bucket padding was written past the prompt's blocks"
+
+
+@pytest.mark.parametrize("prompt_len", [128, 256, 1024])
+def test_a_prompt_filling_the_whole_rope_table_does_not_free_it(generator, prompt_len):
+    """A power-of-two prompt must not leave the model's cos/sin tables freed.
+
+    ``ensure_rope_capacity`` rounds up to a power of two, so at these lengths
+    the window ``prefill_hidden`` wants is the *entire* table -- and
+    ``ttnn.slice`` returns a view, not a copy, when asked for the whole tensor.
+    Deallocating that view frees ``cos_table`` itself, and the failure lands on
+    the *next* prefill as "Input Tensor is not allocated", nowhere near the
+    cause. Two prefills, because one alone cannot see it.
+    """
+    tokens = torch.arange(1000, 1000 + prompt_len, dtype=torch.long).unsqueeze(0)
+    for _ in range(2):
+        generator.reset()
+        logits = generator.prefill_forward(
+            tokens,
+            page_table=generator.make_page_table([prompt_len + 1]),
+            kv_cache=generator._ensure_kv_cache(),
+            prompt_lens=[prompt_len],
+            sampling_mode="host",
+        )
+        assert torch.isfinite(logits).all()
+    assert generator.model.cos_table.is_allocated()
+    assert generator.model.sin_table.is_allocated()

code/models/demos/blackhole/qwen3_coder_30b_a3b/tests/test_moe.py

@@ -0,0 +1,232 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+"""TTNN MoE block vs the HuggingFace reference. Router first, experts after.
+
+The router is tested on its own because it is the one place in this model where
+a *discrete* disagreement is possible. Everything else degrades smoothly with
+precision; top-k selection does not. HF softmaxes over all 128 experts in fp32
+and we run bf16, so two experts with near-equal probability can swap places.
+When that happens the token is routed to a genuinely different expert and the
+MoE output for that token is unrelated to the reference -- a handful of such
+tokens drags whole-block PCC down in a way that looks like a numerics problem
+but is actually a selection problem.
+
+``test_router_selection_matches`` separates the two by comparing the chosen
+expert *sets* directly, so a later PCC dip can be attributed rather than
+guessed at.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+from loguru import logger
+
+import ttnn
+from models.common.utility_functions import comp_allclose, comp_pcc
+
+from ..tt.functional_decoder import (
+    MoEConfig,
+    build_expert_sparsity,
+    moe_prefill,
+    router_forward,
+    upload_expert_weights,
+    upload_router_weight,
+)
+from ..tt.weight_mapping import convert_moe_weights
+from .reference import build_reference_layer, layer_state_dict
+
+LAYER_IDX = 0
+PCC_REQUIRED = 0.99
+
+
+@pytest.fixture(scope="module")
+def reference():
+    return build_reference_layer(LAYER_IDX)
+
+
+@pytest.fixture(scope="module")
+def torch_weights():
+    return convert_moe_weights(layer_state_dict(LAYER_IDX), n_experts=128)
+
+
+def _hidden(config, seq_len, seed=0):
+    torch.manual_seed(seed)
+    return torch.randn(1, seq_len, config.hidden_size, dtype=torch.float32) * 0.02
+
+
+def _reference_router(layer, hidden):
+    """Return ``(dense [S, E], indices [S, k])`` from the reference router."""
+    flat = hidden.view(-1, hidden.shape[-1])
+    with torch.no_grad():
+        _, scores, indices = layer.mlp.gate(flat)
+    dense = torch.zeros(flat.shape[0], 128, dtype=scores.dtype)
+    dense.scatter_(-1, indices, scores)
+    return dense, indices
+
+
+def _run_router(mesh_device, hf_config, torch_weights, hidden):
+    w = upload_router_weight(torch_weights["router"], mesh_device)
+    tt_in = ttnn.from_torch(
+        hidden.unsqueeze(0),
+        dtype=ttnn.bfloat16,
+        layout=ttnn.TILE_LAYOUT,
+        device=mesh_device,
+        memory_config=ttnn.DRAM_MEMORY_CONFIG,
+    )
+    out = router_forward(tt_in, w, MoEConfig.from_hf(hf_config))
+    return ttnn.to_torch(out).reshape(-1, 128).float()
+
+
+@pytest.mark.parametrize("mesh_device", [(1, 1)], ids=["1x1"], indirect=True)
+@pytest.mark.parametrize("seq_len", [32, 128], ids=["s32", "s128"])
+def test_router_dense_weights_vs_reference(mesh_device, reference, torch_weights, seq_len):
+    layer, hf_config = reference
+    hidden = _hidden(hf_config, seq_len)
+
+    ref_dense, _ = _reference_router(layer, hidden)
+    tt_dense = _run_router(mesh_device, hf_config, torch_weights, hidden)
+
+    passing, pcc_message = comp_pcc(ref_dense, tt_dense, PCC_REQUIRED)
+    logger.info(comp_allclose(ref_dense, tt_dense))
+    logger.info(f"router dense seq={seq_len}: {pcc_message}")
+    assert passing, f"router dense weights (seq={seq_len}) below {PCC_REQUIRED}: {pcc_message}"
+
+
+@pytest.mark.parametrize("mesh_device", [(1, 1)], ids=["1x1"], indirect=True)
+def test_router_selection_matches(mesh_device, reference, torch_weights):
+    """How many tokens pick a different set of 8 experts than the reference.
+
+    Reported explicitly rather than silently folded into PCC, because a
+    selection flip and a numerics error need completely different fixes.
+
+    The bound is measured, not guessed. Holding the softmax in fp32 and letting
+    only the projection run in bf16 costs 5/128 tokens on this checkpoint (a
+    host-side simulation of the same arithmetic agrees), so the floor is ~4%.
+    The threshold sits at 10% to leave room for a different activation sample
+    while still catching the regression that matters: dropping the softmax to
+    bf16 sends this straight to ~65%.
+    """
+    seq_len = 128
+    layer, hf_config = reference
+    hidden = _hidden(hf_config, seq_len)
+
+    ref_dense, ref_indices = _reference_router(layer, hidden)
+    tt_dense = _run_router(mesh_device, hf_config, torch_weights, hidden)
+    tt_indices = tt_dense.topk(8, dim=-1).indices
+
+    mismatched = 0
+    worst_missed_weight = 0.0
+    for token in range(seq_len):
+        ref_set = set(ref_indices[token].tolist())
+        tt_set = set(tt_indices[token].tolist())
+        if ref_set != tt_set:
+            mismatched += 1
+            # How much routing weight the reference put on experts we skipped.
+            # Near-ties sit at the bottom of the top-8, so this should be small;
+            # a large value means a genuinely wrong selection, not rounding.
+            for e in ref_set - tt_set:
+                worst_missed_weight = max(worst_missed_weight, float(ref_dense[token, e]))
+
+    logger.info(
+        f"router selection: {mismatched}/{seq_len} tokens differ from the fp32 reference; "
+        f"largest missed routing weight {worst_missed_weight:.4f}"
+    )
+    assert mismatched <= seq_len * 0.10, (
+        f"{mismatched}/{seq_len} tokens routed to a different expert set -- "
+        "far above the ~4% bf16-projection floor. Check that the router softmax "
+        "and topk are still running in fp32 before suspecting anything else."
+    )
+
+
+def _upload_dense(dense: torch.Tensor, mesh_device):
+    return ttnn.from_torch(
+        dense.reshape(1, 1, *dense.shape).float(),
+        dtype=ttnn.bfloat16,
+        layout=ttnn.TILE_LAYOUT,
+        device=mesh_device,
+        memory_config=ttnn.DRAM_MEMORY_CONFIG,
+    )
+
+
+@pytest.mark.parametrize("mesh_device", [(1, 1)], ids=["1x1"], indirect=True)
+@pytest.mark.parametrize("seq_len", [32, 128], ids=["s32", "s128"])
+def test_experts_with_reference_routing(mesh_device, reference, torch_weights, seq_len):
+    """Expert math alone, with the router's own selection taken out of play.
+
+    The reference's exact fp32 routing weights are fed straight to our experts,
+    so any shortfall here is the sparse_matmul / SwiGLU / reduce path and not
+    the ~5/128 near-tie selection differences the router legitimately has.
+    Isolating the two is what makes the end-to-end number interpretable.
+    """
+    layer, hf_config = reference
+    config = MoEConfig.from_hf(hf_config)
+    hidden = _hidden(hf_config, seq_len)
+
+    with torch.no_grad():
+        ref_out = layer.mlp(hidden)
+
+    ref_dense, _ = _reference_router(layer, hidden)
+    weights = upload_expert_weights(torch_weights, mesh_device, config)
+    sparsity = build_expert_sparsity(mesh_device, config.num_experts)
+
+    tt_in = ttnn.from_torch(
+        hidden.unsqueeze(0),
+        dtype=ttnn.bfloat16,
+        layout=ttnn.TILE_LAYOUT,
+        device=mesh_device,
+        memory_config=ttnn.DRAM_MEMORY_CONFIG,
+    )
+    tt_out = moe_prefill(tt_in, _upload_dense(ref_dense, mesh_device), weights, config, sparsity)
+    tt_out_torch = ttnn.to_torch(tt_out).squeeze(0)
+
+    passing, pcc_message = comp_pcc(ref_out, tt_out_torch, PCC_REQUIRED)
+    logger.info(comp_allclose(ref_out, tt_out_torch))
+    logger.info(f"experts (reference routing) seq={seq_len}: {pcc_message}")
+    assert passing, f"experts (seq={seq_len}) below {PCC_REQUIRED}: {pcc_message}"
+
+
+@pytest.mark.parametrize("mesh_device", [(1, 1)], ids=["1x1"], indirect=True)
+@pytest.mark.parametrize("seq_len", [32, 128], ids=["s32", "s128"])
+def test_moe_block_end_to_end(mesh_device, reference, torch_weights, seq_len):
+    """Our router driving our experts, against the reference MoE block."""
+    layer, hf_config = reference
+    config = MoEConfig.from_hf(hf_config)
+    hidden = _hidden(hf_config, seq_len)
+
+    with torch.no_grad():
+        ref_out = layer.mlp(hidden)
+
+    w_router = upload_router_weight(torch_weights["router"], mesh_device)
+    weights = upload_expert_weights(torch_weights, mesh_device, config)
+    sparsity = build_expert_sparsity(mesh_device, config.num_experts)
+
+    tt_in = ttnn.from_torch(
+        hidden.unsqueeze(0),
+        dtype=ttnn.bfloat16,
+        layout=ttnn.TILE_LAYOUT,
+        device=mesh_device,
+        memory_config=ttnn.DRAM_MEMORY_CONFIG,
+    )
+    routing = router_forward(tt_in, w_router, config)
+    tt_out = moe_prefill(tt_in, routing, weights, config, sparsity)
+    tt_out_torch = ttnn.to_torch(tt_out).squeeze(0)
+
+    passing, pcc_message = comp_pcc(ref_out, tt_out_torch, PCC_REQUIRED)
+    logger.info(comp_allclose(ref_out, tt_out_torch))
+    logger.info(f"MoE block end-to-end seq={seq_len}: {pcc_message}")
+    assert passing, f"MoE block (seq={seq_len}) below {PCC_REQUIRED}: {pcc_message}"
+
+
+@pytest.mark.parametrize("mesh_device", [(1, 1)], ids=["1x1"], indirect=True)
+def test_router_weights_sum_to_one(mesh_device, reference, torch_weights):
+    """norm_topk_prob=True, so each token's 8 weights must renormalise to 1."""
+    layer, hf_config = reference
+    hidden = _hidden(hf_config, 32)
+    tt_dense = _run_router(mesh_device, hf_config, torch_weights, hidden)
+
+    sums = tt_dense.sum(dim=-1)
+    nonzero = (tt_dense > 0).sum(dim=-1)
+    assert torch.allclose(sums, torch.ones_like(sums), atol=2e-2), f"router weights not normalised: {sums[:4]}"
+    assert (nonzero == 8).all(), f"expected exactly 8 active experts per token, got {nonzero.unique().tolist()}"

code/models/demos/blackhole/qwen3_coder_30b_a3b/tests/test_multichip_decoder.py

@@ -0,0 +1,1057 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+"""Correctness of the multichip decoder on the full 4-die P300_X2 mesh.
+
+The reference these tests compare against is the **single-chip TTNN optimized
+decoder**, not HuggingFace, and it is run *on the same mesh with every tensor
+replicated*. That is the whole trick of this file: a mesh op is SPMD, so
+uploading the unsharded stage-02 weights with ``ReplicateTensorToMesh`` makes
+each of the four dies independently compute the exact single-chip answer, in the
+same process, from the same host tensors, with the same program cache. The
+comparison then isolates sharding and collective bugs from every source of
+numerical difference that HF-vs-TTNN would drag in.
+
+``test_baseline_upload_is_actually_replicated`` is what stops that reference
+quietly becoming meaningless -- if ``from_torch`` ever stopped replicating, the
+baseline would still produce *a* number and every PCC below would still pass.
+
+Two tests here are load-bearing in a way their size does not suggest:
+
+* ``test_topk_is_identical_across_dies``. The expert-parallel scheme assumes the
+  four dies agree on the global top-8 from bit-identical replicated logits, so
+  that the four 32-expert windows partition it. If ``ttnn.topk`` ever broke a
+  tie differently on one die the layer would be **silently wrong** -- no shape
+  error, no assert, just PCC drift -- so the property is asserted directly
+  rather than argued from "same program, same input".
+
+* ``test_expert_window_can_be_empty``. Under EP the locally-live expert count is
+  data-dependent in 0..8, which is why decode must pass ``nnz=None``. The zero
+  case is the one that never happens by accident in a random test and is exactly
+  where an uninitialised output buffer would leak a NaN into the all-reduce.
+
+Every test opens the mesh with ``fabric_config=FABRIC_1D_RING``. Without it the
+collectives have no fabric to run on; with ``FABRIC_1D`` they would run but on
+the linear topology the CCL sweep measured 1.2-1.8x slower.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+from loguru import logger
+
+import ttnn
+from models.common.modules.tt_ccl import default_topology
+from models.common.utility_functions import comp_pcc
+
+from ..tt import functional_decoder as F
+from ..tt import multichip_decoder as MC
+from ..tt import optimized_decoder as O
+from ..tt.weight_mapping import convert_layer_weights
+from .reference import build_reference_layer, layer_state_dict, rotary_embeddings
+
+LAYER_IDX = 0
+# Against the replicated single-chip baseline only sharding and the collectives
+# differ, so 0.99 -- the threshold used against HF, where dtype and kernel choice
+# differ too -- was two orders of magnitude looser than the measured margin. The
+# actuals span 0.99945 (two stacked layers, the only one below 0.9996) to
+# 1.0; 0.999 sits below the worst of them and still catches anything that would
+# make the sharding wrong rather than merely different.
+PCC_VS_SINGLE_CHIP = 0.999
+PCC_VS_HF = 0.995
+MAX_SEQ = 1024
+BLOCK_SIZE = 32
+TRACE_REGION_SIZE = 90000000
+
+# Ring fabric must be configured before the mesh is opened, which is what this
+# indirect parametrisation does (conftest.set_fabric runs ahead of the open).
+MESH_PARAMS = {"trace_region_size": TRACE_REGION_SIZE, "fabric_config": ttnn.FabricConfig.FABRIC_1D_RING}
+mesh_4 = pytest.mark.parametrize("mesh_device", [MC.MESH_SHAPE], ids=["1x4"], indirect=True)
+ring_fabric = pytest.mark.parametrize("device_params", [MESH_PARAMS], indirect=True)
+
+
+@pytest.fixture(scope="module")
+def reference():
+    return build_reference_layer(LAYER_IDX)
+
+
+@pytest.fixture(scope="module")
+def torch_weights(reference):
+    _, hf_config = reference
+    return convert_layer_weights(layer_state_dict(LAYER_IDX), hf_config)
+
+
+def _hidden(hf_config, seq_len, seed=0):
+    torch.manual_seed(seed)
+    return torch.randn(1, 1, seq_len, hf_config.hidden_size, dtype=torch.float32) * 0.02
+
+
+def _reference_layer(layer, hf_config, hidden):
+    """HF layer output for a ``[1, 1, S, H]`` input, returned as ``[1, S, H]``."""
+    hidden = hidden.reshape(1, -1, hf_config.hidden_size)
+    seq_len = hidden.shape[1]
+    cos, sin = rotary_embeddings(hf_config, seq_len)
+    mask = torch.full((seq_len, seq_len), float("-inf")).triu(1).reshape(1, 1, seq_len, seq_len)
+    with torch.no_grad():
+        out = layer(hidden, position_embeddings=(cos, sin), attention_mask=mask)
+    return out[0] if isinstance(out, tuple) else out
+
+
+def _replicate(t, mesh_device, dtype=ttnn.bfloat16):
+    return ttnn.from_torch(
+        t,
+        dtype=dtype,
+        layout=ttnn.TILE_LAYOUT,
+        device=mesh_device,
+        memory_config=ttnn.DRAM_MEMORY_CONFIG,
+        mesh_mapper=ttnn.ReplicateTensorToMesh(mesh_device),
+    )
+
+
+def _per_die(t, mesh_device, dim: int = 0) -> torch.Tensor:
+    """All four dies' copies of a tensor, concatenated along ``dim``.
+
+    Every activation in this layer is ``[1, 1, ., .]``, so concatenating on dim 0
+    puts die *d* at index *d* and reads like a stack -- but it *is* a
+    concatenation, which matters for the KV cache, whose dim 0 is the block index
+    and which therefore has to be reassembled on the head axis instead.
+    """
+    return ttnn.to_torch(t, mesh_composer=ttnn.ConcatMeshToTensor(mesh_device, dim=dim))
+
+
+class Fixture:
+    """Both paths uploaded side by side onto the same mesh.
+
+    ``baseline`` is the stage-02 optimized decoder with every tensor replicated,
+    so each die computes the full single-chip layer; ``multichip`` is the sharded
+    stage-03 path. They share the router weight tensor, the RoPE caches and the
+    host weights, so the only difference between them is the parallelisation.
+    """
+
+    def __init__(self, mesh_device, hf_config, torch_weights):
+        self.mesh = mesh_device
+        self.hf = hf_config
+        self.config = MC.MeshDecoderConfig.from_hf(hf_config)
+        self.ctx = MC.mesh_context(mesh_device)
+        self.torch_router = torch_weights["router"]
+        self.multichip = MC.upload_multichip_weights(torch_weights, mesh_device, self.config)
+        self.baseline_experts = O.upload_optimized_weights(torch_weights, mesh_device, self.config.global_config.moe)
+        self.baseline = F.DecoderLayerWeights(
+            input_layernorm=self.multichip.input_layernorm,
+            post_attention_layernorm=self.multichip.post_attention_layernorm,
+            attention=None,  # the optimized path reads OptimizedWeights.attention
+            router=self.multichip.router,
+            experts=None,
+        )
+        self.cos, self.sin = F.build_rope_cache(hf_config, MAX_SEQ, mesh_device)
+        self.baseline_sparsity = F.build_expert_sparsity(mesh_device, self.config.global_config.moe.num_experts)
+        self.sparsity = MC.build_local_sparsity(mesh_device, self.config.local_moe)
+
+    def rep(self, t):
+        return _replicate(t, self.mesh)
+
+    def dies(self, t, dim: int = 0):
+        return _per_die(t, self.mesh, dim)
+
+    def baseline_prefill(self, x, kv_cache=None, user_id=0):
+        return O.decoder_layer_prefill_optimized(
+            self.rep(x),
+            self.baseline,
+            self.config.global_config,
+            self.cos,
+            self.sin,
+            self.baseline_sparsity,
+            self.baseline_experts,
+            kv_cache=kv_cache,
+            user_id=user_id,
+        )
+
+    def multichip_prefill(self, x, kv_cache=None, user_id=0):
+        return MC.decoder_layer_prefill_multichip(
+            self.rep(x),
+            self.multichip,
+            self.config,
+            self.ctx,
+            self.cos,
+            self.sin,
+            self.sparsity,
+            kv_cache=kv_cache,
+            user_id=user_id,
+        )
+
+
+@pytest.fixture
+def fixture(mesh_device, reference, torch_weights):
+    _, hf_config = reference
+    return Fixture(mesh_device, hf_config, torch_weights)
+
+
+# --- host-side weight transforms ---------------------------------------------
+
+
+def test_wqkv_column_split_is_head_interleaved(reference, torch_weights):
+    """Die *d* must get Q heads 8d..8d+7 plus K head d and V head d.
+
+    A contiguous 4-way split of the checkpoint's ``[Wq | Wk | Wv]`` gives die 0
+    nothing but Q heads and die 3 nothing but K and V, and produces **no shape
+    error** -- ``nlp_create_qkv_heads_decode(num_heads=8, num_kv_heads=1)``
+    accepts 1280 columns whatever is in them. This is a host-only test because
+    that is where the bug would live and where it is cheapest to catch.
+    """
+    _, hf_config = reference
+    cfg = MC.MeshDecoderConfig.from_hf(hf_config)
+    a = cfg.global_config.attention
+    n, hd = cfg.num_devices, a.head_dim
+    full = torch_weights["wqkv"].reshape(a.hidden_size, -1)
+    permuted = MC.head_interleaved_wqkv(full, a, n)
+
+    q_end = a.num_attention_heads * hd
+    k_end = q_end + a.num_key_value_heads * hd
+    per_die = permuted.shape[-1] // n
+    q_per = a.num_attention_heads // n
+
+    for d in range(n):
+        shard = permuted[:, d * per_die : (d + 1) * per_die]
+        assert shard.shape[-1] == q_per * hd + 2 * hd, shard.shape
+        expect_q = full[:, d * q_per * hd : (d + 1) * q_per * hd]
+        expect_k = full[:, q_end + d * hd : q_end + (d + 1) * hd]
+        expect_v = full[:, k_end + d * hd : k_end + (d + 1) * hd]
+        assert torch.equal(shard[:, : q_per * hd], expect_q), f"die {d} Q heads"
+        assert torch.equal(shard[:, q_per * hd : q_per * hd + hd], expect_k), f"die {d} K head"
+        assert torch.equal(shard[:, q_per * hd + hd :], expect_v), f"die {d} V head"
+
+    # And it is a permutation, not a rewrite: every column survives exactly once.
+    # Compared as sorted multisets rather than row sums -- the reordering changes
+    # float addition order, so ``sum`` differs in the last bits even when nothing
+    # has been lost, and asserting on it fails for the wrong reason.
+    assert torch.equal(permuted.sort(dim=-1).values, full.sort(dim=-1).values)
+    logger.info(f"wqkv head-interleaved split verified for {n} dies, per-die N = {per_die}")
+
+
+def test_repo_default_topology_is_wrong_for_this_mesh():
+    """Documents *why* Ring is passed explicitly, so it cannot be "simplified" away.
+
+    ``tt_ccl.default_topology()`` only returns ``Ring`` for 8-device T3K and
+    Galaxy; for this 4-device Blackhole mesh it returns ``Linear``, which the CCL
+    sweep measured at 1.21x slower at decode size and 1.79x at 2 MB. The module
+    constant must therefore disagree with the helper.
+    """
+    assert MC.TOPOLOGY is ttnn.Topology.Ring
+    assert MC.NUM_LINKS == 2
+    logger.info(f"multichip_decoder overrides default_topology (callable: {default_topology.__name__})")
+
+
+# --- the replicated baseline itself ------------------------------------------
+
+
+@ring_fabric
+@mesh_4
+def test_baseline_upload_is_actually_replicated(fixture):
+    """The single-chip reference is only a reference if all four dies agree.
+
+    Everything else in this file divides by this. If ``ReplicateTensorToMesh``
+    ever stopped replicating, or the mesh stopped being SPMD, the baseline would
+    still return numbers and every PCC below would still pass against them.
+    """
+    out = fixture.dies(fixture.baseline_prefill(_hidden(fixture.hf, 128)))
+    spread = (out - out[0:1]).abs().max().item()
+    logger.info(f"replicated single-chip baseline: max spread across 4 dies = {spread:.3e}")
+    assert spread == 0.0, f"replicated baseline differs across dies by {spread}; it is not a valid reference"
+
+
+# --- the determinism assumption the whole scheme rests on --------------------
+
+
+@ring_fabric
+@mesh_4
+@pytest.mark.parametrize("seed", [0, 1, 2, 3, 4, 5, 6, 7])
+def test_topk_is_identical_across_dies(fixture, seed):
+    """All four dies must select the same 8 experts from the replicated logits.
+
+    Expert parallelism partitions the 128 experts into four 32-wide windows and
+    each die keeps only the winners inside its own. That is the global top-8 only
+    if the four dies agree; if they disagree the layer double-counts some experts
+    and drops others, with no error of any kind. Random inputs are checked, and
+    so is an all-zero input, where every logit is the router's bias-free
+    projection of zero and the top-8 is decided **entirely by tie-breaking** --
+    the degenerate case an ordinary test never reaches.
+    """
+    hidden = _hidden(fixture.hf, 128, seed=seed) if seed else torch.zeros(1, 1, 128, fixture.hf.hidden_size)
+    logits = ttnn.linear(
+        fixture.rep(hidden), fixture.multichip.router, dtype=ttnn.float32, memory_config=ttnn.DRAM_MEMORY_CONFIG
+    )
+    _, indices = ttnn.topk(logits, k=fixture.config.global_config.moe.num_experts_per_tok, dim=-1, sorted=True)
+    per_die = fixture.dies(indices)
+    for d in range(1, fixture.config.num_devices):
+        assert torch.equal(per_die[0], per_die[d]), (
+            f"seed={seed}: die {d} selected different experts than die 0 "
+            f"({int((per_die[0] != per_die[d]).sum())} of {per_die[0].numel()} slots differ). "
+            "The four expert windows are no longer a partition of the global top-8."
+        )
+    logger.info(f"topk seed={seed}: 4 dies bit-identical over {per_die[0].numel()} selections")
+
+
+@ring_fabric
+@mesh_4
+@pytest.mark.parametrize("seq_len", [1, 33, 128], ids=["decode", "s33", "s128"])
+def test_router_windows_partition_global_routing(fixture, seq_len):
+    """Concatenating the four local windows must reproduce the global dense routing.
+
+    This is the direct statement of the EP contract: the multichip router returns
+    ``[1, 1, S, 32]`` per die, and stitching them in device order must equal what
+    the single-chip router returns as one ``[1, 1, S, 128]`` row -- same experts,
+    same weights, normalised by the same global denominator.
+    """
+    x = fixture.rep(_hidden(fixture.hf, seq_len))
+    moe = fixture.config.global_config.moe
+    global_dense = fixture.dies(O.router_forward_optimized(x, fixture.multichip.router, moe))[0].float()
+    local = fixture.dies(
+        MC.router_forward_multichip(
+            x, fixture.multichip.router, fixture.multichip.expert_window, moe, fixture.config.local_moe
+        )
+    ).float()
+
+    n_local = fixture.config.local_moe.num_experts
+    stitched = torch.cat([local[d].reshape(-1, n_local) for d in range(fixture.config.num_devices)], dim=-1)
+    reference = global_dense.reshape(-1, moe.num_experts)
+    delta = (stitched - reference).abs().max().item()
+    logger.info(f"router windows seq={seq_len}: max |stitched - global| = {delta:.3e}")
+    assert torch.equal(stitched > 0, reference > 0), "the four windows do not select the global top-8"
+    assert delta == 0.0, f"routing weights differ by {delta}; the window matmul is not exact"
+    assert ((stitched > 0).sum(dim=-1) == moe.num_experts_per_tok).all()
+
+
+# --- prefill ------------------------------------------------------------------
+
+
+@ring_fabric
+@mesh_4
+@pytest.mark.parametrize("seq_len", [32, 128, 512, 33, 100, 257], ids=["s32", "s128", "s512", "s33", "s100", "s257"])
+def test_multichip_prefill_vs_single_chip(fixture, seq_len):
+    """Prefill against the single-chip TTNN baseline, aligned and non-aligned.
+
+    The non-aligned lengths are the point of the parametrisation: nothing in the
+    multichip path may turn a decoder that accepted any prompt length into one
+    that only accepts multiples of a chunk, tile, page or collective block. The
+    collectives scatter on dim 3 (hidden, 2048), which is independent of S.
+    """
+    hidden = _hidden(fixture.hf, seq_len)
+    base = fixture.dies(fixture.baseline_prefill(hidden))[0:1].float()
+    multi = fixture.dies(fixture.multichip_prefill(hidden))
+
+    spread = (multi - multi[0:1]).abs().max().item()
+    assert spread == 0.0, f"S={seq_len}: layer output differs across dies by {spread}; the all-reduce is not complete"
+    assert tuple(multi.shape) == (fixture.config.num_devices, 1, seq_len, fixture.hf.hidden_size), (
+        f"S={seq_len}: four dies of [1,1,S,H] concatenated on dim 0 came back {tuple(multi.shape)}; "
+        "the replicated layer contract is not intact"
+    )
+
+    passing, message = comp_pcc(base, multi[0:1].float(), PCC_VS_SINGLE_CHIP)
+    logger.info(f"multichip prefill S={seq_len} vs single-chip TTNN: {message}")
+    assert passing, f"multichip prefill S={seq_len} below {PCC_VS_SINGLE_CHIP} vs single-chip: {message}"
+
+
+@ring_fabric
+@mesh_4
+@pytest.mark.parametrize("seq_len", [128, 33], ids=["s128", "s33"])
+def test_multichip_prefill_vs_hf(fixture, reference, seq_len):
+    """The end-to-end bar: the same 0.995 PCC against HF the single chip clears."""
+    layer, hf_config = reference
+    hidden = _hidden(hf_config, seq_len)
+    ref = _reference_layer(layer, hf_config, hidden)
+    multi = fixture.dies(fixture.multichip_prefill(hidden))[0].reshape(1, seq_len, hf_config.hidden_size)
+    passing, message = comp_pcc(ref, multi.float(), PCC_VS_HF)
+    logger.info(f"multichip prefill S={seq_len} vs HF: {message}")
+    assert passing, f"multichip prefill S={seq_len} below {PCC_VS_HF} vs HF: {message}"
+
+
+@ring_fabric
+@mesh_4
+def test_multichip_prefill_is_deterministic(fixture):
+    """Bitwise repeatability, including through two collectives per layer."""
+    hidden = _hidden(fixture.hf, 128)
+    outs = [fixture.dies(fixture.multichip_prefill(hidden)).clone() for _ in range(3)]
+    assert torch.equal(outs[0], outs[1]), "multichip prefill run 1 != run 2 (bitwise)"
+    assert torch.equal(outs[0], outs[2]), "multichip prefill run 1 != run 3 (bitwise)"
+    logger.info("multichip prefill: 3 runs bit-identical on all 4 dies")
+
+
+# --- KV cache and decode ------------------------------------------------------
+
+
+@ring_fabric
+@mesh_4
+def test_local_kv_cache_layout(fixture):
+    """Each die owns exactly one KV head, and the four together hold the whole cache.
+
+    This is the memory half of the TP decision: 512 B per token per layer per die
+    instead of 2048. The test does not merely check the *shape* -- it prefills
+    both paths from the same prompt and asserts that stacking the four dies' K
+    caches on the head axis reproduces the single-chip cache, which is what
+    proves the head *assignment* matches the wqkv column split rather than just
+    the head count.
+    """
+    cfg = fixture.config
+    base_kv = F.create_kv_cache(fixture.mesh, cfg.global_config.attention, 1, 128, block_size=BLOCK_SIZE)
+    mc_kv = MC.create_mesh_kv_cache(fixture.mesh, cfg, 1, 128, block_size=BLOCK_SIZE)
+
+    assert mc_kv.k.shape[1] == 1, f"per-die KV cache has {mc_kv.k.shape[1]} heads, expected 1"
+    assert base_kv.k.shape[1] == cfg.global_config.attention.num_key_value_heads
+    assert mc_kv.is_paged and base_kv.is_paged
+
+    hidden = _hidden(fixture.hf, 64)
+    fixture.baseline_prefill(hidden, kv_cache=base_kv)
+    fixture.multichip_prefill(hidden, kv_cache=mc_kv)
+
+    # The cache's dim 0 is the physical block index, so the four dies are
+    # reassembled on the *head* axis -- which is also exactly what makes this a
+    # test of head ownership rather than of head count.
+    n_kv = cfg.global_config.attention.num_key_value_heads
+    # The baseline cache is replicated, so any one die's copy is the reference;
+    # take die 0's four heads out of the 16 the concat produces.
+    base_k = fixture.dies(base_kv.k, dim=1)[:, :n_kv].float()  # [blocks, 4, block, head_dim]
+    stitched = fixture.dies(mc_kv.k, dim=1).float()  # 4 x [blocks, 1, block, head_dim]
+    assert stitched.shape == base_k.shape, (stitched.shape, base_k.shape)
+    passing, message = comp_pcc(base_k, stitched, 0.999)
+    logger.info(f"local KV head layout: 4x[.,1,.,.] stitched vs single-chip [.,4,.,.]: {message}")
+    assert passing, f"per-die KV heads are not the single-chip heads in device order: {message}"
+
+    bytes_per_token = mc_kv.k.shape[1] * cfg.local_attention.head_dim * 2 * 2
+    logger.info(f"per-die KV: {bytes_per_token} B/token/layer (single chip: {bytes_per_token * 4})")
+    assert bytes_per_token == 512
+
+
+@ring_fabric
+@mesh_4
+@pytest.mark.parametrize("block_size", [None, 32], ids=["contiguous", "paged32"])
+def test_multichip_decode_vs_single_chip(fixture, block_size):
+    """One decode step against the single-chip baseline, both cache modes.
+
+    Both paths are prefilled with the same prompt into their own caches, so this
+    covers the paged write path (``paged_fill_cache``), the paged update
+    (``paged_update_cache`` at 1 KV head, which the design phase flagged as
+    unexercised), the page table, and ``cur_pos_tensor``, not just the matmuls.
+    """
+    cfg = fixture.config
+    prompt = 32
+    full = _hidden(fixture.hf, prompt + 1)
+    base_kv = F.create_kv_cache(fixture.mesh, cfg.global_config.attention, 1, MAX_SEQ, block_size=block_size)
+    mc_kv = MC.create_mesh_kv_cache(fixture.mesh, cfg, 1, MAX_SEQ, block_size=block_size)
+    fixture.baseline_prefill(full[:, :, :prompt, :], kv_cache=base_kv)
+    fixture.multichip_prefill(full[:, :, :prompt, :], kv_cache=mc_kv)
+
+    pos = ttnn.from_torch(
+        torch.tensor([prompt], dtype=torch.int32),
+        dtype=ttnn.int32,
+        device=fixture.mesh,
+        mesh_mapper=ttnn.ReplicateTensorToMesh(fixture.mesh),
+    )
+    token = full[:, :, prompt : prompt + 1, :]
+
+    base = fixture.dies(
+        O.decoder_layer_decode_optimized(
+            fixture.rep(token),
+            fixture.baseline,
+            cfg.global_config,
+            fixture.cos,
+            fixture.sin,
+            base_kv,
+            pos,
+            prompt,
+            packed_experts=fixture.baseline_experts,
+        )
+    )[0:1].float()
+    multi = fixture.dies(
+        MC.decoder_layer_decode_multichip(
+            fixture.rep(token), fixture.multichip, cfg, fixture.ctx, fixture.cos, fixture.sin, mc_kv, pos, prompt
+        )
+    )
+    spread = (multi - multi[0:1]).abs().max().item()
+    assert spread == 0.0, f"decode output differs across dies by {spread}"
+    kind = "contiguous" if block_size is None else f"paged({block_size})"
+    passing, message = comp_pcc(base, multi[0:1].float(), PCC_VS_SINGLE_CHIP)
+    logger.info(f"multichip decode [{kind}] vs single-chip TTNN: {message}")
+    assert passing, f"multichip decode [{kind}] below {PCC_VS_SINGLE_CHIP}: {message}"
+
+
+@ring_fabric
+@mesh_4
+def test_multichip_decode_contiguous_batch8(fixture):
+    """The contiguous-cache SDPA workaround at batch > 1, against the single chip.
+
+    ``_sdpa_program_config`` is the layer's one hand-written program config and
+    the only place the multichip path departs from stage 02's tuning. It exists
+    because at TP=4 the contiguous cache has 1 KV head per die, so at batch 1
+    SDPA-decode asks for all 110 worker cores on that head and
+    ``sdpa_decode_program_factory.cpp:245`` refuses anything over 64. But
+    ``num_cores_per_head`` divides by the batch, so at batch 8 the op would have
+    asked for 13 and been legal *without* the config -- and the config is
+    supplied unconditionally on the contiguous path.
+
+    That makes batch > 1 the case that actually tests the workaround rather than
+    the failure it works around: here the cap is not rescuing anything, it is
+    only constraining, and so are the ``q_chunk_size``/``k_chunk_size`` of 32
+    that come with it (the default path picked its own). Every other contiguous
+    test is batch 1 and every batch > 1 test is paged, so without this one both
+    the cap and the chunk sizes are exercised in exactly one configuration.
+    """
+    cfg = fixture.config
+    batch, prompt = 8, 32
+    per_user = [_hidden(fixture.hf, prompt + 1, seed=u) for u in range(batch)]
+
+    base_kv = F.create_kv_cache(fixture.mesh, cfg.global_config.attention, batch, 128, block_size=None)
+    mc_kv = MC.create_mesh_kv_cache(fixture.mesh, cfg, batch, 128, block_size=None)
+    assert not mc_kv.is_paged
+    for user, hidden in enumerate(per_user):
+        fixture.baseline_prefill(hidden[:, :, :prompt, :], kv_cache=base_kv, user_id=user)
+        fixture.multichip_prefill(hidden[:, :, :prompt, :], kv_cache=mc_kv, user_id=user)
+
+    tokens = torch.cat([h[:, :, prompt, :] for h in per_user], dim=1).reshape(1, 1, batch, fixture.hf.hidden_size)
+    pos = ttnn.from_torch(
+        torch.full((batch,), prompt, dtype=torch.int32),
+        dtype=ttnn.int32,
+        device=fixture.mesh,
+        mesh_mapper=ttnn.ReplicateTensorToMesh(fixture.mesh),
+    )
+
+    base = (
+        fixture.dies(
+            O.decoder_layer_decode_optimized(
+                fixture.rep(tokens),
+                fixture.baseline,
+                cfg.global_config,
+                fixture.cos,
+                fixture.sin,
+                base_kv,
+                pos,
+                prompt,
+                packed_experts=fixture.baseline_experts,
+            )
+        )[0]
+        .reshape(-1, fixture.hf.hidden_size)[:batch]
+        .float()
+    )
+    multi_all = fixture.dies(
+        MC.decoder_layer_decode_multichip(
+            fixture.rep(tokens), fixture.multichip, cfg, fixture.ctx, fixture.cos, fixture.sin, mc_kv, pos, prompt
+        )
+    )
+    spread = (multi_all - multi_all[0:1]).abs().max().item()
+    assert spread == 0.0, f"decode output differs across dies by {spread}"
+    multi = multi_all[0].reshape(-1, fixture.hf.hidden_size)[:batch].float()
+
+    assert (multi - multi[0:1]).abs().max().item() > 1e-3, "all users identical (broadcast bug)"
+    for user in range(batch):
+        passing, message = comp_pcc(base[user : user + 1], multi[user : user + 1], PCC_VS_SINGLE_CHIP)
+        logger.info(f"multichip decode [contiguous, batch 8] user {user} vs single-chip TTNN: {message}")
+        assert passing, f"contiguous batch-8 user {user} below {PCC_VS_SINGLE_CHIP}: {message}"
+
+
+@ring_fabric
+@mesh_4
+def test_multichip_multi_step_decode_vs_hf(fixture, reference):
+    """Four consecutive decode steps against HF, each at its own position."""
+    layer, hf_config = reference
+    cfg = fixture.config
+    prompt, steps = 32, 4
+    full = _hidden(hf_config, prompt + steps)
+    ref = _reference_layer(layer, hf_config, full)
+
+    kv = MC.create_mesh_kv_cache(fixture.mesh, cfg, 1, MAX_SEQ, block_size=BLOCK_SIZE)
+    fixture.multichip_prefill(full[:, :, :prompt, :], kv_cache=kv)
+
+    for step in range(steps):
+        p = prompt + step
+        pos = ttnn.from_torch(
+            torch.tensor([p], dtype=torch.int32),
+            dtype=ttnn.int32,
+            device=fixture.mesh,
+            mesh_mapper=ttnn.ReplicateTensorToMesh(fixture.mesh),
+        )
+        out = MC.decoder_layer_decode_multichip(
+            fixture.rep(full[:, :, p : p + 1, :]),
+            fixture.multichip,
+            cfg,
+            fixture.ctx,
+            fixture.cos,
+            fixture.sin,
+            kv,
+            pos,
+            p,
+        )
+        got = fixture.dies(out)[0].reshape(1, -1).float()
+        passing, message = comp_pcc(ref[:, p, :], got, 0.99)
+        logger.info(f"multichip decode step {step} (pos {p}) vs HF: {message}")
+        assert passing, f"multichip decode step {step} below 0.99: {message}"
+
+
+@ring_fabric
+@mesh_4
+@pytest.mark.parametrize("batch", [1, 2, 8, 32], ids=["b1", "b2", "b8", "b32"])
+def test_multichip_decode_batch(fixture, reference, batch):
+    """Multi-user decode, each user against its own HF reference.
+
+    32 is the ceiling and it is a TTNN op limit that TP does not move:
+    ``nlp_create_qkv_heads_decode_device_operation.cpp:51`` asserts
+    ``num_users <= 32``, and that op is on the per-die path too. Per-user
+    references are what prove routing is per-user rather than broadcast, which
+    matters more under EP than on one die -- a die whose window is empty for one
+    user and full for another exercises the dynamic ``nnz`` path in both
+    directions inside a single program.
+    """
+    layer, hf_config = reference
+    cfg = fixture.config
+    prompt = 32
+    kv = MC.create_mesh_kv_cache(fixture.mesh, cfg, batch, 128, block_size=BLOCK_SIZE)
+    per_user = [_hidden(hf_config, prompt + 1, seed=u) for u in range(batch)]
+    for user, hidden in enumerate(per_user):
+        fixture.multichip_prefill(hidden[:, :, :prompt, :], kv_cache=kv, user_id=user)
+
+    tokens = torch.cat([h[:, :, prompt, :] for h in per_user], dim=1).reshape(1, 1, batch, hf_config.hidden_size)
+    pos = ttnn.from_torch(
+        torch.full((batch,), prompt, dtype=torch.int32),
+        dtype=ttnn.int32,
+        device=fixture.mesh,
+        mesh_mapper=ttnn.ReplicateTensorToMesh(fixture.mesh),
+    )
+    out = MC.decoder_layer_decode_multichip(
+        fixture.rep(tokens), fixture.multichip, cfg, fixture.ctx, fixture.cos, fixture.sin, kv, pos, prompt
+    )
+    got = fixture.dies(out)[0].reshape(-1, hf_config.hidden_size)[:batch].float()
+    assert torch.isfinite(got).all(), f"batch={batch} produced non-finite values"
+
+    for user, hidden in enumerate(per_user):
+        ref_user = _reference_layer(layer, hf_config, hidden)[:, prompt, :]
+        passing, message = comp_pcc(ref_user, got[user : user + 1], 0.99)
+        logger.info(f"multichip decode batch={batch} user {user} vs HF: {message}")
+        assert passing, f"batch={batch} user {user} below 0.99: {message}"
+
+    if batch > 1:
+        assert (got - got[0]).abs().max().item() > 1e-3, f"batch={batch}: all users identical (broadcast bug)"
+
+
+# --- stage 04: the layer's own shape and layout contract ----------------------
+
+
+@ring_fabric
+@mesh_4
+@pytest.mark.parametrize("batch", [1, 8], ids=["b1", "b8"])
+def test_decode_output_layout_matches_input(fixture, batch):
+    """The decode layer must return exactly the tensor contract it takes.
+
+    Replicated ``[1, 1, B, 2048]``, bfloat16, TILE, DRAM-interleaved, logical
+    shape included -- that is what lets 48 layers stack with no boundary
+    conversion, and it is the *inter-layer residual layout contract* that
+    ``doc/optimized_multichip_decoder/README.md`` writes down for full-model
+    bringup.
+
+    It is asserted rather than assumed because stage 04's persistent collective
+    buffers can break it silently. A persistent output buffer imposes its own
+    logical shape on the op's result, and the layer's two all-reduces have the
+    same *padded* shape but different *logical* ones -- the attention partial is
+    32 rows out of ``wo``, the expert partial is ``batch``. Keyed on the padded
+    shape alone they collide and the layer returns a 32-row tensor; every test
+    that compares a path against itself still passes.
+    """
+    cfg = fixture.config
+    prompt = 32
+    kv = MC.create_mesh_kv_cache(fixture.mesh, cfg, batch, 128, block_size=BLOCK_SIZE)
+    hidden = _hidden(fixture.hf, prompt + 1)
+    for user in range(batch):
+        fixture.multichip_prefill(hidden[:, :, :prompt, :], kv_cache=kv, user_id=user)
+
+    token = fixture.rep(hidden[:, :, prompt, :].reshape(1, 1, 1, -1).repeat(1, 1, batch, 1))
+    pos = ttnn.from_torch(
+        torch.full((batch,), prompt, dtype=torch.int32),
+        dtype=ttnn.int32,
+        device=fixture.mesh,
+        mesh_mapper=ttnn.ReplicateTensorToMesh(fixture.mesh),
+    )
+    out = MC.decoder_layer_decode_multichip(
+        token, fixture.multichip, cfg, fixture.ctx, fixture.cos, fixture.sin, kv, pos, prompt
+    )
+    assert list(out.shape) == list(token.shape), (
+        f"decode layer changed the logical shape: in {list(token.shape)}, out {list(out.shape)}. "
+        "48 of these stack, so the output contract must equal the input contract."
+    )
+    assert out.dtype == token.dtype, f"dtype changed: {token.dtype} -> {out.dtype}"
+    assert out.layout == token.layout, f"layout changed: {token.layout} -> {out.layout}"
+    assert out.memory_config() == token.memory_config(), (
+        f"memory config changed: {token.memory_config()} -> {out.memory_config()}; "
+        "there must be no inter-layer reshard"
+    )
+    spread = (fixture.dies(out) - fixture.dies(out)[0:1]).abs().max().item()
+    assert spread == 0.0, f"decode output differs across dies by {spread}"
+
+
+# --- the dynamic-nnz hazard ---------------------------------------------------
+
+
+@ring_fabric
+@mesh_4
+def test_expert_window_can_be_empty(fixture):
+    """A die holding none of the global top-8 must contribute an exact zero.
+
+    Constructed rather than hoped for: the router weight keeps amplified real
+    rows for experts 0..31 and zeroed rows for 32..127, so the routing collapses
+    onto the low end of the expert range and at least two dies end up holding
+    none of the global top-8. Measured, the split is **[6, 2, 0, 0]** -- die 0
+    full, die 1 partial, dies 2 and 3 empty -- which is a better test than a
+    clean [8,0,0,0] would have been, because it exercises both hazards at once:
+
+    * ``E_local = 0``, where a host-computed ``nnz`` would be 8 against zero live
+      sparsity entries and where an uninitialised ``sparse_matmul`` output would
+      put a NaN into the all-reduce and poison the layer on *every* die;
+    * ``0 < E_local < top_k``, the ordinary EP case, which is data-dependent and
+      is exactly why no single ``nnz`` can be computed on the host for a program
+      that runs on four dies at once.
+
+    So the assertions are on the *properties*, not on the exact split: the live
+    counts must sum to top-8 (the windows are a partition), at least one die must
+    be empty (the hazard is reached), every empty die must contribute an exact
+    zero, and the whole layer must still match the single-chip baseline.
+    """
+    cfg = fixture.config
+    hf = fixture.hf
+    n_local = cfg.local_moe.num_experts
+
+    # Amplified real router rows for experts 0..31, zeroed rows for 32..127. The
+    # zero rows give those experts a logit of exactly 0, so they only win a slot
+    # when fewer than 8 of the first 32 come out positive -- which is what
+    # produces the [6, 2, 0, 0] split rather than [8, 0, 0, 0].
+    #
+    # **The gain is 4x, and that number is a hazard, not a taste.** A first
+    # version of this test used 10x *synthetic* rows, which -- against the
+    # rms-normed activation, not the 0.02-scaled raw hidden -- produces logits
+    # with a standard deviation near 450 and a top-8 spread past 1000.
+    # ``exp(-1000)`` is exactly zero in bf16, so two of the eight routing weights
+    # underflowed, ``count_nonzero(sparsity)`` fell below the ``nnz = top_k *
+    # batch`` that the *single-chip* baseline passes, and the board deadlocked
+    # exactly as ``sparse_matmul_device_operation.cpp:205-211`` says it will --
+    # it had to be killed and reset. The multichip leg, on ``nnz=None``, was
+    # unaffected. That is this stage's own reproduction of the hazard the design
+    # phase only read about, and it is recorded in ``work_log.md``. At 4x on the
+    # real rows the top-8 spread is a few units and every weight stays normal.
+    forced = torch.zeros(cfg.global_config.moe.num_experts, hf.hidden_size)
+    forced[:n_local] = 4.0 * fixture.torch_router[:n_local]
+    forced_router = _replicate(
+        forced.T.contiguous().reshape(1, 1, hf.hidden_size, cfg.global_config.moe.num_experts), fixture.mesh
+    )
+
+    x = fixture.rep(_hidden(hf, 1))
+    routing = MC.router_forward_multichip(
+        x, forced_router, fixture.multichip.expert_window, cfg.global_config.moe, cfg.local_moe
+    )
+    live = [int((fixture.dies(routing)[d].reshape(-1) > 0).sum()) for d in range(cfg.num_devices)]
+    logger.info(f"forced routing: live experts per die = {live} (sum {sum(live)})")
+    assert sum(live) == cfg.global_config.moe.num_experts_per_tok, (
+        f"the four windows hold {sum(live)} experts, not the global top-"
+        f"{cfg.global_config.moe.num_experts_per_tok}: {live}"
+    )
+    empty = [d for d in range(cfg.num_devices) if live[d] == 0]
+    assert empty, f"the forced routing did not empty any die: {live}"
+
+    partial = fixture.dies(MC.moe_decode_multichip(x, routing, fixture.multichip.experts, cfg.local_moe)).float()
+    assert torch.isfinite(partial).all(), "an empty expert window produced non-finite output"
+    empty_max = partial[empty].abs().max().item()
+    logger.info(f"empty-window partials: max |value| on dies {empty} = {empty_max}")
+    assert empty_max == 0.0, f"a die with no live experts contributed {empty_max}, not an exact zero"
+
+    forced_baseline = F.DecoderLayerWeights(
+        input_layernorm=fixture.multichip.input_layernorm,
+        post_attention_layernorm=fixture.multichip.post_attention_layernorm,
+        attention=None,
+        router=forced_router,
+        experts=None,
+    )
+    forced_multichip = MC.MultichipWeights(
+        input_layernorm=fixture.multichip.input_layernorm,
+        post_attention_layernorm=fixture.multichip.post_attention_layernorm,
+        router=forced_router,
+        expert_window=fixture.multichip.expert_window,
+        experts=fixture.multichip.experts,
+    )
+    hidden = _hidden(hf, 128)
+    base = fixture.dies(
+        O.decoder_layer_prefill_optimized(
+            fixture.rep(hidden),
+            forced_baseline,
+            cfg.global_config,
+            fixture.cos,
+            fixture.sin,
+            fixture.baseline_sparsity,
+            fixture.baseline_experts,
+        )
+    )[0:1].float()
+    multi = fixture.dies(
+        MC.decoder_layer_prefill_multichip(
+            fixture.rep(hidden),
+            forced_multichip,
+            cfg,
+            fixture.ctx,
+            fixture.cos,
+            fixture.sin,
+            fixture.sparsity,
+        )
+    )[0:1].float()
+    passing, message = comp_pcc(base, multi, PCC_VS_SINGLE_CHIP)
+    logger.info(f"layer under maximally unbalanced routing vs single-chip: {message}")
+    assert passing, message
+
+
+# --- stacking and trace -------------------------------------------------------
+
+
+@ring_fabric
+@mesh_4
+def test_stacked_layer_io_contract(fixture):
+    """The layer's output must be usable as its own input, unmodified.
+
+    Stage 04 stacks 48 of these. The contract is a replicated
+    ``[1, 1, B, 2048]`` in DRAM in both directions, so feeding the output
+    straight back in must work with no gather, reshard, layout change or dtype
+    cast in between -- and must still match the single-chip baseline stacked the
+    same way, which is what rules out a per-layer boundary conversion hiding
+    inside the comparison.
+    """
+    hidden = _hidden(fixture.hf, 128)
+
+    base_out = fixture.baseline_prefill(hidden)
+    base_out2 = O.decoder_layer_prefill_optimized(
+        base_out,
+        fixture.baseline,
+        fixture.config.global_config,
+        fixture.cos,
+        fixture.sin,
+        fixture.baseline_sparsity,
+        fixture.baseline_experts,
+    )
+
+    multi_out = fixture.multichip_prefill(hidden)
+    assert multi_out.memory_config() == ttnn.DRAM_MEMORY_CONFIG
+    assert multi_out.layout == ttnn.TILE_LAYOUT and multi_out.dtype == ttnn.bfloat16
+    multi_out2 = MC.decoder_layer_prefill_multichip(
+        multi_out,
+        fixture.multichip,
+        fixture.config,
+        fixture.ctx,
+        fixture.cos,
+        fixture.sin,
+        fixture.sparsity,
+    )
+    assert multi_out2.shape == multi_out.shape
+    assert multi_out2.memory_config() == multi_out.memory_config()
+    assert multi_out2.dtype == multi_out.dtype and multi_out2.layout == multi_out.layout
+
+    passing, message = comp_pcc(
+        fixture.dies(base_out2)[0:1].float(), fixture.dies(multi_out2)[0:1].float(), PCC_VS_SINGLE_CHIP
+    )
+    logger.info(f"two stacked multichip layers vs two stacked single-chip layers: {message}")
+    assert passing, f"stacked layers diverge: {message}"
+
+
+@ring_fabric
+@mesh_4
+def test_multichip_decode_is_traceable(fixture):
+    """Warmed trace capture and replay on the mesh, with a live input buffer.
+
+    Trace capture and CCL interact: global semaphores and any persistent CCL
+    buffer must exist before ``begin_trace_capture`` and nothing may allocate
+    inside it. ``MeshContext`` allocates its semaphores at construction and the
+    ``_ones_column`` constant is populated by the eager warm-up call below, which
+    is what makes the capture legal.
+    """
+    cfg = fixture.config
+    prompt = 32
+    full = _hidden(fixture.hf, prompt + 1)
+    kv = MC.create_mesh_kv_cache(fixture.mesh, cfg, 1, MAX_SEQ, block_size=BLOCK_SIZE)
+    fixture.multichip_prefill(full[:, :, :prompt, :], kv_cache=kv)
+
+    tt_in = fixture.rep(full[:, :, prompt : prompt + 1, :])
+    pos = ttnn.from_torch(
+        torch.tensor([prompt], dtype=torch.int32),
+        dtype=ttnn.int32,
+        device=fixture.mesh,
+        mesh_mapper=ttnn.ReplicateTensorToMesh(fixture.mesh),
+    )
+
+    def step():
+        return MC.decoder_layer_decode_multichip(
+            tt_in, fixture.multichip, cfg, fixture.ctx, fixture.cos, fixture.sin, kv, pos, prompt
+        )
+
+    eager = fixture.dies(step()).clone()
+    ttnn.synchronize_device(fixture.mesh)
+
+    trace_id = ttnn.begin_trace_capture(fixture.mesh, cq_id=0)
+    traced_out = step()
+    ttnn.end_trace_capture(fixture.mesh, trace_id, cq_id=0)
+
+    ttnn.execute_trace(fixture.mesh, trace_id, cq_id=0, blocking=True)
+    replayed = fixture.dies(traced_out).clone()
+    passing, message = comp_pcc(eager.float(), replayed.float(), 0.999)
+    logger.info(f"multichip traced decode vs eager: {message}")
+    assert passing, f"traced replay disagrees with eager: {message}"
+
+    other = torch.randn(1, 1, 1, fixture.hf.hidden_size) * 0.02
+    ttnn.copy_host_to_device_tensor(
+        ttnn.from_torch(
+            other,
+            dtype=ttnn.bfloat16,
+            layout=ttnn.TILE_LAYOUT,
+            mesh_mapper=ttnn.ReplicateTensorToMesh(fixture.mesh),
+        ),
+        tt_in,
+    )
+    ttnn.execute_trace(fixture.mesh, trace_id, cq_id=0, blocking=True)
+    changed = fixture.dies(traced_out).clone()
+    delta = (replayed.float() - changed.float()).abs().max().item()
+    logger.info(f"multichip traced replay delta after input swap = {delta:.6f}")
+    assert delta > 1e-3, "the mesh trace is not reading the live input buffer"
+
+    ttnn.release_trace(fixture.mesh, trace_id)
+
+
+@ring_fabric
+@mesh_4
+def test_multichip_decode_stress_is_deterministic(fixture):
+    """Repeated decode at the same position must be bit-identical, run to run.
+
+    The collectives are the new source of non-determinism here: two async CCLs
+    per layer with cycling semaphores, on a ring where four dies race to the same
+    reduction. Reduction order is fixed by the topology, so the result must be
+    exactly repeatable; a drift would mean the semaphore cycling is letting two
+    collectives overlap.
+    """
+    cfg = fixture.config
+    prompt = 32
+    full = _hidden(fixture.hf, prompt + 1)
+    kv = MC.create_mesh_kv_cache(fixture.mesh, cfg, 1, MAX_SEQ, block_size=BLOCK_SIZE)
+    fixture.multichip_prefill(full[:, :, :prompt, :], kv_cache=kv)
+    pos = ttnn.from_torch(
+        torch.tensor([prompt], dtype=torch.int32),
+        dtype=ttnn.int32,
+        device=fixture.mesh,
+        mesh_mapper=ttnn.ReplicateTensorToMesh(fixture.mesh),
+    )
+    token = fixture.rep(full[:, :, prompt : prompt + 1, :])
+
+    outs = []
+    for _ in range(20):
+        outs.append(
+            fixture.dies(
+                MC.decoder_layer_decode_multichip(
+                    token, fixture.multichip, cfg, fixture.ctx, fixture.cos, fixture.sin, kv, pos, prompt
+                )
+            ).clone()
+        )
+    for i, o in enumerate(outs[1:], start=1):
+        assert torch.equal(outs[0], o), f"decode iteration {i} differs from iteration 0 (bitwise)"
+    logger.info("multichip decode: 20 consecutive steps bit-identical across all 4 dies")
+
+
+# --- runtime fallback audit ---------------------------------------------------
+
+
+@ring_fabric
+@mesh_4
+@pytest.mark.parametrize("batch", [1, 32], ids=["b1", "b32"])
+def test_no_runtime_fallbacks(fixture, batch):
+    """None of the imported single-chip helpers may quietly take a slower path.
+
+    All three of them see different inputs under TP/EP than they were tuned
+    against, and all three fall back *silently* -- a PCC test cannot tell the
+    difference. In particular ``_dram_sharded_ok`` needs both weight dims
+    divisible by ``8 banks x 32 = 256``, and per-die wqkv N is 1280 = 5x256, one
+    factor of two from failing; if it ever did, stage 02's 1.11x DRAM-sharded
+    decode attention would disappear with no error at all.
+    """
+    audit = MC.fallback_audit(fixture.multichip, fixture.config, batch)
+    logger.info(f"multichip fallback audit at batch {batch}: {audit}")
+    assert audit["dram_sharded_taken"], "decode attention fell back to the interleaved path"
+    assert audit["dram_sharded_qkv"] == (2048, 1280), audit
+    assert audit["dram_sharded_wo"] == (1024, 2048), audit
+    # **Literals, not the module constants.** ``EXPERT_IN0_BLOCK_W_*`` are now
+    # derived from ``DEFAULT_PRECISION``, which is the same value the audit
+    # resolves from -- so comparing them was an identity that could not fail and
+    # could no longer catch a width regression. These are the widths stage 07
+    # selected and measured; changing the default must fail here and be
+    # re-measured, not silently ratified. ``test_precision_config.py`` pins the
+    # same two literals through the full construction path.
+    assert audit["gate_up_in0_block_w"] == 64, "gate/up block width moved off the stage-07 selection"
+    assert audit["down_in0_block_w"] == 24, "down block width moved off the stage-07 selection"
+    assert (O.EXPERT_IN0_BLOCK_W_GATE_UP, O.EXPERT_IN0_BLOCK_W_DOWN) == (
+        64,
+        24,
+    ), "the module constants no longer agree with the selected widths"
+    assert audit["local_heads"] == (8, 1) and audit["local_experts"] == 32, audit
+    # Batch 1 is the latency target and must keep the intermediates in L1 -- the
+    # traced A/B says L1 is 7.6% faster there. Batch 32 must not: the allocator
+    # refuses 234.88 MB outright (bank_manager.cpp:462), so a budget that let it
+    # through would be a crash, not a slow path. The inherited 40 MB constant
+    # would have separated these two correctly by luck while silently changing
+    # the answer for batches 2 to 16; the swept 128 MB is chosen against the
+    # measured L1-vs-DRAM crossover. See probes/l1_budget_probe.py.
+    assert audit["expert_intermediate_buffer"] == ("L1" if batch == 1 else "DRAM"), audit
+    # Stage 04. The sharded residual norm writes exactly the L1 shard the
+    # DRAM-sharded qkv projection reads, which is what lets the first norm's
+    # output cross into attention with no conversion at all. If that equality
+    # ever breaks, TTNN inserts a reshard between them and the layer gets slower
+    # with no error -- the same failure mode as the three above.
+    assert audit["norm_shard_feeds_qkv_directly"], (
+        "the sharded norm's output shard no longer matches attention's qkv input shard; "
+        "a silent reshard has been reintroduced between them"
+    )
+    assert audit["norm_shard_cores"] == 8, audit
+
+
+def test_meta_rope_weights_match_hf():
+    """The Meta channel permutation is a *pair*: Q/K rows and the QK-norm vectors.
+
+    Stage 01 chose HF-style RoPE precisely so neither permutation was needed.
+    Stage 04 adopts ``rotary_embedding_llama`` on the decode path, which brings
+    both back, and crossing them "runs fine and silently produces garbage"
+    (``weight_mapping.py``). This asserts the whole convention on the host, with
+    no device, so a mismatch fails here rather than as a PCC that is merely
+    lower.
+    """
+    import torch
+
+    from ..tt.weight_mapping import hf_to_meta_channels, permute_head_vector_to_meta, permute_wqkv_to_meta
+
+    hd, nh, nkv, hidden = 128, 8, 1, 2048
+    perm = hf_to_meta_channels(hd)
+    inv = torch.argsort(perm)
+
+    # 1. The permutation is a permutation, and it is the interleave it claims.
+    assert sorted(perm.tolist()) == list(range(hd))
+    assert perm[0] == 0 and perm[1] == hd // 2 and perm[2] == 1
+
+    # 2. HF rope on HF-ordered channels == Meta rope on Meta-ordered channels.
+    torch.manual_seed(0)
+    x = torch.randn(4, hd)
+    c, s = torch.randn(hd // 2).abs(), torch.randn(hd // 2)
+    cos_hf, sin_hf = torch.cat([c, c]), torch.cat([s, s])
+    hf = x * cos_hf + torch.cat([-x[:, hd // 2 :], x[:, : hd // 2]], dim=-1) * sin_hf
+    xm = x[:, perm]
+    cos_m, sin_m = cos_hf[perm], sin_hf[perm]
+    rot = torch.stack([-xm[:, 1::2], xm[:, 0::2]], dim=-1).reshape(xm.shape)
+    meta = xm * cos_m + rot * sin_m
+    assert torch.allclose(meta[:, inv], hf, atol=1e-6), (meta[:, inv] - hf).abs().max()
+
+    # 3. permute_wqkv_to_meta touches Q and K and leaves V alone.
+    wqkv = torch.randn(1, 1, hidden, (nh + 2 * nkv) * hd)
+    out = permute_wqkv_to_meta(wqkv, n_heads=nh, n_kv_heads=nkv, head_dim=hd)
+    assert out.shape == wqkv.shape
+    v0 = (nh + nkv) * hd
+    assert torch.equal(out[..., v0:], wqkv[..., v0:]), "V was permuted"
+    for h in range(nh + nkv):
+        lo = h * hd
+        assert torch.equal(out[..., lo : lo + hd], wqkv[..., lo : lo + hd][..., perm])
+    assert not torch.equal(out[..., :hd], wqkv[..., :hd]), "Q was not permuted"
+
+    # 4. Applying it twice is not identity -- i.e. forgetting it is detectable.
+    vec = torch.randn(hd)
+    assert not torch.equal(permute_head_vector_to_meta(vec, head_dim=hd), vec)
+    assert torch.equal(permute_head_vector_to_meta(vec, head_dim=hd)[inv], vec)

code/models/demos/blackhole/qwen3_coder_30b_a3b/tests/test_optimized_decoder.py

@@ -0,0 +1,503 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+"""Correctness of the optimized decoder — every stage-01 guarantee, re-checked.
+
+Optimization here changed program configs and packed two weights together; it
+did not change the maths. So these tests deliberately re-run the *same*
+contracts stage 01 established rather than a reduced subset, because the way an
+optimization usually breaks a model is by quietly narrowing what still works:
+a program config legal only at tile-aligned lengths, a packed weight whose
+halves are swapped, a trace that captured a stale buffer.
+
+``test_optimized_vs_functional_precision_delta`` is the sharpest of these.
+Packing gate/up and widening ``in0_block_w`` are value-preserving, but the
+optimized path also holds expert weights in bfloat4_b and attention projections
+in bfloat8_b, so bit-identity is not available -- and asserting it would be
+asserting the optimization away. Instead it bounds the gap between the two
+implementations at 0.999 PCC. That is far tighter than either one's distance to
+HF, so a real defect (swapped packed halves, a mis-sliced block) still cannot
+hide inside it, while the quantisation that was measured and accepted can.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+from loguru import logger
+
+import ttnn
+from models.common.utility_functions import comp_allclose, comp_pcc
+
+from ..tt import functional_decoder as F
+from ..tt import optimized_decoder as O
+from ..tt.weight_mapping import convert_layer_weights
+from .reference import build_reference_layer, layer_state_dict, rotary_embeddings
+
+LAYER_IDX = 0
+PCC_REQUIRED = 0.995  # same bar as the functional decoder
+MAX_SEQ = 1024
+BLOCK_SIZE = 32
+TRACE_REGION_SIZE = 50331648
+
+
+@pytest.fixture(scope="module")
+def reference():
+    return build_reference_layer(LAYER_IDX)
+
+
+@pytest.fixture(scope="module")
+def torch_weights(reference):
+    _, hf_config = reference
+    return convert_layer_weights(layer_state_dict(LAYER_IDX), hf_config)
+
+
+def _hidden(hf_config, seq_len, seed=0):
+    torch.manual_seed(seed)
+    return torch.randn(1, seq_len, hf_config.hidden_size, dtype=torch.float32) * 0.02
+
+
+def _reference_layer(layer, hf_config, hidden):
+    seq_len = hidden.shape[1]
+    cos, sin = rotary_embeddings(hf_config, seq_len)
+    mask = torch.full((seq_len, seq_len), float("-inf")).triu(1).reshape(1, 1, seq_len, seq_len)
+    with torch.no_grad():
+        out = layer(hidden, position_embeddings=(cos, sin), attention_mask=mask)
+    return out[0] if isinstance(out, tuple) else out
+
+
+def _to_device(t, mesh_device):
+    return ttnn.from_torch(
+        t, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=mesh_device, memory_config=ttnn.DRAM_MEMORY_CONFIG
+    )
+
+
+def _norm(t, mesh_device):
+    return ttnn.from_torch(
+        t.reshape(1, 1, 1, -1).float(),
+        dtype=ttnn.bfloat16,
+        layout=ttnn.TILE_LAYOUT,
+        device=mesh_device,
+        memory_config=ttnn.DRAM_MEMORY_CONFIG,
+    )
+
+
+def _build(mesh_device, hf_config, torch_weights, *, functional_weights=False):
+    """Upload only what the path under test actually reads.
+
+    ``decoder_layer_prefill_optimized`` / ``decoder_layer_decode_optimized``
+    read exactly three tensors off ``DecoderLayerWeights`` -- the two norms and
+    the router -- and take every projection and expert weight from
+    ``OptimizedWeights``. Calling ``F.upload_layer_weights`` here as well
+    uploaded ~1.2 GB of bf16 experts plus a third copy of wqkv/wo that nothing
+    on the optimized path ever touches. That, not a model or op limit, is what
+    used to exhaust DRAM at batch 8, so the batch coverage below was capped by
+    the harness measuring its own waste.
+
+    ``functional_weights=True`` restores the full set, needed only by the test
+    that runs the *functional* layer side by side for a precision delta.
+    """
+    config = F.DecoderLayerConfig.from_hf(hf_config)
+    if functional_weights:
+        weights = F.upload_layer_weights(torch_weights, mesh_device, config)
+    else:
+        weights = F.DecoderLayerWeights(
+            input_layernorm=_norm(torch_weights["input_layernorm"], mesh_device),
+            post_attention_layernorm=_norm(torch_weights["post_attention_layernorm"], mesh_device),
+            attention=None,  # optimized path uses OptimizedWeights.attention
+            router=F.upload_router_weight(torch_weights["router"], mesh_device),
+            experts=None,  # optimized path uses OptimizedWeights.gate_up_proj/down_proj
+        )
+    packed = O.upload_packed_expert_weights(torch_weights, mesh_device, config.moe)
+    cos_cache, sin_cache = F.build_rope_cache(hf_config, MAX_SEQ, mesh_device)
+    sparsity = F.build_expert_sparsity(mesh_device, config.moe.num_experts)
+    return config, weights, packed, cos_cache, sin_cache, sparsity
+
+
+@pytest.mark.parametrize("mesh_device", [(1, 1)], ids=["1x1"], indirect=True)
+@pytest.mark.parametrize("seq_len", [32, 128, 512, 33, 100, 257], ids=["s32", "s128", "s512", "s33", "s100", "s257"])
+def test_optimized_prefill_vs_reference(mesh_device, reference, torch_weights, seq_len):
+    """Aligned and non-aligned lengths must both survive the new program configs."""
+    layer, hf_config = reference
+    config, weights, packed, cos, sin, sparsity = _build(mesh_device, hf_config, torch_weights)
+    hidden = _hidden(hf_config, seq_len)
+
+    ref_out = _reference_layer(layer, hf_config, hidden)
+    tt_out = ttnn.to_torch(
+        O.decoder_layer_prefill_optimized(
+            _to_device(hidden.unsqueeze(0), mesh_device), weights, config, cos, sin, sparsity, packed
+        )
+    ).squeeze(0)
+
+    passing, pcc_message = comp_pcc(ref_out, tt_out, PCC_REQUIRED)
+    logger.info(f"optimized prefill seq={seq_len}: {pcc_message}")
+    assert passing, f"optimized prefill (seq={seq_len}) below {PCC_REQUIRED}: {pcc_message}"
+
+
+@pytest.mark.parametrize("mesh_device", [(1, 1)], ids=["1x1"], indirect=True)
+@pytest.mark.parametrize("seq_len", [128, 33], ids=["s128", "s33"])
+def test_optimized_vs_functional_precision_delta(mesh_device, reference, torch_weights, seq_len):
+    """Bound how far the optimized path may drift from the functional one.
+
+    The two are deliberately *not* bit-identical: the optimized path holds
+    expert weights in bfloat4_b and attention projections in bfloat8_b, which is
+    what makes it fast. So this asserts a tight bound on the gap rather than
+    equality -- close enough that a real defect (swapped packed halves, a
+    mis-sliced block) still cannot hide, but loose enough to admit the
+    quantisation that was measured and accepted.
+
+    The bound is empirical and is budgeted, not guessed. Against HF the
+    functional layer scores ~0.9995 and the optimized one ~0.9990, so the two
+    can differ by ~0.0013 purely from quantisation; measured, they differ by
+    0.00102 at S=128 and 0.00133 at S=33. 0.998 sits just outside that and well
+    inside anything a structural bug would produce -- swapping the packed
+    gate/up halves, for instance, drops this to below 0.5.
+    """
+    _, hf_config = reference
+    config, weights, packed, cos, sin, sparsity = _build(mesh_device, hf_config, torch_weights, functional_weights=True)
+    hidden = _hidden(hf_config, seq_len)
+
+    base = ttnn.to_torch(
+        F.decoder_layer_prefill(_to_device(hidden.unsqueeze(0), mesh_device), weights, config, cos, sin, sparsity)
+    )
+    opt = ttnn.to_torch(
+        O.decoder_layer_prefill_optimized(
+            _to_device(hidden.unsqueeze(0), mesh_device), weights, config, cos, sin, sparsity, packed
+        )
+    )
+
+    passing, message = comp_pcc(base, opt, 0.998)
+    logger.info(comp_allclose(base, opt))
+    logger.info(f"seq={seq_len} optimized vs functional (bf16 vs {O.EXPERT_WEIGHT_DTYPE}): {message}")
+    assert passing, f"optimized diverges from functional beyond expert quantisation: {message}"
+
+
+@pytest.mark.parametrize("mesh_device", [(1, 1)], ids=["1x1"], indirect=True)
+@pytest.mark.parametrize("seq_len", [1, 128, 33], ids=["decode", "s128", "s33"])
+def test_optimized_router_matches_functional(mesh_device, reference, torch_weights, seq_len):
+    """The optimized router must route *identically*, not merely closely.
+
+    ``router_forward_optimized`` removes both keepdim reductions -- the max
+    becomes column 0 of the sorted top-k, and the sum moves after the scatter
+    and becomes a matmul. Neither is supposed to change the answer, so the test
+    asserts the strong property: the same 8 experts, in the same slots, with
+    weights equal to the functional router's within bf16 representation error.
+    A weaker PCC bound would let a genuinely different routing pass, which is
+    the failure mode that matters -- misrouting a token replaces its experts
+    outright rather than perturbing its output.
+    """
+    _, hf_config = reference
+    config = F.DecoderLayerConfig.from_hf(hf_config)
+    w_router = F.upload_router_weight(torch_weights["router"], mesh_device)
+    x = _to_device(_hidden(hf_config, seq_len).unsqueeze(0), mesh_device)
+
+    base = ttnn.to_torch(F.router_forward(x, w_router, config.moe)).float().reshape(seq_len, -1)
+    opt = ttnn.to_torch(O.router_forward_optimized(x, w_router, config.moe)).float().reshape(seq_len, -1)
+
+    assert torch.equal(base > 0, opt > 0), (
+        f"seq={seq_len}: optimized router selected different experts "
+        f"({(( base > 0) != (opt > 0)).sum().item()} slots differ)"
+    )
+    assert ((base > 0).sum(dim=-1) == config.moe.num_experts_per_tok).all()
+    delta = (base - opt).abs().max().item()
+    sums = opt.sum(dim=-1)
+    logger.info(
+        f"router seq={seq_len}: max |functional - optimized| = {delta:.3e}, weight sums in "
+        f"[{sums.min():.5f}, {sums.max():.5f}]"
+    )
+    assert delta < 5e-3, f"seq={seq_len}: routing weights differ by {delta}"
+    assert torch.allclose(sums, torch.ones_like(sums), atol=2e-2), f"seq={seq_len}: weights not normalised"
+
+
+@pytest.mark.parametrize("mesh_device", [(1, 1)], ids=["1x1"], indirect=True)
+@pytest.mark.parametrize("seq_len", [1, 33, 100, 128], ids=["decode", "s33", "s100", "s128"])
+def test_optimized_router_padding_is_zero(mesh_device, reference, torch_weights, seq_len):
+    """The routing tensor's tile row-padding must be exact zero, not +inf.
+
+    ``router_forward_optimized`` moved the sum after the scatter, so the divide
+    runs over whole tiles. Rows ``seq_len``..``ceil(seq_len/32)*32`` have a zero
+    numerator *and* a zero denominator, and unguarded ``ttnn.div`` returns
+    **+inf** there -- which the functional router, dividing before the scatter,
+    never did. No consumer was found that observes it (``to_torch`` returns the
+    logical shape, the sparsity path drops the padding when it converts to
+    ROW_MAJOR, and the scale multiply, ``rms_norm`` and ``fast_reduce_nc``
+    reduce along axes that are tile-aligned or not the padded one), so this is
+    a latent hazard rather than a live bug -- which is exactly the kind that a
+    later consumer turns into a silent NaN. The divisor is clamped in
+    ``router_forward_optimized``; this test is what stops the clamp being
+    optimized back out.
+
+    ``to_torch_with_padded_shape`` is the point of the test: ``ttnn.to_torch``
+    slices to the logical shape and would pass no matter what the padding held.
+    """
+    _, hf_config = reference
+    config = F.DecoderLayerConfig.from_hf(hf_config)
+    w_router = F.upload_router_weight(torch_weights["router"], mesh_device)
+    x = _to_device(_hidden(hf_config, seq_len).unsqueeze(0), mesh_device)
+
+    out = O.router_forward_optimized(x, w_router, config.moe)
+    padded = out.cpu().to_torch_with_padded_shape().float()
+    assert torch.isfinite(padded).all(), (
+        f"seq={seq_len}: routing tensor has {int((~torch.isfinite(padded)).sum())} non-finite "
+        f"entries in its padded shape {tuple(padded.shape)}"
+    )
+    pad = padded[..., seq_len:, :]
+    logger.info(
+        f"router seq={seq_len}: padded {tuple(padded.shape)}, {pad.numel()} padding entries, "
+        f"max |pad| = {pad.abs().max().item() if pad.numel() else 0.0}"
+    )
+    assert (pad == 0).all(), (
+        f"seq={seq_len}: {int((pad != 0).sum())} of {pad.numel()} padding entries are non-zero "
+        f"(max {pad.abs().max().item()})"
+    )
+
+
+@pytest.mark.parametrize("mesh_device", [(1, 1)], ids=["1x1"], indirect=True)
+@pytest.mark.parametrize("block_size", [None, 32], ids=["contiguous", "paged32"])
+def test_optimized_decode_matches_prefill(mesh_device, reference, torch_weights, block_size):
+    """Paged and contiguous KV caches both still work on the optimized path."""
+    layer, hf_config = reference
+    config, weights, packed, cos, sin, sparsity = _build(mesh_device, hf_config, torch_weights)
+    prompt_len = 32
+    hidden_full = _hidden(hf_config, prompt_len + 1)
+    ref_out = _reference_layer(layer, hf_config, hidden_full)[:, prompt_len, :]
+
+    kv_cache = F.create_kv_cache(mesh_device, config.attention, max_batch=1, max_seq_len=MAX_SEQ, block_size=block_size)
+    O.decoder_layer_prefill_optimized(
+        _to_device(hidden_full[:, :prompt_len, :].unsqueeze(0), mesh_device),
+        weights,
+        config,
+        cos,
+        sin,
+        sparsity,
+        packed,
+        kv_cache=kv_cache,
+    )
+    current_pos = ttnn.from_torch(torch.tensor([prompt_len], dtype=torch.int32), dtype=ttnn.int32, device=mesh_device)
+    out = O.decoder_layer_decode_optimized(
+        _to_device(hidden_full[:, prompt_len, :].reshape(1, 1, 1, hf_config.hidden_size), mesh_device),
+        weights,
+        config,
+        cos,
+        sin,
+        kv_cache,
+        current_pos,
+        prompt_len,
+        packed_experts=packed,
+    )
+    tt_out = ttnn.to_torch(out).reshape(1, hf_config.hidden_size)
+
+    passing, pcc_message = comp_pcc(ref_out, tt_out, 0.99)
+    kind = "contiguous" if block_size is None else f"paged({block_size})"
+    logger.info(f"optimized decode [{kind}]: {pcc_message}")
+    assert passing, f"optimized decode [{kind}] below 0.99: {pcc_message}"
+
+
+@pytest.mark.parametrize("mesh_device", [(1, 1)], ids=["1x1"], indirect=True)
+def test_optimized_multi_step_decode(mesh_device, reference, torch_weights):
+    """Several steps against a paged cache, each checked at its own position."""
+    layer, hf_config = reference
+    config, weights, packed, cos, sin, sparsity = _build(mesh_device, hf_config, torch_weights)
+    prompt_len, steps = 32, 4
+    hidden_full = _hidden(hf_config, prompt_len + steps)
+    ref_out = _reference_layer(layer, hf_config, hidden_full)
+
+    kv_cache = F.create_kv_cache(mesh_device, config.attention, max_batch=1, max_seq_len=MAX_SEQ, block_size=BLOCK_SIZE)
+    O.decoder_layer_prefill_optimized(
+        _to_device(hidden_full[:, :prompt_len, :].unsqueeze(0), mesh_device),
+        weights,
+        config,
+        cos,
+        sin,
+        sparsity,
+        packed,
+        kv_cache=kv_cache,
+    )
+
+    for step in range(steps):
+        pos = prompt_len + step
+        current_pos = ttnn.from_torch(torch.tensor([pos], dtype=torch.int32), dtype=ttnn.int32, device=mesh_device)
+        out = O.decoder_layer_decode_optimized(
+            _to_device(hidden_full[:, pos, :].reshape(1, 1, 1, hf_config.hidden_size), mesh_device),
+            weights,
+            config,
+            cos,
+            sin,
+            kv_cache,
+            current_pos,
+            pos,
+            packed_experts=packed,
+        )
+        passing, pcc_message = comp_pcc(ref_out[:, pos, :], ttnn.to_torch(out).reshape(1, -1), 0.99)
+        logger.info(f"optimized decode step {step} (pos {pos}): {pcc_message}")
+        assert passing, f"optimized decode step {step} below 0.99: {pcc_message}"
+
+
+@pytest.mark.parametrize("mesh_device", [(1, 1)], ids=["1x1"], indirect=True)
+def test_optimized_prefill_is_deterministic(mesh_device, reference, torch_weights):
+    """Bitwise repeatability, as required of the functional path."""
+    _, hf_config = reference
+    config, weights, packed, cos, sin, sparsity = _build(mesh_device, hf_config, torch_weights)
+    hidden = _hidden(hf_config, 128)
+
+    outs = [
+        ttnn.to_torch(
+            O.decoder_layer_prefill_optimized(
+                _to_device(hidden.unsqueeze(0), mesh_device), weights, config, cos, sin, sparsity, packed
+            )
+        ).clone()
+        for _ in range(3)
+    ]
+    assert torch.equal(outs[0], outs[1]), "optimized prefill run 1 != run 2 (bitwise)"
+    assert torch.equal(outs[0], outs[2]), "optimized prefill run 1 != run 3 (bitwise)"
+    logger.info("optimized prefill: 3 runs bit-identical")
+
+
+@pytest.mark.parametrize("device_params", [{"trace_region_size": TRACE_REGION_SIZE}], indirect=True)
+@pytest.mark.parametrize("mesh_device", [(1, 1)], ids=["1x1"], indirect=True)
+def test_optimized_decode_is_traceable(mesh_device, reference, torch_weights):
+    """Trace capture, bit-exact replay, and a live input buffer."""
+    _, hf_config = reference
+    config, weights, packed, cos, sin, sparsity = _build(mesh_device, hf_config, torch_weights)
+    prompt_len = 32
+    hidden_full = _hidden(hf_config, prompt_len + 1)
+
+    kv_cache = F.create_kv_cache(mesh_device, config.attention, max_batch=1, max_seq_len=MAX_SEQ, block_size=BLOCK_SIZE)
+    O.decoder_layer_prefill_optimized(
+        _to_device(hidden_full[:, :prompt_len, :].unsqueeze(0), mesh_device),
+        weights,
+        config,
+        cos,
+        sin,
+        sparsity,
+        packed,
+        kv_cache=kv_cache,
+    )
+
+    tt_in = _to_device(hidden_full[:, prompt_len, :].reshape(1, 1, 1, hf_config.hidden_size), mesh_device)
+    current_pos = ttnn.from_torch(torch.tensor([prompt_len], dtype=torch.int32), dtype=ttnn.int32, device=mesh_device)
+
+    def step():
+        return O.decoder_layer_decode_optimized(
+            tt_in, weights, config, cos, sin, kv_cache, current_pos, prompt_len, packed_experts=packed
+        )
+
+    eager = ttnn.to_torch(step()).clone()
+    trace_id = ttnn.begin_trace_capture(mesh_device, cq_id=0)
+    traced_out = step()
+    ttnn.end_trace_capture(mesh_device, trace_id, cq_id=0)
+
+    ttnn.execute_trace(mesh_device, trace_id, cq_id=0, blocking=True)
+    replayed = ttnn.to_torch(traced_out).clone()
+    passing, message = comp_pcc(eager, replayed, 0.999)
+    logger.info(f"optimized traced vs eager: {message}")
+    assert passing, f"optimized traced replay disagrees with eager: {message}"
+
+    # The trace must read the live buffer, not a value captured at record time.
+    other = (_hidden(hf_config, 1, seed=99)).reshape(1, 1, 1, hf_config.hidden_size)
+    ttnn.copy_host_to_device_tensor(ttnn.from_torch(other, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT), tt_in)
+    ttnn.execute_trace(mesh_device, trace_id, cq_id=0, blocking=True)
+    changed = ttnn.to_torch(traced_out).clone()
+    delta = (replayed.float() - changed.float()).abs().max().item()
+    logger.info(f"optimized traced replay delta after input swap = {delta:.6f}")
+    assert delta > 1e-3, "optimized trace is not reading the live input buffer"
+
+    ttnn.release_trace(mesh_device, trace_id)
+
+
+@pytest.mark.parametrize("mesh_device", [(1, 1)], ids=["1x1"], indirect=True)
+@pytest.mark.parametrize("batch", [1, 2, 8, 32], ids=["b1", "b2", "b8", "b32"])
+def test_optimized_decode_batch(mesh_device, reference, torch_weights, batch):
+    """Multi-user decode. Batch 1 is the latency target, but capability must hold.
+
+    This did not work in stage 01 at any batch above 1: ``sparse_matmul``
+    resolves ``batch_length`` differently depending on which operand is flagged
+    sparse, and the down projection landed on a branch that ignores the batch
+    dimension entirely. Each user is given a different prompt so a broadcast bug
+    -- every row returning user 0's answer -- cannot pass.
+
+    Writing this test surfaced a second defect beyond the sparsity one:
+    ``attention_prefill`` hardcoded ``user_id=0``, so every user's prompt
+    overwrote slot 0 and the other slots stayed empty. Per-user PCC failed at
+    0.92-0.93 until ``user_id`` was threaded through prefill. Both halves were
+    needed for multi-user decode to work end to end.
+
+    Coverage used to stop at 2, blamed on a harness limit: every
+    parametrisation re-uploaded **1.63 GB** of weights (the bf16 functional set
+    *and* the optimized set) without reclaiming the previous ones, so b8
+    exhausted DRAM in ``bank_manager.cpp`` before the layer ran. That was
+    self-inflicted -- the optimized path never reads the functional experts or
+    the functional attention copy. ``_build`` now uploads only what this path
+    touches, **0.38 GB**, and b8 and b32 run.
+
+    Both figures are derived from the shipped uploads rather than recalled:
+    functional experts 3 x 128 x 768 x 2048 elem at bf16 = 1.208 GB, plus
+    wqkv+wo bf16 = 37.75 MB; optimized experts (128 x 1536 x 2048 packed gate/up
+    plus 128 x 768 x 2048 down) at bfloat4_b 0.5625 B/elem = 339.7 MB, plus two
+    copies of wqkv+wo at bfloat8_b 1.0625 B/elem = 40.11 MB. Earlier revisions
+    of this docstring said ~2.4 GB and ~0.63 GB, and neither was derived.
+
+    32 is the real ceiling, and it is a **TTNN op limit**, not this layer's
+    shape choice: ``nlp_create_qkv_heads_decode_device_operation.cpp:51``
+    asserts ``num_users <= num_users_supported`` with ``num_users_supported =
+    32`` hardcoded at line 45 of that file, and that op is on the interleaved
+    attention path as well as the DRAM-sharded one. ``_dram_sharded_usable``
+    does refuse the sharded projections past B=32 -- ``_width_sharded_l1``
+    hardcodes a 32-row shard and ``_dram_sharded_program_config`` sets
+    ``per_core_M=1`` -- but the interleaved fallback it selects then fails in
+    ``nlp_create_qkv_heads_decode`` too. The guard buys a comprehensible
+    failure, not a working larger batch.
+    """
+    layer, hf_config = reference
+    config, weights, packed, cos, sin, sparsity = _build(mesh_device, hf_config, torch_weights)
+    prompt_len = 32
+
+    # Small cache on purpose: this test is about multi-user expert routing, and
+    # a full-length cache per user would exhaust DRAM before reaching the point.
+    kv_cache = F.create_kv_cache(mesh_device, config.attention, max_batch=batch, max_seq_len=128, block_size=BLOCK_SIZE)
+    per_user = [_hidden(hf_config, prompt_len + 1, seed=u) for u in range(batch)]
+    # Each user's prompt must actually land in the cache, or decode attends
+    # zeros and the test proves nothing about multi-user routing.
+    for user, hidden_full in enumerate(per_user):
+        O.decoder_layer_prefill_optimized(
+            _to_device(hidden_full[:, :prompt_len, :].unsqueeze(0), mesh_device),
+            weights,
+            config,
+            cos,
+            sin,
+            sparsity,
+            packed,
+            kv_cache=kv_cache,
+            user_id=user,
+        )
+
+    tokens = torch.cat([h[:, prompt_len, :] for h in per_user], dim=0)  # [batch, hidden]
+    current_pos = ttnn.from_torch(
+        torch.full((batch,), prompt_len, dtype=torch.int32), dtype=ttnn.int32, device=mesh_device
+    )
+    out = O.decoder_layer_decode_optimized(
+        _to_device(tokens.reshape(1, 1, batch, hf_config.hidden_size), mesh_device),
+        weights,
+        config,
+        cos,
+        sin,
+        kv_cache,
+        current_pos,
+        prompt_len,
+        packed_experts=packed,
+    )
+    tt_out = ttnn.to_torch(out).reshape(-1, hf_config.hidden_size)[:batch].float()
+
+    assert torch.isfinite(tt_out).all(), f"batch={batch} produced non-finite values"
+
+    # Every user is checked against its own HF reference. This is what proves
+    # the routing is per-user rather than broadcast.
+    for user, hidden_full in enumerate(per_user):
+        ref_user = _reference_layer(layer, hf_config, hidden_full)[:, prompt_len, :]
+        passing, message = comp_pcc(ref_user, tt_out[user : user + 1], 0.99)
+        logger.info(f"optimized decode batch={batch} user {user} vs HF: {message}")
+        assert passing, f"batch={batch} user {user} decode below 0.99: {message}"
+
+    if batch > 1:
+        spread = (tt_out - tt_out[0]).abs().max().item()
+        assert spread > 1e-3, f"batch={batch}: all users returned identical output (broadcast bug)"

code/models/demos/blackhole/qwen3_coder_30b_a3b/tests/test_perf.py

@@ -0,0 +1,611 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+"""Warmed prefill and traced warmed decode latency for the functional decoder.
+
+These numbers are the stage-01 baseline that the optimized decoder has to beat,
+so the measurement conditions matter as much as the values:
+
+* **Warmed.** The first call to any shape compiles kernels and populates the
+  program cache. Timing that measures the compiler. Every configuration runs
+  warmup iterations that are discarded.
+* **Traced decode.** Decode is short enough that host dispatch overhead is a
+  large share of wall time, so an eager measurement mostly reports Python.
+  Replaying a captured trace is what a serving stack actually does.
+* **Device-synchronised.** Dispatch is asynchronous; without an explicit
+  synchronise the host would time enqueue calls rather than execution.
+* **Median, not mean.** One descheduled iteration should not move the number.
+
+Results are written to ``doc/functional_decoder/`` as CSV so later stages can
+diff against them rather than re-deriving a baseline.
+
+Two harness rules the CSVs depend on:
+
+* **Every test here is marked ``models_performance_bare_metal``.** These tests
+  overwrite the published CSVs, and ``TT_METAL_WATCHER=10`` inflates device
+  timings roughly 8x, so a watcher run over the whole suite silently replaced
+  the prefill baseline with 4358 us/token. Watcher runs must deselect them:
+  ``TT_METAL_WATCHER=10 pytest ... -m "not models_performance_bare_metal"``.
+* **The decode CSVs are rewritten whole, not appended to.** The decode test is
+  parametrised over context length, so each parametrisation contributes one
+  row; rows accumulate in ``_DECODE_ROWS`` for the life of the process and the
+  file is rewritten from scratch each time. Appending stacked eight interleaved
+  runs into ``doc/functional_decoder/perf_decode.csv`` before this was fixed.
+"""
+
+from __future__ import annotations
+
+import csv
+import statistics
+import time
+from pathlib import Path
+
+import pytest
+import torch
+from loguru import logger
+
+import ttnn
+
+from ..tt.functional_decoder import (
+    DecoderLayerConfig,
+    build_expert_sparsity,
+    build_rope_cache,
+    create_kv_cache,
+    decoder_layer_decode,
+    decoder_layer_prefill,
+    upload_layer_weights,
+)
+from ..tt.weight_mapping import convert_layer_weights
+from .reference import build_reference_layer, layer_state_dict
+
+LAYER_IDX = 0
+MAX_SEQ = 4096
+BLOCK_SIZE = 32
+TRACE_REGION_SIZE = 50331648
+
+PREFILL_LENGTHS = [128, 512, 1024, 2048]
+PREFILL_WARMUP, PREFILL_ITERS = 1, 5
+DECODE_WARMUP, DECODE_ITERS = 10, 100
+
+DOC_DIR = Path(__file__).resolve().parents[1] / "doc" / "functional_decoder"
+
+DECODE_FIELDS = ["context_len", "median_ms", "min_ms", "max_ms", "tok_per_s_per_layer", "iters"]
+
+# {csv path: [row, ...]} for this process only, so a rerun truncates.
+_DECODE_ROWS: dict[Path, list[dict]] = {}
+
+
+def _write_decode_row(path: Path, row: dict) -> None:
+    """Add a row and rewrite the whole file (see module docstring)."""
+    path.parent.mkdir(parents=True, exist_ok=True)
+    rows = _DECODE_ROWS.setdefault(path, [])
+    rows[:] = [r for r in rows if r["context_len"] != row["context_len"]] + [row]
+    rows.sort(key=lambda r: r["context_len"])
+    with path.open("w", newline="") as fh:
+        writer = csv.DictWriter(fh, fieldnames=DECODE_FIELDS)
+        writer.writeheader()
+        writer.writerows(rows)
+
+
+@pytest.fixture(scope="module")
+def reference():
+    return build_reference_layer(LAYER_IDX)
+
+
+@pytest.fixture(scope="module")
+def torch_weights(reference):
+    _, hf_config = reference
+    return convert_layer_weights(layer_state_dict(LAYER_IDX), hf_config)
+
+
+def _to_device(t, mesh_device):
+    return ttnn.from_torch(
+        t, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=mesh_device, memory_config=ttnn.DRAM_MEMORY_CONFIG
+    )
+
+
+def _write_csv(name: str, fieldnames: list[str], rows: list[dict]) -> Path:
+    DOC_DIR.mkdir(parents=True, exist_ok=True)
+    path = DOC_DIR / name
+    with path.open("w", newline="") as fh:
+        writer = csv.DictWriter(fh, fieldnames=fieldnames)
+        writer.writeheader()
+        writer.writerows(rows)
+    logger.info(f"wrote {path}")
+    return path
+
+
+@pytest.mark.models_performance_bare_metal
+@pytest.mark.parametrize("device_params", [{"trace_region_size": TRACE_REGION_SIZE}], indirect=True)
+@pytest.mark.parametrize("mesh_device", [(1, 1)], ids=["1x1"], indirect=True)
+def test_prefill_latency(mesh_device, reference, torch_weights):
+    _, hf_config = reference
+    config = DecoderLayerConfig.from_hf(hf_config)
+    weights = upload_layer_weights(torch_weights, mesh_device, config)
+    cos_cache, sin_cache = build_rope_cache(hf_config, MAX_SEQ, mesh_device)
+    sparsity = build_expert_sparsity(mesh_device, config.moe.num_experts)
+
+    rows = []
+    for seq_len in PREFILL_LENGTHS:
+        torch.manual_seed(0)
+        hidden = torch.randn(1, 1, seq_len, hf_config.hidden_size) * 0.02
+        tt_in = _to_device(hidden, mesh_device)
+
+        def once():
+            out = decoder_layer_prefill(tt_in, weights, config, cos_cache, sin_cache, sparsity)
+            ttnn.synchronize_device(mesh_device)
+            ttnn.deallocate(out)
+
+        for _ in range(PREFILL_WARMUP):
+            once()
+
+        samples = []
+        for _ in range(PREFILL_ITERS):
+            t0 = time.perf_counter()
+            once()
+            samples.append((time.perf_counter() - t0) * 1e3)
+
+        median = statistics.median(samples)
+        per_tok = median / seq_len * 1e3  # us/token
+        logger.info(
+            f"prefill S={seq_len:>5}: median {median:8.2f} ms  "
+            f"min {min(samples):8.2f}  max {max(samples):8.2f}  ({per_tok:6.1f} us/token)"
+        )
+        rows.append(
+            {
+                "seq_len": seq_len,
+                "median_ms": round(median, 3),
+                "min_ms": round(min(samples), 3),
+                "max_ms": round(max(samples), 3),
+                "us_per_token": round(per_tok, 2),
+                "iters": PREFILL_ITERS,
+            }
+        )
+        ttnn.deallocate(tt_in)
+
+    _write_csv(
+        "perf_prefill.csv",
+        ["seq_len", "median_ms", "min_ms", "max_ms", "us_per_token", "iters"],
+        rows,
+    )
+    assert all(r["median_ms"] > 0 for r in rows)
+
+
+@pytest.mark.models_performance_bare_metal
+@pytest.mark.parametrize("device_params", [{"trace_region_size": TRACE_REGION_SIZE}], indirect=True)
+@pytest.mark.parametrize("mesh_device", [(1, 1)], ids=["1x1"], indirect=True)
+@pytest.mark.parametrize("context_len", [128, 1024, 4096], ids=["ctx128", "ctx1k", "ctx4k"])
+def test_decode_latency_traced(mesh_device, reference, torch_weights, context_len):
+    """Traced single-token decode latency at several cache depths.
+
+    Swept over context because decode cost is dominated by the SDPA read over
+    the cache, so a single depth would not say whether latency is flat or grows
+    with the conversation.
+    """
+    _, hf_config = reference
+    config = DecoderLayerConfig.from_hf(hf_config)
+    weights = upload_layer_weights(torch_weights, mesh_device, config)
+    cos_cache, sin_cache = build_rope_cache(hf_config, MAX_SEQ, mesh_device)
+    kv_cache = create_kv_cache(mesh_device, config.attention, max_batch=1, max_seq_len=MAX_SEQ, block_size=BLOCK_SIZE)
+
+    pos = context_len - 1
+    torch.manual_seed(0)
+    tt_in = _to_device(torch.randn(1, 1, 1, hf_config.hidden_size) * 0.02, mesh_device)
+    current_pos = ttnn.from_torch(torch.tensor([pos], dtype=torch.int32), dtype=ttnn.int32, device=mesh_device)
+
+    def step():
+        return decoder_layer_decode(
+            tt_in, weights, config, cos_cache, sin_cache, kv_cache, current_pos, token_index=pos
+        )
+
+    step()  # compile outside the capture
+    ttnn.synchronize_device(mesh_device)
+
+    trace_id = ttnn.begin_trace_capture(mesh_device, cq_id=0)
+    step()
+    ttnn.end_trace_capture(mesh_device, trace_id, cq_id=0)
+
+    for _ in range(DECODE_WARMUP):
+        ttnn.execute_trace(mesh_device, trace_id, cq_id=0, blocking=True)
+
+    samples = []
+    for _ in range(DECODE_ITERS):
+        t0 = time.perf_counter()
+        ttnn.execute_trace(mesh_device, trace_id, cq_id=0, blocking=True)
+        samples.append((time.perf_counter() - t0) * 1e3)
+
+    median = statistics.median(samples)
+    logger.info(
+        f"traced decode ctx={context_len:>5}: median {median:7.3f} ms  "
+        f"min {min(samples):7.3f}  max {max(samples):7.3f}  "
+        f"({1e3 / median:7.1f} tok/s/layer)"
+    )
+
+    _write_decode_row(
+        DOC_DIR / "perf_decode.csv",
+        {
+            "context_len": context_len,
+            "median_ms": round(median, 4),
+            "min_ms": round(min(samples), 4),
+            "max_ms": round(max(samples), 4),
+            "tok_per_s_per_layer": round(1e3 / median, 1),
+            "iters": DECODE_ITERS,
+        },
+    )
+
+    ttnn.release_trace(mesh_device, trace_id)
+    assert median > 0
+
+
+# --- optimized path -----------------------------------------------------------
+# Same harness, same conditions, writing to doc/optimized_decoder/ so the two
+# stages' CSVs are directly diffable rather than needing re-derivation.
+
+OPT_DOC_DIR = Path(__file__).resolve().parents[1] / "doc" / "optimized_decoder"
+
+
+@pytest.mark.models_performance_bare_metal
+@pytest.mark.parametrize("device_params", [{"trace_region_size": TRACE_REGION_SIZE}], indirect=True)
+@pytest.mark.parametrize("mesh_device", [(1, 1)], ids=["1x1"], indirect=True)
+def test_optimized_prefill_latency(mesh_device, reference, torch_weights):
+    from ..tt import optimized_decoder as O
+
+    _, hf_config = reference
+    config = DecoderLayerConfig.from_hf(hf_config)
+    weights = upload_layer_weights(torch_weights, mesh_device, config)
+    packed = O.upload_packed_expert_weights(torch_weights, mesh_device, config.moe)
+    cos_cache, sin_cache = build_rope_cache(hf_config, MAX_SEQ, mesh_device)
+    sparsity = build_expert_sparsity(mesh_device, config.moe.num_experts)
+
+    rows = []
+    for seq_len in PREFILL_LENGTHS:
+        torch.manual_seed(0)
+        tt_in = _to_device(torch.randn(1, 1, seq_len, hf_config.hidden_size) * 0.02, mesh_device)
+
+        def once():
+            out = O.decoder_layer_prefill_optimized(tt_in, weights, config, cos_cache, sin_cache, sparsity, packed)
+            ttnn.synchronize_device(mesh_device)
+            ttnn.deallocate(out)
+
+        for _ in range(PREFILL_WARMUP):
+            once()
+        samples = []
+        for _ in range(PREFILL_ITERS):
+            t0 = time.perf_counter()
+            once()
+            samples.append((time.perf_counter() - t0) * 1e3)
+
+        median = statistics.median(samples)
+        logger.info(
+            f"OPTIMIZED prefill S={seq_len:>5}: median {median:8.2f} ms ({median / seq_len * 1e3:6.1f} us/token)"
+        )
+        rows.append(
+            {
+                "seq_len": seq_len,
+                "median_ms": round(median, 3),
+                "min_ms": round(min(samples), 3),
+                "max_ms": round(max(samples), 3),
+                "us_per_token": round(median / seq_len * 1e3, 2),
+                "iters": PREFILL_ITERS,
+            }
+        )
+        ttnn.deallocate(tt_in)
+
+    OPT_DOC_DIR.mkdir(parents=True, exist_ok=True)
+    with (OPT_DOC_DIR / "perf_prefill.csv").open("w", newline="") as fh:
+        wr = csv.DictWriter(fh, fieldnames=["seq_len", "median_ms", "min_ms", "max_ms", "us_per_token", "iters"])
+        wr.writeheader()
+        wr.writerows(rows)
+    assert all(r["median_ms"] > 0 for r in rows)
+
+
+@pytest.mark.models_performance_bare_metal
+@pytest.mark.parametrize("device_params", [{"trace_region_size": TRACE_REGION_SIZE}], indirect=True)
+@pytest.mark.parametrize("mesh_device", [(1, 1)], ids=["1x1"], indirect=True)
+@pytest.mark.parametrize("context_len", [128, 1024, 4096], ids=["ctx128", "ctx1k", "ctx4k"])
+def test_optimized_decode_latency_traced(mesh_device, reference, torch_weights, context_len):
+    from ..tt import optimized_decoder as O
+
+    _, hf_config = reference
+    config = DecoderLayerConfig.from_hf(hf_config)
+    weights = upload_layer_weights(torch_weights, mesh_device, config)
+    packed = O.upload_packed_expert_weights(torch_weights, mesh_device, config.moe)
+    cos_cache, sin_cache = build_rope_cache(hf_config, MAX_SEQ, mesh_device)
+    kv_cache = create_kv_cache(mesh_device, config.attention, max_batch=1, max_seq_len=MAX_SEQ, block_size=BLOCK_SIZE)
+
+    pos = context_len - 1
+    torch.manual_seed(0)
+    tt_in = _to_device(torch.randn(1, 1, 1, hf_config.hidden_size) * 0.02, mesh_device)
+    current_pos = ttnn.from_torch(torch.tensor([pos], dtype=torch.int32), dtype=ttnn.int32, device=mesh_device)
+
+    def step():
+        return O.decoder_layer_decode_optimized(
+            tt_in, weights, config, cos_cache, sin_cache, kv_cache, current_pos, pos, packed_experts=packed
+        )
+
+    step()
+    ttnn.synchronize_device(mesh_device)
+    trace_id = ttnn.begin_trace_capture(mesh_device, cq_id=0)
+    step()
+    ttnn.end_trace_capture(mesh_device, trace_id, cq_id=0)
+    for _ in range(DECODE_WARMUP):
+        ttnn.execute_trace(mesh_device, trace_id, cq_id=0, blocking=True)
+    samples = []
+    for _ in range(DECODE_ITERS):
+        t0 = time.perf_counter()
+        ttnn.execute_trace(mesh_device, trace_id, cq_id=0, blocking=True)
+        samples.append((time.perf_counter() - t0) * 1e3)
+
+    median = statistics.median(samples)
+    logger.info(
+        f"OPTIMIZED traced decode ctx={context_len:>5}: median {median:7.3f} ms ({1e3 / median:7.1f} tok/s/layer)"
+    )
+
+    _write_decode_row(
+        OPT_DOC_DIR / "perf_decode.csv",
+        {
+            "context_len": context_len,
+            "median_ms": round(median, 4),
+            "min_ms": round(min(samples), 4),
+            "max_ms": round(max(samples), 4),
+            "tok_per_s_per_layer": round(1e3 / median, 1),
+            "iters": DECODE_ITERS,
+        },
+    )
+    ttnn.release_trace(mesh_device, trace_id)
+    assert median > 0
+
+
+# --- multichip path -----------------------------------------------------------
+# Stage 03, on the full 4-die P300_X2 mesh. Three tests, and the first of them is
+# the one that makes the other two mean anything:
+#
+#   test_multichip_baseline_1x1_*  re-measures the *single-chip* optimized layer,
+#       with this harness, in this tree, on one die, and writes it into
+#       doc/multichip_decoder/. Stage 02's CSVs already hold a warmed single-chip
+#       baseline, but they are also the artifact its own README quotes cell by
+#       cell, and re-running them here would move their third significant figure
+#       and silently invalidate that document. A stage-owned copy of the baseline
+#       is cheaper than a cross-stage prose/artifact mismatch.
+#   test_multichip_prefill_latency / test_multichip_decode_latency_traced
+#       measure the same lengths on the mesh, so speedup is one CSV cell divided
+#       by another rather than a number quoted from anywhere.
+#
+# All three are marked models_performance_bare_metal: they rewrite published
+# CSVs and TT_METAL_WATCHER inflates device timings, so a watcher run must
+# deselect them.
+
+MC_DOC_DIR = Path(__file__).resolve().parents[1] / "doc" / "optimized_multichip_decoder"
+# Stage 04 note. ``tt/multichip_decoder.py`` is optimized **in place**, so the
+# four tests below now measure the stage-04 path. They therefore write into
+# ``doc/optimized_multichip_decoder/``; ``doc/multichip_decoder/perf_*.csv`` are
+# stage 03's frozen *before* numbers and are deliberately never regenerated --
+# re-pointing this constant back would overwrite the baseline half of every
+# before/after table in both READMEs. The stage-04 decode before/after is also
+# measured in one process by
+# ``doc/optimized_multichip_decoder/probes/layer_levers.py``, whose "stage 03"
+# leg is a verbatim copy of the committed stage-03 layer body.
+
+# Ring fabric must be set before the mesh opens; the conftest device_params hook
+# does that, which is why it is spelled here rather than with set_fabric_config.
+MC_DEVICE_PARAMS = {"trace_region_size": TRACE_REGION_SIZE, "fabric_config": ttnn.FabricConfig.FABRIC_1D_RING}
+MC_MESH = (1, 4)
+
+PREFILL_FIELDS = ["seq_len", "median_ms", "min_ms", "max_ms", "us_per_token", "iters"]
+
+
+def _prefill_sweep(once_factory, lengths=PREFILL_LENGTHS, label="") -> list[dict]:
+    rows = []
+    for seq_len in lengths:
+        once = once_factory(seq_len)
+        for _ in range(PREFILL_WARMUP):
+            once()
+        samples = []
+        for _ in range(PREFILL_ITERS):
+            t0 = time.perf_counter()
+            once()
+            samples.append((time.perf_counter() - t0) * 1e3)
+        median = statistics.median(samples)
+        logger.info(f"{label} prefill S={seq_len:>5}: median {median:8.2f} ms ({median / seq_len * 1e3:6.1f} us/token)")
+        rows.append(
+            {
+                "seq_len": seq_len,
+                "median_ms": round(median, 3),
+                "min_ms": round(min(samples), 3),
+                "max_ms": round(max(samples), 3),
+                "us_per_token": round(median / seq_len * 1e3, 2),
+                "iters": PREFILL_ITERS,
+            }
+        )
+    return rows
+
+
+def _write_rows(path: Path, fieldnames: list[str], rows: list[dict]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with path.open("w", newline="") as fh:
+        wr = csv.DictWriter(fh, fieldnames=fieldnames)
+        wr.writeheader()
+        wr.writerows(rows)
+    logger.info(f"wrote {path}")
+
+
+def _traced_decode_median(mesh_device, step) -> tuple[float, float, float]:
+    step()
+    ttnn.synchronize_device(mesh_device)
+    trace_id = ttnn.begin_trace_capture(mesh_device, cq_id=0)
+    step()
+    ttnn.end_trace_capture(mesh_device, trace_id, cq_id=0)
+    for _ in range(DECODE_WARMUP):
+        ttnn.execute_trace(mesh_device, trace_id, cq_id=0, blocking=True)
+    samples = []
+    for _ in range(DECODE_ITERS):
+        t0 = time.perf_counter()
+        ttnn.execute_trace(mesh_device, trace_id, cq_id=0, blocking=True)
+        samples.append((time.perf_counter() - t0) * 1e3)
+    ttnn.release_trace(mesh_device, trace_id)
+    return statistics.median(samples), min(samples), max(samples)
+
+
+@pytest.mark.models_performance_bare_metal
+@pytest.mark.parametrize("device_params", [{"trace_region_size": TRACE_REGION_SIZE}], indirect=True)
+@pytest.mark.parametrize("mesh_device", [(1, 1)], ids=["1x1"], indirect=True)
+def test_optimized_multichip_baseline_1x1_prefill(mesh_device, reference, torch_weights):
+    """Stage-04's own copy of the warmed single-chip prefill baseline.
+
+    ``optimized_decoder.py`` is untouched by stage 04, so this re-measures the
+    same code stage 03 did; it is re-run rather than quoted so the speedup
+    columns in this stage's README are one CSV cell divided by another taken in
+    the same session."""
+    from ..tt import optimized_decoder as O
+
+    _, hf_config = reference
+    config = DecoderLayerConfig.from_hf(hf_config)
+    weights = upload_layer_weights(torch_weights, mesh_device, config)
+    packed = O.upload_packed_expert_weights(torch_weights, mesh_device, config.moe)
+    cos_cache, sin_cache = build_rope_cache(hf_config, MAX_SEQ, mesh_device)
+    sparsity = build_expert_sparsity(mesh_device, config.moe.num_experts)
+
+    def factory(seq_len):
+        torch.manual_seed(0)
+        tt_in = _to_device(torch.randn(1, 1, seq_len, hf_config.hidden_size) * 0.02, mesh_device)
+
+        def once():
+            out = O.decoder_layer_prefill_optimized(tt_in, weights, config, cos_cache, sin_cache, sparsity, packed)
+            ttnn.synchronize_device(mesh_device)
+            ttnn.deallocate(out)
+
+        return once
+
+    rows = _prefill_sweep(factory, label="BASELINE 1x1")
+    _write_rows(MC_DOC_DIR / "perf_baseline_1x1_prefill.csv", PREFILL_FIELDS, rows)
+    assert all(r["median_ms"] > 0 for r in rows)
+
+
+@pytest.mark.models_performance_bare_metal
+@pytest.mark.parametrize("device_params", [{"trace_region_size": TRACE_REGION_SIZE}], indirect=True)
+@pytest.mark.parametrize("mesh_device", [(1, 1)], ids=["1x1"], indirect=True)
+@pytest.mark.parametrize("context_len", [128, 1024, 4096], ids=["ctx128", "ctx1k", "ctx4k"])
+def test_optimized_multichip_baseline_1x1_decode(mesh_device, reference, torch_weights, context_len):
+    """Stage-04's own copy of the warmed single-chip traced decode baseline."""
+    from ..tt import optimized_decoder as O
+
+    _, hf_config = reference
+    config = DecoderLayerConfig.from_hf(hf_config)
+    weights = upload_layer_weights(torch_weights, mesh_device, config)
+    packed = O.upload_packed_expert_weights(torch_weights, mesh_device, config.moe)
+    cos_cache, sin_cache = build_rope_cache(hf_config, MAX_SEQ, mesh_device)
+    kv_cache = create_kv_cache(mesh_device, config.attention, max_batch=1, max_seq_len=MAX_SEQ, block_size=BLOCK_SIZE)
+
+    pos = context_len - 1
+    torch.manual_seed(0)
+    tt_in = _to_device(torch.randn(1, 1, 1, hf_config.hidden_size) * 0.02, mesh_device)
+    current_pos = ttnn.from_torch(torch.tensor([pos], dtype=torch.int32), dtype=ttnn.int32, device=mesh_device)
+
+    median, lo, hi = _traced_decode_median(
+        mesh_device,
+        lambda: O.decoder_layer_decode_optimized(
+            tt_in, weights, config, cos_cache, sin_cache, kv_cache, current_pos, pos, packed_experts=packed
+        ),
+    )
+    logger.info(f"BASELINE 1x1 traced decode ctx={context_len:>5}: median {median:7.4f} ms")
+    _write_decode_row(
+        MC_DOC_DIR / "perf_baseline_1x1_decode.csv",
+        {
+            "context_len": context_len,
+            "median_ms": round(median, 4),
+            "min_ms": round(lo, 4),
+            "max_ms": round(hi, 4),
+            "tok_per_s_per_layer": round(1e3 / median, 1),
+            "iters": DECODE_ITERS,
+        },
+    )
+    assert median > 0
+
+
+@pytest.mark.models_performance_bare_metal
+@pytest.mark.parametrize("device_params", [MC_DEVICE_PARAMS], indirect=True)
+@pytest.mark.parametrize("mesh_device", [MC_MESH], ids=["1x4"], indirect=True)
+def test_optimized_multichip_prefill_latency(mesh_device, reference, torch_weights):
+    from ..tt import multichip_decoder as MC
+
+    _, hf_config = reference
+    config = MC.MeshDecoderConfig.from_hf(hf_config)
+    ctx = MC.mesh_context(mesh_device)
+    weights = MC.upload_multichip_weights(torch_weights, mesh_device, config)
+    cos_cache, sin_cache = build_rope_cache(hf_config, MAX_SEQ, mesh_device)
+    sparsity = MC.build_local_sparsity(mesh_device, config.local_moe)
+
+    def factory(seq_len):
+        torch.manual_seed(0)
+        tt_in = ttnn.from_torch(
+            torch.randn(1, 1, seq_len, hf_config.hidden_size) * 0.02,
+            dtype=ttnn.bfloat16,
+            layout=ttnn.TILE_LAYOUT,
+            device=mesh_device,
+            memory_config=ttnn.DRAM_MEMORY_CONFIG,
+            mesh_mapper=ttnn.ReplicateTensorToMesh(mesh_device),
+        )
+
+        def once():
+            out = MC.decoder_layer_prefill_multichip(tt_in, weights, config, ctx, cos_cache, sin_cache, sparsity)
+            ttnn.synchronize_device(mesh_device)
+            ttnn.deallocate(out)
+
+        return once
+
+    rows = _prefill_sweep(factory, label="MULTICHIP 1x4")
+    _write_rows(MC_DOC_DIR / "perf_prefill.csv", PREFILL_FIELDS, rows)
+    assert all(r["median_ms"] > 0 for r in rows)
+
+
+@pytest.mark.models_performance_bare_metal
+@pytest.mark.parametrize("device_params", [MC_DEVICE_PARAMS], indirect=True)
+@pytest.mark.parametrize("mesh_device", [MC_MESH], ids=["1x4"], indirect=True)
+@pytest.mark.parametrize("context_len", [128, 1024, 4096], ids=["ctx128", "ctx1k", "ctx4k"])
+def test_optimized_multichip_decode_latency_traced(mesh_device, reference, torch_weights, context_len):
+    """Warmed trace replay on the mesh. Same harness as the 1x1 baseline above."""
+    from ..tt import multichip_decoder as MC
+
+    _, hf_config = reference
+    config = MC.MeshDecoderConfig.from_hf(hf_config)
+    ctx = MC.mesh_context(mesh_device)
+    weights = MC.upload_multichip_weights(torch_weights, mesh_device, config)
+    cos_cache, sin_cache = build_rope_cache(hf_config, MAX_SEQ, mesh_device)
+    kv_cache = MC.create_mesh_kv_cache(mesh_device, config, 1, MAX_SEQ, block_size=BLOCK_SIZE)
+
+    pos = context_len - 1
+    torch.manual_seed(0)
+    tt_in = ttnn.from_torch(
+        torch.randn(1, 1, 1, hf_config.hidden_size) * 0.02,
+        dtype=ttnn.bfloat16,
+        layout=ttnn.TILE_LAYOUT,
+        device=mesh_device,
+        memory_config=ttnn.DRAM_MEMORY_CONFIG,
+        mesh_mapper=ttnn.ReplicateTensorToMesh(mesh_device),
+    )
+    current_pos = ttnn.from_torch(
+        torch.tensor([pos], dtype=torch.int32),
+        dtype=ttnn.int32,
+        device=mesh_device,
+        mesh_mapper=ttnn.ReplicateTensorToMesh(mesh_device),
+    )
+
+    median, lo, hi = _traced_decode_median(
+        mesh_device,
+        lambda: MC.decoder_layer_decode_multichip(
+            tt_in, weights, config, ctx, cos_cache, sin_cache, kv_cache, current_pos, pos
+        ),
+    )
+    logger.info(f"MULTICHIP 1x4 traced decode ctx={context_len:>5}: median {median:7.4f} ms")
+    _write_decode_row(
+        MC_DOC_DIR / "perf_decode.csv",
+        {
+            "context_len": context_len,
+            "median_ms": round(median, 4),
+            "min_ms": round(lo, 4),
+            "max_ms": round(hi, 4),
+            "tok_per_s_per_layer": round(1e3 / median, 1),
+            "iters": DECODE_ITERS,
+        },
+    )
+    assert median > 0

code/models/demos/blackhole/qwen3_coder_30b_a3b/tests/test_precision_config.py

@@ -0,0 +1,369 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+"""Gates for ``tt/precision.py`` -- the precision config stage 07 sweeps.
+
+The stage-07 goal asks for a selected precision config that "later
+full-model/vLLM construction paths actually consume by default", and says
+explicitly that **a JSON field ignored by hard-coded model code does not satisfy
+this requirement**. So this file is arranged around that sentence:
+
+* :func:`test_default_is_the_shipped_policy` and the alias tests pin that the
+  *default* changed nothing -- every constant stages 02-06 measured at is still
+  what a no-argument construction produces;
+* :func:`test_non_default_precision_reaches_the_device` builds a real model on
+  the mesh at a **non-default** value and asserts on what the device actually
+  holds: a different weight dtype, a different program-config block width, a
+  different compute-kernel fidelity, and a smaller per-die expert allocation.
+  That is the assertion the goal is really asking for;
+* the round-trip tests pin that config -> JSON -> config is lossless and that
+  the JSON carries every field.
+
+The host-only tests need no device. The two device tests build **two-layer**
+models -- the observable is per-layer and a 48-layer load is three minutes --
+on a module-scoped mesh, exactly as ``test_full_model.py`` does.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import json
+
+import pytest
+
+import ttnn
+
+# **Absolute imports, deliberately.** ``tt/generator.py`` imports ``tt.model``
+# by absolute path, and under this repo's ``--import-mode=importlib`` a relative
+# ``from ..tt import model`` here resolves to a *second* copy of the module (no
+# ``models/__init__.py``, so pytest roots the package at this directory).
+# The identity assertions below would then be comparing two different classes,
+# and the device tests would be inspecting a model built by the other copy.
+from models.demos.blackhole.qwen3_coder_30b_a3b.tt import model as M
+from models.demos.blackhole.qwen3_coder_30b_a3b.tt import multichip_decoder as MC
+from models.demos.blackhole.qwen3_coder_30b_a3b.tt import optimized_decoder as O
+from models.demos.blackhole.qwen3_coder_30b_a3b.tt.generator import build_generator
+from models.demos.blackhole.qwen3_coder_30b_a3b.tt.model import DEFAULT_TRACE_REGION_SIZE
+from models.demos.blackhole.qwen3_coder_30b_a3b.tt.precision import DEFAULT_PRECISION, PrecisionConfig
+
+MODEL_DIR = "models/demos/blackhole/qwen3_coder_30b_a3b"
+
+#: Every field the stage-07 goal enumerates, mapped to the config field(s) that
+#: carry it. The test below fails if any of these stops being serialised, which
+#: is the cheap way to notice a field being dropped from the artifact.
+GOAL_FIELDS = {
+    "experts gate_up weight dtype": ["experts_gate_up_dtype"],
+    "experts down weight dtype": ["experts_down_dtype"],
+    "attention qkv weight dtype": ["attention_qkv_dtype"],
+    "attention wo weight dtype": ["attention_wo_dtype"],
+    "lm_head weight dtype": ["lm_head_dtype"],
+    "router weight dtype": ["router_dtype"],
+    "embedding weight dtype": ["embedding_dtype"],
+    "per-group compute fidelity": [
+        "experts_fidelity",
+        "attention_fidelity",
+        "router_window_fidelity",
+        "lm_head_fidelity",
+        "norm_fidelity",
+    ],
+    "activation/residual dtype": ["activation_dtype"],
+    "CCL dtype": ["ccl_dtype"],
+    "KV-cache dtype": ["kv_cache_dtype"],
+    "logits/sampling dtype": ["logits_dtype", "sampling_dtype"],
+}
+
+
+# -- the default is the shipped policy ----------------------------------------
+
+
+def test_default_is_the_shipped_policy():
+    """The literal values stages 02-06 measured at, re-asserted here.
+
+    Written out rather than compared against the module constants, which are now
+    *derived* from this config -- comparing them to each other would be a
+    tautology. If a shipped value is ever changed, this test is the thing that
+    has to be changed with it, deliberately.
+    """
+    p = DEFAULT_PRECISION
+    assert p.experts_gate_up_dtype is ttnn.bfloat4_b
+    assert p.experts_down_dtype is ttnn.bfloat4_b
+    # Stage 07: retuned to the full-K ceilings the 48-layer sweep measured.
+    assert p.experts_gate_up_in0_block_w == 64
+    assert p.experts_down_in0_block_w == 24
+    assert p.experts_fidelity is ttnn.MathFidelity.LoFi
+    assert p.attention_qkv_dtype is ttnn.bfloat8_b
+    assert p.attention_wo_dtype is ttnn.bfloat8_b
+    assert p.attention_fidelity is None, "the projections take the op default; see _attention_compute_kernel_config"
+    assert p.lm_head_dtype is ttnn.bfloat8_b
+    assert p.lm_head_fidelity is ttnn.MathFidelity.HiFi2
+    assert p.router_dtype is ttnn.bfloat16
+    assert p.router_window_fidelity is ttnn.MathFidelity.HiFi4, "the one-hot window matmul must select, not approximate"
+    assert p.embedding_dtype is ttnn.bfloat16
+    assert p.norm_weight_dtype is ttnn.bfloat16
+    assert p.norm_fidelity is ttnn.MathFidelity.HiFi4
+    assert p.activation_dtype is ttnn.bfloat16
+    assert p.ccl_dtype is None and p.effective_ccl_dtype is ttnn.bfloat16
+    assert p.kv_cache_dtype is ttnn.bfloat16
+    assert p.logits_dtype is ttnn.bfloat16
+    assert p.sampling_dtype is ttnn.bfloat16
+
+
+def test_module_constants_still_resolve_to_the_default():
+    """The stage-02/04 names are aliases now; they must still read the same.
+
+    Probes under ``doc/`` and several stage-02 tests import these, and the point
+    of keeping them was that nothing outside this file had to change.
+    """
+    assert O.EXPERT_WEIGHT_DTYPE is DEFAULT_PRECISION.experts_gate_up_dtype
+    assert O.EXPERT_IN0_BLOCK_W_GATE_UP == DEFAULT_PRECISION.experts_gate_up_in0_block_w
+    assert O.EXPERT_IN0_BLOCK_W_DOWN == DEFAULT_PRECISION.experts_down_in0_block_w
+    assert O.EXPERT_MATH_FIDELITY is DEFAULT_PRECISION.experts_fidelity
+    assert O.ATTENTION_WEIGHT_DTYPE is DEFAULT_PRECISION.attention_qkv_dtype
+    assert M.LM_HEAD_WEIGHT_DTYPE is DEFAULT_PRECISION.lm_head_dtype
+    assert M.EMBED_WEIGHT_DTYPE is DEFAULT_PRECISION.embedding_dtype
+
+
+def test_default_construction_paths_all_resolve_to_the_same_object():
+    """``None`` means the shipped policy on every entry point that takes one."""
+    assert M._resolve_precision(None) is DEFAULT_PRECISION
+    assert M._resolve_precision(DEFAULT_PRECISION) is DEFAULT_PRECISION
+    assert M._resolve_precision(DEFAULT_PRECISION.to_dict()) == DEFAULT_PRECISION
+
+
+# -- serialisation -------------------------------------------------------------
+
+
+def test_json_round_trip_is_lossless():
+    for config in (
+        DEFAULT_PRECISION,
+        DEFAULT_PRECISION.with_overrides(experts_gate_up_dtype="bfloat8_b", experts_gate_up_in0_block_w=32),
+        DEFAULT_PRECISION.with_overrides(attention_fidelity="HiFi4", ccl_dtype="bfloat8_b"),
+    ):
+        assert PrecisionConfig.from_json(config.to_json()) == config
+        # and again through a second hop, so an asymmetric coercion cannot hide
+        assert PrecisionConfig.from_json(PrecisionConfig.from_json(config.to_json()).to_json()) == config
+
+
+def test_json_carries_every_field_the_goal_lists():
+    payload = json.loads(DEFAULT_PRECISION.to_json())
+    declared = {f.name for f in dataclasses.fields(PrecisionConfig)}
+    assert set(payload) == declared, "to_dict() must emit exactly the dataclass fields"
+    for description, names in GOAL_FIELDS.items():
+        for name in names:
+            assert name in payload, f"{description} is missing from the serialised config"
+
+
+def test_json_is_plain_names_not_repr():
+    """The artifact has to be readable and diffable, not ``DataType.BFLOAT4_B``."""
+    payload = json.loads(DEFAULT_PRECISION.to_json())
+    assert payload["experts_gate_up_dtype"] == "bfloat4_b"
+    assert payload["experts_fidelity"] == "LoFi"
+    assert payload["attention_fidelity"] is None
+    assert payload["ccl_dtype"] is None
+    assert payload["experts_gate_up_in0_block_w"] == 64
+
+
+def test_write_and_read_json_file(tmp_path):
+    config = DEFAULT_PRECISION.with_overrides(lm_head_dtype="bfloat16")
+    path = config.write_json(tmp_path / "nested" / "selected_precision_config.json")
+    assert PrecisionConfig.read_json(path) == config
+    # the file form is what a construction path is handed
+    assert M._resolve_precision(str(path)) == config
+
+
+def test_unknown_names_are_rejected(expect_error):
+    with expect_error(ValueError, "unknown dtype"):
+        PrecisionConfig(experts_gate_up_dtype="bfloat3_b")
+    with expect_error(ValueError, "unknown math fidelity"):
+        PrecisionConfig(experts_fidelity="HiFi9")
+    with expect_error(ValueError, "unknown precision fields"):
+        PrecisionConfig.from_dict({**DEFAULT_PRECISION.to_dict(), "expert_dtype": "bfloat8_b"})
+    with expect_error(ValueError, "unknown precision fields"):
+        DEFAULT_PRECISION.with_overrides(expert_weight_dtype="bfloat8_b")
+    with expect_error(ValueError, "may not be None"):
+        PrecisionConfig(activation_dtype=None)
+
+
+def test_config_is_frozen(expect_error):
+    # the fixture requires a match string; frozen dataclasses name the field
+    with expect_error(dataclasses.FrozenInstanceError, "experts_gate_up_dtype"):
+        DEFAULT_PRECISION.experts_gate_up_dtype = ttnn.bfloat16
+
+
+# -- the config is actually consumed (device) ----------------------------------
+
+
+@pytest.fixture(scope="module")
+def mesh_device():
+    ttnn.set_fabric_config(ttnn.FabricConfig.FABRIC_1D_RING)
+    mesh = ttnn.open_mesh_device(mesh_shape=ttnn.MeshShape(*MC.MESH_SHAPE), trace_region_size=DEFAULT_TRACE_REGION_SIZE)
+    yield mesh
+    ttnn.close_mesh_device(mesh)
+    ttnn.set_fabric_config(ttnn.FabricConfig.DISABLED)
+
+
+#: A precision that differs from the shipped one in four independently
+#: observable ways: a wider expert weight, a different block width for it, a
+#: different expert fidelity, and a narrower lm_head.
+NON_DEFAULT = DEFAULT_PRECISION.with_overrides(
+    experts_gate_up_dtype="bfloat8_b",
+    experts_gate_up_in0_block_w=32,  # also a divisor of 2048/32 = 64
+    experts_fidelity="HiFi4",
+    lm_head_dtype="bfloat4_b",
+    attention_wo_dtype="bfloat16",
+)
+
+
+def _build(mesh_device, precision):
+    return build_generator(
+        MODEL_DIR,
+        mesh_device,
+        override_num_layers=2,
+        max_context_len=1024,
+        max_batch_size=1,
+        precision=precision,
+    )
+
+
+def test_non_default_precision_reaches_the_device(mesh_device):
+    """Construct at ``NON_DEFAULT`` and assert on what the *device* holds.
+
+    This is the goal's "a JSON field ignored by hard-coded model code does not
+    satisfy this requirement" test. Nothing here reads ``model.precision``: every
+    assertion is against a dtype read back off an uploaded tensor, a block width
+    resolved by ``_tuned_sparse_matmul_config``, or a byte count computed from
+    the allocated shape.
+    """
+    default_gen = _build(mesh_device, None)
+    try:
+        base = default_gen.model.runtime_fallback_audit()
+        base_lm_head = str(default_gen.model.lm_head.dtype)
+    finally:
+        default_gen.teardown()
+
+    gen = _build(mesh_device, NON_DEFAULT)
+    try:
+        audit = gen.model.runtime_fallback_audit()
+        lm_head_dtype = str(gen.model.lm_head.dtype)
+
+        # 1. a different weight dtype reached the device
+        assert base["device_experts_gate_up_dtype"] == str(ttnn.bfloat4_b)
+        assert audit["device_experts_gate_up_dtype"] == str(ttnn.bfloat8_b)
+        assert audit["device_attention_wo_dtype"] == str(ttnn.bfloat16)
+        assert base["device_attention_wo_dtype"] == str(ttnn.bfloat8_b)
+        assert base_lm_head == str(ttnn.bfloat8_b)
+        assert lm_head_dtype == str(ttnn.bfloat4_b)
+
+        # 2. gate/up moved but down did **not** -- the two expert groups are
+        #    genuinely separate fields, not one knob wearing two names
+        assert audit["device_experts_down_dtype"] == str(ttnn.bfloat4_b)
+
+        # 3. a different block width in the resolved program config
+        assert base["gate_up_in0_block_w"] == 64
+        assert audit["gate_up_in0_block_w"] == 32
+        assert audit["down_in0_block_w"] == 24, "down's width was not overridden and must not move"
+
+        # 4. a different fidelity in the compute kernel config
+        assert base["expert_math_fidelity"] == str(ttnn.MathFidelity.LoFi)
+        assert audit["expert_math_fidelity"] == str(ttnn.MathFidelity.HiFi4)
+
+        # 5. an allocation-size change: bfloat4_b -> bfloat8_b on gate/up
+        #    roughly doubles the gate/up half of the per-die expert footprint
+        assert audit["device_expert_bytes_per_die"] > base["device_expert_bytes_per_die"]
+        grew = audit["device_expert_bytes_per_die"] - base["device_expert_bytes_per_die"]
+        assert grew > 20 * 1024 * 1024, f"expected tens of MB per die, got {grew}"
+
+        # 6. and it still runs -- a config that reaches the device but wedges it
+        #    would not be sweepable
+        ids = gen.tokenizer("def fib(n):", add_special_tokens=False)["input_ids"]
+        out = gen.generate(ids, 4, enable_trace=True, sampling_mode="device", top_k=1)
+        assert len(out) == 4
+    finally:
+        gen.teardown()
+
+
+def test_default_construction_audit_matches_the_shipped_values(mesh_device):
+    """The default path puts the shipped dtypes and widths on the device."""
+    gen = _build(mesh_device, None)
+    try:
+        audit = gen.model.runtime_fallback_audit()
+        assert audit["device_experts_gate_up_dtype"] == str(ttnn.bfloat4_b)
+        assert audit["device_experts_down_dtype"] == str(ttnn.bfloat4_b)
+        assert audit["device_attention_qkv_dtype"] == str(ttnn.bfloat8_b)
+        assert audit["device_attention_wo_dtype"] == str(ttnn.bfloat8_b)
+        assert audit["device_attention_qkv_decode_dtype"] == str(ttnn.bfloat8_b)
+        assert audit["device_router_dtype"] == str(ttnn.bfloat16)
+        assert audit["device_norm_weight_dtype"] == str(ttnn.bfloat16)
+        assert audit["gate_up_in0_block_w"] == 64
+        assert audit["down_in0_block_w"] == 24
+        assert audit["expert_math_fidelity"] == str(ttnn.MathFidelity.LoFi)
+        assert audit["attention_math_fidelity"] is None
+        assert audit["router_window_math_fidelity"] == str(ttnn.MathFidelity.HiFi4)
+        assert audit["ccl_dtype"] == str(ttnn.bfloat16)
+        assert audit["activation_dtype"] == str(ttnn.bfloat16)
+        assert str(gen.model.ensure_internal_kv_cache()[0].k.dtype) == str(ttnn.bfloat16)
+        # the audit also reports the config itself, which is what a sweep row
+        # records alongside its measurement
+        assert audit["precision"] == DEFAULT_PRECISION.to_dict()
+    finally:
+        gen.teardown()
+
+
+#: The four fields stage 07's original selection proof could not see.
+#:
+#: They had **no audit entry at all**, which made "this lever does nothing" and
+#: "this lever is not wired up" indistinguishable -- three sweep rows produced
+#: ``device_audit`` blocks byte-identical to the baseline's. For
+#: ``norm_fidelity`` it was the second: ``decode_residual_norm`` built its
+#: compute config from the module default and never saw ``self.precision``, so
+#: the field was a documented knob with no effect and ``R21_norm_hifi2``
+#: measured nothing. This config moves all four away from their defaults.
+TERMINAL_AND_NORM = DEFAULT_PRECISION.with_overrides(
+    norm_fidelity="HiFi2",
+    lm_head_fidelity="LoFi",
+    logits_dtype="bfloat8_b",
+    sampling_dtype="bfloat8_b",
+)
+
+
+def test_fidelity_and_terminal_dtypes_reach_the_device(mesh_device):
+    """The four fields that used to change nothing observable must change it now.
+
+    Regression test for a dead config field. Every assertion is against
+    something the *device* or the *ops* hold: the fidelities come off the
+    ``compute_kernel_config`` objects the norm and lm_head are handed, the two
+    dtypes off the tensors the terminal path actually produced during a real
+    traced decode. A default-vs-override diff is asserted for each, so a field
+    that silently stops being threaded fails here rather than in a sweep row
+    three stages later.
+    """
+    gen = _build(mesh_device, None)
+    try:
+        ids = gen.tokenizer("def fib(n):", add_special_tokens=False)["input_ids"]
+        gen.generate(ids, 4, enable_trace=True, sampling_mode="device", top_k=1)
+        base = gen.model.runtime_fallback_audit()
+    finally:
+        gen.teardown()
+
+    assert base["norm_math_fidelity"] == str(ttnn.MathFidelity.HiFi4)
+    assert base["lm_head_math_fidelity"] == str(ttnn.MathFidelity.HiFi2)
+    assert base["logits_dtype_observed"] == "bfloat16"
+    assert base["sampling_dtype_observed"] == "bfloat16"
+    assert base["terminal_dtype_source"] == "device_readback"
+
+    gen = _build(mesh_device, TERMINAL_AND_NORM)
+    try:
+        ids = gen.tokenizer("def fib(n):", add_special_tokens=False)["input_ids"]
+        out = gen.generate(ids, 4, enable_trace=True, sampling_mode="device", top_k=1)
+        assert len(out) == 4, "a config that reaches the device but wedges it is not sweepable"
+        audit = gen.model.runtime_fallback_audit()
+    finally:
+        gen.teardown()
+
+    # norm_fidelity: the field that was NOT threaded. It reaches only the decode
+    # residual norms -- the prefill norms pass no compute config at all -- which
+    # is why the audit name is norm_math_fidelity rather than something global.
+    assert audit["norm_math_fidelity"] == str(ttnn.MathFidelity.HiFi2)
+    assert audit["lm_head_math_fidelity"] == str(ttnn.MathFidelity.LoFi)
+    assert audit["logits_dtype_observed"] == "bfloat8_b"
+    assert audit["sampling_dtype_observed"] == "bfloat8_b"

code/models/demos/blackhole/qwen3_coder_30b_a3b/tests/test_reference.py

@@ -0,0 +1,123 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+"""Validate the layer-only reference before any TTNN work is built against it.
+
+The reference does one thing the checkpoint does not: it fuses each expert's
+``gate_proj``/``up_proj`` into a single ``gate_up_proj`` and stacks all experts.
+If that fusion were reversed, the layer would still run and produce
+plausible-looking numbers -- it would just be wrong. Everything downstream is
+compared against this reference, so an error here is invisible forever after.
+
+``test_moe_matches_unfused_reimplementation`` therefore recomputes the MoE block
+from the raw per-expert checkpoint tensors, following
+``Qwen3MoeSparseMoeBlock.forward`` literally, and requires the two to agree.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from models.common.utility_functions import comp_pcc
+
+from .reference import build_reference_layer, layer_state_dict, rotary_embeddings
+
+LAYER_IDX = 0
+SEQ_LEN = 32
+
+
+@pytest.fixture(scope="module")
+def reference():
+    layer, config = build_reference_layer(LAYER_IDX)
+    return layer, config
+
+
+def _hidden(config, seq_len=SEQ_LEN, seed=0):
+    """Activations roughly matching what reaches a decoder layer post-embedding."""
+    torch.manual_seed(seed)
+    return torch.randn(1, seq_len, config.hidden_size, dtype=torch.float32) * 0.02
+
+
+def test_layer_forward_is_finite_and_non_degenerate(reference):
+    layer, config = reference
+    hidden = _hidden(config)
+    cos, sin = rotary_embeddings(config, SEQ_LEN)
+
+    with torch.no_grad():
+        out = layer(hidden, position_embeddings=(cos, sin), attention_mask=None)
+    out = out[0] if isinstance(out, tuple) else out
+
+    assert out.shape == hidden.shape
+    assert torch.isfinite(out).all()
+    assert out.std() > 1e-6, "output is constant -- layer is not doing anything"
+    assert not torch.allclose(out, hidden), "output equals input -- residual-only path"
+
+
+def test_layer_forward_is_deterministic(reference):
+    layer, config = reference
+    hidden = _hidden(config)
+    cos, sin = rotary_embeddings(config, SEQ_LEN)
+
+    with torch.no_grad():
+        a = layer(hidden, position_embeddings=(cos, sin), attention_mask=None)
+        b = layer(hidden, position_embeddings=(cos, sin), attention_mask=None)
+    a = a[0] if isinstance(a, tuple) else a
+    b = b[0] if isinstance(b, tuple) else b
+
+    assert torch.equal(a, b), "reference is not deterministic; PCC comparisons would be unstable"
+
+
+def test_moe_matches_unfused_reimplementation(reference):
+    """Recompute the MoE block from raw per-expert tensors and require agreement.
+
+    This is the guard on the gate/up fusion order and on the router's
+    softmax -> top-k -> renormalise ordering.
+    """
+    layer, config = reference
+    # build_reference_layer upcasts the bf16 checkpoint into the fp32 module, so
+    # cast here too -- this test is about the fusion order, not about dtype.
+    sd = {k: v.float() for k, v in layer_state_dict(LAYER_IDX).items()}
+    hidden = _hidden(config)
+    flat = hidden.view(-1, config.hidden_size)
+
+    with torch.no_grad():
+        fused_out = layer.mlp(hidden).view(-1, config.hidden_size)
+
+    # --- independent implementation, straight from Qwen3MoeTopKRouter.forward ---
+    logits = torch.nn.functional.linear(flat, sd["mlp.gate.weight"])
+    probs = torch.softmax(logits, dim=-1, dtype=torch.float)  # over ALL experts, fp32
+    top_w, top_i = torch.topk(probs, config.num_experts_per_tok, dim=-1)
+    if config.norm_topk_prob:
+        top_w = top_w / top_w.sum(dim=-1, keepdim=True)
+    top_w = top_w.to(logits.dtype)
+
+    # --- and from Qwen3MoeExperts.forward, but with unfused checkpoint tensors ---
+    manual = torch.zeros_like(flat)
+    for token in range(flat.shape[0]):
+        for slot in range(config.num_experts_per_tok):
+            e = int(top_i[token, slot])
+            x = flat[token]
+            gate = torch.nn.functional.linear(x, sd[f"mlp.experts.{e}.gate_proj.weight"])
+            up = torch.nn.functional.linear(x, sd[f"mlp.experts.{e}.up_proj.weight"])
+            h = torch.nn.functional.silu(gate) * up
+            manual[token] += torch.nn.functional.linear(h, sd[f"mlp.experts.{e}.down_proj.weight"]) * top_w[token, slot]
+
+    passing, message = comp_pcc(manual, fused_out, pcc=0.9999)
+    assert passing, f"fused MoE disagrees with unfused checkpoint math: {message}"
+
+
+def test_router_selects_expected_expert_count(reference):
+    """top_k experts per token, weights renormalised to 1."""
+    layer, config = reference
+    hidden = _hidden(config)
+    flat = hidden.view(-1, config.hidden_size)
+
+    with torch.no_grad():
+        _, scores, indices = layer.mlp.gate(flat)
+
+    assert indices.shape == (flat.shape[0], config.num_experts_per_tok)
+    assert indices.min() >= 0 and indices.max() < config.num_experts
+    if config.norm_topk_prob:
+        sums = scores.float().sum(dim=-1)
+        assert torch.allclose(sums, torch.ones_like(sums), atol=1e-5), f"router weights not normalised: {sums[:4]}"

code/models/demos/blackhole/qwen3_coder_30b_a3b/tests/test_rmsnorm.py

@@ -0,0 +1,99 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+"""First on-device module for Qwen3-Coder-30B-A3B: RMSNorm vs the HF reference.
+
+RMSNorm is the smallest piece of the decoder layer, so it is brought up first --
+it proves the whole harness (open a 1x1 mesh, upload real checkpoint weights,
+run, compare PCC against the layer-only reference) before GQA, RoPE and the MoE
+block are layered on top. When attention PCC later misbehaves, normalisation is
+already ruled out.
+
+Two Qwen3-specific details are asserted rather than assumed:
+  * eps is 1e-6, not the module default of 1e-5 -- so the config path is used.
+  * the norm is the plain variant; ``add_unit_offset`` stays False. Qwen3.5/3.6
+    use a zero-centred RMSNorm that folds a "+1" into the weight; this
+    checkpoint does not.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import pytest
+import torch
+from loguru import logger
+
+import ttnn
+from models.common.auto_compose import to_torch_auto_compose
+from models.common.modules.lazy_weight import LazyWeight
+from models.common.modules.rmsnorm.rmsnorm_1d import RMSNorm1D, RMSNorm1DConfig
+from models.common.utility_functions import comp_allclose, comp_pcc
+
+from .reference import build_reference_layer
+
+LAYER_IDX = 0
+PCC_REQUIRED = 0.999  # normalisation alone should be near-exact; the layer bar is 0.995
+
+
+@pytest.fixture(scope="module")
+def reference():
+    return build_reference_layer(LAYER_IDX)
+
+
+@pytest.mark.parametrize("mesh_device", [(1, 1)], ids=["1x1"], indirect=True)
+@pytest.mark.parametrize(
+    "norm_name,seq_len,mode",
+    [
+        ("input_layernorm", 32, "prefill"),
+        ("input_layernorm", 128, "prefill"),
+        ("post_attention_layernorm", 32, "prefill"),
+        ("input_layernorm", 1, "decode"),
+    ],
+    ids=["input_s32", "input_s128", "postattn_s32", "input_decode"],
+)
+def test_rmsnorm_vs_reference(
+    mesh_device: ttnn.MeshDevice,
+    reference,
+    norm_name: str,
+    seq_len: int,
+    mode: str,
+):
+    layer, config = reference
+    torch.manual_seed(0)
+
+    ref_norm = getattr(layer, norm_name)
+    dim = config.hidden_size
+    assert config.rms_norm_eps == 1e-6, f"unexpected eps {config.rms_norm_eps}"
+
+    # Activations scaled like what actually reaches a decoder layer, per the
+    # skill's guidance -- not arbitrary large randoms.
+    torch_input = (torch.randn(1, 1, seq_len, dim, dtype=torch.float32) * 0.02).to(torch.bfloat16)
+
+    ttnn.SetDefaultDevice(mesh_device)
+    try:
+        cache_dir = Path("model_cache/qwen3_coder_30b_a3b/rmsnorm")
+        lazy_weight = LazyWeight(
+            source=ref_norm.weight.data.clone(),
+            dtype=ttnn.bfloat16,
+            cache_dir_weight_name=(cache_dir, f"{norm_name}_L{LAYER_IDX}"),
+        )
+        tt_model = RMSNorm1D.from_config(
+            RMSNorm1DConfig(
+                weight=lazy_weight,
+                eps=config.rms_norm_eps,
+                add_unit_offset=False,  # plain RMSNorm, not the zero-centred Qwen3.5/3.6 variant
+            )
+        )
+        tt_out = tt_model.forward(LazyWeight(source=torch_input, dtype=ttnn.bfloat16), mode=mode)
+        tt_out_torch = to_torch_auto_compose(tt_out)
+    finally:
+        ttnn.SetDefaultDevice(None)
+
+    with torch.no_grad():
+        ref_out = ref_norm(torch_input.to(torch.float32)).to(torch.bfloat16)
+
+    passing, pcc_message = comp_pcc(ref_out, tt_out_torch, PCC_REQUIRED)
+    logger.info(comp_allclose(ref_out, tt_out_torch))
+    logger.info(f"RMSNorm[{norm_name}] {mode} seq={seq_len}: {pcc_message}")
+    assert passing, f"RMSNorm {norm_name} ({mode}, seq={seq_len}) below {PCC_REQUIRED}: {pcc_message}"

code/models/demos/blackhole/qwen3_coder_30b_a3b/tests/test_trace.py

@@ -0,0 +1,200 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+"""Capture the decode step as a trace and replay it.
+
+Tracing is the real test of every "on-device, trace-compatible" claim made
+while building this layer. A trace records device commands once and replays
+them, so anything that is not a pure device op -- a host round-trip, a
+Python-side branch on tensor *values*, a shape that depends on the data --
+either fails to capture or replays stale results. Notably, this is what makes
+the router's design load-bearing: it keeps top-k selection and the scatter on
+device precisely so this step is possible.
+
+Two properties are checked, and the second is the one that catches real bugs:
+
+1. the traced output matches the eager output for the same input;
+2. replaying with *different* input produces *different*, still-correct output.
+
+Property 2 is essential. A trace whose input tensor was captured by value
+rather than written in place replays the original activations forever and
+therefore passes property 1 perfectly, every time.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+from loguru import logger
+
+import ttnn
+from models.common.utility_functions import comp_allclose, comp_pcc
+
+from ..tt.functional_decoder import (
+    DecoderLayerConfig,
+    build_expert_sparsity,
+    build_rope_cache,
+    create_kv_cache,
+    decoder_layer_decode,
+    decoder_layer_prefill,
+    upload_layer_weights,
+)
+from ..tt.weight_mapping import convert_layer_weights
+from .reference import build_reference_layer, layer_state_dict, rotary_embeddings
+
+LAYER_IDX = 0
+PCC_REQUIRED = 0.99
+MAX_SEQ = 256
+PROMPT_LEN = 32
+# Reserved at device open; the capture fails outright if the graph needs more.
+TRACE_REGION_SIZE = 50331648
+
+
+@pytest.fixture(scope="module")
+def reference():
+    return build_reference_layer(LAYER_IDX)
+
+
+@pytest.fixture(scope="module")
+def torch_weights(reference):
+    _, hf_config = reference
+    return convert_layer_weights(layer_state_dict(LAYER_IDX), hf_config)
+
+
+def _hidden(hf_config, seq_len, seed=0):
+    torch.manual_seed(seed)
+    return torch.randn(1, seq_len, hf_config.hidden_size, dtype=torch.float32) * 0.02
+
+
+def _reference_layer(layer, hf_config, hidden):
+    seq_len = hidden.shape[1]
+    cos, sin = rotary_embeddings(hf_config, seq_len)
+    mask = torch.full((seq_len, seq_len), float("-inf")).triu(1).reshape(1, 1, seq_len, seq_len)
+    with torch.no_grad():
+        out = layer(hidden, position_embeddings=(cos, sin), attention_mask=mask)
+    return out[0] if isinstance(out, tuple) else out
+
+
+def _to_device(t, mesh_device):
+    return ttnn.from_torch(
+        t, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=mesh_device, memory_config=ttnn.DRAM_MEMORY_CONFIG
+    )
+
+
+@pytest.mark.parametrize("device_params", [{"trace_region_size": TRACE_REGION_SIZE}], indirect=True)
+@pytest.mark.parametrize("mesh_device", [(1, 1)], ids=["1x1"], indirect=True)
+def test_decode_step_is_traceable(mesh_device, reference, torch_weights):
+    layer, hf_config = reference
+    config = DecoderLayerConfig.from_hf(hf_config)
+
+    weights = upload_layer_weights(torch_weights, mesh_device, config)
+    cos_cache, sin_cache = build_rope_cache(hf_config, MAX_SEQ, mesh_device)
+    sparsity = build_expert_sparsity(mesh_device, config.moe.num_experts)
+    kv_cache = create_kv_cache(mesh_device, config.attention, max_batch=1, max_seq_len=MAX_SEQ)
+
+    hidden_full = _hidden(hf_config, PROMPT_LEN + 2)
+    ref_out = _reference_layer(layer, hf_config, hidden_full)
+
+    decoder_layer_prefill(
+        _to_device(hidden_full[:, :PROMPT_LEN, :].unsqueeze(0), mesh_device),
+        weights,
+        config,
+        cos_cache,
+        sin_cache,
+        sparsity,
+        kv_cache=kv_cache,
+    )
+
+    # Persistent input buffers: a trace replays writes to the *same* addresses,
+    # so inputs must be updated in place rather than rebound each step.
+    tt_in = _to_device(hidden_full[:, PROMPT_LEN, :].reshape(1, 1, 1, hf_config.hidden_size), mesh_device)
+    current_pos = ttnn.from_torch(torch.tensor([PROMPT_LEN], dtype=torch.int32), dtype=ttnn.int32, device=mesh_device)
+
+    def step():
+        return decoder_layer_decode(
+            tt_in, weights, config, cos_cache, sin_cache, kv_cache, current_pos, token_index=PROMPT_LEN
+        )
+
+    # Warm up so program compilation happens outside the capture.
+    eager_out = ttnn.to_torch(step()).reshape(1, hf_config.hidden_size)
+
+    trace_id = ttnn.begin_trace_capture(mesh_device, cq_id=0)
+    traced_out = step()
+    ttnn.end_trace_capture(mesh_device, trace_id, cq_id=0)
+
+    ttnn.execute_trace(mesh_device, trace_id, cq_id=0, blocking=True)
+    replayed = ttnn.to_torch(traced_out).reshape(1, hf_config.hidden_size)
+
+    passing, pcc_message = comp_pcc(eager_out, replayed, 0.999)
+    logger.info(comp_allclose(eager_out, replayed))
+    logger.info(f"traced vs eager: {pcc_message}")
+    assert passing, f"traced replay disagrees with eager execution: {pcc_message}"
+
+    passing, pcc_message = comp_pcc(ref_out[:, PROMPT_LEN, :], replayed, PCC_REQUIRED)
+    logger.info(f"traced vs reference: {pcc_message}")
+    assert passing, f"traced output below {PCC_REQUIRED} vs reference: {pcc_message}"
+
+    ttnn.release_trace(mesh_device, trace_id)
+
+
+@pytest.mark.parametrize("device_params", [{"trace_region_size": TRACE_REGION_SIZE}], indirect=True)
+@pytest.mark.parametrize("mesh_device", [(1, 1)], ids=["1x1"], indirect=True)
+def test_traced_replay_follows_new_input(mesh_device, reference, torch_weights):
+    """Writing a new token into the input buffer must change the traced output.
+
+    Guards the failure mode a same-input trace test cannot see: if the capture
+    bound the input by value, replay reproduces the first token's result
+    forever and every equality check still passes.
+    """
+    layer, hf_config = reference
+    config = DecoderLayerConfig.from_hf(hf_config)
+
+    weights = upload_layer_weights(torch_weights, mesh_device, config)
+    cos_cache, sin_cache = build_rope_cache(hf_config, MAX_SEQ, mesh_device)
+    sparsity = build_expert_sparsity(mesh_device, config.moe.num_experts)
+    kv_cache = create_kv_cache(mesh_device, config.attention, max_batch=1, max_seq_len=MAX_SEQ)
+
+    hidden_full = _hidden(hf_config, PROMPT_LEN + 1)
+    decoder_layer_prefill(
+        _to_device(hidden_full[:, :PROMPT_LEN, :].unsqueeze(0), mesh_device),
+        weights,
+        config,
+        cos_cache,
+        sin_cache,
+        sparsity,
+        kv_cache=kv_cache,
+    )
+
+    token_a = hidden_full[:, PROMPT_LEN, :].reshape(1, 1, 1, hf_config.hidden_size)
+    token_b = (_hidden(hf_config, 1, seed=99)).reshape(1, 1, 1, hf_config.hidden_size)
+
+    tt_in = _to_device(token_a, mesh_device)
+    current_pos = ttnn.from_torch(torch.tensor([PROMPT_LEN], dtype=torch.int32), dtype=ttnn.int32, device=mesh_device)
+
+    def step():
+        return decoder_layer_decode(
+            tt_in, weights, config, cos_cache, sin_cache, kv_cache, current_pos, token_index=PROMPT_LEN
+        )
+
+    step()  # warm up / compile
+
+    trace_id = ttnn.begin_trace_capture(mesh_device, cq_id=0)
+    traced_out = step()
+    ttnn.end_trace_capture(mesh_device, trace_id, cq_id=0)
+
+    ttnn.execute_trace(mesh_device, trace_id, cq_id=0, blocking=True)
+    out_a = ttnn.to_torch(traced_out).reshape(-1).float().clone()
+
+    # Overwrite the captured input buffer in place, then replay.
+    ttnn.copy_host_to_device_tensor(ttnn.from_torch(token_b, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT), tt_in)
+    ttnn.execute_trace(mesh_device, trace_id, cq_id=0, blocking=True)
+    out_b = ttnn.to_torch(traced_out).reshape(-1).float().clone()
+
+    delta = (out_a - out_b).abs().max().item()
+    logger.info(f"max|out_a - out_b| after swapping the input token = {delta:.6f}")
+    assert delta > 1e-3, (
+        "traced replay produced identical output for a different input token -- "
+        "the trace is not reading the live input buffer"
+    )
+
+    ttnn.release_trace(mesh_device, trace_id)

code/models/demos/blackhole/qwen3_coder_30b_a3b/tt-model-localplugin.yaml

@@ -0,0 +1,152 @@
+# A/B VARIANT of tt-model.yaml: identical EXCEPT that the vLLM plugin comes from a local
+# checkout rather than a cloned ref. Published as a separate repo so both packaging paths
+# can be exercised end to end.
+#
+# This one file is the whole authoring interface:
+#   tt-model package --container models/demos/blackhole/qwen3_coder_30b_a3b/tt-model.yaml
+#
+# It lives next to the model on purpose: the serving recipe then travels through review
+# in the same PR as the code it serves.
+schema: "5.1"
+
+repo: raahemnabeel/qwen3-coder-30b-a3b-localplugin
+name: qwen3-coder-30b-a3b-localplugin
+weights: Qwen/Qwen3-Coder-30B-A3B-Instruct   # a POINTER — 57 GB, never baked into the image
+kind: vllm-plugin                            # stock vLLM + the standalone tenstorrent/vllm-tt-plugin
+arch: blackhole
+
+source:
+  tt_metal: /home/raahem/metal-publish
+
+  # EXACTLY what ships — an allowlist, never a denylist. The image's tt-metal tree holds
+  # NO models/ except these, so an under-specified list fails the image's own build-time
+  # import check on the author's machine, not on a consumer's first boot.
+  code:
+    # tt/ imports models.common.modules.sampling.sampling_1d,
+    # models.common.modules.tt_ccl and models.common.readiness_check.contract
+    - models/common
+    # the whole model, 1.6 MB: tt/, vllm_bundle/, config/, tests/, README.md
+    - models/demos/blackhole/qwen3_coder_30b_a3b
+
+  ubuntu: "24.04"
+  python: "3.12"
+
+runtime:
+  # The plugin monkeypatches vLLM internals, so this pin is load-bearing, not cosmetic.
+  # This box serves with 0.24.0 (reported 0.24.0+empty — the local tag comes from the
+  # VLLM_TARGET_DEVICE=empty sdist build, which is what the image does too).
+  vllm: {version: "0.24.0"}
+
+  # A/B VARIANT: the author's LOCAL plugin checkout, staged into the build context the
+  # same way source.tt_metal is — uncommitted work included, nothing fetched. The sibling
+  # manifest (tt-model.yaml) pins a pushed SHA and clones it during the build instead.
+  # Both produce an image whose venv holds the plugin; serving cannot tell them apart.
+  plugin:
+    path: /home/raahem/vllm-tt-plugin
+
+  # The directory the plugin SCANS — it walks the CHILDREN of this dir for
+  # vllm_metadata.json. Registration chain:
+  #   register_tt_models() -> _register_models_from_extra_dir() -> this dir ->
+  #   vllm_bundle/qwen3_coder_30b_a3b_instruct/vllm_metadata.json ->
+  #   Qwen3MoeForCausalLM registered as TTQwen3MoeForCausalLM ->
+  #   "tt_qwen3_coder_30b_a3b_instruct:Qwen3CoderForCausalLM" -> the shim ->
+  #   models.demos.blackhole.qwen3_coder_30b_a3b.tt.generator_vllm:Qwen3CoderForCausalLM
+  # The shim's parents[6] still resolves in-container: the tree lands at
+  # /opt/tt-metal/models/demos/blackhole/<model>/vllm_bundle/<bundle>/, so six levels up
+  # is /opt/tt-metal, which is also PYTHONPATH. Setting this also makes the build pass
+  # TT_VLLM_BUILTIN_MODELS=0, so registration comes solely from here.
+  extra_models_dir: models/demos/blackhole/qwen3_coder_30b_a3b/vllm_bundle
+
+  # Uncomment once requirements.lock is committed next to this file, after the first
+  # successful build — later builds then resolve nothing at all.
+  # lock: requirements.lock
+
+serve:
+  port: 8000
+
+  # 256000, NOT the 262144 HF advertises. config/context_contract.json records
+  # current_supported_context: 256000 with capability_reduction: true, and
+  # tt/generator_vllm.py reads that file as the source of truth and REFUSES to start
+  # above it. (Reason on record: an unresolved prefill cliff in the last ~3,100 tokens,
+  # and 17.5 min TTFT at that length even on the path that works.)
+  max_model_len: 256000
+
+  # 64, from the validated tt-inference-server P300X2 spec, which is authoritative here.
+  # config/context_contract.json's prose says "paging is unchanged (block_size 32)"; that
+  # text is older than the serving config and the spec supersedes it.
+  block_size: 64
+
+  # One configuration, measured. max_num_seqs only bounds what vLLM batches IN FLIGHT, and
+  # with the stage-11 decode-width ladder now the model default, batch-1 decode is already
+  # efficient at this cap — so a lower cap buys nothing. Measured on this box, 128 output
+  # tokens, tt-model 0.1.0:
+  #
+  #                          seqs=1        seqs=32
+  #   single-stream TTFT     0.232 s       0.219 s
+  #   single-stream decode   50.6 tok/s    49.2 tok/s
+  #   8 concurrent, total    46.7 tok/s    71.4 tok/s
+  #   8 concurrent, /user     5.8 tok/s     8.9 tok/s
+  #   slowest of the 8       21.9 s        14.4 s
+  #
+  # seqs=32 is identical for one user and strictly better for several, so there is no
+  # latency-vs-capacity trade to expose as separate profiles. Untested above 8 concurrent,
+  # and TTFT-under-load is unmeasured; if either turns out to matter, `serve_profiles:`
+  # is how you would split this again.
+  hardware: p300x2
+  mesh_device: P300x2
+  max_num_seqs: 32
+
+  # Structured, not a JSON string: tt-model renders it into --additional-config.
+  additional_config:
+    tt:
+      sample_on_device_mode: all
+      trace_region_size: 50331648
+      fabric_config: FABRIC_1D_RING
+
+  # Tool calling. tt-model emits --enable-auto-tool-choice alongside the parser, because
+  # vLLM hard-errors on the parser flag without it.
+  capabilities:
+    tool_parser: qwen3_coder
+
+  # Runtime environment, from the validated tt-inference-server P300X2 spec.
+  # MESH_DEVICE is NOT listed: it is derived from each profile's mesh_device.
+  # VLLM_TARGET_DEVICE is NOT listed either: it is a BUILD-time variable for the vLLM
+  # fork, and this image builds stock vLLM with VLLM_TARGET_DEVICE=empty plus the plugin.
+  env:
+    ARCH_NAME: blackhole
+    VLLM_CONFIGURE_LOGGING: "1"
+    VLLM_RPC_TIMEOUT: "900000"
+    VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
+    TORCHDYNAMO_DISABLE: "1"
+    TT_METAL_OPERATION_TIMEOUT_SECONDS: "120.0"
+    # The variable-width decode ladder (stage 11, doc/batch_scaling). This is the model's
+    # default now, so setting it changes nothing today — it is kept to pin the behaviour
+    # explicitly, so a future change to that default cannot silently move this model.
+    QWEN3_DECODE_WIDTHS: "1,2,4,8,16,32"
+
+  # Anything without a named field goes here, verbatim.
+  args:
+    - [--max-num-batched-tokens, "256000"]
+    - [--max-log-len, "32"]
+    - [--generation-config, vllm]
+    - [--seed, "9472"]
+
+
+# Build-time assertions, run INSIDE the finished image after the USER switch. These make
+# the code/ allowlist and the tree prune safe: an under-shipped image fails HERE.
+verify:
+  # config/ is RUNTIME DATA, not documentation. tt/generator_vllm.py reads both files;
+  # _supported_context() swallows OSError and silently falls back to in-code defaults, so
+  # a missing file would serve at the wrong precision/context with no error at all.
+  - "from pathlib import Path; p = Path('/opt/tt-metal/models/demos/blackhole/qwen3_coder_30b_a3b'); assert (p/'config'/'selected_precision_config.json').is_file(), 'precision config missing — the model would silently use in-code defaults'; assert (p/'config'/'context_contract.json').is_file(), 'context contract missing — _supported_context() would silently fall back'"
+  # the real adapter. (The bundle shim is NOT asserted here: it is importable only once
+  # something has put its folder on sys.path, which is what the plugin does. tt-model
+  # verifies that resolution generically for every model — see launchers.py's
+  # RESOLVE_EXTRA_MODELS — so a per-model assertion would only duplicate it, badly.)
+  - "from models.demos.blackhole.qwen3_coder_30b_a3b.tt.generator_vllm import Qwen3CoderForCausalLM; assert Qwen3CoderForCausalLM"
+
+card:
+  quickstart: |
+    ### Use it
+    An OpenAI-compatible server on `http://127.0.0.1:8000`. Point any OpenAI client at it
+    with model id `Qwen/Qwen3-Coder-30B-A3B-Instruct`.

code/models/demos/blackhole/qwen3_coder_30b_a3b/tt-model.yaml

@@ -0,0 +1,242 @@
+# tt-model container package (v5.1) — qwen3-coder-30b-a3b on QB2 (2x P300, 4 chips).
+#
+# This one file is the whole authoring interface:
+#   tt-model package --container models/demos/blackhole/qwen3_coder_30b_a3b/tt-model.yaml
+#
+# It lives next to the model on purpose: the serving recipe then travels through review
+# in the same PR as the code it serves.
+schema: "5.1"
+
+repo: raahemnabeel/qwen3-coder-30b-a3b
+name: qwen3-coder-30b-a3b
+weights: Qwen/Qwen3-Coder-30B-A3B-Instruct   # a POINTER — 57 GB, never baked into the image
+kind: vllm-plugin                            # stock vLLM + the standalone tenstorrent/vllm-tt-plugin
+arch: blackhole
+
+source:
+  tt_metal: /home/raahem/metal-publish
+
+  # EXACTLY what ships — an allowlist, never a denylist. The image's tt-metal tree holds
+  # NO models/ except these, so an under-specified list fails the image's own build-time
+  # import check on the author's machine, not on a consumer's first boot.
+  code:
+    # tt/ imports models.common.modules.sampling.sampling_1d,
+    # models.common.modules.tt_ccl and models.common.readiness_check.contract
+    - models/common
+    # the whole model, 1.6 MB: tt/, vllm_bundle/, config/, tests/, README.md
+    - models/demos/blackhole/qwen3_coder_30b_a3b
+
+  ubuntu: "24.04"
+  python: "3.12"
+
+runtime:
+  # The plugin monkeypatches vLLM internals, so this pin is load-bearing, not cosmetic.
+  # This box serves with 0.24.0 (reported 0.24.0+empty — the local tag comes from the
+  # VLLM_TARGET_DEVICE=empty sdist build, which is what the image does too).
+  vllm: {version: "0.24.0"}
+
+  # THE PLUGIN, bundled from your own checkout — the same hermetic treatment
+  # source.tt_metal gets. Whatever is in this directory ships, uncommitted work
+  # included; nothing is fetched at build time and nothing is resolved at serve time.
+  # `package` records its HEAD sha and flags the tree dirty if it is.
+  #
+  # Alternatives, when a local checkout is not the right input:
+  #   plugin: {repo: https://github.com/tenstorrent/vllm-tt-plugin, ref: <pushed sha>}
+  #   plugin: {version: "0.1.0"}
+  plugin:
+    path: /home/raahem/vllm-tt-plugin
+
+  # The directory the plugin SCANS — it walks the CHILDREN of this dir for
+  # vllm_metadata.json. Registration chain:
+  #   register_tt_models() -> _register_models_from_extra_dir() -> this dir ->
+  #   vllm_bundle/qwen3_coder_30b_a3b_instruct/vllm_metadata.json ->
+  #   Qwen3MoeForCausalLM registered as TTQwen3MoeForCausalLM ->
+  #   "tt_qwen3_coder_30b_a3b_instruct:Qwen3CoderForCausalLM" -> the shim ->
+  #   models.demos.blackhole.qwen3_coder_30b_a3b.tt.generator_vllm:Qwen3CoderForCausalLM
+  # The shim's parents[6] still resolves in-container: the tree lands at
+  # /opt/tt-metal/models/demos/blackhole/<model>/vllm_bundle/<bundle>/, so six levels up
+  # is /opt/tt-metal, which is also PYTHONPATH. Setting this also makes the build pass
+  # TT_VLLM_BUILTIN_MODELS=0, so registration comes solely from here.
+  extra_models_dir: models/demos/blackhole/qwen3_coder_30b_a3b/vllm_bundle
+
+  # Uncomment once requirements.lock is committed next to this file, after the first
+  # successful build — later builds then resolve nothing at all.
+  # lock: requirements.lock
+
+serve:
+  # The port the image binds under a bare `docker run`. `tt-model serve` does NOT use this
+  # as its seed — it opens 20000 and walks past busy ports — so leave it at vLLM's default.
+  port: 8000
+
+  # 256000, NOT the 262144 HF advertises. config/context_contract.json records
+  # current_supported_context: 256000 with capability_reduction: true, and
+  # tt/generator_vllm.py reads that file as the source of truth and REFUSES to start
+  # above it. (Reason on record: an unresolved prefill cliff in the last ~3,100 tokens,
+  # and 17.5 min TTFT at that length even on the path that works.)
+  max_model_len: 256000
+
+  # 64, from the validated tt-inference-server P300X2 spec, which is authoritative here.
+  # config/context_contract.json's prose says "paging is unchanged (block_size 32)"; that
+  # text is older than the serving config and the spec supersedes it.
+  block_size: 64
+
+  # One configuration, measured. max_num_seqs only bounds what vLLM batches IN FLIGHT, and
+  # with the stage-11 decode-width ladder now the model default, batch-1 decode is already
+  # efficient at this cap — so a lower cap buys nothing. Measured on this box, 128 output
+  # tokens, tt-model 0.1.0:
+  #
+  # (All figures below are p300x2. The p150x4 profile is validated by the author but its
+  # throughput was not separately measured, so do not read these numbers as applying to it.)
+  #
+  #                          seqs=1        seqs=32
+  #   single-stream TTFT     0.232 s       0.219 s
+  #   single-stream decode   50.6 tok/s    49.2 tok/s
+  #   8 concurrent, total    46.7 tok/s    71.4 tok/s
+  #   8 concurrent, /user     5.8 tok/s     8.9 tok/s
+  #   slowest of the 8       21.9 s        14.4 s
+  #
+  # seqs=32 is identical for one user and strictly better for several, so there is no
+  # latency-vs-capacity trade to expose as separate profiles. Untested above 8 concurrent,
+  # and TTFT-under-load is unmeasured; if either turns out to matter, `serve_profiles:`
+  # is how you would split this again.
+  #
+  # Those decode figures are the ARGMAX path, which the --override-generation-config below
+  # now makes the default request — so they describe what an unparameterised client gets.
+  # A client that asks for sampling gets 38.6 tok/s instead; that gap is the strategy, not
+  # the configuration. Re-measured on this box after the warmup change: unparameterised
+  # 49.147 tok/s at 0.304 s TTFT first request, 8 concurrent 77.55 tok/s aggregate
+  # (9.69 tok/s/u), max TTFT 0.90 s.
+  # hardware/mesh_device are the ONLY per-profile fields — see serve_profiles below.
+  max_num_seqs: 32
+
+  # Structured, not a JSON string: tt-model renders it into --additional-config.
+  additional_config:
+    tt:
+      sample_on_device_mode: all
+      trace_region_size: 50331648
+      fabric_config: FABRIC_1D_RING
+
+  # Tool calling. tt-model emits --enable-auto-tool-choice alongside the parser, because
+  # vLLM hard-errors on the parser flag without it.
+  capabilities:
+    tool_parser: qwen3_coder
+
+  # Runtime environment, from the validated tt-inference-server P300X2 spec.
+  # MESH_DEVICE is NOT listed: it is derived from each profile's mesh_device.
+  # VLLM_TARGET_DEVICE is NOT listed either: it is a BUILD-time variable for the vLLM
+  # fork, and this image builds stock vLLM with VLLM_TARGET_DEVICE=empty plus the plugin.
+  env:
+    ARCH_NAME: blackhole
+    VLLM_CONFIGURE_LOGGING: "1"
+    VLLM_RPC_TIMEOUT: "900000"
+    VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
+    TORCHDYNAMO_DISABLE: "1"
+    TT_METAL_OPERATION_TIMEOUT_SECONDS: "120.0"
+    # The variable-width decode ladder (stage 11, doc/batch_scaling). This is the model's
+    # default now, so setting it changes nothing today — it is kept to pin the behaviour
+    # explicitly, so a future change to that default cannot silently move this model.
+    QWEN3_DECODE_WIDTHS: "1,2,4,8,16,32"
+    # Warm a decode graph for every ladder width in BOTH sampling strategies at startup,
+    # rather than compiling one on whichever request first needs it. Costs 9.5 s of boot
+    # (decode warmup 2.41 s -> 11.89 s) and 24.8 MB/die; buys 1.64 s off the first
+    # request's TTFT (1.9416 s -> 0.2989 s) and removes the same stall from the first
+    # request at any other width or strategy. "both" is the model default; it is pinned
+    # here for the same reason QWEN3_DECODE_WIDTHS is.
+    QWEN3_WARMUP_SAMPLING: "both"
+    # Prefill bucketing (doc/prefill_buckets). A prefill program is compiled per exact
+    # sequence length, so without this every new prompt length paid a fresh compile —
+    # measured on 4 dies, cold cache: 4–11 s on top of the 0.89 ms/token floor, and
+    # recurring, not a startup cost (2443 tok cost 22.87 s, then 2444 tok cost 4.29 s
+    # again). Bucketing rounds each prefill to a ladder rung so the shape space is
+    # finite and warmable. "proportional" bounds the padding waste at 1/8th of the
+    # prefill at every length; measured steady-state tax is 2–3%, against 26% for a
+    # half-power-of-two ladder and ~111 s worst case for plain powers of two.
+    QWEN3_PREFILL_BUCKETS: "proportional"
+    # Warm the ladder once per HOST, not once per boot. The rungs and the cached-suffix
+    # shapes cost 217.9 s to compile on a cold cache and ~68 k tokens of prefill to run,
+    # so "auto" runs them only when its marker is absent from TT_METAL_CACHE — which
+    # tt-model mounts per model and keeps across container removal on purpose. Clearing
+    # that cache clears the marker, and the next boot re-warms. "full" ignores the
+    # marker; "off" skips it and leaves every shape to compile on a real request.
+    QWEN3_PREFILL_WARMUP: "auto"
+    # Warm every rung up to and including the one that covers 8192 tokens (23 rungs).
+    # Longer prompts still compile lazily, once each, and persist in the same cache.
+    # Raising this lengthens only the first boot on a fresh host; "0" warms all 52.
+    QWEN3_PREFILL_WARMUP_MAX: "8192"
+    # The cached-suffix (prefix-cache hit) path is a separate program set, keyed on the
+    # chunked-SDPA chunk size AND the suffix length. Reachable chunk sizes are just
+    # {64, 128, 256} at block_size 64, so this caps the second axis: suffix rungs up to
+    # 1024, which is where a hit's new tail actually lands.
+    QWEN3_PREFILL_WARMUP_SUFFIX_MAX: "1024"
+
+  # Anything without a named field goes here, verbatim.
+  args:
+    - [--max-num-batched-tokens, "256000"]
+    - [--max-log-len, "32"]
+    - [--generation-config, vllm]
+    # A GREEDY DEFAULT, measured. `--generation-config vllm` alone means an
+    # unparameterised request arrives at vLLM's own temperature=1.0 / top_p=1.0, and the
+    # model routes anything with k>1 or p>0 to split sampling: 25.88 ms/token (38.6 t/s/u)
+    # against 20.35 ms (49.1) for argmax, on 4 dies, 24-token completions. For a coding
+    # model that is the wrong default three times over — 21 % slower, non-deterministic
+    # across identical prompts, and full-distribution sampling at temperature 1.0 is the
+    # worst of the three available regimes for code. Overriding temperature to 0 makes the
+    # default request greedy (verified: 49.147 t/s/u unparameterised) while leaving
+    # repetition_penalty unset, so the decode graph stays on penalty mode 0. A client that
+    # wants sampling still sends temperature itself and gets the split path, warm — the
+    # model warms both strategies at startup (QWEN3_WARMUP_SAMPLING below).
+    #
+    # QUOTING IS LOAD-BEARING. launchers.py joins these with spaces into ONE
+    # Written BARE. The extra single quotes this used to carry were a workaround for the
+    # vllm-fork path, which joins args into one --additional-server-args string that the
+    # runner shlex-splits: bare JSON lost its double quotes there and reached vLLM as
+    # {temperature: 0}. This model is vllm-plugin, which passes argv loose -- compose_run
+    # hands docker the list and the entrypoint ends in `exec "$@"`, so nothing consumes a
+    # shell quote and the literal ' characters reached json.loads instead. tt-model now
+    # shlex-quotes the fork path itself, so bare is correct for both kinds.
+    - [--override-generation-config, '{"temperature": 0}']
+    - [--seed, "9472"]
+
+
+# TWO BOARDS, ONE CONFIGURATION. p300x2 (2x P300) and p150x4 (4x P150) are both four
+# Blackhole chips on a (1, 4) mesh — tt-model derives 4 from either label and
+# parse_mesh_device resolves both SKUs to the same (1, 4) — so nothing above this line
+# differs between them: same context bound, same block size, same slot count, same
+# trace region, same fabric, same env, same args. Both were tested on hardware.
+#
+# They are separate profiles rather than one because `hardware` is what tt-model uses to
+# state a package's device requirement, and a p150x4 box should not have to read "p300x2"
+# and infer that the chip counts happen to match.
+#
+# p300x2 stays the default: it is the box every number in this file was measured on.
+# `tt-model serve --profile p150x4` selects the other; `tt-model profiles` lists both.
+default_profile: p300x2
+serve_profiles:
+  - name: p300x2
+    hardware: p300x2
+    mesh_device: P300x2
+  - name: p150x4
+    hardware: p150x4
+    mesh_device: P150x4
+
+# Build-time assertions, run INSIDE the finished image after the USER switch. These make
+# the code/ allowlist and the tree prune safe: an under-shipped image fails HERE.
+verify:
+  # config/ is RUNTIME DATA, not documentation. tt/generator_vllm.py reads both files;
+  # _supported_context() swallows OSError and silently falls back to in-code defaults, so
+  # a missing file would serve at the wrong precision/context with no error at all.
+  - "from pathlib import Path; p = Path('/opt/tt-metal/models/demos/blackhole/qwen3_coder_30b_a3b'); assert (p/'config'/'selected_precision_config.json').is_file(), 'precision config missing — the model would silently use in-code defaults'; assert (p/'config'/'context_contract.json').is_file(), 'context contract missing — _supported_context() would silently fall back'"
+  # the real adapter. (The bundle shim is NOT asserted here: it is importable only once
+  # something has put its folder on sys.path, which is what the plugin does. tt-model
+  # verifies that resolution generically for every model — see launchers.py's
+  # RESOLVE_EXTRA_MODELS — so a per-model assertion would only duplicate it, badly.)
+  - "from models.demos.blackhole.qwen3_coder_30b_a3b.tt.generator_vllm import Qwen3CoderForCausalLM; assert Qwen3CoderForCausalLM"
+
+card:
+  description: >
+    Qwen3-Coder-30B-A3B-Instruct — a 30B mixture-of-experts coding model for agentic
+    development work, with tool calling and a 256K context, served on Blackhole via vLLM.
+  quickstart: |
+    ### Use it
+    Point any OpenAI client at the address `tt-model serve` prints (port 20000 unless it
+    was busy), with model id `Qwen/Qwen3-Coder-30B-A3B-Instruct`.

code/models/demos/blackhole/qwen3_coder_30b_a3b/tt/functional_decoder.py

@@ -0,0 +1,1169 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+"""TTNN decoder layer for Qwen3-Coder-30B-A3B-Instruct.
+
+Built bottom-up and validated piece by piece against the layer-only HuggingFace
+reference in ``tests/reference.py``. Attention lands first; the MoE block and
+the composed layer follow.
+
+Shape of this model's attention
+-------------------------------
+    hidden 2048 | 32 Q heads | 4 KV heads (8:1 GQA) | head_dim 128
+
+``32 * 128 = 4096 != hidden``, so ``o_proj`` is ``[2048, 4096]`` -- head_dim is
+an independent config field here, not ``hidden / n_heads``. There is no
+attention bias and no sliding window, so every layer is plain causal attention.
+
+Two Qwen3-specific things the more common Llama-shaped ports do not have:
+
+  * **QK-norm.** An RMSNorm over ``head_dim`` is applied per head to Q and to K,
+    after the head split and *before* RoPE. Note it is Q and K only -- gemma4,
+    which this file follows closely, additionally norms V because Gemma has a
+    ``v_norm``. Qwen3 does not; adding it there would silently corrupt V.
+
+  * **RoPE theta lives in ``rope_scaling``.** ``config.rope_theta`` does not
+    exist on this checkpoint; the value (1e7, the long-context setting) is
+    nested at ``config.rope_scaling["rope_theta"]``. Rather than reach in and
+    risk a stale default, the cos/sin cache is produced by
+    ``Qwen3MoeRotaryEmbedding`` itself, which is correct by construction.
+
+RoPE convention is HF-style throughout -- see the note in ``weight_mapping.py``.
+Weights therefore keep their checkpoint channel order.
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+
+import torch
+
+import ttnn
+
+
+@dataclass(frozen=True)
+class AttentionConfig:
+    hidden_size: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    head_dim: int
+    rms_norm_eps: float
+
+    @classmethod
+    def from_hf(cls, config) -> "AttentionConfig":
+        assert not getattr(config, "attention_bias", False), "attention bias is not wired up"
+        assert not getattr(config, "use_sliding_window", False), "sliding-window attention is not wired up"
+        return cls(
+            hidden_size=config.hidden_size,
+            num_attention_heads=config.num_attention_heads,
+            num_key_value_heads=config.num_key_value_heads,
+            head_dim=config.head_dim,
+            rms_norm_eps=config.rms_norm_eps,
+        )
+
+
+@dataclass
+class AttentionWeights:
+    """Device-resident attention weights, as produced by ``weight_mapping``."""
+
+    wqkv: ttnn.Tensor  # [1, 1, hidden, (n_heads + 2*n_kv_heads) * head_dim]
+    wo: ttnn.Tensor  # [1, 1, n_heads * head_dim, hidden]
+    q_norm: ttnn.Tensor  # [1, 1, 1, head_dim]
+    k_norm: ttnn.Tensor  # [1, 1, 1, head_dim]
+
+
+def build_rope_cache(hf_config, max_seq_len: int, device) -> tuple[ttnn.Tensor, ttnn.Tensor]:
+    """Upload a ``[1, 1, max_seq_len, head_dim]`` cos/sin pair for HF-style RoPE.
+
+    Generated by the reference rotary module so ``rope_theta`` (and any future
+    scaling) can never drift from the checkpoint.
+    """
+    from ..tests.reference import rotary_embeddings
+
+    cos, sin = rotary_embeddings(hf_config, max_seq_len)  # [1, S, head_dim]
+    out = []
+    for t in (cos, sin):
+        out.append(
+            ttnn.from_torch(
+                t.unsqueeze(0).float(),
+                dtype=ttnn.bfloat16,
+                layout=ttnn.TILE_LAYOUT,
+                device=device,
+                memory_config=ttnn.DRAM_MEMORY_CONFIG,
+            )
+        )
+    return out[0], out[1]
+
+
+def _per_head_rms_norm(tensor: ttnn.Tensor, weight: ttnn.Tensor, eps: float) -> ttnn.Tensor:
+    """RMSNorm over the last dim, applied independently per head.
+
+    ``ttnn.rms_norm`` normalises the last dimension of a 2D-ish tile grid, so
+    the heads are folded into the row dimension and unfolded afterwards.
+    Input/output ``[1, n_heads, seq, head_dim]``.
+    """
+    shape = tensor.shape
+    flat = ttnn.reshape(tensor, (1, 1, shape[1] * shape[2], shape[3]))
+    normed = ttnn.rms_norm(flat, weight=weight, epsilon=eps)
+    return ttnn.reshape(normed, shape)
+
+
+def _apply_rope(tensor: ttnn.Tensor, cos_cache, sin_cache, token_index=None) -> ttnn.Tensor:
+    """HF-style rotary embedding, restoring the logical dim-2 length afterwards.
+
+    The op pads dim 2 up to a tile multiple, and dim 2 means different things in
+    the two modes -- which makes this a hazard twice over:
+
+    * **decode**: dim 2 is the head count. Q already has 32 heads so nothing
+      happens, but K has only ``num_kv_heads`` = 4 and comes back 32-deep. The
+      reshape declares logical-vs-padded before slicing, following the
+      tt_transformers ``_hf_rope_decode`` pattern.
+
+    * **prefill**: dim 2 is the sequence. At a non-tile-aligned length (say 33)
+      Q and K return padded to 64 while **V never passes through RoPE** and
+      stays 33, so SDPA rejects them with "K and V sequence length must match".
+      Slicing back is what keeps odd prompt lengths working.
+    """
+    orig = tensor.shape
+    out = ttnn.experimental.rotary_embedding(tensor, cos_cache, sin_cache, token_index)
+    if out.shape[2] != orig[2]:
+        if token_index is not None:
+            out = ttnn.reshape(out, (orig[0], orig[1], orig[2], orig[3]), (orig[0], orig[1], 32, orig[3]))
+            out = out[:, :, : orig[2]]
+        else:
+            out = ttnn.slice(out, [0, 0, 0, 0], [orig[0], orig[1], orig[2], orig[3]])
+    return out
+
+
+def rope_transformation_matrix() -> "torch.Tensor":
+    """The 32x32 matrix ``rotary_embedding_llama`` rotates a tile row with.
+
+    ``+1`` at ``(2i, 2i+1)`` and ``-1`` at ``(2i+1, 2i)`` -- i.e. the Meta
+    channel pairing, expressed as a matmul so the kernel needs no gather. This
+    is what makes the llama op cheaper than the HF one at the same core count:
+    the HF op reads a cos/sin row out of a DRAM cache per call, this one
+    multiplies by a resident 32x32 tile.
+    """
+    import torch
+
+    d = 32
+    m = torch.zeros(1, 1, d, d)
+    m[..., torch.arange(0, d, 2), torch.arange(1, d, 2)] = 1
+    m[..., torch.arange(1, d, 2), torch.arange(0, d, 2)] = -1
+    return m
+
+
+def apply_rope_llama(tensor: ttnn.Tensor, cos_shard, sin_shard, trans_mat) -> ttnn.Tensor:
+    """Meta-style rotary embedding, decode mode.
+
+    ``tensor``, ``cos_shard`` and ``sin_shard`` must all be the height-sharded
+    ``[1, batch, 32, head_dim]`` L1 config ``nlp_create_qkv_heads_decode``
+    produces, and ``tensor`` must be in **Meta** channel order -- which is a
+    property of the *weights*, established once at upload by
+    ``weight_mapping.permute_wqkv_to_meta``, not of anything done here.
+
+    Measured against ``_apply_rope`` at the shipped per-die decode shape:
+    3.84 -> 1.26 us, ``max|diff|`` exactly 0.0 and PCC 1.0000000
+    (``doc/optimized_multichip_decoder/probes/rope_probe.py``).
+
+    Unlike ``_apply_rope`` this needs **no** reshape-and-slice afterwards: the
+    op is shape-preserving in decode mode, so the dim-2 padding hazard that
+    ``_apply_rope`` documents does not arise.
+    """
+    return ttnn.experimental.rotary_embedding_llama(tensor, cos_shard, sin_shard, trans_mat, is_decode_mode=True)
+
+
+def _concat_heads_decode(attn: ttnn.Tensor, config: AttentionConfig) -> ttnn.Tensor:
+    """Merge heads in decode mode. ``attn`` ``[1, batch, num_heads, head_dim]``.
+
+    ``nlp_concat_heads_decode`` is the multi-core decode variant and requires a
+    height-sharded input with one core per user, whereas SDPA-decode hands back
+    a DRAM interleaved tensor -- hence the reshard. The op also pads batch up to
+    a full tile, so the result is sliced back to the logical batch.
+    """
+    batch = attn.shape[1]
+
+    grid_x = min(batch, 8)
+    while batch % grid_x:
+        grid_x -= 1
+    grid_y = batch // grid_x
+    core_grid = ttnn.CoreRangeSet({ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(grid_x - 1, grid_y - 1))})
+    shard_cfg = ttnn.create_sharded_memory_config(
+        shape=(ttnn.TILE_SIZE, config.head_dim),
+        core_grid=core_grid,
+        strategy=ttnn.ShardStrategy.HEIGHT,
+        orientation=ttnn.ShardOrientation.ROW_MAJOR,
+        use_height_and_width_as_shard_shape=True,
+    )
+
+    sharded = ttnn.to_memory_config(attn, shard_cfg)
+    out = ttnn.experimental.nlp_concat_heads_decode(sharded, num_heads=config.num_attention_heads)
+    ttnn.deallocate(sharded)
+    out = ttnn.sharded_to_interleaved(out, ttnn.DRAM_MEMORY_CONFIG)
+    if out.shape[2] != batch:
+        out = out[:, :, :batch, :]
+    return out
+
+
+@dataclass
+class KVCache:
+    """K/V cache, either contiguous per user or paged through a block table.
+
+    Paged mode is what a serving stack actually uses: logical positions are
+    mapped to physical blocks by ``page_table``, so users can share a block pool
+    instead of each reserving ``max_seq_len``. Both modes are kept because they
+    exercise different kernels -- ``paged_*`` ops versus their contiguous
+    counterparts -- and the contiguous path is the simpler thing to bisect
+    against when a paged result looks wrong.
+    """
+
+    k: ttnn.Tensor
+    v: ttnn.Tensor
+    page_table: ttnn.Tensor | None = None
+    block_size: int = 0
+
+    @property
+    def is_paged(self) -> bool:
+        return self.page_table is not None
+
+
+def create_kv_cache(
+    device,
+    config: AttentionConfig,
+    max_batch: int,
+    max_seq_len: int,
+    block_size: int | None = None,
+) -> KVCache:
+    """Allocate a KV cache.
+
+    ``block_size=None`` gives a contiguous cache of
+    ``[max_batch, n_kv_heads, max_seq_len, head_dim]``. Passing a block size
+    switches to a paged cache of ``[num_blocks, n_kv_heads, block_size,
+    head_dim]`` with an identity page table -- block ``b`` of user ``u`` lives
+    at physical block ``u * blocks_per_seq + b``. A real scheduler would hand
+    out blocks from a free pool; the identity mapping keeps the plumbing honest
+    (every op still goes through the table) without pulling a block allocator
+    into the decoder.
+    """
+    n_kv, head_dim = config.num_key_value_heads, config.head_dim
+
+    if block_size is None:
+        shape = (max_batch, n_kv, max_seq_len, head_dim)
+        page_table = None
+    else:
+        blocks_per_seq = math.ceil(max_seq_len / block_size)
+        shape = (max_batch * blocks_per_seq, n_kv, block_size, head_dim)
+        page_table = ttnn.from_torch(
+            torch.arange(max_batch * blocks_per_seq, dtype=torch.int32).reshape(max_batch, blocks_per_seq),
+            dtype=ttnn.int32,
+            layout=ttnn.ROW_MAJOR_LAYOUT,
+            device=device,
+        )
+
+    k, v = (
+        ttnn.from_torch(
+            torch.zeros(shape),
+            dtype=ttnn.bfloat16,
+            layout=ttnn.TILE_LAYOUT,
+            device=device,
+            memory_config=ttnn.DRAM_MEMORY_CONFIG,
+        )
+        for _ in range(2)
+    )
+    return KVCache(k=k, v=v, page_table=page_table, block_size=block_size or 0)
+
+
+def match_cache_dtype(cache: ttnn.Tensor, x: ttnn.Tensor) -> ttnn.Tensor:
+    """Return ``x`` in the cache tensor's dtype, casting only if they differ.
+
+    **For the fill (prefill) writers only.** The two cache writers have
+    *opposite* dtype contracts, and neither says so out loud, which is how
+    stage 07's ``R19_kv_bfp8`` came to score at chance.
+    ``doc/datatype_sweep/probes/kv_bfp8_diagnosis.py`` measures both at the op
+    level, with no model:
+
+    ==================  ============  ============  =============================
+    op                  cache         input         round-trip PCC
+    ==================  ============  ============  =============================
+    paged_fill_cache    bfloat16      bfloat16      1.0 (control)
+    paged_fill_cache    bfloat8_b     bfloat16      **NaN**
+    paged_fill_cache    bfloat8_b     bfloat8_b     1.0
+    paged_update_cache  bfloat16      bfloat16      1.0 (control)
+    paged_update_cache  bfloat8_b     bfloat16      0.999969
+    paged_update_cache  bfloat8_b     bfloat8_b     **rejected by the op**
+    ==================  ============  ============  =============================
+
+    So ``paged_fill_cache`` needs the input cast **to** the cache dtype -- it
+    validates the input against a permissive ``OR`` that a mismatch satisfies,
+    and then writes NaN -- while ``paged_update_cache`` needs the input left
+    **alone**: it converts into the cache itself and hard-rejects a block-float
+    input (``paged_update_cache_device_operation.cpp:296``, *"Data type of input
+    tensor for update cache must be FLOAT32 or BFLOAT16"*). Casting at the
+    decode writer would turn silent corruption into a hard crash, which is why
+    this helper is applied at the fill sites and deliberately not at the update
+    sites.
+
+    Taking the dtype off the **cache tensor itself** rather than off
+    ``precision.kv_cache_dtype`` keeps the guarantee true by construction: the
+    thing the write must agree with is the allocated cache, and this cannot
+    drift from it even if a caller allocates a cache some other way.
+
+    The cast is a no-op in the shipped configuration (``kv_cache_dtype ==
+    activation_dtype == bfloat16``), so it costs nothing unless the cache dtype
+    is actually moved.
+    """
+    if x.dtype == cache.dtype:
+        return x
+    return ttnn.typecast(x, cache.dtype, memory_config=x.memory_config())
+
+
+#: Which prefill attention branch ran, and at what chunk size. Instrumentation
+#: only -- nothing reads it in production. It exists because a PCC of ~1.0 on a
+#: split prefill has two explanations, "the arithmetic coincided" and "the
+#: chunked branch never executed", and they are indistinguishable from the PCC
+#: alone. A probe can assert `chunked > 0` and make the question un-askable.
+PREFILL_ATTENTION_BRANCHES = {"standard": 0, "chunked": 0, "chunk_sizes": []}
+
+
+def sdpa_chunk_size(chunk_start_idx: int) -> int:
+    """q_chunk_size == k_chunk_size for a chunked-prefill offset.
+
+    ``chunk_start_idx`` must be a multiple of BOTH q_chunk_size and k_chunk_size
+    (sdpa_nanobind.cpp:487-493), so the largest legal power of two is the offset's
+    own lowest set bit. Capped at 256: 512 overflows L1 -- measured
+    (1760704 B against a 1572864 B limit) and already recorded at
+    multichip_decoder.py:972 as rejected at every length.
+
+    Derivation copied from models/tt_transformers model_config.py:1552-1582
+    rather than hardcoded, so it tracks upstream.
+    """
+    return min(256, chunk_start_idx & -chunk_start_idx)
+
+
+def _fill_cache(kv_cache: KVCache, k: ttnn.Tensor, v: ttnn.Tensor, user_id: int, fill_page_table=None) -> None:
+    """Write a whole prompt's K/V into the cache.
+
+    The paged kernel writes block-at-a-time, so a prompt that does not fill its
+    last block is zero-padded up to a block boundary first. Those trailing
+    positions are never read: decode passes ``cur_pos``, and SDPA only attends
+    up to it.
+
+    K/V are cast to the cache's own dtype before the write -- see
+    :func:`match_cache_dtype` for why the ops will not do it for us. The cast
+    goes *after* the pad, because ``ttnn.pad`` is a bfloat16/float32 op.
+    """
+    if not kv_cache.is_paged:
+        ttnn.fill_cache(kv_cache.k, match_cache_dtype(kv_cache.k, k), user_id)
+        ttnn.fill_cache(kv_cache.v, match_cache_dtype(kv_cache.v, v), user_id)
+        return
+
+    seq_len = k.shape[2]
+    padded = math.ceil(seq_len / kv_cache.block_size) * kv_cache.block_size
+    if padded != seq_len:
+        pad = [(0, 0), (0, 0), (0, padded - seq_len), (0, 0)]
+        k = ttnn.pad(k, pad, value=0.0)
+        v = ttnn.pad(v, pad, value=0.0)
+
+    # ``fill_page_table`` is the suffix write of a split prefill: a SINGLE-ROW
+    # table already sliced to the blocks the suffix occupies, so the op's
+    # block-0-relative write lands at the right absolute offset. batch_idx is 0
+    # because the row has already been selected. None => the shipped whole-prompt
+    # write, unchanged.
+    table = kv_cache.page_table if fill_page_table is None else fill_page_table
+    batch_idx = user_id if fill_page_table is None else 0
+    ttnn.experimental.paged_fill_cache(kv_cache.k, match_cache_dtype(kv_cache.k, k), table, batch_idx=batch_idx)
+    ttnn.experimental.paged_fill_cache(kv_cache.v, match_cache_dtype(kv_cache.v, v), table, batch_idx=batch_idx)
+
+
+def attention_decode(
+    x: ttnn.Tensor,
+    weights: AttentionWeights,
+    config: AttentionConfig,
+    cos_cache: ttnn.Tensor,
+    sin_cache: ttnn.Tensor,
+    kv_cache: KVCache,
+    current_pos: ttnn.Tensor,
+    token_index: int,
+    compute_kernel_config=None,
+) -> ttnn.Tensor:
+    """Single-token attention against the KV cache.
+
+    ``x`` is ``[1, 1, batch, hidden]`` and the return matches. ``current_pos``
+    is an int32 tensor of shape ``[batch]`` holding each user's write position;
+    ``token_index`` is the same value as a Python int, needed because the
+    rotary op takes a scalar rather than a tensor.
+    """
+    k_cache, v_cache, page_table = kv_cache.k, kv_cache.v, kv_cache.page_table
+
+    xqkv = ttnn.linear(
+        x,
+        weights.wqkv,
+        dtype=ttnn.bfloat16,
+        memory_config=ttnn.DRAM_MEMORY_CONFIG,
+        compute_kernel_config=compute_kernel_config,
+    )
+
+    # Blackhole: nlp_create_qkv_heads_decode's interleaved reader zeroes
+    # odd-indexed Q rows when the fused input sits in DRAM, due to a NoC
+    # DRAM-read alignment violation (tt-metal #16667). Staging through L1 takes
+    # a different code path and is unaffected. No-op on Wormhole.
+    if xqkv.memory_config().buffer_type == ttnn.BufferType.DRAM:
+        xqkv = ttnn.to_memory_config(xqkv, ttnn.L1_MEMORY_CONFIG)
+
+    q, k, v = ttnn.experimental.nlp_create_qkv_heads_decode(
+        xqkv,
+        num_heads=config.num_attention_heads,
+        num_kv_heads=config.num_key_value_heads,
+        memory_config=ttnn.L1_HEIGHT_SHARDED_MEMORY_CONFIG,
+    )
+    ttnn.deallocate(xqkv)
+
+    # The decode ops disagree about layout: rms_norm wants interleaved DRAM,
+    # while paged_update_cache requires a *sharded* update tensor
+    # (paged_update_cache_device_operation.cpp:255). So remember the sharded
+    # config the split produced and restore it once the norm and RoPE are done.
+    kv_sharded_mem = k.memory_config()
+
+    q = _per_head_rms_norm(ttnn.to_memory_config(q, ttnn.DRAM_MEMORY_CONFIG), weights.q_norm, config.rms_norm_eps)
+    k = _per_head_rms_norm(ttnn.to_memory_config(k, ttnn.DRAM_MEMORY_CONFIG), weights.k_norm, config.rms_norm_eps)
+
+    q = _apply_rope(q, cos_cache, sin_cache, token_index=token_index)
+    k = _apply_rope(k, cos_cache, sin_cache, token_index=token_index)
+
+    k = ttnn.to_memory_config(k, kv_sharded_mem)  # v never left the sharded layout
+    # page_table=None is the contiguous path; the same op serves both.
+    # NOT cast to the cache dtype -- unlike the fill writers above.
+    # ``paged_update_cache`` requires a FLOAT32/BFLOAT16 update and converts into
+    # the cache itself; handing it a block-float input is rejected outright
+    # (``paged_update_cache_device_operation.cpp:296``). Measured both ways in
+    # ``doc/datatype_sweep/probes/kv_bfp8_diagnosis.json``. See
+    # :func:`match_cache_dtype`.
+    ttnn.experimental.paged_update_cache(k_cache, k, update_idxs_tensor=current_pos, page_table=page_table)
+    ttnn.experimental.paged_update_cache(v_cache, v, update_idxs_tensor=current_pos, page_table=page_table)
+    ttnn.deallocate(k)
+    ttnn.deallocate(v)
+
+    if kv_cache.is_paged:
+        attn = ttnn.transformer.paged_scaled_dot_product_attention_decode(
+            q,
+            k_cache,
+            v_cache,
+            page_table_tensor=page_table,
+            cur_pos_tensor=current_pos,
+            scale=config.head_dim**-0.5,
+            memory_config=ttnn.DRAM_MEMORY_CONFIG,
+        )
+    else:
+        attn = ttnn.transformer.scaled_dot_product_attention_decode(
+            q,
+            k_cache,
+            v_cache,
+            cur_pos_tensor=current_pos,
+            scale=config.head_dim**-0.5,
+            memory_config=ttnn.DRAM_MEMORY_CONFIG,
+        )
+    ttnn.deallocate(q)
+
+    attn = _concat_heads_decode(attn, config)
+    out = ttnn.linear(
+        attn,
+        weights.wo,
+        dtype=ttnn.bfloat16,
+        memory_config=ttnn.DRAM_MEMORY_CONFIG,
+        compute_kernel_config=compute_kernel_config,
+    )
+    ttnn.deallocate(attn)
+    return out
+
+
+def attention_prefill(
+    x: ttnn.Tensor,
+    weights: AttentionWeights,
+    config: AttentionConfig,
+    cos_cache: ttnn.Tensor,
+    sin_cache: ttnn.Tensor,
+    kv_cache: KVCache | None = None,
+    user_id: int = 0,
+    compute_kernel_config=None,
+    sdpa_program_config=None,
+    activation_dtype=ttnn.bfloat16,
+    start_pos: int = 0,
+    chunk_page_table=None,
+    fill_page_table=None,
+    fill_len: int | None = None,
+) -> ttnn.Tensor:
+    """Causal self-attention over a full sequence. ``x``/return ``[1, 1, S, hidden]``.
+
+    When ``kv_cache`` is given, the post-RoPE K and V for the whole sequence are
+    written into it so a decode pass can continue from position S.
+
+    ``sdpa_program_config`` is passed straight through to the prefill SDPA op and
+    defaults to ``None`` -- the op default, which is what **every** caller
+    currently uses, including the multichip path, and what every number in this
+    file was measured at. It is a seam, exactly like
+    ``attention_decode_optimized``'s: stage 06 used it to build and measure a
+    length-dependent chunking worth 6.3-6.8x on this op at S >= 4096, and then
+    **did not adopt it** because it costs a top-1 point on ``run_teacher_forcing``
+    and buys nothing at the 158-token prompt that gate uses. See
+    ``multichip_decoder._sdpa_prefill_program_config`` for the numbers and for
+    what it would take to adopt.
+
+    ``activation_dtype`` is the dtype the two projections emit. It defaults to
+    ``ttnn.bfloat16``, which is the literal that used to be written here, so
+    every existing caller is unchanged; the multichip prefill layer passes
+    ``precision.activation_dtype`` so that field reaches prefill attention and
+    not only decode.
+
+    ``fill_len`` is how many of ``S`` rows are REAL tokens, and exists for
+    bucketed prefill: the caller pads the prompt up to a bucket so the matmuls
+    and SDPA see one of a handful of shapes, but the cache write must stay at
+    the true length. vLLM allocates ``ceil(real/block_size)`` blocks and not one
+    more, so writing the padded length would run off the end of the user's page
+    table and into another request's pages. ``None`` means "all of ``S`` is
+    real" -- the unbucketed path, byte-identical to before this seam existed.
+    """
+    xqkv = ttnn.linear(
+        x,
+        weights.wqkv,
+        dtype=activation_dtype,
+        memory_config=ttnn.DRAM_MEMORY_CONFIG,
+        compute_kernel_config=compute_kernel_config,
+    )
+
+    q, k, v = ttnn.experimental.nlp_create_qkv_heads(
+        xqkv,
+        num_heads=config.num_attention_heads,
+        num_kv_heads=config.num_key_value_heads,
+        transpose_k_heads=False,  # SDPA wants [.., S, head_dim], not a pre-transposed K
+        memory_config=ttnn.DRAM_MEMORY_CONFIG,
+    )
+    ttnn.deallocate(xqkv)
+
+    # QK-norm, then RoPE -- this order is what Qwen3MoeAttention.forward does.
+    q = _per_head_rms_norm(q, weights.q_norm, config.rms_norm_eps)
+    k = _per_head_rms_norm(k, weights.k_norm, config.rms_norm_eps)
+
+    q = _apply_rope(q, cos_cache, sin_cache)
+    k = _apply_rope(k, cos_cache, sin_cache)
+
+    # Seed the cache with the prompt's post-RoPE K/V so decode can continue.
+    if kv_cache is not None:
+        if fill_len is not None and int(fill_len) < int(k.shape[2]):
+            # Bucketed prefill: drop the padding rows before the write. SDPA
+            # above already ran at the padded length -- which is the point, it
+            # is the shape-hungry op -- and causality means the padding rows
+            # could not have influenced any real row's output.
+            real = int(fill_len)
+            k_fill = ttnn.slice(k, [0, 0, 0, 0], [k.shape[0], k.shape[1], real, k.shape[3]])
+            v_fill = ttnn.slice(v, [0, 0, 0, 0], [v.shape[0], v.shape[1], real, v.shape[3]])
+            _fill_cache(kv_cache, k_fill, v_fill, user_id=user_id, fill_page_table=fill_page_table)
+            ttnn.deallocate(k_fill)
+            ttnn.deallocate(v_fill)
+        else:
+            _fill_cache(kv_cache, k, v, user_id=user_id, fill_page_table=fill_page_table)
+
+    # GQA is handled inside SDPA: it broadcasts the 4 KV heads across 32 Q heads.
+    # Default scale is head_dim ** -0.5, which is what Qwen3 uses.
+    if start_pos > 0:
+        # Split prefill: this chunk's Q attends to the WHOLE cached prefix, which
+        # lives in the paged cache -- so the read goes through the paged kernel
+        # with the user's full page-table row, not the local k/v above.
+        chunk = sdpa_chunk_size(start_pos)
+        PREFILL_ATTENTION_BRANCHES["chunked"] += 1
+        PREFILL_ATTENTION_BRANCHES["chunk_sizes"].append(chunk)
+        prog = ttnn.SDPAProgramConfig(
+            compute_with_storage_grid_size=q.device().compute_with_storage_grid_size(),
+            q_chunk_size=chunk,
+            k_chunk_size=chunk,
+            exp_approx_mode=True,
+        )
+        attn = ttnn.transformer.chunked_scaled_dot_product_attention(
+            q,
+            kv_cache.k,
+            kv_cache.v,
+            chunk_page_table,
+            chunk_start_idx=start_pos,
+            program_config=prog,
+            compute_kernel_config=compute_kernel_config,
+        )
+    else:
+        PREFILL_ATTENTION_BRANCHES["standard"] += 1
+        attn = ttnn.transformer.scaled_dot_product_attention(
+            q, k, v, is_causal=True, program_config=sdpa_program_config
+        )
+    for t in (q, k, v):
+        ttnn.deallocate(t)
+
+    attn = ttnn.experimental.nlp_concat_heads(attn, memory_config=ttnn.DRAM_MEMORY_CONFIG)
+    out = ttnn.linear(
+        attn,
+        weights.wo,
+        dtype=activation_dtype,
+        memory_config=ttnn.DRAM_MEMORY_CONFIG,
+        compute_kernel_config=compute_kernel_config,
+    )
+    ttnn.deallocate(attn)
+    return out
+
+
+@dataclass(frozen=True)
+class MoEConfig:
+    hidden_size: int
+    num_experts: int
+    num_experts_per_tok: int
+    moe_intermediate_size: int
+    norm_topk_prob: bool
+
+    @classmethod
+    def from_hf(cls, config) -> "MoEConfig":
+        assert config.hidden_act == "silu", f"expected silu, got {config.hidden_act}"
+        assert config.decoder_sparse_step == 1 and not config.mlp_only_layers, "every layer is expected to be MoE"
+        return cls(
+            hidden_size=config.hidden_size,
+            num_experts=config.num_experts,
+            num_experts_per_tok=config.num_experts_per_tok,
+            moe_intermediate_size=config.moe_intermediate_size,
+            norm_topk_prob=config.norm_topk_prob,
+        )
+
+
+def router_forward(x: ttnn.Tensor, w_router: ttnn.Tensor, config: MoEConfig) -> ttnn.Tensor:
+    """Dense routing weights ``[1, 1, S, num_experts]``: top-k weights, zeros elsewhere.
+
+    Returning a *dense* tensor rather than (values, indices) keeps everything on
+    device and trace-compatible -- the experts consume it as a sparsity pattern
+    and as the post-hoc scaling factor.
+
+    Selection happens on the **raw logits**, not on softmax probabilities, and
+    the softmax is taken over the 8 survivors only. That is algebraically the
+    same answer HF computes, and numerically a far better one.
+
+    Why it is the same answer. Softmax is monotonic, so the top-k of the
+    probabilities is the top-k of the logits -- the 128-wide softmax cannot
+    change *who* wins. And with ``norm_topk_prob`` the shared denominator
+    cancels out of the renormalisation::
+
+        w_i = [exp(x_i)/Z] / sum_{j in top8}[exp(x_j)/Z]
+            =  exp(x_i)     / sum_{j in top8} exp(x_j)
+
+    so ``Z``, the sum over all 128 experts, is never needed. (This cancellation
+    is what makes the rewrite legal; it does **not** hold if
+    ``norm_topk_prob`` is False, hence the assert.)
+
+    Why it is a better answer. Measured on this checkpoint with a 128-token
+    activation sample:
+
+        ttnn.topk on fp32 logits ............ 0/128 tokens misrouted (exact)
+        ttnn.topk on fp32 softmax probs ..... 0/128 (topk itself is fine)
+        ttnn.softmax fp32 vs torch fp32 ..... max abs error 3.3e-4
+
+    The last line is the problem. Softmax over 128 experts leaves the 8th-place
+    probability near 0.008 with only ~1.4e-5 separating it from the 9th, so a
+    3.3e-4 error decides the cut by luck -- and ``ttnn.softmax`` carries that
+    error even when handed an fp32 tensor, so no dtype change rescues it.
+    Routing through the full softmax cost 34/128 misrouted tokens (83/128 with
+    a bf16 softmax, dense PCC 0.88). Comparing logits sidesteps it: the gaps
+    there are ordinary-sized.
+
+    The projection stays bf16 -- it is accurate to PCC 0.9999973 and costs only
+    3/128 on its own -- but it accumulates into fp32 so the comparison is clean.
+    Weights are cast to bf16 for the scatter, which has no fp32 tiled support;
+    that is harmless, since representing a weight to 0.4% is not the same
+    problem as deciding which weights exist.
+    """
+    assert config.norm_topk_prob, (
+        "router selects on raw logits, which relies on the softmax denominator "
+        "cancelling during top-k renormalisation; that only holds when "
+        "norm_topk_prob is True"
+    )
+
+    logits = ttnn.linear(x, w_router, dtype=ttnn.float32, memory_config=ttnn.DRAM_MEMORY_CONFIG)
+
+    top_logits, top_indices = ttnn.topk(logits, k=config.num_experts_per_tok, dim=-1)
+
+    # softmax over the 8 survivors, written out rather than calling
+    # ttnn.softmax so the reduction is the ttnn.sum already exercised by
+    # test_router_weights_sum_to_one. Subtracting the max is for exp() range
+    # only; any shared shift cancels in the division.
+    top_max = ttnn.max(top_logits, dim=-1, keepdim=True)
+    exp_logits = ttnn.exp(ttnn.sub(top_logits, top_max))
+    total = ttnn.sum(exp_logits, dim=-1, keepdim=True)
+    top_values = ttnn.div(exp_logits, total)
+
+    zeros = ttnn.typecast(ttnn.zeros_like(logits), ttnn.bfloat16)
+    dense = ttnn.scatter(
+        zeros,
+        dim=-1,
+        index=top_indices,
+        src=ttnn.typecast(top_values, ttnn.bfloat16),
+    )
+    for t in (logits, top_logits, top_indices, top_max, exp_logits, total, top_values):
+        ttnn.deallocate(t)
+    return dense
+
+
+def upload_router_weight(router: torch.Tensor, device) -> ttnn.Tensor:
+    """``[num_experts, hidden]`` checkpoint tensor -> ``[1, 1, hidden, num_experts]``."""
+    return ttnn.from_torch(
+        router.T.contiguous().reshape(1, 1, router.shape[1], router.shape[0]).float(),
+        dtype=ttnn.bfloat16,
+        layout=ttnn.TILE_LAYOUT,
+        device=device,
+        memory_config=ttnn.DRAM_MEMORY_CONFIG,
+    )
+
+
+@dataclass
+class ExpertWeights:
+    gate_proj: ttnn.Tensor  # [1, num_experts, hidden, intermediate]
+    up_proj: ttnn.Tensor  # [1, num_experts, hidden, intermediate]
+    down_proj: ttnn.Tensor  # [1, num_experts, intermediate, hidden]
+
+
+def upload_expert_weights(torch_weights: dict[str, torch.Tensor], device, config: MoEConfig) -> ExpertWeights:
+    """Unfuse and transpose the batched expert tensors into sparse_matmul layout.
+
+    ``weight_mapping`` keeps the checkpoint's fused ``[E, 2I, H]`` form because
+    that is what the HF module holds and what the reference test validates
+    against. ``sparse_matmul`` wants them separate and transposed, so the split
+    happens here rather than at conversion time -- the fusion order stays
+    guarded by ``test_moe_matches_unfused_reimplementation``.
+    """
+    fused = torch_weights["experts_gate_up"]  # [E, 2I, H], gate first
+    inter = config.moe_intermediate_size
+
+    # bf16 while correctness is being established, so a PCC miss means a bug
+    # rather than quantisation. bfloat8_b is the intended production dtype
+    # (gemma4 and gpt_oss both ship it) and is a later, measurable step -- one
+    # layer's experts are ~604M params, 1.2GB at bf16, which fits comfortably.
+    def up(t: torch.Tensor) -> ttnn.Tensor:
+        return ttnn.from_torch(
+            t.contiguous().float(),
+            dtype=ttnn.bfloat16,
+            layout=ttnn.TILE_LAYOUT,
+            device=device,
+            memory_config=ttnn.DRAM_MEMORY_CONFIG,
+        )
+
+    return ExpertWeights(
+        gate_proj=up(fused[:, :inter, :].transpose(-2, -1).unsqueeze(0)),
+        up_proj=up(fused[:, inter:, :].transpose(-2, -1).unsqueeze(0)),
+        # checkpoint down_proj is [E, H, I]; sparse_matmul wants [1, E, I, H]
+        down_proj=up(torch_weights["experts_down"].transpose(-2, -1).unsqueeze(0)),
+    )
+
+
+def build_expert_sparsity(device, num_experts: int) -> ttnn.Tensor:
+    """All-ones sparsity mask ``[1, 1, 1, E]``, ROW_MAJOR.
+
+    Every expert is computed for every token group; the routing weights zero
+    out the inactive ones after the down projection. That is the established
+    gpt_oss/gemma4 pattern -- it trades compute for a static, trace-friendly
+    shape rather than gathering per-token expert assignments.
+    """
+    return ttnn.from_torch(
+        torch.ones(1, 1, 1, num_experts, dtype=torch.bfloat16),
+        layout=ttnn.ROW_MAJOR_LAYOUT,
+        dtype=ttnn.bfloat16,
+        device=device,
+    )
+
+
+def _expert_compute_kernel_config(device):
+    """HiFi4, and ``fp32_dest_acc_en`` deliberately OFF.
+
+    The functional path holds expert weights in **bf16** (see
+    ``upload_expert_weights``), so the accumulation over hidden/intermediate has
+    16 bits of mantissa to preserve and LoFi -- the matmul default, which keeps
+    the top 5 -- is not enough. Enabling fp32 dest accumulation looks like the
+    natural next lever but must not be used here: it halves the matmul dest from
+    8 tiles to 4, which corrupts expert output on Blackhole (tt-metal #49068,
+    hit on BH-QB-2). HiFi4 alone provides the accuracy.
+
+    The optimized decoder quantises the experts to bfloat4_b, where there is no
+    longer any low-order mantissa for HiFi4 to resolve, and measured LoFi as
+    both faster and marginally more accurate -- so it overrides this with its
+    own config rather than reusing it. See
+    ``optimized_decoder.EXPERT_MATH_FIDELITY``.
+    """
+    return ttnn.init_device_compute_kernel_config(
+        device.arch(),
+        math_fidelity=ttnn.MathFidelity.HiFi4,
+        math_approx_mode=False,
+        fp32_dest_acc_en=False,
+        packer_l1_acc=False,
+    )
+
+
+def _sparse_matmul_config(m: int, n: int, in0_block_w: int = 1):
+    """Spread the N dimension over the largest usable slice of an 8x8 grid."""
+    n_tiles = math.ceil(n / 32)
+    best_cores, best_cx, best_cy = 1, 1, 1
+    for num_cores in range(1, min(65, n_tiles + 1)):
+        if n_tiles % num_cores:
+            continue
+        for cy in range(1, 9):
+            if num_cores % cy == 0:
+                cx = num_cores // cy
+                if cx <= 8 and num_cores > best_cores:
+                    best_cores, best_cx, best_cy = num_cores, cx, cy
+                    break
+    per_core_n = n_tiles // best_cores
+    return ttnn.MatmulMultiCoreReuseMultiCast1DProgramConfig(
+        compute_with_storage_grid_size=ttnn.CoreCoord(best_cx, best_cy),
+        in0_block_w=in0_block_w,
+        out_subblock_h=1,
+        out_subblock_w=1,
+        out_block_h=1,
+        out_block_w=per_core_n,
+        per_core_M=max(32, m) // 32,
+        per_core_N=per_core_n,
+        fuse_batch=False,
+        fused_activation=None,
+        mcast_in0=True,
+    )
+
+
+# sparse_matmul folds the group dimension (chunk_len / 32) into M, which grows
+# num_blocks_y and can overflow the core grid. Chunking at exactly one tile
+# keeps group_size == 1 so all blocking comes from N.
+EXPERT_CHUNK_SIZE = 32
+
+
+def _experts_chunk(
+    hidden: ttnn.Tensor,
+    routing: ttnn.Tensor,
+    weights: ExpertWeights,
+    config: MoEConfig,
+    sparsity_base: ttnn.Tensor,
+) -> ttnn.Tensor:
+    """One 32-token chunk through all experts. ``hidden`` ``[1, 1, 32, H]``."""
+    chunk_len = hidden.shape[2]
+    n_experts = config.num_experts
+    hidden_size = config.hidden_size
+    group_size = chunk_len // EXPERT_CHUNK_SIZE
+
+    device = hidden.device()
+    compute_config = _expert_compute_kernel_config(device)
+    output_tile = ttnn.Tile([32, 32])
+    gate_up_config = _sparse_matmul_config(EXPERT_CHUNK_SIZE, config.moe_intermediate_size)
+    down_config = _sparse_matmul_config(EXPERT_CHUNK_SIZE, hidden_size)
+
+    hidden_grouped = ttnn.reshape(hidden, (1, group_size, EXPERT_CHUNK_SIZE, hidden_size))
+    sparsity = ttnn.repeat(sparsity_base, (1, 1, group_size, 1))
+    nnz = n_experts * group_size
+
+    def project(weight):
+        out = ttnn.sparse_matmul(
+            hidden_grouped,
+            weight,
+            sparsity=sparsity,
+            nnz=nnz,
+            memory_config=ttnn.DRAM_MEMORY_CONFIG,
+            output_tile=output_tile,
+            program_config=gate_up_config,
+            compute_kernel_config=compute_config,
+            dtype=ttnn.bfloat16,
+        )
+        inter = out.shape[-1]
+        return ttnn.reshape(ttnn.transpose(out, 1, 3), (1, n_experts, chunk_len, inter)), inter
+
+    gate, inter = project(weights.gate_proj)
+    up, _ = project(weights.up_proj)
+    ttnn.deallocate(hidden_grouped)
+
+    # SwiGLU -- Qwen3 is hidden_act="silu". gemma4, which this follows, uses
+    # GeGLU; swapping the activation runs fine and returns wrong numbers.
+    down_input = ttnn.reshape(ttnn.mul(ttnn.silu(gate), up), (1, n_experts, chunk_len, inter))
+    ttnn.deallocate(gate)
+    ttnn.deallocate(up)
+
+    down = ttnn.sparse_matmul(
+        down_input,
+        weights.down_proj,
+        sparsity=sparsity_base,
+        nnz=n_experts,
+        memory_config=ttnn.DRAM_MEMORY_CONFIG,
+        output_tile=output_tile,
+        program_config=down_config,
+        is_input_a_sparse=True,
+        compute_kernel_config=compute_config,
+        dtype=ttnn.bfloat16,
+    )
+    ttnn.deallocate(down_input)
+
+    # Scale each expert's contribution by its routing weight (zero for the 120
+    # experts this token did not select), then sum over the expert dimension.
+    states = ttnn.reshape(down, (1, n_experts, chunk_len, hidden_size))
+    states = ttnn.mul(states, ttnn.permute(routing, (0, 3, 2, 1)))  # [1, E, S, 1]
+    states = ttnn.unsqueeze_to_4D(ttnn.experimental.fast_reduce_nc(states, dims=[1]))
+    return ttnn.reshape(states, (1, 1, chunk_len, hidden_size))
+
+
+def moe_prefill(
+    x: ttnn.Tensor,
+    routing: ttnn.Tensor,
+    weights: ExpertWeights,
+    config: MoEConfig,
+    sparsity_base: ttnn.Tensor,
+) -> ttnn.Tensor:
+    """Full MoE expert pass over a sequence. ``x`` ``[1, 1, S, H]``, any S.
+
+    ``sparse_matmul`` works a tile at a time, so a sequence that is not a
+    multiple of 32 is zero-padded up to one and the extra rows are dropped from
+    the result. Padded rows carry an all-zero routing vector, so every expert's
+    contribution to them is scaled by zero -- they cost a little compute and
+    cannot perturb the real tokens.
+    """
+    seq_len = x.shape[2]
+    padded_len = math.ceil(seq_len / EXPERT_CHUNK_SIZE) * EXPERT_CHUNK_SIZE
+
+    if padded_len != seq_len:
+        pad = [(0, 0), (0, 0), (0, padded_len - seq_len), (0, 0)]
+        x = ttnn.pad(x, pad, value=0.0)
+        routing = ttnn.pad(routing, pad, value=0.0)
+
+    outputs = []
+    for start in range(0, padded_len, EXPERT_CHUNK_SIZE):
+        end = start + EXPERT_CHUNK_SIZE
+        outputs.append(
+            _experts_chunk(
+                ttnn.slice(x, [0, 0, start, 0], [1, 1, end, config.hidden_size]),
+                ttnn.slice(routing, [0, 0, start, 0], [1, 1, end, config.num_experts]),
+                weights,
+                config,
+                sparsity_base,
+            )
+        )
+    out = outputs[0] if len(outputs) == 1 else ttnn.concat(outputs, dim=2)
+    if padded_len != seq_len:
+        out = ttnn.slice(out, [0, 0, 0, 0], [1, 1, seq_len, config.hidden_size])
+    return out
+
+
+def moe_decode(
+    x: ttnn.Tensor,
+    routing: ttnn.Tensor,
+    weights: ExpertWeights,
+    config: MoEConfig,
+) -> ttnn.Tensor:
+    """MoE for a single token per user. ``x`` ``[1, 1, batch, H]``.
+
+    Unlike the prefill path this uses **real** sparsity: the dense routing
+    tensor is handed to ``sparse_matmul`` directly, so only the selected experts
+    are computed rather than all 128. Prefill cannot do that -- across 32 tokens
+    the union of selected experts approaches the full set, so an all-ones mask
+    with a post-hoc mask-out is both simpler and no more expensive. In decode
+    the working set really is ``batch * top_k``, and skipping the other 120
+    experts per token is the entire reason decode is affordable.
+    """
+    batch = x.shape[2]
+    n_experts = config.num_experts
+    hidden_size = config.hidden_size
+    # True non-zero count of the sparsity tensor. gemma4 passes top_k because
+    # its decode is single-user; with several users each contributing top_k
+    # entries the kernel needs the full count or it under-sizes the work.
+    nnz = config.num_experts_per_tok * batch
+
+    sparsity = ttnn.to_layout(routing, ttnn.ROW_MAJOR_LAYOUT)
+    output_tile = ttnn.Tile([32, 32])
+    compute_config = _expert_compute_kernel_config(x.device())
+    gate_up_config = _sparse_matmul_config(batch, config.moe_intermediate_size)
+    down_config = _sparse_matmul_config(batch, hidden_size)
+
+    def project(weight):
+        out = ttnn.sparse_matmul(
+            x,
+            weight,
+            sparsity=sparsity,
+            nnz=nnz,
+            memory_config=ttnn.L1_MEMORY_CONFIG,
+            output_tile=output_tile,
+            program_config=gate_up_config,
+            compute_kernel_config=compute_config,
+            dtype=ttnn.bfloat16,
+        )
+        inter = out.shape[-1]
+        out = ttnn.transpose(ttnn.reshape(out, (batch, n_experts, 1, inter)), 1, 2)
+        return ttnn.reshape(out, (batch, n_experts, inter)), inter
+
+    gate, inter = project(weights.gate_proj)
+    up, _ = project(weights.up_proj)
+
+    down_input = ttnn.mul(ttnn.silu(gate), up)  # SwiGLU, not GeGLU
+    ttnn.deallocate(gate)
+    ttnn.deallocate(up)
+    down_input = ttnn.reshape(ttnn.transpose(down_input, 1, 0), (1, n_experts, batch, inter))
+
+    down = ttnn.sparse_matmul(
+        down_input,
+        weights.down_proj,
+        sparsity=sparsity,
+        nnz=nnz,
+        memory_config=ttnn.L1_MEMORY_CONFIG,
+        output_tile=output_tile,
+        program_config=down_config,
+        is_input_a_sparse=True,
+        compute_kernel_config=compute_config,
+        dtype=ttnn.bfloat16,
+    )
+    ttnn.deallocate(down_input)
+
+    states = ttnn.reshape(ttnn.permute(down, (0, 2, 1, 3)), (batch, n_experts, hidden_size))
+    states = ttnn.mul(states, ttnn.reshape(routing, (batch, n_experts, 1)))
+    states = ttnn.unsqueeze_to_4D(ttnn.sum(states, dim=1))
+    return ttnn.reshape(states, (1, 1, batch, hidden_size), (1, 1, max(32, batch), hidden_size))
+
+
+@dataclass(frozen=True)
+class DecoderLayerConfig:
+    attention: AttentionConfig
+    moe: MoEConfig
+    rms_norm_eps: float
+
+    @classmethod
+    def from_hf(cls, config) -> "DecoderLayerConfig":
+        return cls(
+            attention=AttentionConfig.from_hf(config),
+            moe=MoEConfig.from_hf(config),
+            rms_norm_eps=config.rms_norm_eps,
+        )
+
+
+@dataclass
+class DecoderLayerWeights:
+    input_layernorm: ttnn.Tensor
+    post_attention_layernorm: ttnn.Tensor
+    attention: AttentionWeights
+    router: ttnn.Tensor
+    experts: ExpertWeights
+
+
+def decoder_layer_prefill(
+    x: ttnn.Tensor,
+    weights: DecoderLayerWeights,
+    config: DecoderLayerConfig,
+    cos_cache: ttnn.Tensor,
+    sin_cache: ttnn.Tensor,
+    sparsity: ttnn.Tensor,
+    kv_cache: KVCache | None = None,
+    user_id: int = 0,
+) -> ttnn.Tensor:
+    """One full decoder layer. ``x`` / return ``[1, 1, S, hidden]``.
+
+    Passing ``kv_cache`` seeds it with the prompt's K/V so decode can continue
+    from position S.
+
+    Pre-norm, matching ``Qwen3MoeDecoderLayer.forward``::
+
+        h = x + attn(norm1(x))
+        out = h + moe(norm2(h))
+
+    Note the MoE consumes a single normed tensor for *both* the router and the
+    experts. Some ports (gemma4) thread separate router/expert inputs because
+    their router applies its own normalisation; Qwen3's does not, and feeding
+    the router the un-normed residual instead would change every routing
+    decision.
+    """
+    eps = config.rms_norm_eps
+
+    normed = ttnn.rms_norm(x, weight=weights.input_layernorm, epsilon=eps)
+    attn_out = attention_prefill(normed, weights.attention, config.attention, cos_cache, sin_cache, kv_cache, user_id)
+    ttnn.deallocate(normed)
+    hidden = ttnn.add(x, attn_out)
+    ttnn.deallocate(attn_out)
+
+    normed = ttnn.rms_norm(hidden, weight=weights.post_attention_layernorm, epsilon=eps)
+    routing = router_forward(normed, weights.router, config.moe)
+    moe_out = moe_prefill(normed, routing, weights.experts, config.moe, sparsity)
+    ttnn.deallocate(normed)
+    ttnn.deallocate(routing)
+
+    out = ttnn.add(hidden, moe_out)
+    ttnn.deallocate(hidden)
+    ttnn.deallocate(moe_out)
+    return out
+
+
+def decoder_layer_decode(
+    x: ttnn.Tensor,
+    weights: DecoderLayerWeights,
+    config: DecoderLayerConfig,
+    cos_cache: ttnn.Tensor,
+    sin_cache: ttnn.Tensor,
+    kv_cache: KVCache,
+    current_pos: ttnn.Tensor,
+    token_index: int,
+) -> ttnn.Tensor:
+    """One decoder layer, single token per user. ``x`` / return ``[1, 1, batch, hidden]``.
+
+    Same graph as ``decoder_layer_prefill`` -- only the attention and expert
+    kernels differ, because decode attends against the KV cache and can exploit
+    real routing sparsity.
+    """
+    eps = config.rms_norm_eps
+
+    normed = ttnn.rms_norm(x, weight=weights.input_layernorm, epsilon=eps)
+    attn_out = attention_decode(
+        normed, weights.attention, config.attention, cos_cache, sin_cache, kv_cache, current_pos, token_index
+    )
+    ttnn.deallocate(normed)
+    hidden = ttnn.add(x, attn_out)
+    ttnn.deallocate(attn_out)
+
+    normed = ttnn.rms_norm(hidden, weight=weights.post_attention_layernorm, epsilon=eps)
+    routing = router_forward(normed, weights.router, config.moe)
+    moe_out = moe_decode(normed, routing, weights.experts, config.moe)
+    ttnn.deallocate(normed)
+    ttnn.deallocate(routing)
+
+    out = ttnn.add(hidden, moe_out)
+    ttnn.deallocate(hidden)
+    ttnn.deallocate(moe_out)
+    return out
+
+
+def upload_layer_weights(torch_weights: dict[str, torch.Tensor], device, config: DecoderLayerConfig):
+    """Everything one decoder layer needs, from ``convert_layer_weights`` output."""
+
+    def norm(t: torch.Tensor) -> ttnn.Tensor:
+        return ttnn.from_torch(
+            t.reshape(1, 1, 1, -1).float(),
+            dtype=ttnn.bfloat16,
+            layout=ttnn.TILE_LAYOUT,
+            device=device,
+            memory_config=ttnn.DRAM_MEMORY_CONFIG,
+        )
+
+    return DecoderLayerWeights(
+        input_layernorm=norm(torch_weights["input_layernorm"]),
+        post_attention_layernorm=norm(torch_weights["post_attention_layernorm"]),
+        attention=upload_attention_weights(torch_weights, device),
+        router=upload_router_weight(torch_weights["router"], device),
+        experts=upload_expert_weights(torch_weights, device, config.moe),
+    )
+
+
+def upload_attention_weights(torch_weights: dict[str, torch.Tensor], device) -> AttentionWeights:
+    """Move the host-side tensors from ``weight_mapping`` onto the device."""
+
+    def up(t: torch.Tensor, pad_to_4d: bool = False) -> ttnn.Tensor:
+        if pad_to_4d:
+            t = t.reshape(1, 1, 1, -1)
+        while t.dim() < 4:
+            t = t.unsqueeze(0)
+        return ttnn.from_torch(
+            t.float(),
+            dtype=ttnn.bfloat16,
+            layout=ttnn.TILE_LAYOUT,
+            device=device,
+            memory_config=ttnn.DRAM_MEMORY_CONFIG,
+        )
+
+    return AttentionWeights(
+        wqkv=up(torch_weights["wqkv"]),
+        wo=up(torch_weights["wo"]),
+        q_norm=up(torch_weights["q_norm"], pad_to_4d=True),
+        k_norm=up(torch_weights["k_norm"], pad_to_4d=True),
+    )

code/models/demos/blackhole/qwen3_coder_30b_a3b/tt/generator.py

@@ -0,0 +1,1637 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+"""Metal-readiness generator for the 4-die Qwen3-Coder-30B-A3B-Instruct model.
+
+Two API levels, as the readiness contract requires:
+
+* **low level** -- ``prefill_forward`` / ``decode_forward``. The caller owns the
+  KV cache, the page table, the per-user prompt lengths and the per-user decode
+  positions, and threads them through each call. Mixed-length prompts, fixed
+  request slots and inactive rows (position ``-1``) are all expressible here.
+  This is the surface a serving adapter drives.
+* **high level** -- ``generate``. Owns the cache and page table and loops
+  deterministically over the low-level calls.
+
+The measured token-out path is **entirely on device**: the model trace produces
+sampler-ready per-die logits, a second trace runs the split sampler, the sampled
+token is written straight into the persistent decode token input through
+``tt_out_tok``, and the trace advances both position tensors itself with
+``ttnn.plus_one``. Between two steady-state tokens the host does exactly one
+thing -- replay two traces -- plus whatever readback the caller asked for.
+``sampling_mode="host"`` is the explicit compatibility mode for tests that need
+host sampling and is never used to produce a performance number.
+"""
+
+from __future__ import annotations
+
+import bisect
+import contextlib
+import math
+import os
+from pathlib import Path
+from typing import Any, Optional, Sequence
+
+import torch
+from transformers import AutoTokenizer
+
+import ttnn
+from models.common.readiness_check.contract import Generator, NextInputFn
+from models.demos.blackhole.qwen3_coder_30b_a3b.tt.model import (
+    HF_MODEL_ID,
+    HF_REVISION,
+    MAX_CONTEXT,
+    NUM_LAYERS,
+    Qwen3CoderModel,
+)
+
+#: ``ttnn.sampling``/``nlp_create_qkv_heads_decode`` both work in 32-slot units.
+SAMPLING_SLOTS = 32
+
+
+#: Prefill bucket ladders, keyed by the value of ``QWEN3_PREFILL_BUCKETS``.
+#:
+#: A prefill program is compiled per *exact* sequence length: the model
+#: deliberately runs each user at its own logical length so nothing needs a
+#: mask (see :meth:`Generator.prefill_forward`). That is the right trade for
+#: throughput and the wrong one for first-touch latency, because the shape
+#: space is then 1..max_cache_len and a warmup cannot enumerate it -- every
+#: new prompt length pays a fresh ~11 s compile.
+#:
+#: Bucketing collapses that space: the prompt is zero-padded up to the next
+#: rung, the extra rows are dropped from the logits (they are past the selected
+#: row) and from the cache write (``fill_len``), and causality means they could
+#: not have influenced any real row. This is exactly what
+#: ``models/tt_transformers`` does with ``get_padded_prefill_len`` -- powers of
+#: two, warmed by ``warmup_model_prefill`` -- and the ladder below is that idea
+#: with a rung added between each power so the padding tax is halved.
+#:
+#: The cost is wasted compute on the padding: at the measured 0.96 ms/token, a
+#: worst-case ``pow2`` prompt pays ~100% of its prefill again, ``pow2_half``
+#: ~50%, ``1k`` at most ~1 s. The benefit is that the whole ladder is finite
+#: and therefore warmable.
+DEFAULT_PREFILL_BUCKETS = "proportional"
+
+#: The ``proportional`` ladder's gap, as a fraction of the rung it follows --
+#: and therefore the worst-case padding tax at any length. An eighth was chosen
+#: to sit well under the ~24% measured on the ``pow2_half`` ladder while
+#: keeping the rung count near 50, which is small enough that the persistent
+#: kernel cache saturates over a real session.
+PROPORTIONAL_LADDER_STEP = 0.125
+
+
+def prefill_bucket_ladder(name: str, max_len: int) -> tuple[int, ...]:
+    """Every prefill length the model will run under ladder ``name``.
+
+    Returned ascending, and the last rung is always ``max_len`` so that every
+    admissible prompt maps onto some rung. ``exact`` returns ``()`` -- the
+    unbucketed behaviour this model shipped with.
+    """
+    max_len = int(max_len)
+    if max_len < 1:
+        raise ValueError("max_len must be positive")
+    if name == "exact":
+        return ()
+    # Below 128 the projections are tile-bound anyway, so one small rung buys
+    # the whole short-prompt range; 512 keeps a 200-token prompt off the 1024
+    # rung, which is the common case for a chat turn. ``proportional`` derives
+    # its own bottom rungs and starts from 128 alone.
+    lengths = [128] if name == "proportional" else [128, 512]
+    if name == "pow2":
+        step = 1024
+        while step < max_len:
+            lengths.append(step)
+            step *= 2
+    elif name == "pow2_half":
+        step = 1024
+        while step < max_len:
+            lengths.append(step)
+            lengths.append(step + step // 2)
+            step *= 2
+    elif name == "1k":
+        step = 1024
+        while step < max_len:
+            lengths.append(step)
+            step += 1024
+    elif name == "proportional":
+        # Gaps that grow with the rung, so the padding tax is bounded as a
+        # FRACTION of the prefill rather than as a token count. That is the
+        # invariant a user actually feels, and neither fixed-step nor doubling
+        # ladders have it: ``1k`` wastes 0.9 s on a 200-token prompt whose real
+        # prefill is 0.18 s, while ``pow2`` charges a prompt just past 131072
+        # about 111 s of padding. This caps the overhead at
+        # ``PROPORTIONAL_LADDER_STEP`` of the prefill at *every* length, for 52
+        # rungs -- few enough that the persistent kernel cache saturates over a
+        # session, and the low rungs (the ones a warmup can afford) are the
+        # ones a coding client spends most of its prompts on.
+        #
+        # The 128 floor keeps the count finite near the bottom, where an eighth
+        # of the length is a handful of tokens; it also means the shortest
+        # prompts pay at most 127 padded tokens, ~0.11 s.
+        step = 128
+        while step < max_len:
+            lengths.append(step)
+            grow = math.ceil(step * PROPORTIONAL_LADDER_STEP / 128) * 128
+            step += max(128, grow)
+    else:
+        raise ValueError("QWEN3_PREFILL_BUCKETS must be exact, proportional, pow2, pow2_half or 1k; " f"got {name!r}")
+    lengths.append(max_len)
+    return tuple(sorted({v for v in lengths if v <= max_len}))
+
+
+def _first_device_to_torch(tensor) -> torch.Tensor:
+    shards = ttnn.get_device_tensors(tensor)
+    return ttnn.to_torch(shards[0] if shards else tensor)
+
+
+class Qwen3CoderGenerator(Generator):
+    """Caller-owned cache/page-table state plus traced on-device token feedback."""
+
+    def __init__(self, model: Qwen3CoderModel, tokenizer):
+        self.model = model
+        self.mesh_device = model.mesh_device
+        self.tokenizer = tokenizer
+        self.batch = model.max_batch_size
+        self.page_block_size = model.page_block_size
+        self.pages_per_user = math.ceil(model.max_cache_len / self.page_block_size)
+        self.num_blocks = self.batch * self.pages_per_user
+
+        self._kv_cache: list | None = None
+        self._trace_model_id = None
+        #: Resolved on first use from ``QWEN3_PREFILL_BUCKETS``.
+        self._prefill_buckets: tuple[int, ...] | None = None
+        self._trace_sampling_id = None
+        self._trace_inputs = None
+        self._trace_logits = None
+        self._trace_sampled = None
+        self._trace_kv_cache = None
+        self._trace_page_table_snapshot: torch.Tensor | None = None
+        self._trace_active_batch = None
+        #: Width of the *captured* decode graph. Equal to ``self.batch`` unless
+        #: the caller asked for a narrower one via ``decode_forward(graph_width=)``.
+        self._trace_graph_width = None
+        #: Highest rotary position the *next* trace replay will gather at. The
+        #: trace advances ``rotary_position`` on device with ``ttnn.plus_one``
+        #: and nothing on device clamps it, so this host-side mirror is the only
+        #: thing that can tell a replay it is about to index past the cos/sin
+        #: table. See ``decode_forward``.
+        self._trace_rotary_position: int | None = None
+        self._decode_warm_key = None
+        #: Decode graph keys whose programs are already in the program cache.
+        #: Unlike ``_decode_warm_key`` this **survives a trace release**: the
+        #: eager warm pass exists to get every program compiled before capture,
+        #: and a program stays compiled after the trace that used it is freed.
+        #: Serving releases and re-captures the decode traces on every prefill
+        #: (a new request is admitted while other slots decode), so without this
+        #: each admission would pay a full eager decode forward it does not need.
+        #: The key is ``_decode_graph_key``, which includes ``rope_cache_len``:
+        #: growing the rotary tables changes the graph's shapes, so those
+        #: programs are *not* already compiled and the warm pass must run.
+        self._decode_compiled_keys: set = set()
+        self._sampling_params = None
+        self._sampling_snapshot = None
+        self._sampling_stochastic = False
+
+        #: Sampling-penalty state. ``_penalty_mode`` is a *graph* property (see
+        #: ``_WatcherCleanSampling1D``'s penalty section): 0 means the penalty ops
+        #: are not in the captured decode trace at all, so an unpenalised request
+        #: pays nothing. ``_penalty_host`` are persistent full-vocabulary staging
+        #: buffers; ``_penalty_prev_*`` remember which columns each row last wrote
+        #: so a step resets only those instead of the whole 151936-wide row.
+        self._penalty_mode = 0
+        self._penalty_host = None
+        self._penalty_local_vocab = None
+        self._penalty_prev_add: list = []
+        self._penalty_prev_rep: list = []
+
+        #: Steady-state host-work counters. Everything except ``replays`` and
+        #: ``caller_token_readbacks`` must stay flat while tokens are produced.
+        self.trace_stats = {
+            "captures": 0,
+            "replays": 0,
+            "releases": 0,
+            "decode_warmups": 0,
+            "token_host_copies": 0,
+            "token_device_copies": 0,
+            "position_host_copies": 0,
+            "rotary_position_host_copies": 0,
+            "page_table_host_copies": 0,
+            "sampling_param_host_copies": 0,
+            "penalty_host_copies": 0,
+            "caller_token_readbacks": 0,
+            "explicit_synchronizations": 0,
+            "resets": 0,
+        }
+        self._allocate_persistent_inputs()
+
+    # -- persistent device state ---------------------------------------------
+
+    def _replicated_host_tensor(self, host: torch.Tensor, *, dtype):
+        return ttnn.from_torch(
+            host.contiguous(),
+            device=None,
+            dtype=dtype,
+            layout=ttnn.ROW_MAJOR_LAYOUT,
+            mesh_mapper=ttnn.ReplicateTensorToMesh(self.mesh_device),
+        )
+
+    def _replicated_device_tensor(self, host: torch.Tensor, *, dtype):
+        return ttnn.from_torch(
+            host.contiguous(),
+            device=self.mesh_device,
+            dtype=dtype,
+            layout=ttnn.ROW_MAJOR_LAYOUT,
+            mesh_mapper=ttnn.ReplicateTensorToMesh(self.mesh_device),
+        )
+
+    def _copy_host(self, host: torch.Tensor, device, *, dtype) -> None:
+        ttnn.copy_host_to_device_tensor(self._replicated_host_tensor(host, dtype=dtype), device)
+
+    def _allocate_persistent_inputs(self) -> None:
+        """Allocate every stable decode input **before** any trace is captured."""
+        self._prefill_page_table = self._replicated_device_tensor(
+            torch.full((self.batch, self.pages_per_user), -1, dtype=torch.int32), dtype=ttnn.int32
+        )
+        self._prefill_sampled = self._replicated_device_tensor(
+            torch.zeros((1, 1, 1, SAMPLING_SLOTS), dtype=torch.int32), dtype=ttnn.uint32
+        )
+        self._width_pools = {}
+        self._decode_trace_input_pool = self._decode_input_pool(self.batch)
+
+    def _decode_input_pool(self, width: int) -> tuple:
+        """The four persistent decode inputs, ``width`` rows wide.
+
+        One pool per captured graph width. The **token** tensor is
+        ``[1,1,1,32]`` at every width because that is ``tt_out_tok``'s shape --
+        the sampler always addresses 32 slots and ``embed_decode`` slices the
+        embedding down to ``model.decode_width``. The other three are the only
+        things that bind a request to a row, and they are exactly what
+        compaction permutes.
+        """
+        width = int(width)
+        pool = self._width_pools.get(width)
+        if pool is not None:
+            return pool
+        pool = (
+            # token: [1,1,1,32] uint32, the tensor ``tt_out_tok`` writes into
+            self._replicated_device_tensor(
+                torch.zeros((1, 1, 1, SAMPLING_SLOTS), dtype=torch.int32), dtype=ttnn.uint32
+            ),
+            # current_pos: [width] int32, consumed by paged_update_cache and SDPA
+            self._replicated_device_tensor(torch.full((width,), -1, dtype=torch.int32), dtype=ttnn.int32),
+            # rotary_position: [1, width] uint32, the cos/sin gather index
+            self._replicated_device_tensor(torch.zeros((1, width), dtype=torch.int32), dtype=ttnn.uint32),
+            # page_table: [width, pages_per_user] int32
+            self._replicated_device_tensor(
+                torch.full((width, self.pages_per_user), -1, dtype=torch.int32), dtype=ttnn.int32
+            ),
+        )
+        self._width_pools[width] = pool
+        return pool
+
+    def _ensure_kv_cache(self):
+        if self._kv_cache is None:
+            self._kv_cache = self.model.allocate_kv_cache(num_blocks=self.num_blocks)
+        return self._kv_cache
+
+    def configure_paging(self, *, page_block_size: int, pages_per_user: int, num_blocks: int) -> None:
+        """Adopt a **caller-owned** paging geometry (the vLLM serving mode).
+
+        Standalone mode derives ``page_block_size`` / ``pages_per_user`` /
+        ``num_blocks`` from the model, allocates its own cache and builds its own
+        page tables. Serving inverts that: vLLM picks the cache block size and
+        the block count, and every page table it hands over is
+        ``[batch, max_num_blocks_per_req]`` at *its* width. The two persistent
+        page-table tensors were sized for the standalone geometry in
+        ``_allocate_persistent_inputs``, so adopting vLLM's means reallocating
+        them -- which is only safe before any trace exists, hence the guard.
+
+        ``tt/generator_vllm.py`` calls this from ``allocate_kv_cache``, which the
+        TT plugin invokes once, before warmup and before any forward.
+        """
+        page_block_size = int(page_block_size)
+        pages_per_user = int(pages_per_user)
+        num_blocks = int(num_blocks)
+        if min(page_block_size, pages_per_user, num_blocks) < 1:
+            raise ValueError("paging geometry must be positive")
+        if self._trace_model_id is not None or self._trace_sampling_id is not None:
+            raise RuntimeError("configure_paging must run before any decode trace is captured")
+        if self._kv_cache is not None:
+            raise RuntimeError("configure_paging must run before the generator allocates its own cache")
+        if (page_block_size, pages_per_user, num_blocks) == (
+            self.page_block_size,
+            self.pages_per_user,
+            self.num_blocks,
+        ):
+            return
+        self.page_block_size = page_block_size
+        self.model.page_block_size = page_block_size
+        self.pages_per_user = pages_per_user
+        self.num_blocks = num_blocks
+        self._allocate_persistent_inputs()
+
+    def decode_device_state(self) -> dict[str, torch.Tensor] | None:
+        """The authoritative per-slot decode state that lives **on device**.
+
+        The traced decode path writes the sampled token straight into the
+        persistent token input and advances ``current_pos`` with
+        ``ttnn.plus_one``, so after step *N* the device -- not the host -- holds
+        the token and position step *N+1* must use. A serving scheduler under
+        async scheduling can be a step behind that, and re-installing its host
+        view would re-decode a position or feed a stale token. This exposes the
+        device view (plus the page table the live trace was captured against) so
+        ``tt/generator_vllm.py`` can keep it for slots that are simply
+        continuing and take the host's only for slots that changed hands.
+
+        Returns ``None`` when no trace is live. Costs two small device reads and
+        is called only on scheduler-layout changes, never per token.
+        """
+        if self._trace_model_id is None or self._trace_inputs is None:
+            return None
+        token, current_pos, _rotary, _page_table = self._trace_inputs
+        # The live trace may be **narrower** than the configured slot count, so
+        # everything is reported at the graph's width and ``width`` says what
+        # that is. Row *i* here is graph row *i*, not necessarily vLLM slot *i* --
+        # the caller owns the mapping (``Qwen3CoderForCausalLM._compaction``).
+        width = self._trace_graph_width or self.batch
+        return {
+            "width": width,
+            "tokens": _first_device_to_torch(token).reshape(-1)[:width].to(torch.int64),
+            "positions": _first_device_to_torch(current_pos).reshape(-1)[:width].to(torch.int64),
+            "page_table": (
+                None if self._trace_page_table_snapshot is None else self._trace_page_table_snapshot.clone()
+            ),
+        }
+
+    def read_sampled_tokens(self, sampled, count: int | None = None) -> torch.Tensor:
+        """Host copy of a sampled-token tensor. The only readback on the token path."""
+        tokens = self._sampled_to_torch(sampled)
+        return tokens if count is None else tokens[: int(count)]
+
+    def _synchronize(self) -> None:
+        ttnn.synchronize_device(self.mesh_device)
+        self.trace_stats["explicit_synchronizations"] += 1
+
+    # -- page tables ----------------------------------------------------------
+
+    def _page_table_to_torch(self, page_table) -> torch.Tensor:
+        if isinstance(page_table, torch.Tensor):
+            host = page_table.detach().cpu().to(torch.int32)
+        elif isinstance(page_table, ttnn.Tensor):
+            host = _first_device_to_torch(page_table).to(torch.int32)
+        else:
+            raise TypeError("page_table must be a torch or TTNN tensor")
+        if host.ndim != 2:
+            raise ValueError(f"page_table must be rank two, got {tuple(host.shape)}")
+        return host
+
+    def _normalise_page_table(self, page_table, active_batch: int, width: int | None = None) -> torch.Tensor:
+        """Trim/pad a caller's block table to ``width`` rows x ``pages_per_user``.
+
+        ``width`` defaults to the configured slot count and is the *graph* width
+        -- the number of rows the captured decode trace has. A narrow graph is
+        handed the first ``width`` rows, which is why the caller must compact its
+        live requests into them first.
+        """
+        width = self.batch if width is None else int(width)
+        host = self._page_table_to_torch(page_table)
+        if host.shape[0] < active_batch or host.shape[0] > self.batch:
+            raise ValueError("page table does not match the configured/active batch")
+        if host.shape[1] < self.pages_per_user:
+            host = torch.nn.functional.pad(host, (0, self.pages_per_user - host.shape[1]), value=-1)
+        elif host.shape[1] > self.pages_per_user:
+            host = host[:, : self.pages_per_user]
+        if host.shape[0] < width:
+            host = torch.nn.functional.pad(host, (0, 0, 0, width - host.shape[0]), value=-1)
+        elif host.shape[0] > width:
+            host = host[:width]
+        return host.contiguous()
+
+    def _sdpa_rounded_page_count(self, token_count: int) -> int:
+        """Physical pages the paged decode SDPA kernel actually reads.
+
+        The kernel rounds a short sequence up to a power-of-two tile count and a
+        long one up to a multiple of eight, and it reads the whole rounded
+        window before causal masking. Every rounded tail page therefore needs a
+        valid mapping even though it holds no live token yet -- allocating only
+        ``ceil(len/block)`` pages produces top-k misses that cliff at exactly
+        those boundaries and look like dtype drift.
+        """
+        if token_count < 1 or token_count > self.model.max_cache_len:
+            raise ValueError("SDPA token count is outside the supported context")
+        logical_pages = math.ceil(token_count / self.page_block_size)
+        if logical_pages <= 8:
+            return 1 << (logical_pages - 1).bit_length()
+        return 8 * math.ceil(logical_pages / 8)
+
+    def make_page_table(self, lengths: Sequence[int]) -> torch.Tensor:
+        """A disjoint physical-block assignment covering each user's horizon."""
+        if len(lengths) > self.batch:
+            raise ValueError(f"{len(lengths)} prompts exceed configured batch {self.batch}")
+        table = torch.full((self.batch, self.pages_per_user), -1, dtype=torch.int32)
+        next_block = 0
+        for user, length in enumerate(lengths):
+            blocks = self._sdpa_rounded_page_count(int(length))
+            if blocks > self.pages_per_user or next_block + blocks > self.num_blocks:
+                raise ValueError("paged KV-cache capacity is insufficient for the requested prompts")
+            table[user, :blocks] = torch.arange(next_block, next_block + blocks, dtype=torch.int32)
+            next_block += blocks
+        return table
+
+    def _validate_page_coverage(self, page_table: torch.Tensor, positions: torch.Tensor, active_batch: int) -> None:
+        assigned: set[int] = set()
+        for slot, position in enumerate(positions.reshape(-1).tolist()[:active_batch]):
+            if position < 0:  # inactive row
+                continue
+            logical_pages = math.ceil((int(position) + 1) / self.page_block_size)
+            rounded_pages = self._sdpa_rounded_page_count(int(position) + 1)
+            if rounded_pages > page_table.shape[1]:
+                raise ValueError(f"slot {slot} page table is too narrow for decode position {position}")
+            physical = [int(v) for v in page_table[slot, :rounded_pages].tolist()]
+            if any(v < 0 or v >= self.num_blocks for v in physical):
+                raise ValueError(f"slot {slot} lacks valid physical pages for the rounded SDPA read at {position}")
+            live = physical[:logical_pages]
+            if len(set(live)) != len(live) or assigned.intersection(live):
+                raise ValueError("active page-table rows must map disjoint physical cache pages")
+            assigned.update(live)
+
+    # -- prefill --------------------------------------------------------------
+
+    @property
+    def prefill_buckets(self) -> tuple[int, ...]:
+        """The prefill length ladder in force, ascending. ``()`` means exact."""
+        if self._prefill_buckets is None:
+            name = os.getenv("QWEN3_PREFILL_BUCKETS", DEFAULT_PREFILL_BUCKETS).strip().lower()
+            self._prefill_buckets = prefill_bucket_ladder(name, self.model.max_cache_len)
+        return self._prefill_buckets
+
+    def prefill_padded_len(self, real_len: int, *, start: int = 0) -> int:
+        """Round a suffix length up to its bucket, or return it unchanged.
+
+        Clamped so ``start + result`` never leaves the supported context: near
+        the very top of the window there may be no rung left to round to, and
+        such a prompt simply keeps its exact length and pays a compile. Every
+        other length lands on a rung.
+        """
+        real_len = int(real_len)
+        ladder = self.prefill_buckets
+        if not ladder:
+            return real_len
+        index = bisect.bisect_left(ladder, real_len)
+        if index >= len(ladder):
+            return real_len
+        return max(real_len, min(ladder[index], self.model.max_cache_len - int(start)))
+
+    @staticmethod
+    def _prefill_starts(active_batch: int, start_pos) -> list[int]:
+        """How many tokens of each row are already in the cache. None/0 => all new."""
+        if start_pos is None:
+            return [0] * active_batch
+        if isinstance(start_pos, int):
+            return [int(start_pos)] * active_batch
+        return [int(v) for v in start_pos]
+
+    def _prefill_rope_horizon(self, prompt_lens, start_pos) -> int:
+        """The highest absolute position any row of this prefill will rotate at.
+
+        Bucketing pads past the prompt, so the RoPE tables have to cover the
+        PADDED horizon. Getting this wrong is not a slow path but a crash:
+        ``prefill_forward`` releases the decode traces when the tables grow,
+        precisely because a captured trace holds the old tables' identities,
+        and a growth that happened later -- inside ``prefill_hidden``, under a
+        live trace -- would slip past that guard.
+        """
+        starts = self._prefill_starts(len(prompt_lens), start_pos)
+        return max(
+            start + self.prefill_padded_len(int(length) - start, start=start)
+            for length, start in zip(prompt_lens, starts)
+        )
+
+    def _release_decode_traces_before_allocating(self) -> None:
+        """Prefill is eager and allocates; a live trace makes that unsafe."""
+        if self._trace_model_id is None and self._trace_sampling_id is None:
+            return
+        self._synchronize()
+        self._release_decode_traces()
+
+    def prefill_forward(
+        self,
+        tokens: torch.Tensor,
+        *,
+        page_table,
+        kv_cache: Any,
+        prompt_lens: Sequence[int],
+        return_all_logits: bool = False,
+        sampling_mode: str = "host",
+        preserve_decode_traces: bool = False,
+        start_pos: Sequence[int] | int | None = None,
+        **kwargs: Any,
+    ):
+        """Prefill arbitrary logical lengths, one user at a time into the cache.
+
+        ``tokens`` is ``[active_batch, width]`` and ``prompt_lens`` gives each
+        row's **real** length. Rows may differ in length, and no mask is needed
+        at any length: nothing is padded to a chunk, tile or page boundary, and
+        the returned logits are always sliced back to the logical prompt length.
+
+        Each row is prefilled at its bucket rather than at its exact length --
+        see :func:`prefill_bucket_ladder`, and ``QWEN3_PREFILL_BUCKETS=exact``
+        to turn that off. Bucketing is invisible here: the padding rows are
+        zeros appended *after* the prompt, so causality keeps them out of every
+        real row's attention, ``fill_len`` keeps them out of the KV cache, and
+        the row this method selects is the real last token either way.
+
+        ``preserve_decode_traces`` keeps a captured decode trace alive across
+        this prefill. Standalone callers never need it -- ``generate`` prefills
+        once, before any decode trace exists. **Serving does**: vLLM admits a new
+        request by prefilling it while other slots are mid-decode, and releasing
+        the decode traces there would re-capture them on the very next token,
+        putting a multi-second stall inside the measured inter-token latency of
+        every other in-flight request. It is safe because prefill's allocations
+        never touch the trace region and every tensor a captured trace holds --
+        ``_decode_trace_input_pool``, ``_trace_logits``, ``_trace_sampled`` --
+        is owned by this object and therefore never freed underneath it. The
+        page table is the one shared binding, and this method rebinds the cache
+        back to the live trace's page-table tensor before it returns.
+        """
+        if sampling_mode not in {"host", "device"}:
+            raise ValueError("sampling_mode must be 'host' or 'device'")
+        if sampling_mode == "device" and return_all_logits:
+            raise ValueError("return_all_logits is incompatible with device sampling")
+        if tokens.ndim != 2:
+            raise ValueError(f"tokens must be [batch,seq], got {tuple(tokens.shape)}")
+        active_batch, logical_width = int(tokens.shape[0]), int(tokens.shape[1])
+        if not 1 <= active_batch <= self.batch:
+            raise ValueError(f"active batch must be in [1,{self.batch}]")
+        if len(prompt_lens) != active_batch or any(not 1 <= int(n) <= logical_width for n in prompt_lens):
+            raise ValueError("prompt_lens must contain one valid logical length per input row")
+        if max(prompt_lens) > self.model.max_cache_len:
+            raise ValueError("prompt exceeds the supported context")
+
+        if preserve_decode_traces:
+            if self._trace_model_id is not None:
+                # The replayed trace is asynchronous; prefill is eager. Let the
+                # queue drain before eager work reads or writes the same cache.
+                self._synchronize()
+            if self.model.ensure_rope_capacity(self._prefill_rope_horizon(prompt_lens, start_pos)):
+                # Growing the tables moves them, and a captured trace holds the
+                # old identities. Nothing can preserve a trace across that.
+                self._release_decode_traces()
+        else:
+            self._release_decode_traces_before_allocating()
+            self.model.ensure_rope_capacity(self._prefill_rope_horizon(prompt_lens, start_pos))
+        caches = self._ensure_kv_cache() if kv_cache is None else kv_cache
+        page_host = self._normalise_page_table(page_table, active_batch)
+        self._copy_host(page_host, self._prefill_page_table, dtype=ttnn.int32)
+        self.model.bind_page_table(caches, self._prefill_page_table)
+        try:
+            return self._prefill_body(
+                tokens,
+                caches,
+                active_batch=active_batch,
+                logical_width=logical_width,
+                prompt_lens=prompt_lens,
+                return_all_logits=return_all_logits,
+                sampling_mode=sampling_mode,
+                start_pos=start_pos,
+                page_table=page_host,
+            )
+        finally:
+            if self._trace_inputs is not None:
+                # Hand the cache back to the tensor the live decode trace was
+                # captured against, so the next replay writes through the page
+                # table the scheduler owns rather than the prefill scratch one.
+                self.model.bind_page_table(caches, self._trace_inputs[3])
+
+    def _prefill_body(
+        self,
+        tokens: torch.Tensor,
+        caches,
+        *,
+        active_batch: int,
+        logical_width: int,
+        prompt_lens: Sequence[int],
+        return_all_logits: bool,
+        sampling_mode: str,
+        start_pos: Sequence[int] | int | None = None,
+        page_table=None,
+    ):
+        # ``start_pos`` is how many tokens of each row are ALREADY in the cache.
+        # None/0 is the shipped whole-prompt prefill and takes an identical path.
+        starts = self._prefill_starts(active_batch, start_pos)
+        per_user_logits: list[torch.Tensor] = []
+        selected_rows = []
+        for user in range(active_batch):
+            prompt_len = int(prompt_lens[user])
+            start = starts[user]
+            # Bucketed prefill: pad the suffix up to its rung so the shape-hungry
+            # ops (QKV/wo projections, MoE, SDPA) see one of a handful of
+            # lengths instead of one per prompt. ``fill_len`` keeps the cache
+            # write at the real length, ``select_prefill_rows`` indexes the real
+            # last row, and the padding rows are causally invisible to it.
+            token_host = tokens[user : user + 1, start:prompt_len].to(torch.int32)
+            real_len = int(token_host.shape[1])
+            padded_len = self.prefill_padded_len(real_len, start=start)
+            fill_len = None
+            if padded_len > real_len:
+                token_host = torch.nn.functional.pad(token_host, (0, padded_len - real_len))
+                fill_len = real_len
+
+            chunk_pt = fill_pt = None
+            if start:
+                block = int(caches[0].block_size)
+                if start % block:
+                    raise ValueError(f"start_pos {start} is not a multiple of the block size {block}")
+                if not 0 < start < prompt_len:
+                    raise ValueError(f"start_pos {start} must be inside (0, prompt_len={prompt_len})")
+                if page_table is None:
+                    raise ValueError("a split prefill needs a page table")
+                row = torch.as_tensor(page_table)[user : user + 1].to(torch.int32)
+                # Two different tables, and the difference is the whole trick:
+                #   chunk_pt -- the user's FULL row, so chunked SDPA can read the
+                #               cached prefix from absolute block 0;
+                #   fill_pt  -- a window over the suffix's blocks, because
+                #               paged_fill_cache writes relative to block 0 of the
+                #               table it is handed.
+                chunk_pt = ttnn.from_torch(
+                    row,
+                    device=self.mesh_device,
+                    dtype=ttnn.int32,
+                    layout=ttnn.ROW_MAJOR_LAYOUT,
+                    mesh_mapper=ttnn.ReplicateTensorToMesh(self.mesh_device),
+                )
+                # WIDTH IS DERIVED FROM THE BUCKET, NOT FROM ``start``. The
+                # obvious ``row[:, start // block :]`` is one column narrower
+                # per cached block, so its shape -- and therefore the
+                # ``paged_fill_cache`` program -- varies with the length of the
+                # prefix, which is unbounded. Measured: an unwarmed prefix
+                # offset cost 0.976 s against 0.178 s for a warmed one, on
+                # identical suffix and chunk sizes. Sizing the window by the
+                # suffix's rung instead leaves one program per rung.
+                #
+                # Trailing ``-1`` entries are never read: ``k`` is sliced to
+                # ``fill_len`` above, so the op writes only
+                # ``ceil(fill_len / block)`` blocks from the front of this table.
+                fill_blocks = math.ceil(padded_len / block)
+                fill_window = row[:, start // block : start // block + fill_blocks]
+                if fill_window.shape[1] < fill_blocks:
+                    fill_window = torch.nn.functional.pad(
+                        fill_window, (0, fill_blocks - fill_window.shape[1]), value=-1
+                    )
+                fill_pt = ttnn.from_torch(
+                    fill_window.contiguous(),
+                    device=self.mesh_device,
+                    dtype=ttnn.int32,
+                    layout=ttnn.ROW_MAJOR_LAYOUT,
+                    mesh_mapper=ttnn.ReplicateTensorToMesh(self.mesh_device),
+                )
+            token_device = ttnn.from_torch(
+                token_host,
+                device=self.mesh_device,
+                dtype=ttnn.uint32,
+                layout=ttnn.ROW_MAJOR_LAYOUT,
+                mesh_mapper=ttnn.ReplicateTensorToMesh(self.mesh_device),
+            )
+            hidden = self.model.prefill_hidden(
+                token_device,
+                kv_cache=caches,
+                user_id=user,
+                start_pos=start,
+                chunk_page_table=chunk_pt,
+                fill_page_table=fill_pt,
+                fill_len=fill_len,
+            )
+            ttnn.deallocate(token_device, True)
+            for t in (chunk_pt, fill_pt):
+                if t is not None:
+                    ttnn.deallocate(t, True)
+            if return_all_logits:
+                normed = self.model.prefill_norm(hidden)
+                ttnn.deallocate(hidden, True)
+                local = self.model.local_logits(normed)
+                ttnn.deallocate(normed, True)
+                host = self.model.gather_logits_to_torch(local)[0, 0, : prompt_len - start, :]  # padding rows dropped
+                ttnn.deallocate(local, True)
+                per_user_logits.append(
+                    torch.nn.functional.pad(host, (0, 0, 0, logical_width - (prompt_len - start))).unsqueeze(0)
+                )
+            else:
+                # The last row of THIS chunk: absolute prompt_len-1 is row
+                # prompt_len-1-start within the suffix.
+                selected_rows.append(self.model.select_prefill_rows(hidden, [prompt_len - 1 - start]))
+                ttnn.deallocate(hidden, True)
+
+        if return_all_logits:
+            return torch.cat(per_user_logits, dim=0)[:, :logical_width]
+
+        selected = (
+            selected_rows[0]
+            if len(selected_rows) == 1
+            else ttnn.concat(selected_rows, dim=2, memory_config=ttnn.DRAM_MEMORY_CONFIG)
+        )
+        normed = self.model.prefill_norm(selected)
+        if selected is not selected_rows[0] or len(selected_rows) > 1:
+            for row in selected_rows:
+                ttnn.deallocate(row, True)
+        else:
+            ttnn.deallocate(selected, True)
+        if sampling_mode == "device":
+            padded = self._pad_rows_to_sampling_slots(normed, active_batch)
+            local = self.model.local_logits(padded)
+            ttnn.deallocate(padded, True)
+            with self._penalties_suspended():
+                sampled = self._sample_device(local, tt_out_tok=self._prefill_sampled)
+            ttnn.deallocate(local, True)
+            return sampled
+        local = self.model.local_logits(normed)
+        ttnn.deallocate(normed, True)
+        host = self.model.gather_logits_to_torch(local, valid_rows=active_batch)[0, 0]
+        ttnn.deallocate(local, True)
+        return host.unsqueeze(1)
+
+    def _pad_rows_to_sampling_slots(self, normed, active_batch: int):
+        """``ttnn.sampling`` works in 32 fixed slots; pad the selected rows up."""
+        rows = int(normed.shape[-2])
+        if rows >= SAMPLING_SLOTS:
+            return normed
+        padded = ttnn.pad(
+            normed,
+            [(0, 0), (0, 0), (0, SAMPLING_SLOTS - rows), (0, 0)],
+            value=0.0,
+            memory_config=ttnn.DRAM_MEMORY_CONFIG,
+        )
+        ttnn.deallocate(normed, True)
+        return padded
+
+    # -- sampling parameters --------------------------------------------------
+
+    def _ensure_sampling_params(self):
+        if self._sampling_params is None:
+            self._sampling_params = (
+                self._replicated_device_tensor(torch.ones(SAMPLING_SLOTS, dtype=torch.int32), dtype=ttnn.uint32),
+                self._replicated_device_tensor(torch.zeros(SAMPLING_SLOTS, dtype=torch.bfloat16), dtype=ttnn.bfloat16),
+                self._replicated_device_tensor(torch.ones(SAMPLING_SLOTS, dtype=torch.bfloat16), dtype=ttnn.bfloat16),
+            )
+            self._sampling_snapshot = ((1,) * SAMPLING_SLOTS, (0.0,) * SAMPLING_SLOTS, (1.0,) * SAMPLING_SLOTS)
+        return self._sampling_params
+
+    @staticmethod
+    def _expand(value, *, active_batch: int, inactive, name: str):
+        active = [value] * active_batch if isinstance(value, (int, float)) else list(value)
+        if len(active) != active_batch:
+            raise ValueError(f"{name} must be scalar or contain {active_batch} values")
+        return active + [inactive] * (SAMPLING_SLOTS - active_batch)
+
+    def set_sampling_params(self, *, top_k=1, top_p=0.0, temperature=1.0, active_batch: int = 1) -> None:
+        """Set per-slot ``(k, p, temperature)``. ``k=1`` is exactly greedy."""
+        if not 1 <= active_batch <= self.batch:
+            raise ValueError(f"active_batch must be in [1,{self.batch}]")
+        k = [int(v) for v in self._expand(top_k, active_batch=active_batch, inactive=1, name="top_k")]
+        p = [float(v) for v in self._expand(top_p, active_batch=active_batch, inactive=0.0, name="top_p")]
+        temp = [
+            float(v) for v in self._expand(temperature, active_batch=active_batch, inactive=0.0, name="temperature")
+        ]
+        if any(not 0.0 <= v <= 1.0 for v in p[:active_batch]):
+            raise ValueError("top_p must be in [0,1]")
+        if any(v < 0.0 for v in temp[:active_batch]):
+            raise ValueError("temperature must be non-negative")
+        for slot in range(active_batch):
+            # A serving stack spells greedy as temperature=0 / top_k=0.
+            if temp[slot] == 0.0 and k[slot] == 0:
+                k[slot] = 1
+        if any(not 1 <= v <= SAMPLING_SLOTS for v in k[:active_batch]):
+            raise ValueError("top_k must be in [1,32] (0 accepted only with temperature=0)")
+        device_temp = []
+        for slot, value in enumerate(temp):
+            if value == 0.0:
+                k[slot], p[slot], value = 1, 0.0, 1.0
+            device_temp.append(1.0 / value)
+        stochastic = any(v > 1 for v in k[:active_batch]) or any(v > 0.0 for v in p[:active_batch])
+        if stochastic != self._sampling_stochastic and self._trace_model_id is not None:
+            self._release_decode_traces()
+        self._sampling_stochastic = stochastic
+        params = self._ensure_sampling_params()
+        snapshot = (tuple(k), tuple(p), tuple(device_temp))
+        if snapshot == self._sampling_snapshot:
+            return
+        for host, device, dtype in (
+            (torch.tensor(k, dtype=torch.int32), params[0], ttnn.uint32),
+            (torch.tensor(p, dtype=torch.bfloat16), params[1], ttnn.bfloat16),
+            (torch.tensor(device_temp, dtype=torch.bfloat16), params[2], ttnn.bfloat16),
+        ):
+            self._copy_host(host, device, dtype=dtype)
+            self.trace_stats["sampling_param_host_copies"] += 1
+        self._sampling_snapshot = snapshot
+
+    # -- sampling penalties ---------------------------------------------------
+
+    @staticmethod
+    def _row_token_ids(history, row: int) -> torch.Tensor:
+        """Row ``row`` of a vLLM ``[rows, L]`` history tensor, -1 padding dropped.
+
+        vLLM pads both ``prompt_tokens`` and ``output_tokens`` with **-1**
+        (``input_batch.make_prompt_token_ids_tensor``: "TT device sampling relies
+        on -1 as the padding sentinel"), and pads the *batch* to ``max_num_reqs``
+        with all-(-1) rows. Dropping every negative entry handles both.
+        """
+        if history is None:
+            return torch.empty(0, dtype=torch.int64)
+        tensor = torch.as_tensor(history)
+        if tensor.ndim == 1:
+            tensor = tensor.reshape(1, -1)
+        if row >= tensor.shape[0]:
+            return torch.empty(0, dtype=torch.int64)
+        ids = tensor[row].reshape(-1).to(torch.int64)
+        return ids[ids >= 0]
+
+    def _ensure_penalty_host(self, slots: int, vocab: int) -> dict:
+        """Per-die staging buffers, **already contiguous in the shard layout**.
+
+        Not one ``[1,1,32,151936]`` tensor. Handing a full-width host tensor to
+        ``ttnn.ShardTensorToMesh(dim=-1)`` makes it re-slice a strided view into
+        four contiguous copies on every decode step, and that reshard -- not
+        tilization, not the wire -- was **6.601 ms of a 6.897 ms** upload.
+        Keeping the four ``[1,1,32,37984]`` buffers contiguous from the start and
+        assembling them with ``ttnn.from_host_shards`` costs **2.049 ms**
+        end to end, 3.4x less.
+
+        The trade is that the global -> (die, local) split now happens here, in
+        host Python, instead of being implied by the mesh mapper. That is the
+        one piece of index arithmetic in this feature, so it is checked rather
+        than trusted: ``penalty_shard_boundary_probe.py``'s
+        ``fast_staging_matches_shard_mapper`` leg builds the same operand both
+        ways and requires the two device tensors to be bit-identical, and its
+        cross-die/boundary legs would fail first if the split were wrong.
+        """
+        if self._penalty_host is None:
+            devices, local = self.model.sampler.penalty_shard_geometry()
+            if devices * local != vocab:
+                raise RuntimeError(f"penalty shard geometry {devices}x{local} does not cover {vocab}")
+            self._penalty_local_vocab = local
+            self._penalty_host = {
+                "rep_neg": [torch.ones((1, 1, slots, local), dtype=torch.bfloat16) for _ in range(devices)],
+                "add": [torch.zeros((1, 1, slots, local), dtype=torch.bfloat16) for _ in range(devices)],
+            }
+            self._penalty_prev_add = [None] * slots
+            self._penalty_prev_rep = [None] * slots
+        return self._penalty_host
+
+    def _penalty_split(self, ids: torch.Tensor):
+        """Global token ids -> ``[(die, local_ids_on_that_die, selector), ...]``.
+
+        ``die = t // local_vocab``, ``local = t % local_vocab`` -- the same
+        contiguous ascending decomposition ``_dist_die_offset`` is built from,
+        with ``local_vocab`` read off the sampler rather than re-derived here.
+        """
+        local_vocab = self._penalty_local_vocab
+        die = torch.div(ids, local_vocab, rounding_mode="floor")
+        local = ids - die * local_vocab
+        out = []
+        for index in range(len(self._penalty_host["rep_neg"])):
+            selector = die == index
+            if bool(selector.any()):
+                out.append((index, local[selector], selector))
+        return out
+
+    def set_penalty_params(
+        self,
+        *,
+        presence=None,
+        frequency=None,
+        repetition=None,
+        prompt_tokens=None,
+        output_tokens=None,
+        active_batch: int = 1,
+    ) -> tuple[bool, bool]:
+        """Stage the three vLLM sampling penalties for the next decode step.
+
+        Returns ``(live, graph_changed)``: whether the penalty stage runs this
+        step, and whether the decode graph changed shape -- the caller must
+        reinstall the trace when it did, because a mode change releases it. Everything
+        here is per **global** token id; the global -> die mapping is done by the
+        same ``ShardTensorToMesh(dim=-1)`` split the logits themselves live under,
+        so no index arithmetic can put a penalty on the wrong die. The argument is
+        in ``_WatcherCleanSampling1D``'s penalty section.
+
+        Semantics are vLLM's ``model_executor/layers/utils.py::apply_penalties``,
+        including its order: repetition (over prompt+output) multiplies the raw
+        logit, then frequency (output counts) and presence (output mask) subtract.
+        """
+        sampler = self.model.sampler
+        slots, vocab = sampler.penalty_buffer_shape()
+        rows = max(0, min(int(active_batch), slots))
+
+        def _row_values(values, neutral):
+            if values is None:
+                return [neutral] * rows
+            if isinstance(values, (int, float)):
+                return [float(values)] * rows
+            listed = [float(v) for v in list(values)[:rows]]
+            return listed + [neutral] * (rows - len(listed))
+
+        presence = _row_values(presence, 0.0)
+        frequency = _row_values(frequency, 0.0)
+        repetition = _row_values(repetition, 1.0)
+
+        rep_rows = [r for r in range(rows) if repetition[r] != 1.0]
+        add_rows = [r for r in range(rows) if presence[r] != 0.0 or frequency[r] != 0.0]
+        if (rep_rows or add_rows) and prompt_tokens is None and output_tokens is None:
+            # vLLM only sends the history when a penalty is live (and only on
+            # decode). Without it there is nothing to key a penalty on; run the
+            # unpenalised graph rather than invent one.
+            rep_rows, add_rows = [], []
+        mode = (1 if rep_rows else 0) | (2 if add_rows else 0)
+
+        if mode != self._penalty_mode:
+            # A graph change, exactly like the argmax/split flip: the ops either
+            # are or are not in the captured trace, so the trace must go.
+            if self._trace_model_id is not None or self._trace_sampling_id is not None:
+                self._release_decode_traces()
+            self._decode_warm_key = None
+            sampler.allocate_penalty_buffers(mode)
+            self._penalty_mode = mode
+            graph_changed = True
+        else:
+            graph_changed = False
+        if mode == 0:
+            return False, graph_changed
+
+        host = self._ensure_penalty_host(slots, vocab)
+        add, rep_neg = host["add"], host["rep_neg"]
+
+        # Reset only what this row wrote last step, not the whole 151936-wide
+        # row: the history is at most the context length and is usually far
+        # shorter, so this is O(history) rather than O(vocabulary).
+        for row in range(slots):
+            if mode & 2:
+                previous = self._penalty_prev_add[row]
+                if previous is not None:
+                    for die, local, _ in self._penalty_split(previous):
+                        add[die][0, 0, row].index_fill_(0, local, 0.0)
+                    self._penalty_prev_add[row] = None
+            if mode & 1:
+                previous = self._penalty_prev_rep[row]
+                if previous is not None:
+                    for die, local, _ in self._penalty_split(previous):
+                        rep_neg[die][0, 0, row].index_fill_(0, local, 1.0)
+                    self._penalty_prev_rep[row] = None
+
+        for row in add_rows:
+            out_ids = self._row_token_ids(output_tokens, row)
+            if out_ids.numel() == 0:
+                continue
+            unique, counts = torch.unique(out_ids, return_counts=True)
+            # f * count(t in output) + q * (count > 0), summed on the host so the
+            # device sees one additive tensor rather than two.
+            values = (counts.to(torch.float32) * frequency[row] + presence[row]).to(torch.bfloat16)
+            for die, local, selector in self._penalty_split(unique):
+                add[die][0, 0, row].index_copy_(0, local, values[selector])
+            self._penalty_prev_add[row] = unique
+
+        for row in rep_rows:
+            ids = torch.cat((self._row_token_ids(prompt_tokens, row), self._row_token_ids(output_tokens, row)))
+            if ids.numel() == 0:
+                continue
+            unique = torch.unique(ids)
+            # Only ``p`` is staged; ``1/p - p`` is derived on device from it. See
+            # ``_WatcherCleanSampling1D._apply_penalties``.
+            for die, local, _ in self._penalty_split(unique):
+                rep_neg[die][0, 0, row].index_fill_(0, local, repetition[row])
+            self._penalty_prev_rep[row] = unique
+
+        buffers = sampler.penalty_device_buffers()
+        for name in (("rep_neg",) if mode & 1 else ()) + (("add",) if mode & 2 else ()):
+            self._upload_penalty_tensor(host[name], buffers[name])
+            self.trace_stats["penalty_host_copies"] += 1
+        return True, graph_changed
+
+    def _upload_penalty_tensor(self, shards: list, device) -> None:
+        """Four contiguous ``[1,1,32,37984]`` host buffers -> the four die shards.
+
+        ``ttnn.from_host_shards`` assembles them into the multi-device host
+        tensor directly, so nothing re-slices a 9.7 MB strided view per step.
+        Shard ``d`` is die ``d``'s columns by the ordering
+        ``ShardTensorToMesh(dim=-1)`` uses, which the probe pins by building the
+        same operand both ways and requiring bit-identical device tensors.
+        """
+        ttnn.copy_host_to_device_tensor(
+            ttnn.from_host_shards(
+                [ttnn.from_torch(shard, device=None, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT) for shard in shards],
+                self.mesh_device.shape,
+            ),
+            device,
+        )
+
+    @contextlib.contextmanager
+    def _penalties_suspended(self):
+        """Run the enclosed sampling without the penalty stage.
+
+        Prefill and the eager host/device decode compatibility paths sample rows
+        that are **not** the decode trace's slots (a prefill's row *i* is the
+        *i*-th admitted request, not slot *i*), and vLLM does not send a token
+        history for them -- it populates ``prompt_tokens``/``output_tokens``
+        "if penalties are needed (decode only)". Applying another slot's staged
+        penalty row to them would penalise the wrong tokens, so the stage is off.
+        """
+        sampler = self.model.sampler
+        saved = sampler._penalty_mode
+        sampler._penalty_mode = 0
+        try:
+            yield
+        finally:
+            sampler._penalty_mode = saved
+
+    def _sample_device(self, logits, *, tt_out_tok=None):
+        """Greedy takes the argmax strategy, anything sampled takes the split one.
+
+        Both are ``Sampling1D``, both traced, both write ``tt_out_tok``. The
+        split is by *request*, not by convenience: greedy is exactly top-1, and
+        at this vocabulary the argmax strategy computes it 6.6x faster
+        (0.928 ms against 6.155 ms, ``doc/optimized_full_model/README.md``,
+        "The sampler comparison"). Changing
+        between the two releases the decode traces, which ``set_sampling_params``
+        already does when ``_sampling_stochastic`` flips.
+        """
+        k, p, temp = self._ensure_sampling_params()
+        if not self._sampling_stochastic:
+            return self.model.sample_greedy_argmax(logits, tt_out_tok=tt_out_tok)
+        return self.model.sample_split(logits, k=k, p=p, temp=temp, tt_out_tok=tt_out_tok)
+
+    def _sampled_to_torch(self, sampled) -> torch.Tensor:
+        self.trace_stats["caller_token_readbacks"] += 1
+        return _first_device_to_torch(sampled).reshape(-1)[: self.batch].to(torch.long)
+
+    # -- decode trace ---------------------------------------------------------
+
+    def _prepare_decode_host_inputs(
+        self, tokens: torch.Tensor, positions: torch.Tensor, page_table: torch.Tensor, width: int | None = None
+    ):
+        width = self.batch if width is None else int(width)
+        tokens = tokens.reshape(-1).to(torch.int64)
+        positions = positions.reshape(-1).to(torch.int64)
+        if tokens.numel() > width or positions.numel() > width:
+            raise ValueError("decode batch exceeds the graph width")
+        padded_tokens = torch.zeros(SAMPLING_SLOTS, dtype=torch.int32)
+        padded_tokens[: tokens.numel()] = tokens.to(torch.int32)
+        padded_positions = torch.full((width,), -1, dtype=torch.int32)
+        padded_positions[: positions.numel()] = positions.to(torch.int32)
+        rotary = torch.clamp(padded_positions, min=0).reshape(1, width)
+        return (
+            self._replicated_host_tensor(padded_tokens.reshape(1, 1, 1, SAMPLING_SLOTS), dtype=ttnn.uint32),
+            self._replicated_host_tensor(padded_positions, dtype=ttnn.int32),
+            self._replicated_host_tensor(rotary, dtype=ttnn.uint32),
+            self._replicated_host_tensor(page_table, dtype=ttnn.int32),
+        )
+
+    def _restore_trace_inputs(self, host_inputs, *, include_page_table: bool, token_device=None) -> None:
+        count = 4 if include_page_table else 3
+        start = 0
+        if token_device is not None:
+            ttnn.copy(token_device, self._trace_inputs[0])
+            self.trace_stats["token_device_copies"] += 1
+            start = 1
+        for index in range(start, count):
+            ttnn.copy_host_to_device_tensor(host_inputs[index], self._trace_inputs[index])
+        if token_device is None:
+            self.trace_stats["token_host_copies"] += 1
+        self.trace_stats["position_host_copies"] += 1
+        self.trace_stats["rotary_position_host_copies"] += 1
+        if include_page_table:
+            self.trace_stats["page_table_host_copies"] += 1
+
+    def _decode_graph_key(self, kv_cache, graph_width: int) -> tuple:
+        """Everything that changes which programs the decode graph needs.
+
+        ``rope_cache_len`` is part of it because ``_ensure_decode_rope_capacity``
+        reallocates the cos/sin tables at a *new length* when the horizon grows,
+        which changes the shapes ``ttnn.embedding`` and the untilize behind it
+        run at. Those programs are not in the cache yet, and a trace capture
+        cannot compile them ("Cannot load new binaries during trace capture").
+        Without the length in the key, ``_decode_compiled_keys`` would claim the
+        graph was already warm and skip the eager pass that compiles them.
+        """
+        return (
+            id(kv_cache),
+            graph_width,
+            self._sampling_stochastic,
+            self._penalty_mode,
+            self.model.rope_cache_len,
+            # ``Qwen3CoderModel.active_row_gating`` adds four small ops per step
+            # and one broadcast multiply per layer; flipping it is a different
+            # program set, so a stale "already compiled" claim would try to load
+            # binaries inside an open capture. Same failure mode
+            # ``rope_cache_len`` was added for -- see the docstring above and
+            # ``doc/vllm_integration/work_log.md`` §12.
+            self.model.active_row_gating,
+        )
+
+    def _warm_decode_graphs(self, host_inputs, kv_cache, *, graph_width: int, initial_token_device=None) -> None:
+        """Compile every program once eagerly.
+
+        Non-negotiable rather than merely tidy: ``_decode_ccl_buffers``
+        allocates the two persistent collective buffers on the first call at a
+        shape, and ``ttnn.from_torch`` inside ``begin_trace_capture`` raises and
+        leaves the capture open -- a hung mesh (stage-04 ``work_log.md`` §6).
+        """
+        self._trace_inputs = self._decode_input_pool(graph_width)
+        self._restore_trace_inputs(host_inputs, include_page_table=True, token_device=initial_token_device)
+        token, current_pos, rotary_pos, page_table = self._trace_inputs
+        self.model.bind_page_table(kv_cache, page_table)
+        with self.model.decode_width_scope(graph_width):
+            logits = self.model.decode_forward_from_ttnn_inputs(
+                # ``advance_position=True`` here as well as in the capture: every
+                # op the traced graph contains must already be in the program
+                # cache, and that includes the two ``ttnn.plus_one`` calls. The
+                # positions this leaves behind are overwritten by the restore below.
+                token,
+                current_pos,
+                rotary_position=rotary_pos,
+                kv_cache=kv_cache,
+                advance_position=True,
+            )
+        self._sample_device(logits, tt_out_tok=token)
+        ttnn.deallocate(logits, True)
+        self._synchronize()
+        self._restore_trace_inputs(host_inputs, include_page_table=True, token_device=initial_token_device)
+        self._synchronize()
+        self._decode_warm_key = self._decode_graph_key(kv_cache, graph_width)
+        self._decode_compiled_keys.add(self._decode_warm_key)
+        self.trace_stats["decode_warmups"] += 1
+
+    def _capture_decode_traces(
+        self, host_inputs, kv_cache, *, graph_width: int, active_batch: int, initial_token_device=None
+    ) -> None:
+        self._trace_inputs = self._decode_input_pool(graph_width)
+        model_trace_id = sampling_trace_id = None
+        model_open = sampling_open = False
+        try:
+            warm_key = self._decode_graph_key(kv_cache, graph_width)
+            if self._decode_warm_key != warm_key and warm_key not in self._decode_compiled_keys:
+                self._warm_decode_graphs(
+                    host_inputs, kv_cache, graph_width=graph_width, initial_token_device=initial_token_device
+                )
+            self._restore_trace_inputs(host_inputs, include_page_table=True, token_device=initial_token_device)
+            self._synchronize()
+            token, current_pos, rotary_pos, page_table = self._trace_inputs
+            self.model.bind_page_table(kv_cache, page_table)
+
+            model_trace_id = ttnn.begin_trace_capture(self.mesh_device, cq_id=0)
+            model_open = True
+            with self.model.decode_width_scope(graph_width):
+                logits = self.model.decode_forward_from_ttnn_inputs(
+                    token, current_pos, rotary_position=rotary_pos, kv_cache=kv_cache, advance_position=True
+                )
+            ttnn.end_trace_capture(self.mesh_device, model_trace_id, cq_id=0)
+            model_open = False
+            self._synchronize()
+
+            sampling_trace_id = ttnn.begin_trace_capture(self.mesh_device, cq_id=0)
+            sampling_open = True
+            sampled = self._sample_device(logits, tt_out_tok=token)
+            ttnn.end_trace_capture(self.mesh_device, sampling_trace_id, cq_id=0)
+            sampling_open = False
+            self._synchronize()
+        except Exception:
+            if sampling_open:
+                ttnn.end_trace_capture(self.mesh_device, sampling_trace_id, cq_id=0)
+            if model_open:
+                ttnn.end_trace_capture(self.mesh_device, model_trace_id, cq_id=0)
+            for trace_id in (sampling_trace_id, model_trace_id):
+                if trace_id is not None:
+                    try:
+                        ttnn.release_trace(self.mesh_device, trace_id)
+                    except Exception:
+                        pass
+            raise
+
+        self._trace_model_id = model_trace_id
+        self._trace_sampling_id = sampling_trace_id
+        self._trace_logits = logits
+        self._trace_sampled = sampled
+        self._trace_kv_cache = kv_cache
+        self._trace_page_table_snapshot = self._page_table_to_torch(host_inputs[3]).clone()
+        self._trace_active_batch = active_batch
+        self._trace_graph_width = graph_width
+        self.trace_stats["captures"] += 1
+        self._restore_trace_inputs(host_inputs, include_page_table=True, token_device=initial_token_device)
+        self._synchronize()
+
+    def _refresh_trace_state(
+        self, host_inputs, kv_cache, *, graph_width: int, active_batch: int, initial_token_device=None
+    ) -> None:
+        new_page_table = self._page_table_to_torch(host_inputs[3])
+        shape_changed = (
+            self._trace_page_table_snapshot is not None
+            and new_page_table.shape != self._trace_page_table_snapshot.shape
+        )
+        if self._trace_model_id is not None and (
+            kv_cache is not self._trace_kv_cache
+            or graph_width != self._trace_graph_width
+            or active_batch != self._trace_active_batch
+            or shape_changed
+        ):
+            self._release_decode_traces()
+        if self._trace_model_id is None:
+            self._capture_decode_traces(
+                host_inputs,
+                kv_cache,
+                graph_width=graph_width,
+                active_batch=active_batch,
+                initial_token_device=initial_token_device,
+            )
+            return
+        self._restore_trace_inputs(host_inputs, include_page_table=False, token_device=initial_token_device)
+        if not torch.equal(new_page_table, self._trace_page_table_snapshot):
+            ttnn.copy_host_to_device_tensor(host_inputs[3], self._trace_inputs[3])
+            self._trace_page_table_snapshot = new_page_table.clone()
+            self.trace_stats["page_table_host_copies"] += 1
+
+    def _refresh_persistent_page_table(self, page_table, kv_cache, *, active_batch: int) -> None:
+        if self._trace_model_id is None:
+            raise RuntimeError("decode trace is not initialized")
+        if kv_cache is not self._trace_kv_cache:
+            raise RuntimeError("KV-cache identity changed; initialize a new trace")
+        if active_batch != self._trace_active_batch:
+            raise RuntimeError("fixed active slots changed; initialize a new trace")
+        if page_table is None:
+            return
+        new_page_table = self._normalise_page_table(page_table, active_batch, width=self._trace_graph_width)
+        if torch.equal(new_page_table, self._trace_page_table_snapshot):
+            return  # unchanged page table costs zero host copies
+        ttnn.copy_host_to_device_tensor(
+            self._replicated_host_tensor(new_page_table, dtype=ttnn.int32), self._trace_inputs[3]
+        )
+        self._trace_page_table_snapshot = new_page_table.clone()
+        self.trace_stats["page_table_host_copies"] += 1
+
+    def _copy_forced_tokens(self, tokens: torch.Tensor) -> None:
+        """Teacher forcing: overwrite the fed-back token, keep everything else."""
+        values = tokens.reshape(-1).to(torch.int64)
+        if values.numel() != self._trace_active_batch:
+            raise ValueError(f"expected {self._trace_active_batch} forced tokens, got {values.numel()}")
+        host = torch.zeros(SAMPLING_SLOTS, dtype=torch.int32)
+        host[: values.numel()] = values.to(torch.int32)
+        ttnn.copy_host_to_device_tensor(
+            self._replicated_host_tensor(host.reshape(1, 1, 1, SAMPLING_SLOTS), dtype=ttnn.uint32),
+            self._trace_inputs[0],
+        )
+        self.trace_stats["token_host_copies"] += 1
+
+    def _replay_split_sampling(self):
+        ttnn.execute_trace(self.mesh_device, self._trace_model_id, cq_id=0, blocking=False)
+        ttnn.execute_trace(self.mesh_device, self._trace_sampling_id, cq_id=0, blocking=False)
+        self.trace_stats["replays"] += 1
+        return self._trace_sampled
+
+    def _release_decode_traces(self) -> None:
+        released = self._trace_model_id is not None or self._trace_sampling_id is not None
+        for trace_id in (self._trace_model_id, self._trace_sampling_id):
+            if trace_id is not None:
+                ttnn.release_trace(self.mesh_device, trace_id)
+        if released:
+            self.trace_stats["releases"] += 1
+        self._trace_model_id = None
+        self._trace_sampling_id = None
+        self._trace_inputs = None
+        self._trace_logits = None
+        self._trace_sampled = None
+        self._trace_kv_cache = None
+        self._trace_page_table_snapshot = None
+        self._trace_active_batch = None
+        self._trace_graph_width = None
+        self._trace_rotary_position = None
+        self._decode_warm_key = None
+
+    def _ensure_decode_rope_capacity(self, required_len: int) -> None:
+        """Grow the cos/sin tables for decode, releasing traces if they move.
+
+        ``Qwen3CoderModel.ensure_rope_capacity`` reallocates ``cos_table`` and
+        ``sin_table`` when it grows them, and a captured trace holds the *old*
+        tensor identities -- replaying it afterwards would gather from freed
+        DRAM. ``prefill_forward`` and ``generate`` are safe because both release
+        the decode traces before they call it; the low-level ``decode_forward``
+        has no such release, so it does one here and only when the tables
+        actually moved.
+        """
+        if required_len > self.model.max_cache_len:
+            raise ValueError(f"decode horizon {required_len} exceeds the supported context {self.model.max_cache_len}")
+        if self.model.ensure_rope_capacity(required_len):
+            self._release_decode_traces()
+
+    def decode_forward(
+        self,
+        tokens: torch.Tensor | None,
+        start_pos: torch.Tensor | None,
+        *,
+        page_table,
+        kv_cache: Any,
+        sampling_mode: str = "host",
+        enable_trace: bool = False,
+        active_batch: int | None = None,
+        graph_width: int | None = None,
+        decode_horizon: int | None = None,
+        validate_page_coverage: bool = True,
+        **kwargs: Any,
+    ):
+        """One decode step.
+
+        With ``enable_trace=True, sampling_mode="device"`` this is the delivered
+        path: pass ``start_pos``/``page_table`` on the first step to install the
+        trace, then call with ``tokens=None, start_pos=None, page_table=None``
+        and the traces replay over persistent state -- the sampled token from
+        step *N* is already the token input of step *N+1*, and both position
+        tensors were advanced on device inside the model trace.
+
+        **Rotary capacity.** The cos/sin tables are sized lazily, and the traced
+        loop advances ``rotary_position`` with ``ttnn.plus_one`` with nothing on
+        device to clamp it, so a replay past the table length would gather
+        out of range and silently rotate at the wrong position. Pass
+        ``decode_horizon`` -- the highest position this trace will ever decode
+        at, i.e. ``prompt_len + max_new_tokens - 1`` -- on the installing call
+        and the tables are grown once, up front, to cover the whole run. Without
+        it the tables are sized only for ``start_pos`` and a replay that would
+        step past them raises instead of returning a wrong answer.
+        ``generate`` sizes for its own horizon and never hits either path.
+        """
+        if sampling_mode not in {"host", "device"}:
+            raise ValueError("sampling_mode must be 'host' or 'device'")
+        caches = self._ensure_kv_cache() if kv_cache is None else kv_cache
+        inferred = self._trace_active_batch if tokens is None else int(tokens.numel())
+        active_batch = inferred if active_batch is None else int(active_batch)
+        if active_batch is None or not 1 <= active_batch <= self.batch:
+            raise ValueError(f"active_batch must be in [1,{self.batch}]")
+        if tokens is not None and tokens.numel() != active_batch:
+            raise ValueError("tokens do not match active_batch")
+        if start_pos is not None and start_pos.numel() != active_batch:
+            raise ValueError("start_pos does not match active_batch")
+        # ``graph_width`` is how many rows the captured decode graph has;
+        # ``active_batch`` is how many of them the caller is filling. They are
+        # the same on the shipped path. A caller that has compacted its live
+        # requests into rows ``0..active_batch-1`` may ask for a narrower graph,
+        # which is the whole point of ``doc/batch_scaling``: expert, router and
+        # SDPA cost is paid per row *configured*, so the only way to stop paying
+        # for 32 rows when one is live is to capture a graph that has fewer.
+        if graph_width is None:
+            graph_width = self._trace_graph_width if start_pos is None and self._trace_graph_width else self.batch
+        graph_width = int(graph_width)
+        if not active_batch <= graph_width <= self.batch:
+            raise ValueError(f"graph_width must be in [{active_batch},{self.batch}], got {graph_width}")
+
+        if enable_trace and sampling_mode == "device":
+            if start_pos is not None:
+                if page_table is None:
+                    raise ValueError("initial trace state requires positions and page_table")
+                highest = int(start_pos.reshape(-1).max().item())
+                horizon = highest + 1 if decode_horizon is None else int(decode_horizon)
+                if horizon < highest + 1:
+                    raise ValueError("decode_horizon is below the requested start_pos")
+                self._ensure_decode_rope_capacity(horizon)
+                page_host = self._normalise_page_table(page_table, active_batch, width=graph_width)
+                if validate_page_coverage:
+                    self._validate_page_coverage(page_host, start_pos, active_batch)
+                initial_token_device = self._prefill_sampled if tokens is None else None
+                host_tokens = torch.zeros(active_batch, dtype=torch.long) if tokens is None else tokens
+                host_inputs = self._prepare_decode_host_inputs(host_tokens, start_pos, page_host, width=graph_width)
+                self._refresh_trace_state(
+                    host_inputs,
+                    caches,
+                    graph_width=graph_width,
+                    active_batch=active_batch,
+                    initial_token_device=initial_token_device,
+                )
+                # The installing call also replays once, at ``highest``.
+                self._trace_rotary_position = highest
+            else:
+                self._refresh_persistent_page_table(page_table, caches, active_batch=active_batch)
+                if tokens is not None:
+                    self._copy_forced_tokens(tokens)
+                if self._trace_rotary_position is not None:
+                    # ``ttnn.plus_one`` already moved the device tensor on the
+                    # previous replay; this replay gathers at that position.
+                    self._trace_rotary_position += 1
+                    if self._trace_rotary_position >= self.model.rope_cache_len:
+                        raise RuntimeError(
+                            f"decode position {self._trace_rotary_position} is past the rotary table "
+                            f"({self.model.rope_cache_len} entries). ``ttnn.embedding`` would gather out of "
+                            "range inside the replayed trace and rotate at a wrong position without "
+                            "raising. Re-install the trace with decode_horizon= set to the highest "
+                            "position this run will reach."
+                        )
+            return self._replay_split_sampling()
+
+        if tokens is None or start_pos is None or page_table is None:
+            raise ValueError("eager/host decode requires tokens, start_pos and page_table")
+        self._release_decode_traces_before_allocating()
+        # Eager decode gathers cos/sin at ``start_pos`` too, and holds no trace
+        # by this point, so the tables can simply be grown to fit.
+        self._ensure_decode_rope_capacity(int(start_pos.reshape(-1).max().item()) + 1)
+        page_host = self._normalise_page_table(page_table, active_batch)
+        if validate_page_coverage:
+            self._validate_page_coverage(page_host, start_pos, active_batch)
+        host_inputs = self._prepare_decode_host_inputs(tokens, start_pos, page_host)
+        device_inputs = [
+            ttnn.to_device(tensor, self.mesh_device, memory_config=ttnn.DRAM_MEMORY_CONFIG) for tensor in host_inputs
+        ]
+        self.model.bind_page_table(caches, device_inputs[3])
+        logits = self.model.decode_forward_from_ttnn_inputs(
+            device_inputs[0],
+            device_inputs[1],
+            rotary_position=device_inputs[2],
+            kv_cache=caches,
+            advance_position=False,
+        )
+        if sampling_mode == "device":
+            with self._penalties_suspended():
+                sampled = self._sample_device(logits)
+            ttnn.deallocate(logits, True)
+            return sampled
+        host = self.model.gather_logits_to_torch(logits, valid_rows=active_batch)[0, 0]
+        ttnn.deallocate(logits, True)
+        return host
+
+    # -- high level -----------------------------------------------------------
+
+    def _generate_host_compat(
+        self, prompt_token_ids: list[int], max_new_tokens: int, *, next_input: Optional[NextInputFn]
+    ) -> list[int]:
+        """Explicit host-sampling compatibility mode. Never a measured path."""
+        self._release_decode_traces()
+        kv_cache = self._ensure_kv_cache()
+        horizon = len(prompt_token_ids) + max_new_tokens - 1
+        page_host = self.make_page_table([horizon])
+        logits = self.prefill_forward(
+            torch.tensor([prompt_token_ids]),
+            page_table=page_host,
+            kv_cache=kv_cache,
+            prompt_lens=[len(prompt_token_ids)],
+            sampling_mode="host",
+        )
+        predicted = int(logits[0, 0].argmax().item())
+        outputs: list[int] = []
+        for step in range(max_new_tokens):
+            outputs.append(predicted)
+            next_token = next_input(step, predicted) if next_input is not None else predicted
+            if step + 1 == max_new_tokens:
+                break
+            decoded = self.decode_forward(
+                torch.tensor([[next_token]]),
+                torch.tensor([len(prompt_token_ids) + step]),
+                page_table=page_host,
+                kv_cache=kv_cache,
+                sampling_mode="host",
+                enable_trace=False,
+            )
+            predicted = int(decoded[0].argmax().item())
+        return outputs
+
+    def generate(
+        self,
+        prompt_token_ids: list[int],
+        max_new_tokens: int,
+        *,
+        next_input: Optional[NextInputFn] = None,
+        enable_trace: bool = True,
+        sampling_mode: str = "device",
+        stop_on_eos: bool = False,
+        top_k=1,
+        top_p=0.0,
+        temperature=1.0,
+        **kwargs: Any,
+    ) -> list[int]:
+        """Prefill, then loop the traced split-sampling decode path."""
+        if not prompt_token_ids or max_new_tokens < 1:
+            return []
+        horizon = len(prompt_token_ids) + max_new_tokens - 1
+        if horizon > self.model.max_cache_len:
+            raise ValueError("prompt plus requested output exceeds the supported context")
+        self._release_decode_traces_before_allocating()
+        self.model.ensure_rope_capacity(horizon)
+        if sampling_mode == "host":
+            return self._generate_host_compat(prompt_token_ids, max_new_tokens, next_input=next_input)
+        if sampling_mode != "device":
+            raise ValueError("sampling_mode must be 'device' or 'host'")
+        if not enable_trace and max_new_tokens > 1:
+            raise ValueError("the optimized token-out path requires enable_trace=True")
+        self.set_sampling_params(top_k=top_k, top_p=top_p, temperature=temperature, active_batch=1)
+
+        kv_cache = self._ensure_kv_cache()
+        page_host = self.make_page_table([horizon])
+        sampled = self.prefill_forward(
+            torch.tensor([prompt_token_ids]),
+            page_table=page_host,
+            kv_cache=kv_cache,
+            prompt_lens=[len(prompt_token_ids)],
+            sampling_mode="device",
+        )
+        predicted = int(self._sampled_to_torch(sampled)[0].item())
+        outputs: list[int] = []
+        for step in range(max_new_tokens):
+            outputs.append(predicted)
+            forced = next_input(step, predicted) if next_input is not None else predicted
+            if step + 1 == max_new_tokens:
+                break
+            if stop_on_eos and next_input is None and predicted == self.tokenizer.eos_token_id:
+                break
+            initial = step == 0
+            sampled = self.decode_forward(
+                (torch.tensor([[forced]]) if next_input is not None else None),
+                torch.tensor([len(prompt_token_ids)]) if initial else None,
+                page_table=page_host if initial else None,
+                kv_cache=kv_cache,
+                sampling_mode="device",
+                enable_trace=True,
+                active_batch=1,
+                decode_horizon=horizon,
+            )
+            predicted = int(self._sampled_to_torch(sampled)[0].item())
+        return outputs
+
+    def reset(self) -> None:
+        """Wipe per-prompt state, keeping weights, buffers and program cache."""
+        if self._trace_model_id is not None or self._trace_sampling_id is not None:
+            self._synchronize()
+            self._release_decode_traces()
+        if self._kv_cache is not None:
+            self.model.reset_kv_cache(self._kv_cache)
+        empty = torch.full((self.batch, self.pages_per_user), -1, dtype=torch.int32)
+        self._copy_host(empty, self._prefill_page_table, dtype=ttnn.int32)
+        self._copy_host(
+            torch.zeros((1, 1, 1, SAMPLING_SLOTS), dtype=torch.int32), self._prefill_sampled, dtype=ttnn.uint32
+        )
+        for width, pool in self._width_pools.items():
+            token, current_pos, rotary_pos, page_table = pool
+            self._copy_host(torch.zeros((1, 1, 1, SAMPLING_SLOTS), dtype=torch.int32), token, dtype=ttnn.uint32)
+            self._copy_host(torch.full((width,), -1, dtype=torch.int32), current_pos, dtype=ttnn.int32)
+            self._copy_host(torch.zeros((1, width), dtype=torch.int32), rotary_pos, dtype=ttnn.uint32)
+            self._copy_host(empty[:width], page_table, dtype=ttnn.int32)
+        self._trace_rotary_position = None
+        self.trace_stats["resets"] += 1
+        self._synchronize()
+
+    def teardown(self) -> None:
+        self._release_decode_traces()
+        #: Not cleared by ``_release_decode_traces`` on purpose -- a released
+        #: trace leaves its programs compiled, which is the whole point of the
+        #: set. Cleared here because teardown deallocates the KV cache, and the
+        #: key holds ``id(kv_cache)``: a later allocation could land on the same
+        #: address with a different shape and falsely claim to be warm.
+        self._decode_compiled_keys.clear()
+        if self._kv_cache is not None:
+            for cache in self._kv_cache:
+                ttnn.deallocate(cache.k, True)
+                ttnn.deallocate(cache.v, True)
+            self._kv_cache = None
+
+
+def _resolve_snapshot(model_path: str | Path | None = None) -> Path:
+    if model_path is not None:
+        path = Path(model_path)
+        if not path.exists():
+            raise FileNotFoundError(path)
+        return path
+    hf_home = Path(os.getenv("HF_HOME", Path.home() / ".cache" / "huggingface"))
+    snapshot = hf_home / "hub" / "models--Qwen--Qwen3-Coder-30B-A3B-Instruct" / "snapshots" / HF_REVISION
+    if (snapshot / "model.safetensors.index.json").is_file():
+        return snapshot
+    from huggingface_hub import snapshot_download
+
+    return Path(snapshot_download(HF_MODEL_ID, revision=HF_REVISION))
+
+
+def build_generator(model_dir: str | Path, mesh_device, **kwargs) -> Generator:
+    """Readiness discovery factory. See ``models/common/readiness_check/contract.py``.
+
+    **This is the construction path the precision config has to reach.** The
+    readiness runners, the qualitative suite and (later) vLLM all arrive here
+    and none of them can pass a Python object, so ``precision`` is accepted as a
+    kwarg *and* read from ``QWEN3_PRECISION_CONFIG`` in the environment as a
+    path to a ``selected_precision_config.json``. Unset -- which is every run to
+    date -- means ``DEFAULT_PRECISION``, i.e. the shipped policy, so this
+    default is the one the stage-07 goal asks for rather than a JSON field that
+    hard-coded model code ignores.
+    """
+    snapshot = _resolve_snapshot(kwargs.pop("model_path", os.getenv("QWEN3_CODER_30B_MODEL_PATH")))
+    tokenizer = AutoTokenizer.from_pretrained(snapshot)
+    max_batch_size = int(kwargs.pop("max_batch_size", 1))
+    max_context_len = int(kwargs.pop("max_context_len", MAX_CONTEXT))
+    override_num_layers = kwargs.pop("override_num_layers", None)
+    num_layers = NUM_LAYERS if override_num_layers is None else int(override_num_layers)
+    page_block_size = int(kwargs.pop("page_block_size", 32))
+    rope_cache_len = int(kwargs.pop("rope_cache_len", 8192))
+    precision = kwargs.pop("precision", os.getenv("QWEN3_PRECISION_CONFIG") or None)
+    if kwargs:
+        raise TypeError(f"unsupported build_generator kwargs: {sorted(kwargs)}")
+    model = Qwen3CoderModel.from_checkpoint(
+        snapshot,
+        mesh_device=mesh_device,
+        max_batch_size=max_batch_size,
+        max_cache_len=max_context_len,
+        num_layers=num_layers,
+        page_block_size=page_block_size,
+        rope_cache_len=rope_cache_len,
+        precision=precision,
+    )
+    return Qwen3CoderGenerator(model, tokenizer)
+
+
+__all__ = ["Qwen3CoderGenerator", "build_generator"]

code/models/demos/blackhole/qwen3_coder_30b_a3b/tt/generator_vllm.py

@@ -0,0 +1,1568 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+"""vLLM serving adapter for Qwen3-Coder-30B-A3B-Instruct on 4 Blackhole dies.
+
+This file is **translation only**. Every device operation it causes is a call
+into ``tt/generator.py``'s low-level surface -- ``prefill_forward`` /
+``decode_forward`` / ``set_sampling_params`` / ``configure_paging`` -- which is
+the same surface the standalone readiness runners drive. There is no model code
+here, no second sampler, no host argmax, no full-logits readback on the measured
+path and no Python readback/writeback token-feedback loop.
+
+The three things it actually has to reconcile
+---------------------------------------------
+
+**1. Who owns the cache.** Standalone, the generator allocates its own paged
+cache and builds its own page tables. Serving, vLLM owns both: it picks the
+block size and the block count, calls ``allocate_kv_cache`` once, and hands that
+exact cache and a ``[batch, max_num_blocks_per_req]`` block table into every
+forward. ``allocate_kv_cache`` therefore installs vLLM's geometry through
+``Qwen3CoderGenerator.configure_paging`` and allocates the cache through
+``Qwen3CoderModel.allocate_kv_cache(num_blocks=...)``; the generator's own
+``_kv_cache`` is never created, so no standalone-cache assumption can survive.
+
+**2. Who owns the token.** The generator's traced decode path writes the sampled
+token straight back into the persistent decode token input with ``tt_out_tok``
+and advances ``current_pos``/``rotary_position`` on device with
+``ttnn.plus_one``. So after step *N* the **device**, not the host, holds the
+token and position step *N+1* needs. vLLM's scheduler also tracks them, and
+under ``--async-scheduling`` its copy is a step behind. On a steady decode step
+this adapter therefore passes *nothing* -- ``decode_forward(None, None,
+page_table=..., ...)`` replays the two traces over state the device already
+owns. Only when vLLM says the slot layout changed (``reset_batch``, or the step
+right after a prefill) does it reinstall host state, and even then it keeps the
+device's token and position for slots that are merely continuing
+(``_merge_scheduler_view``), so an async-ahead scheduler cannot re-decode a
+position or feed back a stale token.
+
+**3. Who owns sampling.** vLLM would rather hand us logits. It only stops doing
+that when the model declares ``supports_sample_on_device`` and the server runs
+with ``sample_on_device_mode: all`` -- then it sends per-row
+``(temperature, top_k, top_p)`` and expects token ids back. That is the measured
+path and it is exactly the full model's canonical split sampling: greedy takes
+``Qwen3CoderModel.sample_greedy_argmax``, anything sampled takes
+``sample_split``, both are ``_WatcherCleanSampling1D``, both are traced, both
+write ``tt_out_tok``. The TT plugin still routes a few request shapes to host
+sampling on its own (logprobs on a mesh that is not 8 or 32 dies, ``min_p``,
+``bad_words``, ``logit_bias``, structured output) -- for those it passes
+``sampling_params=None`` and wants logits. That is served by the generator's
+pre-existing, explicit ``sampling_mode="host"`` compatibility mode. It is opt-in
+per request by vLLM, never used for a performance number, and it never displaces
+the traced path. The eager decode does release the captured decode traces (it
+allocates, so ``generator.decode_forward`` calls
+``_release_decode_traces_before_allocating``); the adapter therefore sets
+``_needs_decode_install`` and the next device-sampled step re-captures through
+``_refresh_trace_state`` rather than replaying a stale trace.
+"""
+
+from __future__ import annotations
+
+import json
+import math
+import os
+from collections import deque
+from pathlib import Path
+from typing import Any, Optional
+
+import torch
+from loguru import logger
+
+import ttnn
+from models.demos.blackhole.qwen3_coder_30b_a3b.tt.functional_decoder import sdpa_chunk_size
+from models.demos.blackhole.qwen3_coder_30b_a3b.tt.generator import SAMPLING_SLOTS, Qwen3CoderGenerator, build_generator
+from models.demos.blackhole.qwen3_coder_30b_a3b.tt.model import HF_MODEL_ID, MAX_CONTEXT
+
+#: This port's directory. Everything below reads its policy from here, not from
+#: a vLLM flag, so serving cannot silently run a different model than readiness.
+MODEL_DIR = Path(__file__).resolve().parents[1]
+
+#: The datatype-sweep selection (stage 07). ``build_generator`` accepts a path
+#: and threads it into ``Qwen3CoderModel``, so weight groups, activation dtype,
+#: CCL dtype, KV-cache dtype, compute fidelities and layer exceptions all come
+#: from this one file on the serving path exactly as they do on the readiness
+#: path. ``QWEN3_PRECISION_CONFIG`` still overrides it, for sweeps.
+SELECTED_PRECISION_CONFIG = MODEL_DIR / "config" / "selected_precision_config.json"
+
+#: ``config/context_contract.json`` is the single source of truth for served
+#: context. ``get_max_tokens_all_users`` and ``initialize_vllm_model`` both read
+#: it rather than trusting a CLI value, so a ``--max-model-len`` above the
+#: recorded capability fails loudly instead of serving a quietly-clipped model.
+CONTEXT_CONTRACT = MODEL_DIR / "config" / "context_contract.json"
+
+#: Whether a serving prefill may run while the decode traces stay captured.
+#: **Off by default, and that is a measurement, not caution.** Keeping them alive
+#: hangs the NoC: prefill's collectives share the sampler's persistent CCL
+#: buffers and semaphores with the captured graph, and after a few admissions a
+#: replay waits on a semaphore value an eager collective already consumed.
+#: ``doc/vllm_integration/triage/tt-triage-preserve-traces-hang.txt`` is that
+#: hang, caught with ``dump_running_operations`` reporting ``NOC0 CB0..3 active
+#: (0xFFFFFFFF). NoC is likely hung.`` on device 0. Releasing on prefill --
+#: which is what ``Qwen3CoderGenerator.prefill_forward`` has always done -- makes
+#: the next capture re-establish that state, and
+#: ``Qwen3CoderGenerator._decode_compiled_keys`` keeps the re-capture from
+#: paying for a second eager warm pass. Set ``QWEN3_VLLM_PRESERVE_DECODE_TRACES=1``
+#: to reproduce the hang.
+PRESERVE_DECODE_TRACES = os.getenv("QWEN3_VLLM_PRESERVE_DECODE_TRACES", "0") not in ("0", "", "false", "no")
+
+
+#: Prefix caching is ON by default from phase 3. This is no longer a feature gate
+#: but a kill switch: ``QWEN3_PREFIX_CACHING=0`` restores the pre-phase-3 refusal
+#: without touching ``model_capabilities``, so the two can be reverted separately.
+_PREFIX_CACHING_ENABLED = os.getenv("QWEN3_PREFIX_CACHING", "1") not in ("0", "", "false", "no")
+
+
+def _supported_context() -> int:
+    try:
+        contract = json.loads(CONTEXT_CONTRACT.read_text())
+    except (OSError, ValueError):
+        return MAX_CONTEXT
+    return int(contract.get("current_supported_context") or MAX_CONTEXT)
+
+
+def _reorder_history(history, order, picks):
+    """Reorder a vLLM ``[rows, L]`` token history into graph-row order.
+
+    Type-preserving **on purpose**. ``Generator._row_token_ids`` consumes these
+    with ``torch.as_tensor(history)``, which raises ``TypeError: only integer
+    tensors of a single element can be converted to an index`` when handed a
+    *list of 1-D tensors*. Rebuilding a tensor history with a list comprehension
+    therefore turns a working penalised decode into a crash -- and only when the
+    width ladder is active, since with the ladder off ``order`` is ``None`` and
+    the history is passed through untouched. That is exactly the shape of bug
+    that reaches production: invisible on the default path, fatal on the new one.
+
+    Anything that supports fancy indexing (torch tensors, numpy arrays) is
+    indexed so its type survives; genuine python sequences keep the list
+    comprehension, which is already what ``_row_token_ids`` expects from them.
+    """
+    if history is None:
+        return None
+    if isinstance(history, torch.Tensor):
+        return history[order.to(torch.long)]
+    if hasattr(history, "__getitem__") and hasattr(history, "dtype"):  # numpy & friends
+        return history[list(picks)]
+    return [history[i] for i in picks]
+
+
+def _as_int_list(values, length: int, default) -> list:
+    """vLLM hands per-row sampling params over as python lists; normalise them."""
+    if values is None:
+        return [default] * length
+    # vLLM hands prompt_lens/start_pos over as numpy arrays and sampling params
+    # as python lists, and the plugin builds some fields from torch tensors.
+    if hasattr(values, "tolist") and not isinstance(values, (list, tuple)):
+        values = values.tolist()
+    if not isinstance(values, (list, tuple)):
+        values = [values] * length
+    out = list(values)[:length]
+    out.extend([default] * (length - len(out)))
+    return out
+
+
+class Qwen3CoderForCausalLM:
+    """The class vLLM instantiates. Registered as ``TTQwen3MoeForCausalLM``.
+
+    ``Qwen3MoeForCausalLM`` is the architecture string in this checkpoint's
+    ``config.json``; the TT plugin registers every model under a ``TT`` prefix.
+    """
+
+    #: Read off the *class* by ``vllm_tt_plugin.platform.check_and_update_config``
+    #: before anything is instantiated.
+    #:
+    #: * ``supports_sample_on_device`` -- the full model's traced split sampling
+    #:   is the measured token-out path; without this flag the readiness runner's
+    #:   ``sample_on_device_mode: all`` is a hard config error.
+    #: * ``supports_async_decode`` -- ``decode_forward(read_from_device=False)``
+    #:   returns device handles, ``read_decode_output(async_read=True)`` does the
+    #:   deferred read and records an event, ``process_decode_output_host`` does
+    #:   host formatting only. This also gates ``--async-scheduling``, which is
+    #:   safe here because ``_merge_scheduler_view`` prefers the device's token
+    #:   and position over an async-ahead scheduler's.
+    #: * ``supports_prefix_caching`` -- **False**. Phase 3 REVERTED it: with caching
+    #:   on, cold-vs-warm greedy output matched on only 1 of 10 prompts, while the
+    #:   same test with --no-enable-prefix-caching matched 10/10. See
+    #:   doc/prefix_caching/probes/phase3_cold_warm_rate.json and
+    #:   phase3_control_no_prefix_caching.json. The adapter wiring below is correct
+    #:   and stays; the flag must not go back to True until that gap is closed.
+    #:   (Historical note, kept because the wiring depends on it:)
+    #:   vLLM then sends a non-zero ``start_pos`` (= ``num_computed_tokens``) with
+    #:   the FULL prompt and the FULL ``prompt_lens``; the model slices the suffix
+    #:   itself, matching tt_transformers' ``tokens[i, num_cached:seq_len]``.
+    #:   Two vLLM invariants make our generator-side guards exact rather than
+    #:   defensive, both READ OFF vllm rather than assumed:
+    #:     - ``max_cache_hit_length = request.num_tokens - 1`` (kv_cache_manager.py)
+    #:       so a full hit still recomputes the last token: ``start < prompt_len``.
+    #:     - cache hits are whole blocks and ``allocate_slots`` requires
+    #:       block-aligned ``num_computed_tokens``: ``start % 32 == 0``.
+    #:   Chunked prefill is force-disabled by the plugin platform, so prefix
+    #:   caching is the only source of a non-zero ``start_pos``.
+    #:   Kill switch: ``QWEN3_PREFIX_CACHING=0`` restores the old refusal.
+    model_capabilities = {
+        # QUALITY-GATE EDIT (doc/prefix_caching/QUALITY_BAR.md): flipped True to
+        # run the caching-ON arm. Revert this single line to False if the gate
+        # fails. See doc/prefix_caching/quality_gate/.
+        "supports_prefix_caching": True,
+        "supports_async_decode": True,
+        "supports_sample_on_device": True,
+    }
+
+    # -- construction ---------------------------------------------------------
+
+    @classmethod
+    def initialize_vllm_model(
+        cls,
+        hf_config,
+        mesh_device,
+        max_batch_size,
+        max_seq_len: int | None = None,
+        n_layers: int | None = None,
+        tt_data_parallel: int = 1,
+        optimizations: str | None = None,
+        **kwargs: Any,
+    ) -> "Qwen3CoderForCausalLM":
+        if int(tt_data_parallel) != 1:
+            raise ValueError(
+                f"tt_data_parallel={tt_data_parallel} is unsupported: this port occupies the whole "
+                "1x4 mesh with tensor parallelism, so there is no submesh left to replicate onto."
+            )
+        supported = _supported_context()
+        max_seq_len = supported if max_seq_len is None else int(max_seq_len)
+        if max_seq_len > supported:
+            raise ValueError(
+                f"--max-model-len {max_seq_len} exceeds the context recorded in {CONTEXT_CONTRACT} " f"({supported})."
+            )
+        max_batch_size = int(max_batch_size)
+        if not 1 <= max_batch_size <= SAMPLING_SLOTS:
+            raise ValueError(
+                f"max_num_seqs={max_batch_size} is outside [1,{SAMPLING_SLOTS}]. "
+                "nlp_create_qkv_heads_decode and ttnn.sampling both address 32 fixed user slots."
+            )
+
+        # Reduced serving target for the bring-up inner loop only: the same
+        # adapter, generator, registration, cache/page-table shapes, terminal
+        # norm/LM head, sampler and trace behaviour, with fewer copies of the one
+        # layer kind this model has. Never used for accuracy or performance
+        # evidence -- the final run leaves it unset and gets all 48 layers.
+        reduced = os.getenv("QWEN3_VLLM_NUM_LAYERS")
+        if reduced:
+            n_layers = int(reduced)
+            logger.warning(
+                "QWEN3_VLLM_NUM_LAYERS={} -- REDUCED serving target, bring-up inner loop only. "
+                "Do not report accuracy or performance from this server.",
+                n_layers,
+            )
+
+        precision = os.getenv("QWEN3_PRECISION_CONFIG") or str(SELECTED_PRECISION_CONFIG)
+        generator = build_generator(
+            MODEL_DIR,
+            mesh_device,
+            max_batch_size=max_batch_size,
+            max_context_len=max_seq_len,
+            # The traced decode loop advances ``rotary_position`` on device and
+            # nothing on device clamps it, so the cos/sin tables must already
+            # cover every position this server may be asked to serve. Sizing
+            # them here means no serving step can ever grow them -- growing
+            # reallocates, and a captured trace holds the old identities.
+            rope_cache_len=max_seq_len,
+            precision=precision,
+            **({} if n_layers is None else {"override_num_layers": int(n_layers)}),
+        )
+        # Logged after the generator exists so ``active_row_gating`` is read off
+        # the model that was actually built rather than re-parsed from the
+        # environment here. Every leg of an A/B then carries its own
+        # configuration in its own server log, instead of the two legs being
+        # distinguishable only by the work-log prose that says which was which.
+        logger.info(
+            "Qwen3-Coder-30B-A3B vLLM init: max_num_seqs={} max_model_len={} precision={} "
+            "active_row_gating={} optimizations={}",
+            max_batch_size,
+            max_seq_len,
+            precision,
+            generator.model.active_row_gating,
+            optimizations,
+        )
+        return cls(generator, max_model_len=max_seq_len, max_num_seqs=max_batch_size)
+
+    def __init__(self, generator: Qwen3CoderGenerator, *, max_model_len: int, max_num_seqs: int):
+        self.generator = generator
+        self.model = generator.model
+        self.mesh_device = generator.mesh_device
+        self.max_model_len = int(max_model_len)
+        self.max_num_seqs = int(max_num_seqs)
+
+        #: vLLM-owned cache; set by ``allocate_kv_cache`` and never re-created.
+        self.kv_cache: list | None = None
+        #: True until a decode step has installed host state into the trace.
+        #: Every prefill sets it, because prefill admits a new request into a
+        #: slot whose device token/position belong to whoever held it before.
+        self._needs_decode_install = True
+        #: Set by ``warmup_model_prefill`` so the plugin's two-phase warmup does
+        #: not repeat the prefill sweep (the plugin resets this itself).
+        self.already_warmed_up_prefill = False
+        #: Runtime-fallback bookkeeping, reported by ``serving_audit``.
+        self._audit = {
+            "device_sampled_decode_steps": 0,
+            "host_sampled_decode_steps": 0,
+            "device_sampled_prefills": 0,
+            "host_sampled_prefills": 0,
+            "decode_trace_installs": 0,
+            "top_k_clamped_requests": 0,
+            "penalised_decode_steps": 0,
+            "ignored_seed_requests": 0,
+            #: Steps on which vLLM actually took the async split -- i.e. called
+            #: ``read_decode_output(async_read=True)`` rather than reading the
+            #: device handle synchronously inside ``execute_model``. This is what
+            #: makes ``supports_async_decode`` a measurement instead of a claim;
+            #: see the one-time log line in ``read_decode_output``.
+            "async_decode_reads": 0,
+            "sync_decode_reads": 0,
+        }
+        self._warned: set[str] = set()
+
+        #: Decode graph widths this server may capture, ascending. Read from
+        #: ``QWEN3_DECODE_WIDTHS`` (comma-separated, e.g. ``1,8,32``); anything
+        #: above ``max_num_seqs`` is dropped and ``max_num_seqs`` is always
+        #: present, so the default -- unset -- is exactly the shipped single
+        #: fixed-width graph and nothing below changes behaviour.
+        #:
+        #: Why this exists: expert, router and paged-SDPA cost is paid per row
+        #: **configured**, not per row live (``doc/optimized_vllm/README.md``'s
+        #: control curve: 227.9 ms fixed + 1.28 ms x live_rows at 32 slots). The
+        #: only lever that removes the fixed term is a graph with fewer rows.
+        self._decode_widths = self._configured_widths()
+        #: Graph row -> vLLM slot for the live trace, or ``None`` when the graph
+        #: is full width and the mapping is the identity. Rewritten only on an
+        #: install, which is the only step on which the batch layout may change
+        #: (``model_runner.py`` sets ``reset_batch`` from a sticky
+        #: ``_decode_layout_changed_since_last_decode``).
+        self._compaction: torch.Tensor | None = None
+        #: One entry per decode forward that is still awaiting its host read: the
+        #: graph-row -> vLLM-slot mapping **that forward was issued with**.
+        #:
+        #: Why a queue and not just ``_compaction``. The un-permutation happens in
+        #: ``process_decode_output_host``, which under ``--async-scheduling`` does
+        #: not run in the same step as the forward that produced the tokens: the
+        #: forward returns a device handle and the host read happens later. A
+        #: mapping stored on the adapter can therefore be **rewritten by a later
+        #: install before the earlier step's tokens are scattered**, which would
+        #: put every token on the wrong slot -- silently, as a correctness bug.
+        #:
+        #: The plugin happens to order this safely today: only a layout change can
+        #: rewrite the mapping, and ``model_runner.py`` drains pending async decodes
+        #: whenever the layout changed. But that is an invariant of a *different*
+        #: repository, which this one must not modify and cannot pin with a test.
+        #: Pairing each output with the mapping its own forward used replaces
+        #: that dependency with a weaker and detectable one. It no longer relies
+        #: on drain-on-layout-change; it does still assume the plugin finalizes
+        #: decode steps in issue order and exactly once, which is an invariant of
+        #: the same foreign repository. The difference is that a violation now
+        #: raises (tag mismatch, underflow, or depth cap) instead of silently
+        #: scattering a step's tokens through another step's permutation.
+        self._pending_orders: deque = deque()
+        #: Monotonic id of the next decode forward to be issued, and of the next
+        #: output expected. They are compared on every pop: a queue that has
+        #: skipped or reordered an entry shows up as a tag mismatch rather than
+        #: as tokens quietly landing on the wrong requests.
+        self._next_issue_tag = 0
+        self._next_output_tag = 0
+        #: Hard ceiling on outstanding decode forwards. Async scheduling runs at
+        #: most a step or two ahead; anything approaching this is a leak, not
+        #: depth.
+        self._pending_orders_cap = 64
+        self._audit["narrow_decode_installs"] = 0
+        self._audit["decode_graph_width"] = self.max_num_seqs
+        #: Times the output path found no queued mapping. Must stay 0; a nonzero
+        #: value means forwards and host reads are not paired one-to-one. The
+        #: path raises rather than guessing -- applying the adapter's current
+        #: mapping here would be the exact mis-scatter the queue exists to
+        #: prevent -- so this counter records a raise, not a silent fallback.
+        self._audit["compaction_fifo_underflows"] = 0
+        self._audit["compaction_fifo_max_depth"] = 0
+
+    @property
+    def _compaction_enabled(self) -> bool:
+        """Whether the row-mapping queue is in use at all.
+
+        Derived from ``_decode_widths`` rather than cached, because probes and
+        tests rebind that list after construction to switch the ladder on and
+        off; a cached flag would go stale and silently disable the pairing.
+
+        With the ladder disabled there is no permutation to pair, so nothing is
+        pushed or popped and the shipped path gains neither the bookkeeping nor
+        its failure modes -- in particular it cannot raise the errors below.
+        """
+        return len(self._decode_widths) > 1
+
+    def _reset_pending_orders(self) -> None:
+        """Drop queued mappings whose outputs can no longer be read.
+
+        Called wherever the decode traces are released or replaced. A queued
+        entry refers to a forward whose sampled tokens live in the trace's
+        persistent output tensor; once that trace is gone the handle cannot be
+        read at all, so the entry is dead and keeping it would desync every
+        later pop. The tags are realigned rather than zeroed so the invariant
+        ("the n-th output pairs with the n-th forward") survives the reset.
+        """
+        self._pending_orders.clear()
+        self._next_output_tag = self._next_issue_tag
+
+    #: The ladder used when ``QWEN3_DECODE_WIDTHS`` is unset. Powers of two to
+    #: ``max_num_seqs``: each step runs in the narrowest graph that holds the
+    #: live rows, so the cost a user pays tracks occupancy instead of the slot
+    #: count the server was configured with.
+    #:
+    #: On by default because the alternative default is *known wrong*: a
+    #: ``max_num_seqs=32`` server decodes a single user at 4.3464 t/s/u against
+    #: 49.3636 with the ladder, and a deployment that simply does not set an
+    #: environment variable gets the slow one. That is the same failure shape as
+    #: a missing ``sample_on_device_mode`` -- a config key whose absence looks
+    #: like broken hardware rather than a default.
+    #:
+    #: ``QWEN3_DECODE_WIDTHS=32`` (or any single width equal to ``max_num_seqs``)
+    #: restores the previous fixed-width behaviour exactly.
+    DEFAULT_DECODE_WIDTHS = (1, 2, 4, 8, 16, 32)
+
+    def _configured_widths(self) -> list[int]:
+        """Decode graph widths this server may capture, ascending.
+
+        ``max_num_seqs`` is always present -- a graph can never be wider than
+        the slots the caller sends, and the full-width graph must exist as the
+        fallback -- and anything above it is dropped.
+        """
+        raw = os.getenv("QWEN3_DECODE_WIDTHS", "").strip()
+        source = raw.split(",") if raw else [str(w) for w in self.DEFAULT_DECODE_WIDTHS]
+        widths = {self.max_num_seqs}
+        for piece in source:
+            piece = piece.strip()
+            if not piece:
+                continue
+            value = int(piece)
+            if 1 <= value <= self.max_num_seqs:
+                widths.add(value)
+        return sorted(widths)
+
+    def _choose_width(self, live_rows: int) -> int:
+        """Narrowest configured graph that can hold ``live_rows`` requests."""
+        for width in self._decode_widths:
+            if width >= max(1, live_rows):
+                return width
+        return self.max_num_seqs
+
+    @staticmethod
+    def _compaction_order(host_positions: torch.Tensor, width: int, rows: int) -> torch.Tensor:
+        """Graph row -> vLLM slot, live slots first, then spare slots.
+
+        The live slots go to rows ``0..live-1`` in their original order; the
+        remaining graph rows are filled from the *inactive* slots so that every
+        graph row still names a distinct vLLM slot and therefore still carries a
+        real (zero-filled) page-table row. Those rows install ``current_pos =
+        -1``, which is the inactive sentinel the traced graph already relies on.
+
+        This is a permutation of **only** the three per-row inputs -- position,
+        rotary index and page-table row -- plus the token. No KV page moves: the
+        cache is reached exclusively through page-table entries, so a request's
+        pages are wherever its page-table row says they are, in whatever graph
+        row that row is installed.
+        """
+        live = torch.nonzero(host_positions >= 0, as_tuple=False).reshape(-1)
+        spare = torch.nonzero(host_positions < 0, as_tuple=False).reshape(-1)
+        order = torch.cat((live, spare))[:width]
+        if order.numel() < width:  # fewer vLLM slots than graph rows: cannot happen
+            raise RuntimeError(f"cannot fill a {width}-row graph from {rows} slots")
+        return order.to(torch.int64)
+
+    def _warn_once(self, key: str, message: str) -> None:
+        if key not in self._warned:
+            self._warned.add(key)
+            logger.warning(message)
+
+    # -- scheduler sizing -----------------------------------------------------
+
+    @classmethod
+    def get_max_tokens_all_users(
+        cls,
+        model_name: str = "",
+        num_devices: int = 1,
+        tt_data_parallel: int = 1,
+        max_model_len: int | None = None,
+        max_num_seqs: int | None = None,
+        **kwargs: Any,
+    ) -> int:
+        """Total KV tokens vLLM may allocate blocks for, across all users.
+
+        The whole advertised context for one user. ``config/context_contract.json``
+        records 262144 as the supported context and the paged decode probe that
+        reached position 262143; this is what makes vLLM size enough blocks for
+        a single request to actually use it. At this port's 4 dies the KV cost is
+        512 B per token per layer per die over 48 layers -- 24 KiB per token per
+        die, so 262144 tokens is 6.29 GiB of the 34.18 GiB each die reports.
+
+        The worker adds ``block_size * max_num_seqs`` of its own headroom on top
+        and converts to blocks, so nothing here needs to model that.
+        """
+        supported = _supported_context()
+        return supported if max_model_len is None else min(int(max_model_len), supported)
+
+    # -- vLLM-owned KV cache --------------------------------------------------
+
+    def allocate_kv_cache(self, kv_cache_shape, dtype, num_layers: int):
+        """Allocate the attention KV cache **for vLLM**, at vLLM's geometry.
+
+        ``kv_cache_shape`` is ``(num_blocks, num_kv_heads_per_device, block_size,
+        head_dim)``; the plugin has already divided the head count by the mesh
+        size, which for this port's 4 dies and 4 KV heads gives the 1 local head
+        per die the model expects. The block size is vLLM's, so it is installed
+        into the generator here -- before warmup, before any forward, before any
+        trace -- rather than assumed.
+
+        ``dtype`` is vLLM's torch dtype and is deliberately **not** used: the KV
+        dtype is part of the selected precision policy
+        (``kv_cache_dtype`` in ``selected_precision_config.json``) and serving
+        must not silently run a different one than the sweep measured.
+        """
+        num_blocks, kv_heads, block_size, head_dim = (int(v) for v in kv_cache_shape)
+        if num_layers != self.model.num_layers:
+            if not os.getenv("QWEN3_VLLM_NUM_LAYERS"):
+                raise ValueError(f"vLLM asked for {num_layers} attention layers, model has {self.model.num_layers}")
+            # Reduced bring-up target: vLLM sized blocks for the real depth, the
+            # model only has a few layers. Allocating the model's depth is the
+            # right thing -- it is strictly less memory and the page geometry,
+            # which is what this loop is testing, is unchanged.
+            logger.warning(
+                "Reduced target: allocating {} layer caches, vLLM planned for {}",
+                self.model.num_layers,
+                num_layers,
+            )
+        if kv_heads != self.model.config.local_attention.num_key_value_heads:
+            raise ValueError(
+                f"vLLM computed {kv_heads} local KV heads, this port shards to "
+                f"{self.model.config.local_attention.num_key_value_heads} per die"
+            )
+        if head_dim != self.model.head_dim:
+            raise ValueError(f"vLLM head_dim {head_dim} != model head_dim {self.model.head_dim}")
+
+        pages_per_user = min(math.ceil(self.max_model_len / block_size), num_blocks)
+        self.generator.configure_paging(
+            page_block_size=block_size, pages_per_user=pages_per_user, num_blocks=num_blocks
+        )
+        logger.info(
+            "vLLM-owned KV cache: {} blocks x {} tokens = {} tokens, {} local KV heads, dtype {} "
+            "(from the selected precision config; vLLM asked for {})",
+            num_blocks,
+            block_size,
+            num_blocks * block_size,
+            kv_heads,
+            self.model.precision.kv_cache_dtype,
+            dtype,
+        )
+        self.kv_cache = self.model.allocate_kv_cache(num_blocks=num_blocks)
+        return self.kv_cache
+
+    # -- warmup ---------------------------------------------------------------
+
+    def warmup_model_prefill(self, *, kv_cache, can_sample_on_device: bool, enable_trace: bool, **kwargs: Any) -> None:
+        """Compile the prefill programs on the shapes serving will actually use.
+
+        A prefill program is compiled per sequence length, so what "the shapes
+        serving will use" means depends entirely on whether bucketing is on:
+
+        * **Bucketed** (the default, :func:`prefill_bucket_ladder`) the shape
+          space is the ladder -- 19 rungs at ``pow2_half`` -- and it is finite,
+          so it can be enumerated and warmed. That is the whole reason
+          bucketing exists.
+        * **``QWEN3_PREFILL_BUCKETS=exact``** every distinct prompt length is
+          its own program and no warmup can cover them. This falls back to the
+          two lengths it always warmed, one of them (129) deliberately aligned
+          to nothing: not a multiple of the page block, of a tile, or of any
+          power of two, because the serving path must accept such a length.
+
+        Both arms also warm the **cached-suffix** branch, which is a different
+        program set again: at ``start_pos > 0`` attention switches to
+        ``chunked_scaled_dot_product_attention`` at
+        ``q_chunk_size = k_chunk_size = sdpa_chunk_size(start_pos)``, and that
+        is reachable at only three values -- ``min(256, start & -start)`` over
+        block-aligned starts is ``{64, 128, 256}`` at a 64-token block. Three
+        warm passes therefore cover **every** prefix-cache hit; leaving them
+        cold cost a measured 4.7 s on the first request that hit the cache.
+
+        Prefill is eager (``enable_trace`` is accepted and ignored -- this port
+        has no prefill trace), so the plugin's second warmup phase is a no-op.
+
+        **Depth is gated on the kernel cache, not on the boot.** ``tt-model``
+        bind-mounts a per-model host directory as ``TT_METAL_CACHE`` and keeps
+        it across container removal on purpose (``tt_kernel/container.py``:
+        "the ~10-minute cost the mounted TT_METAL_CACHE exists to avoid paying
+        twice"). Compiles therefore persist, but *running* the ladder does not
+        get cheaper -- a rung is a real prefill, and the default ladder up to
+        the cap is ~26 k tokens, ~26 s at 0.96 ms/token. Paying that once per
+        host is a bargain and once per restart is not, so ``auto`` warms only
+        when its marker is absent from the cache directory. The marker lives
+        *inside* that directory, so clearing the cache clears the marker too
+        and the next boot re-warms.
+
+        ``QWEN3_PREFILL_WARMUP=auto|full|off`` (``full`` ignores the marker,
+        ``off`` skips entirely) and ``QWEN3_PREFILL_WARMUP_MAX`` (rungs above
+        it are left to compile lazily; ``0`` means no cap) are the knobs.
+        """
+        # Also a boot-progress landmark. `tt-model serve` (the kernel package
+        # manager, which packages this model) tracks boot phases by matching
+        # server-log lines, and its "warming up the model" phase starts on
+        # ``Warming up prefill`` / ``Starting decode warmup``
+        # (tt_kernel/boot_progress.py::VLLM_PHASES). Logging only completion
+        # lines left that phase dark for the whole warmup.
+        if self.already_warmed_up_prefill or enable_trace:
+            return
+        self.already_warmed_up_prefill = True
+
+        mode = os.getenv("QWEN3_PREFILL_WARMUP", "auto").strip().lower()
+        if mode not in ("auto", "full", "off"):
+            raise ValueError(f"QWEN3_PREFILL_WARMUP must be auto, full or off; got {mode!r}")
+        if mode == "off":
+            logger.info("Prefill warmup skipped (QWEN3_PREFILL_WARMUP=off)")
+            return
+
+        lengths = self._warmup_prefill_lengths()
+        suffixes = self._warmup_suffix_starts()
+        marker = self._prefill_warmup_marker(lengths, suffixes)
+        if mode == "auto" and marker is not None and marker.exists():
+            # The programs are already in the persistent kernel cache; a first
+            # request loads them rather than compiling them.
+            logger.info(
+                "Prefill warmup skipped -- kernel cache already primed for this ladder ({})",
+                marker,
+            )
+            return
+
+        logger.info(
+            "Warming up prefill ({} lengths {}; {} cached-suffix shapes, chunk sizes {} x suffix rungs {})",
+            len(lengths),
+            list(lengths),
+            len(suffixes),
+            sorted({chunk for _, chunk, _ in suffixes}),
+            sorted({suffix for _, _, suffix in suffixes}),
+        )
+        sampling = self._neutral_sampling_params(1) if can_sample_on_device else None
+        for length in lengths:
+            self.prefill_forward(
+                tokens=torch.zeros((1, length), dtype=torch.int32),
+                page_table=self._warmup_page_table(1, length),
+                kv_cache=kv_cache,
+                enable_trace=False,
+                prompt_lens=[length],
+                start_pos=[0],
+                sampling_params=sampling,
+            )
+        for start, _chunk, suffix in suffixes:
+            # ``suffix`` new tokens on top of a ``start``-token "prefix". The
+            # content is irrelevant -- only the shapes are being compiled -- but
+            # the start must be block-aligned and strictly inside the prompt,
+            # which ``_warmup_suffix_starts`` guarantees.
+            total = start + suffix
+            self.prefill_forward(
+                tokens=torch.zeros((1, total), dtype=torch.int32),
+                page_table=self._warmup_page_table(1, total),
+                kv_cache=kv_cache,
+                enable_trace=False,
+                prompt_lens=[total],
+                start_pos=[start],
+                sampling_params=sampling,
+            )
+        if marker is not None:
+            try:
+                marker.parent.mkdir(parents=True, exist_ok=True)
+                marker.write_text("qwen3-coder-30b-a3b prefill warmup complete\n")
+            except OSError as exc:
+                # A read-only or missing cache directory means the next boot
+                # warms again -- slower, never wrong. Not worth failing a boot.
+                logger.warning("Could not write the prefill warmup marker {}: {}", marker, exc)
+        logger.info("Prefill warmup done ({} lengths + {} cached-suffix shapes)", len(lengths), len(suffixes))
+
+    def _warmup_prefill_lengths(self) -> tuple[int, ...]:
+        """Prefill lengths to warm: the bucket ladder, capped.
+
+        Capped twice over -- by ``QWEN3_PREFILL_WARMUP_MAX`` because a rung is a
+        real prefill and the tail of the ladder is enormous (the top rung is the
+        whole 256 k context, ~4 minutes on its own), and by the KV cache,
+        because a single warm row cannot occupy more blocks than one user is
+        allotted or than the cache holds.
+        """
+        ladder = self.generator.prefill_buckets
+        if not ladder:
+            # ``exact``: nothing to enumerate. 129 is aligned to nothing on
+            # purpose -- see the class docstring.
+            return (129, 128)
+        raw = os.getenv("QWEN3_PREFILL_WARMUP_MAX", "8192").strip()
+        cap = int(raw) if raw else 8192
+        block = self.generator.page_block_size
+        affordable = min(self.generator.pages_per_user, self.generator.num_blocks) * block
+        if cap > 0:
+            # The cap is a promise about PROMPTS, so it has to include the rung
+            # that *covers* it, not merely the rungs below it. Measured the
+            # other way round: at a cap of 8192 an 8192-token prompt rounded to
+            # the 8320 rung, which had been left out, and paid 15.088 s against
+            # 7.610 s once compiled -- the cap silently failing at exactly the
+            # length it named.
+            affordable = min(affordable, self.generator.prefill_padded_len(cap))
+        return tuple(v for v in ladder if v <= affordable)
+
+    def _warmup_suffix_starts(self) -> tuple[tuple[int, int, int], ...]:
+        """``(start_pos, chunk_size, suffix_len)`` for every cached-suffix shape warmed.
+
+        The chunked-SDPA program is keyed on *both* axes, so warming one is not
+        enough:
+
+        * ``chunk_size`` is ``sdpa_chunk_size(start) == min(256, start & -start)``,
+          and a split prefill's ``start`` must be a multiple of the KV block
+          size -- so the reachable set is small and closed (``{64, 128, 256}``
+          at a 64-token block). This enumerates block-aligned starts and keeps
+          the cheapest start reaching each distinct chunk size, rather than
+          hardcoding numbers that were only true at one block size.
+        * ``suffix_len`` is the *bucketed* length of the new tokens, so it is a
+          rung of the prefill ladder -- which is exactly why bucketing had to
+          come first. Unbucketed, this axis is unbounded and the cross product
+          is unwarmable.
+
+        Capped by ``QWEN3_PREFILL_WARMUP_SUFFIX_MAX`` (default 1024) because the
+        cross product is chunk sizes x rungs and only the low rungs are common:
+        a prefix-cache hit's suffix is the *new* tail of a conversation, not a
+        whole prompt. A hit with a longer suffix compiles lazily, once, and the
+        persistent kernel cache keeps it.
+        """
+        block = self.generator.page_block_size
+        affordable = min(self.generator.pages_per_user, self.generator.num_blocks) * block
+        raw = os.getenv("QWEN3_PREFILL_WARMUP_SUFFIX_MAX", "1024").strip()
+        suffix_cap = int(raw) if raw else 1024
+        # 512 is past the point of diminishing returns: chunk sizes saturate at
+        # 256, so no larger start reaches a chunk size a smaller one did not.
+        starts: dict[int, int] = {}
+        start = block
+        while start <= 512:
+            starts.setdefault(sdpa_chunk_size(start), start)
+            start += block
+        rungs = [v for v in (self.generator.prefill_buckets or (block,)) if v <= suffix_cap]
+        out = []
+        for chunk, start in sorted(starts.items()):
+            for rung in rungs:
+                if start + rung <= affordable:
+                    out.append((start, chunk, rung))
+        return tuple(out)
+
+    def _prefill_warmup_marker(self, lengths, suffixes) -> Optional[Path]:
+        """Where ``auto`` records that this ladder has been compiled once.
+
+        Inside the kernel cache directory, and named for everything that would
+        invalidate it: the ladder, the block size, and the mesh shape. If no
+        cache directory is configured there is nothing persistent to key off,
+        so return ``None`` and let ``auto`` behave like ``full``.
+        """
+        root = os.getenv("TT_METAL_CACHE", "").strip()
+        if not root:
+            return None
+        try:
+            shape = "x".join(str(int(v)) for v in tuple(self.generator.mesh_device.shape))
+        except Exception:  # pragma: no cover -- naming detail, never worth a failed boot
+            shape = "unknown"
+        stamp = "-".join(
+            (
+                str(len(lengths)),
+                str(lengths[-1] if lengths else 0),
+                str(len(suffixes)),
+                str(self.generator.page_block_size),
+                shape,
+            )
+        )
+        return Path(root) / f".qwen3_coder_30b_a3b.prefill_warm.v1.{stamp}"
+
+    def warmup_model_decode(
+        self,
+        *,
+        kv_cache,
+        max_batch_size: int,
+        num_blocks: int,
+        can_sample_on_device: bool,
+        enable_trace: bool,
+        **kwargs: Any,
+    ) -> None:
+        """Compile and then capture a decode graph for **every** ladder width.
+
+        Capturing here rather than on the first real token matters twice over:
+        it keeps a multi-second trace capture out of the benchmark's
+        inter-token latency, and it is the phase the plugin gives us for exactly
+        that (phase 1 ``enable_trace=False`` compiles, phase 2 captures).
+
+        **Every rung, not just the serving batch.** A width's *first* capture
+        also pays an eager warm pass to compile its programs -- 1.9-4.8 s
+        against the 713 ms a capture costs once they are cached
+        (``doc/batch_scaling/README.md``, "Cost of switching") -- and
+        ``_decode_compiled_keys`` only spares a width it has already seen. Left
+        lazy, that multi-second compile lands on whichever request first needs
+        the rung, which for a single user against a 32-slot server is the very
+        first one. Warming the ladder moves all of it into the phase built for
+        it, for 24.8 MB of per-width persistent inputs across all six -- 0.11 %
+        of the 21.9 GB/die a maximal prefill leaves free, and already measured
+        with every width captured at the shipped 48 MB trace region.
+
+        **Descending, so width 1 is the one left resident.** Only one trace pair
+        is resident at a time -- ``_refresh_trace_state`` releases on a
+        ``graph_width``, ``active_batch``, ``kv_cache`` or page-table-shape
+        change -- so the last rung warmed is the one a first request can replay
+        for free. Ending at the narrowest makes the single-user case, the one
+        the ladder exists for, cost no capture at all. Each rung is warmed at
+        ``active_batch == graph_width`` and at vLLM's own table width, because
+        both are part of what a release is keyed on: warming a width at some
+        other batch would compile the programs but still discard the trace.
+
+        **Both sampling strategies, greedy last.** ``_decode_graph_key`` covers
+        ``_sampling_stochastic`` as well as the width, and the argmax and
+        split-sampling graphs are different program sets --
+        ``set_sampling_params`` releases the traces when the flag flips. A
+        served model does not choose what its callers send, so both are warmed:
+        measured on 4 dies, an unwarmed strategy cost **1.9416 s** of TTFT on a
+        first request against **0.2989 s** warmed, and the warmed cost holds
+        even for a strategy the server's own default never uses (a
+        Qwen-recommended ``temperature=0.7, top_p=0.8, top_k=20`` request
+        against a greedy-default server: 0.2579 s).
+
+        Greedy is warmed last, so it is the resident one. Order is nearly free
+        -- flipping strategy at width 1 measured 0.2434 s TTFT against 0.2430 s
+        steady, i.e. inside noise -- but greedy is both the faster path
+        (20.35 ms/token, 49.1 t/s/u, against 25.88 ms and 38.6) and what
+        ``tt-model.yaml`` makes the default request, via
+        ``--override-generation-config '{"temperature": 0}'`` on top of
+        ``--generation-config vllm``. Without that override an unparameterised
+        request would arrive at vLLM's own ``temperature=1.0, top_p=1.0``, and
+        ``stochastic = any(k > 1) or any(p > 0.0)`` would put the default path
+        on the slower graph. ``QWEN3_WARMUP_SAMPLING`` trims this to ``greedy``
+        or ``stochastic`` alone when a deployment knows it only serves one.
+
+        Penalties are deliberately *not* warmed. ``_penalty_mode`` is in the key
+        too, but the same ``--generation-config vllm`` is what keeps this
+        checkpoint's ``repetition_penalty=1.05`` from reaching every request, so
+        mode 0 is the steady state and warming a penalised graph would spend
+        1.9-4.8 s per width on a path a default deployment never takes. A server
+        run *without* that flag inverts this, and would want mode 1 warmed
+        instead. ``rope_cache_len`` is stable (``decode_horizon`` is the served
+        context, sized at construction) and ``active_row_gating`` is fixed for
+        the process, so neither adds a dimension here.
+        """
+        if not can_sample_on_device:
+            # Host-sampled decode is eager by construction; nothing to capture.
+            return
+        batch = min(int(max_batch_size), self.max_num_seqs)
+        # Descending, and clipped to what the caller can actually send: a graph
+        # wider than ``batch`` is unreachable this run, and warming it would
+        # compile programs no request can replay. ``batch`` itself is always
+        # warmed -- it is the fallback rung ``_configured_widths`` guarantees --
+        # so with the ladder off this is exactly the previous single warmup.
+        widths = sorted({min(w, batch) for w in self._decode_widths}, reverse=True)
+        strategies = self._warmup_sampling_strategies()
+        # Boot-progress landmark as well as a plan: see warmup_model_prefill.
+        logger.info(
+            "Starting decode warmup ({} graphs: widths {} descending x sampling {})",
+            len(widths) * len(strategies),
+            widths,
+            ["stochastic" if st else "greedy" for st in strategies],
+        )
+        for stochastic in strategies:
+            for graph_width in widths:
+                page_table = self._warmup_page_table(graph_width, self.generator.page_block_size, width=int(num_blocks))
+                self.decode_forward(
+                    tokens=torch.zeros((graph_width, 1), dtype=torch.int32),
+                    page_table=page_table,
+                    kv_cache=kv_cache,
+                    start_pos=torch.zeros(graph_width, dtype=torch.int64),
+                    enable_trace=enable_trace,
+                    read_from_device=True,
+                    sampling_params=self._warmup_sampling_params(graph_width, stochastic=stochastic),
+                    reset_batch=True,
+                )
+        # Each warmup step wrote a token at position 0 of every row it used and
+        # advanced the device positions; the first real request must not inherit
+        # that, and must re-decide its own width against a real batch layout.
+        self._needs_decode_install = True
+        self._reset_pending_orders()
+        self._audit["decode_graph_width"] = widths[-1]
+        logger.info(
+            "Decode warmup done ({} widths {} descending x sampling {}, resident {} {}, enable_trace={})",
+            len(widths) * len(strategies),
+            widths,
+            ["stochastic" if st else "greedy" for st in strategies],
+            "stochastic" if strategies[-1] else "greedy",
+            widths[-1],
+            enable_trace,
+        )
+
+    def _warmup_sampling_strategies(self) -> list[bool]:
+        """Sampling strategies to warm, in warm order -- the last stays resident.
+
+        ``False`` is argmax/greedy, ``True`` is split sampling. Both by default,
+        because a served model does not get to choose what its callers send and
+        an uncompiled strategy costs seconds on a real request -- measured on 4
+        dies, an unwarmed strategy put 1.9416 s on a first request's TTFT
+        against 0.2989 s warmed.
+
+        **Greedy is warmed last, so it is the one left resident.** Order is a
+        near-free choice rather than a load-bearing one: switching strategy at
+        width 1 releases and re-captures, and that cost was unmeasurable
+        (0.2434 s TTFT switching against 0.2430 s steady). Greedy gets the slot
+        because it is both the faster path (20.34 ms/token against 25.88 ms,
+        49.1 t/s/u against 38.6) and what ``tt-model.yaml`` makes the default
+        request via ``--override-generation-config '{"temperature": 0}'``.
+
+        ``QWEN3_WARMUP_SAMPLING=greedy|stochastic`` warms one only, for a
+        deployment that knows which it serves; anything else is rejected rather
+        than silently ignored, since a typo would quietly cost that same
+        multi-second compile on a real request.
+        """
+        choice = os.getenv("QWEN3_WARMUP_SAMPLING", "both").strip().lower()
+        if choice in ("", "both"):
+            return [True, False]
+        if choice == "greedy":
+            return [False]
+        if choice == "stochastic":
+            return [True]
+        raise ValueError(f"QWEN3_WARMUP_SAMPLING must be both, greedy or stochastic; got {choice!r}")
+
+    def _warmup_sampling_params(self, rows: int, *, stochastic: bool):
+        """Neutral params for one warm pass, on the requested sampling strategy.
+
+        The stochastic arm mirrors what vLLM sends for an unparameterised
+        request under ``--generation-config vllm`` -- ``temperature=1.0``,
+        ``top_p=1.0`` -- with ``top_k`` already at the device limit rather than
+        vLLM's "disabled" 0, which ``_apply_sampling_params`` would clamp to the
+        same 32 while incrementing ``top_k_clamped_requests`` and warning. The
+        graph depends only on the boolean, so the values beyond that are
+        irrelevant to what gets compiled.
+        """
+        if not stochastic:
+            return self._neutral_sampling_params(rows)
+        from vllm_tt_plugin.model_input import TTSamplingParams
+
+        return TTSamplingParams(
+            temperature=[1.0] * rows,
+            top_k=[SAMPLING_SLOTS] * rows,
+            top_p=[1.0] * rows,
+            seed=[None] * rows,
+        )
+
+    def _warmup_page_table(self, batch: int, token_count: int, *, width: int | None = None) -> torch.Tensor:
+        """A disjoint block assignment for warmup only, at vLLM's table width."""
+        width = self.generator.pages_per_user if width is None else int(width)
+        blocks = max(1, math.ceil(token_count / self.generator.page_block_size))
+        table = torch.zeros((batch, width), dtype=torch.int32)
+        for row in range(batch):
+            span = min(blocks, width)
+            table[row, :span] = torch.arange(row * span, row * span + span, dtype=torch.int32)
+        return table
+
+    def _neutral_sampling_params(self, rows: int):
+        from vllm_tt_plugin.model_input import TTSamplingParams
+
+        return TTSamplingParams(
+            temperature=[0.0] * rows,
+            top_k=[1] * rows,
+            top_p=[1.0] * rows,
+            # The plugin translates its own "no seed" sentinel to ``None`` before
+            # the model sees it; the dataclass default of ``0`` is a real seed.
+            seed=[None] * rows,
+        )
+
+    # -- sampling translation -------------------------------------------------
+
+    def _apply_sampling_params(self, sampling_params, rows: int, *, order=None, graph_rows: int | None = None) -> None:
+        """vLLM's per-row sampling request -> the generator's ``(k, p, temp)``.
+
+        Nothing here samples. It only sets the three persistent device parameter
+        tensors that ``Qwen3CoderModel.sample_split`` reads, and only when they
+        actually changed -- ``set_sampling_params`` no-ops on an identical
+        snapshot, so a steady greedy benchmark costs zero host copies per token.
+        """
+        temps = _as_int_list(getattr(sampling_params, "temperature", None), rows, 0.0)
+        top_ks = _as_int_list(getattr(sampling_params, "top_k", None), rows, 1)
+        top_ps = _as_int_list(getattr(sampling_params, "top_p", None), rows, 1.0)
+        self._audit_unsupported(sampling_params, rows)
+        # ``ttnn.sampling``'s per-slot parameters address the *graph*'s rows, so
+        # a compacted batch must present them in the same order the rows are in.
+        if order is not None:
+            picks = [int(v) for v in order.tolist()]
+            temps = [temps[i] for i in picks]
+            top_ks = [top_ks[i] for i in picks]
+            top_ps = [top_ps[i] for i in picks]
+        rows = rows if graph_rows is None else int(graph_rows)
+
+        k_out: list[int] = []
+        p_out: list[float] = []
+        t_out: list[float] = []
+        for row in range(rows):
+            temperature = float(temps[row])
+            top_k = int(top_ks[row])
+            top_p = float(top_ps[row])
+            if temperature <= 0.0:
+                # Greedy. The generator maps temperature 0 to k=1, p=0 itself and
+                # then routes to the argmax strategy.
+                k_out.append(1)
+                p_out.append(0.0)
+                t_out.append(0.0)
+                continue
+            if top_k <= 0 or top_k > SAMPLING_SLOTS:
+                # vLLM spells "no top-k" as <=0 and allows any k up to the
+                # vocabulary; ``Sampling1DConfig(max_top_k=32)`` is a device
+                # limit, so both collapse to the widest supported candidate set.
+                if top_k > SAMPLING_SLOTS or top_k <= 0:
+                    self._audit["top_k_clamped_requests"] += 1
+                    self._warn_once(
+                        "top_k",
+                        f"top_k={top_k} clamped to {SAMPLING_SLOTS}: the on-device sampler's "
+                        "max_top_k is 32 candidates per die-gathered slot.",
+                    )
+                top_k = SAMPLING_SLOTS
+            k_out.append(top_k)
+            p_out.append(min(max(top_p, 0.0), 1.0))
+            t_out.append(temperature)
+        self.generator.set_sampling_params(top_k=k_out, top_p=p_out, temperature=t_out, active_batch=rows)
+
+    def _apply_penalties(
+        self, sampling_params, rows: int, prompt_tokens, output_tokens, *, order=None, graph_rows: int | None = None
+    ) -> None:
+        """vLLM's three penalties -> the generator's staged on-device penalty stage.
+
+        The plugin packs ``presence_penalty`` / ``frequency_penalty`` /
+        ``repetition_penalty`` into ``TTSamplingParams`` and sends the token
+        history alongside them (``model_runner.py`` populates ``prompt_tokens``
+        and ``output_tokens`` "if penalties are needed (decode only)"), because
+        ``platform.py`` deliberately does **not** route penalised requests to host
+        sampling. This is the model side of that contract; the stage itself is
+        ``_WatcherCleanSampling1D._apply_penalties``.
+
+        Neutral on every row is the fast path: ``set_penalty_params`` returns
+        False, the ops are not in the captured trace at all, and nothing is
+        uploaded.
+        """
+        presence = _as_int_list(getattr(sampling_params, "presence_penalty", None), rows, 0.0)
+        frequency = _as_int_list(getattr(sampling_params, "frequency_penalty", None), rows, 0.0)
+        repetition = _as_int_list(getattr(sampling_params, "repetition_penalty", None), rows, 1.0)
+        if order is not None:
+            picks = [int(v) for v in order.tolist()]
+            presence = [presence[i] for i in picks]
+            frequency = [frequency[i] for i in picks]
+            repetition = [repetition[i] for i in picks]
+            # The staged penalty rows are per *graph* row too, and the history
+            # they are keyed on has to travel with them.
+            prompt_tokens = _reorder_history(prompt_tokens, order, picks)
+            output_tokens = _reorder_history(output_tokens, order, picks)
+        live, graph_changed = self.generator.set_penalty_params(
+            presence=presence,
+            frequency=frequency,
+            repetition=repetition,
+            prompt_tokens=prompt_tokens,
+            output_tokens=output_tokens,
+            active_batch=rows if graph_rows is None else int(graph_rows),
+        )
+        if graph_changed:
+            # The mode flip released the decode traces; the next step must
+            # reinstall host state rather than replay a freed trace.
+            self._needs_decode_install = True
+            # Any queued mapping refers to a forward whose output tensor the
+            # released trace owned, so those outputs can no longer be read.
+            self._reset_pending_orders()
+        if live:
+            self._audit["penalised_decode_steps"] += 1
+
+    def _audit_unsupported(self, sampling_params, rows: int) -> None:
+        """Record -- loudly, once -- the request features this sampler drops."""
+        seeds = _as_int_list(getattr(sampling_params, "seed", None), rows, None)
+        if any(s is not None for s in seeds):
+            self._audit["ignored_seed_requests"] += 1
+            self._warn_once(
+                "seed",
+                "A per-request seed was supplied but this port's sampler draws from its own device "
+                "RNG buffer; sampled output is not reproducible from the request seed. See "
+                "doc/vllm_integration/README.md, Limitations.",
+            )
+
+    # -- prefill --------------------------------------------------------------
+
+    def prefill_forward(
+        self,
+        *,
+        tokens: torch.Tensor,
+        page_table: torch.Tensor,
+        kv_cache,
+        enable_trace: bool = False,
+        prompt_lens=None,
+        start_pos=None,
+        sampling_params=None,
+        empty_slots=None,
+        page_tables_per_layer=None,
+        **kwargs: Any,
+    ):
+        """One serving prefill step, straight into ``generator.prefill_forward``.
+
+        ``tokens`` is ``[num_reqs, max(prompt_lens)]`` with each row's real
+        length in ``prompt_lens`` and *garbage past it* -- vLLM slices a shared
+        buffer. The generator prefills each row at exactly its own logical
+        length (``tokens[user, :prompt_len]``), so a prompt length that is not a
+        multiple of the page block, the tile height or any chunk size needs no
+        special case: nothing rounds it up on the way in and the selected row is
+        ``prompt_len - 1``.
+        """
+        if page_tables_per_layer is not None:
+            raise ValueError("this port has one uniform full-attention KV-cache group; per-layer tables are not used")
+        active = int(tokens.shape[0])
+        starts = _as_int_list(start_pos, active, 0) if start_pos is not None else [0] * active
+        if any(int(p) != 0 for p in starts) and not _PREFIX_CACHING_ENABLED:
+            raise ValueError(
+                "non-zero prefill start_pos means prefix caching, but it has been "
+                "disabled via QWEN3_PREFIX_CACHING=0 while model_capabilities still "
+                "advertises supports_prefix_caching=True. Those two must agree: either "
+                "unset the kill switch or set supports_prefix_caching=False."
+            )
+        lengths = [int(n) for n in _as_int_list(prompt_lens, active, int(tokens.shape[1]))]
+        device_sampling = sampling_params is not None
+        if device_sampling:
+            self._apply_sampling_params(sampling_params, active)
+            self._audit["device_sampled_prefills"] += 1
+        else:
+            self._audit["host_sampled_prefills"] += 1
+
+        out = self.generator.prefill_forward(
+            tokens.to(torch.int64),
+            page_table=self._page_table_for_generator(page_table, active),
+            kv_cache=self._require_cache(kv_cache),
+            prompt_lens=lengths,
+            sampling_mode="device" if device_sampling else "host",
+            # A new request is admitted while other slots are mid-decode; see
+            # the argument in ``Qwen3CoderGenerator.prefill_forward``, and the
+            # measurement behind this default in
+            # ``doc/vllm_integration/work_log.md``.
+            preserve_decode_traces=PRESERVE_DECODE_TRACES,
+            start_pos=starts if _PREFIX_CACHING_ENABLED else None,
+        )
+        # Whoever held these slots before is gone; the next decode must reinstall
+        # host state rather than replay over the device's stale token/position.
+        self._needs_decode_install = True
+        self._reset_pending_orders()
+        if device_sampling:
+            return self.generator.read_sampled_tokens(out, active).reshape(active, 1)
+        # Host-sampling compatibility mode: vLLM wants ``[B, S, vocab]`` and
+        # reads ``[:, -1, :]``. The generator already returns one row per user.
+        return out.reshape(active, 1, -1)
+
+    # -- decode ---------------------------------------------------------------
+
+    def decode_forward(
+        self,
+        *,
+        tokens: torch.Tensor,
+        page_table: torch.Tensor,
+        kv_cache,
+        start_pos,
+        enable_trace: bool = True,
+        read_from_device: bool = True,
+        sampling_params=None,
+        reset_batch: bool | None = None,
+        slot_remap=None,
+        prompt_tokens=None,
+        output_tokens=None,
+        page_tables_per_layer=None,
+        **kwargs: Any,
+    ):
+        """One serving decode step.
+
+        The batch is always the full ``max_num_seqs`` rows -- vLLM pads it so the
+        trace shape is constant -- with inactive slots carrying position ``-1``,
+        which is the same inactive-row convention the generator's low-level API
+        already had.
+
+        Steady state is the whole point: ``tokens``, ``start_pos`` and
+        ``page_table`` are all passed as ``None``/unchanged, the two traces
+        replay non-blocking, the sampled token is fed back on device and both
+        position tensors advance on device. Host work per token is two
+        ``ttnn.execute_trace`` calls and one page-table equality check.
+        """
+        if page_tables_per_layer is not None:
+            raise ValueError("this port has one uniform full-attention KV-cache group; per-layer tables are not used")
+        caches = self._require_cache(kv_cache)
+        rows = int(tokens.shape[0])
+        host_tokens = tokens.reshape(-1).to(torch.int64)
+        host_positions = torch.as_tensor(start_pos).reshape(-1).to(torch.int64)
+
+        if sampling_params is None:
+            # Explicit host-sampling compatibility mode. vLLM routes a request
+            # here on its own (logprobs on a 4-die mesh, min_p, bad_words,
+            # logit_bias, structured output); it is never the measured path.
+            # The eager decode allocates, so ``decode_forward`` releases the
+            # captured decode traces -- hence ``_needs_decode_install`` below,
+            # which makes the next device-sampled step re-capture through
+            # ``_refresh_trace_state`` instead of replaying a released trace.
+            self._audit["host_sampled_decode_steps"] += 1
+            # Make the demotion loud. vLLM decides this per request and logs
+            # nothing, so without this line a served request silently drops from
+            # ~49 t/s/u to ~3.6 and the only visible symptom is that the model
+            # "got slow" -- which is exactly how this port's batch-scaling defect
+            # was first reported. The server-level ``sample_on_device_mode: all``
+            # is still correct and still says nothing about it.
+            self._warn_once(
+                "host_sampled_decode",
+                "This request was routed to HOST sampling by vLLM, so decode runs "
+                "eager with no captured trace and no width compaction: measured "
+                "3.595 t/s/u against 49.345 on the traced path, a ~14x slowdown "
+                "for the affected requests. On this 4-die mesh the usual cause is "
+                "`logprobs` -- ANY value including 0, because "
+                "`model_runner.check_perform_device_sampling` tests "
+                "`max_num_logprobs is not None` and then rejects a mesh that is "
+                "not 8 or 32 dies. Other triggers are min_p, bad_words, "
+                "logit_bias and structured output. Drop the offending parameter "
+                "to stay on the traced path; see doc/batch_scaling/README.md, "
+                "'logprobs silently cost 14x on this mesh'.",
+            )
+            self._needs_decode_install = True
+            self._reset_pending_orders()
+            logits = self.generator.decode_forward(
+                host_tokens,
+                torch.clamp(host_positions, min=0),
+                page_table=self._page_table_for_generator(page_table, rows),
+                kv_cache=caches,
+                sampling_mode="host",
+                enable_trace=False,
+                active_batch=rows,
+                validate_page_coverage=False,
+            )
+            return logits.reshape(rows, 1, -1)
+
+        self._audit["device_sampled_decode_steps"] += 1
+        install = bool(reset_batch) or self._needs_decode_install
+
+        # The graph width, and with it the graph-row -> vLLM-slot mapping, may
+        # only change on an install: that is the one step the plugin guarantees
+        # is not steady-decode eligible, so nothing is in flight against the old
+        # trace. On every other step the previous mapping still describes the
+        # live trace and is reused unchanged.
+        previous_order = self._compaction
+        if install and len(self._decode_widths) > 1:
+            # ``rows`` is the padded decode batch, normally ``max_num_seqs``; a
+            # graph can never be wider than the slots the caller actually sent.
+            chosen = min(self._choose_width(int((host_positions >= 0).sum())), rows)
+            order = self._compaction_order(host_positions, chosen, rows)
+            identity = chosen == rows and bool(torch.equal(order, torch.arange(rows)))
+            self._compaction = None if identity else order
+            if chosen < rows:
+                self._audit["narrow_decode_installs"] += 1
+            self._audit["decode_graph_width"] = chosen
+        width = rows if self._compaction is None else int(self._compaction.numel())
+        order = self._compaction
+        # With no extra widths configured this stays ``None`` and the generator
+        # keeps its own default -- the full configured slot count -- so the
+        # shipped path is untouched down to which graph gets captured.
+        requested_width = None if self._compaction is None and len(self._decode_widths) == 1 else width
+
+        self._apply_sampling_params(sampling_params, rows, order=order, graph_rows=width)
+        # Before the trace is touched: a penalty-mode change releases the decode
+        # traces (the ops either are or are not in the captured graph), and the
+        # buffers it may allocate cannot be allocated during a capture.
+        self._apply_penalties(sampling_params, rows, prompt_tokens, output_tokens, order=order, graph_rows=width)
+        # A penalty-mode change releases the traces, so it forces an install even
+        # when the scheduler layout did not move. Re-read the flag rather than
+        # trusting the value taken before the call; the width decision above is
+        # unaffected, because the batch layout is what picks the width and that
+        # has not changed.
+        install = install or self._needs_decode_install
+
+        if install:
+            merged_tokens, merged_positions = self._merge_scheduler_view(
+                host_tokens, host_positions, page_table, slot_remap, rows, previous_order
+            )
+            if order is not None:
+                merged_tokens = merged_tokens[order]
+                merged_positions = merged_positions[order]
+            sampled = self.generator.decode_forward(
+                merged_tokens,
+                merged_positions,
+                page_table=self._compact_page_table(page_table, rows, order),
+                kv_cache=caches,
+                sampling_mode="device",
+                enable_trace=True,
+                active_batch=width,
+                graph_width=requested_width,
+                # Sized once at construction to the served context, so this only
+                # asserts the horizon rather than growing anything.
+                decode_horizon=self.max_model_len,
+                # vLLM's block tables are its own: rows of an unused slot are
+                # zero-filled rather than -1, so the standalone disjointness
+                # check does not describe them.
+                validate_page_coverage=False,
+            )
+            self._needs_decode_install = False
+            self._audit["decode_trace_installs"] += 1
+        else:
+            sampled = self.generator.decode_forward(
+                None,
+                None,
+                page_table=self._compact_page_table(page_table, rows, order),
+                kv_cache=caches,
+                sampling_mode="device",
+                enable_trace=True,
+                active_batch=width,
+                graph_width=requested_width,
+            )
+
+        # Pair this forward's tokens with the mapping it was issued with, before
+        # anything can read them back. ``order`` is ``None`` at full width, which
+        # is a meaningful entry: it says "this step needs no un-permutation".
+        if self._compaction_enabled:
+            self._pending_orders.append((self._next_issue_tag, order))
+            self._next_issue_tag += 1
+            depth = len(self._pending_orders)
+            self._audit["compaction_fifo_max_depth"] = max(self._audit["compaction_fifo_max_depth"], depth)
+            if depth > self._pending_orders_cap:
+                # Outputs are being issued and never read: the queue is leaking.
+                # Fail here rather than let it grow unbounded and mis-pair later.
+                raise RuntimeError(
+                    f"decode row-mapping queue reached {depth} entries (cap {self._pending_orders_cap}). "
+                    "Decode forwards are being issued without their outputs being read, so the "
+                    "mapping queue no longer tracks in-flight steps. See "
+                    "Qwen3CoderForCausalLM._pending_orders."
+                )
+        if read_from_device:
+            return self.process_decode_output_host(sampled, is_tokens=True)
+        return sampled
+
+    def _merge_scheduler_view(
+        self,
+        host_tokens: torch.Tensor,
+        host_positions: torch.Tensor,
+        page_table: torch.Tensor,
+        slot_remap,
+        rows: int,
+        previous_order: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Reinstall host state, but keep the device's where the device is right.
+
+        Called only on a layout change, never per token. For a slot that is
+        simply continuing, the device already advanced ``current_pos`` past the
+        token it just sampled, so ``device_pos`` equals the scheduler's position
+        (synchronous scheduling) or is one ahead of it (``--async-scheduling``,
+        where vLLM submits step *N+1* before it has applied token *N*). Taking
+        the device's pair in both cases is what makes async scheduling safe:
+        the host's token would be stale and its position would re-decode a
+        position that is already in the cache.
+
+        A slot that changed hands must take the host's pair. Position continuity
+        alone cannot tell the two apart -- a recycled slot can coincidentally
+        land on a matching position -- so this also requires the slot's page-table
+        row to be byte-identical to the one the live trace was captured against.
+        A newly admitted request is given fresh physical blocks, so its row moves.
+        """
+        state = self.generator.decode_device_state()
+        if state is None or state["page_table"] is None:
+            return host_tokens, host_positions
+
+        # The live trace's rows are *graph* rows. ``previous_order`` says which
+        # vLLM slot each one held; scatter them back into slot order before any
+        # comparison with the scheduler's view, and leave slots the narrow graph
+        # did not cover at the inactive sentinel so they are never "continuing".
+        if previous_order is None:
+            device_tokens = state["tokens"][:rows].clone()
+            device_positions = state["positions"][:rows].clone()
+            snapshot = state["page_table"][:rows]
+        else:
+            covered = previous_order[: state["width"]]
+            device_tokens = torch.zeros(rows, dtype=torch.int64)
+            device_positions = torch.full((rows,), -1, dtype=torch.int64)
+            device_tokens[covered] = state["tokens"][: covered.numel()]
+            device_positions[covered] = state["positions"][: covered.numel()]
+            snapshot = torch.zeros((rows, state["page_table"].shape[1]), dtype=state["page_table"].dtype)
+            snapshot[covered] = state["page_table"][: covered.numel()]
+        incoming = torch.as_tensor(page_table).to(torch.int32)[:rows]
+        width = min(snapshot.shape[1], incoming.shape[1])
+
+        if slot_remap is not None:
+            remap = torch.as_tensor(slot_remap).reshape(-1)[:rows].to(torch.int64)
+            device_tokens = device_tokens[remap]
+            device_positions = device_positions[remap]
+            snapshot = snapshot[remap]
+
+        pages_unchanged = torch.all(snapshot[:, :width] == incoming[:, :width], dim=1)
+        continuing = (
+            ((device_positions == host_positions) | (device_positions == host_positions + 1))
+            & (device_positions >= 0)
+            & (host_positions >= 0)
+            & pages_unchanged
+        )
+        merged_tokens = torch.where(continuing, device_tokens, host_tokens)
+        # ``host_positions`` is the scheduler's view, and the plugin pads rows it
+        # is not serving with ``-1`` (``model_runner.py`` pads decode positions
+        # with ``-1`` "to indicate no position"). That ``-1`` is exactly the
+        # inactive sentinel the traced graph relies on: ``ttnn.plus_one(...,
+        # skip_negative_entries=True)`` leaves it alone across replays, and
+        # ``_decode_active_mask`` derives the expert-gating mask from
+        # ``current_pos >= 0``.
+        #
+        # Clamping it to 0 here would install an inactive row as position 0, the
+        # mask would read it as live, and inactive-row expert gating would
+        # silently become a no-op for that slot until the next prefill released
+        # the traces. Single-request runs never expose it -- every row is either
+        # continuing or genuinely live -- but a server churning 4 of 32 slots
+        # would see the gating win appear and disappear with request turnover.
+        # So preserve the sentinel and only clamp what is not already a sentinel.
+        host_positions_kept = torch.where(host_positions < 0, torch.full_like(host_positions, -1), host_positions)
+        merged_positions = torch.where(continuing, device_positions, host_positions_kept)
+        return merged_tokens, merged_positions
+
+    # -- async split ----------------------------------------------------------
+
+    def read_decode_output(self, tt_out, async_read: bool = False):
+        """Deferred, minimal host read of the sampled-token tensor.
+
+        The payload is one ``[1,1,1,32]`` uint32 tensor -- 128 bytes -- because
+        the token was sampled on device. There is no logits readback here and
+        there is nothing else to move.
+        """
+        if isinstance(tt_out, torch.Tensor):
+            # Host-sampling compatibility mode already returned host logits.
+            return tt_out, []
+        if not async_read:
+            return tt_out.cpu()
+        self._audit["async_decode_reads"] += 1
+        self._warn_once(
+            "async_split",
+            "vLLM took the async decode split: read_decode_output(async_read=True) on a device "
+            "handle returned by decode_forward(read_from_device=False). supports_async_decode is "
+            "being exercised, not merely declared.",
+        )
+        host = tt_out.cpu(blocking=False)
+        return host, [ttnn.record_event(self.mesh_device, 0)]
+
+    def process_decode_output_host(self, tt_out, is_tokens: bool = False):
+        """Host formatting only; submits no device work.
+
+        Accepts the device handle (the plugin's synchronous path skips
+        ``read_decode_output`` entirely), the host ttnn tensor from the async
+        path, or an already-host torch tensor.
+        """
+        if isinstance(tt_out, torch.Tensor):
+            return tt_out
+        if not is_tokens:
+            raise ValueError("host-sampled decode already returns torch logits; nothing to format")
+        if ttnn.is_tensor_storage_on_device(tt_out):
+            # Only a *device*-resident handle means the plugin skipped
+            # ``read_decode_output`` -- its synchronous path -- so the readback
+            # happens now, inside ``execute_model``, rather than after the async
+            # boundary. Counted so the async/sync split is evidence, not prose.
+            #
+            # The discriminator has to be device residency, not
+            # ``not isinstance(tt_out, torch.Tensor)``: the torch case already
+            # returned above, so that test was dead and fired on every step, and
+            # the async path's ``read_decode_output`` hands us
+            # ``tt_out.cpu(blocking=False)`` -- a ttnn *host* tensor, not a
+            # ``torch.Tensor`` -- so async reads were being counted as sync.
+            self._audit["sync_decode_reads"] += 1
+        tokens = self.generator.read_sampled_tokens(tt_out, self.max_num_seqs)
+        # Take the mapping belonging to *this* output, not whatever the adapter
+        # currently holds -- see ``_pending_orders``. FIFO is the right pairing
+        # because decode forwards are finalized in issue order.
+        order = None
+        if self._compaction_enabled:
+            if not self._pending_orders:
+                # There is no safe answer here. Falling back to the adapter's
+                # current mapping is precisely the bug the queue exists to
+                # prevent, and it would be applied in the one state where the
+                # pairing is known to be broken -- every token would go to the
+                # wrong request, silently. A crash is strictly better.
+                self._audit["compaction_fifo_underflows"] += 1
+                raise RuntimeError(
+                    "decode output arrived with no queued row mapping. Forwards and host reads are "
+                    "no longer paired one-to-one, so the sampled tokens cannot be attributed to "
+                    "requests. Refusing to scatter them through a mapping that is not theirs. See "
+                    "Qwen3CoderForCausalLM._pending_orders."
+                )
+            tag, order = self._pending_orders.popleft()
+            if tag != self._next_output_tag:
+                raise RuntimeError(
+                    f"decode row-mapping queue is out of step: popped tag {tag}, expected "
+                    f"{self._next_output_tag}. An output has been skipped or read twice, so this "
+                    "mapping does not belong to these tokens. See "
+                    "Qwen3CoderForCausalLM._pending_orders."
+                )
+            self._next_output_tag += 1
+        if order is not None:
+            # Graph row *i* sampled for vLLM slot ``order[i]``. Scatter back;
+            # slots the narrow graph did not cover hold no live request and vLLM
+            # discards whatever is there.
+            restored = torch.zeros(self.max_num_seqs, dtype=tokens.dtype)
+            restored[order] = tokens[: order.numel()]
+            tokens = restored
+        return tokens.reshape(-1, 1)
+
+    # -- helpers --------------------------------------------------------------
+
+    def _require_cache(self, kv_cache):
+        """The cache vLLM allocated, and only that one.
+
+        The generator would happily allocate its own on a ``None``; in serving
+        that would be a silent second cache that vLLM's block manager knows
+        nothing about, so it is an error instead.
+        """
+        cache = self.kv_cache if kv_cache is None else kv_cache
+        if cache is None:
+            raise RuntimeError("vLLM has not called allocate_kv_cache; there is no serving cache to use")
+        if isinstance(cache, (list, tuple)) and cache and isinstance(cache[0], (list, tuple)):
+            raise ValueError("this port is single-submesh; a per-submesh cache list is not expected")
+        return cache
+
+    def _compact_page_table(self, page_table, rows: int, order) -> torch.Tensor:
+        """vLLM's block table, reordered into graph-row order.
+
+        The page table is the *only* thing that ties a request to its KV pages,
+        so permuting its rows is what moves a request between graph rows -- and
+        it is why nothing in the cache has to move.
+        """
+        table = self._page_table_for_generator(page_table, rows)
+        return table if order is None else table[order].contiguous()
+
+    def _page_table_for_generator(self, page_table, rows: int) -> torch.Tensor:
+        """vLLM's block table at the generator's table width.
+
+        vLLM sizes its table to ``max_num_blocks_per_req``; ``configure_paging``
+        already made that the generator's width, so this is normally a no-op.
+        Where it is not, pad with **0** rather than the generator's standalone
+        ``-1``: the paged decode SDPA kernel rounds its read up to a tile/eight
+        page boundary and dereferences every page in the rounded window before
+        causal masking, so a tail page must map somewhere valid. vLLM pads its
+        own unused entries with 0 for the same reason.
+        """
+        table = torch.as_tensor(page_table).to(torch.int32)
+        if table.ndim != 2:
+            raise ValueError(f"page_table must be rank two, got {tuple(table.shape)}")
+        target = self.generator.pages_per_user
+        if table.shape[1] < target:
+            table = torch.nn.functional.pad(table, (0, target - table.shape[1]), value=0)
+        elif table.shape[1] > target:
+            table = table[:, :target]
+        if table.shape[0] < rows:
+            table = torch.nn.functional.pad(table, (0, 0, 0, rows - table.shape[0]), value=0)
+        return table.contiguous()
+
+    # -- audit ----------------------------------------------------------------
+
+    def serving_audit(self) -> dict:
+        """What the serving path actually did, for the stage's fallback audit."""
+        audit = dict(self._audit)
+        audit["trace_stats"] = dict(self.generator.trace_stats)
+        audit["precision_config"] = str(SELECTED_PRECISION_CONFIG)
+        audit["max_model_len"] = self.max_model_len
+        audit["max_num_seqs"] = self.max_num_seqs
+        audit["page_block_size"] = self.generator.page_block_size
+        audit["pages_per_user"] = self.generator.pages_per_user
+        audit["kv_cache_blocks"] = self.generator.num_blocks
+        audit["model_runtime_fallbacks"] = self.model.runtime_fallback_audit(self.max_num_seqs)
+        return audit
+
+
+#: The architecture string in this checkpoint's ``config.json`` is
+#: ``Qwen3MoeForCausalLM``; the TT plugin registers models ``TT``-prefixed.
+HF_ARCHITECTURE = "Qwen3MoeForCausalLM"
+VLLM_ARCHITECTURE = "TT" + HF_ARCHITECTURE
+
+__all__ = ["Qwen3CoderForCausalLM", "HF_ARCHITECTURE", "VLLM_ARCHITECTURE", "HF_MODEL_ID"]

code/models/demos/blackhole/qwen3_coder_30b_a3b/tt/model.py

@@ -0,0 +1,1723 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+"""Full 48-layer Qwen3-Coder-30B-A3B-Instruct on the 4-die P300_X2 mesh.
+
+Stage 05. This module is the *wrapper* around the stage-04 optimized multichip
+decoder layer and it deliberately changes nothing about that layer's strategy:
+
+* attention TP=4 (8 Q heads, 1 K head, 1 V head per die), experts EP=4 (32 of
+  128 per die), router and both residual RMSNorms and the residual replicated;
+* two all-reduces per layer, ``FABRIC_1D_RING``, 2 links prefill / 1 decode;
+* expert weights ``bfloat4_b`` at LoFi with ``in0_block_w`` 16/12, attention
+  projections ``bfloat8_b`` DRAM-sharded, paged KV cache ``bfloat16``;
+* router top-k in fp32 logit space;
+* **the inter-layer residual layout contract**: every layer takes and returns a
+  replicated ``[1, 1, B, 2048]`` bfloat16 ``TILE`` ``DRAM_MEMORY_CONFIG``
+  tensor, and there is no collective, gather, reshard or layout conversion
+  between layers. ``prefill_hidden`` and ``decode_hidden`` below are literally a
+  ``for`` loop over 48 layers with the residual threaded straight through.
+
+What the wrapper adds, and where each new boundary lives:
+
+``embed_tokens``
+    **Replicated**, bf16, so the embedding output *is* the residual contract
+    with no collective at all. A hidden-sharded embedding would be 4x smaller
+    per die but would owe an all-gather on every prefill chunk and every decode
+    token; at 0.622 GB/die against 22.35 GB of measured headroom
+    (``config/context_contract.json``) the replicated table is free and the
+    collective is not. This is also the shape the stage-03 footprint probe
+    allocated, so the published capacity numbers describe what actually runs.
+
+``model.norm`` (final RMSNorm)
+    Replicated, and shares the layer code: decode uses
+    ``multichip_decoder.decode_residual_norm`` (width-sharded over 8 L1 cores,
+    the same kernel and compute config as the two residual norms), prefill uses
+    the interleaved ``ttnn.rms_norm``.
+
+``lm_head``
+    **Column-parallel over the vocabulary**: die *d* owns columns
+    ``37984*d .. 37984*d+37983`` of ``[2048, 151936]``. 151936 = 4 * 37984 and
+    37984 = 32 * 1187, so the split is exact and needs no vocabulary padding.
+    **Logits never reach the host on the token-out path, and neither strategy
+    all-gathers them.** Both reduce first and gather the survivors: greedy takes
+    a per-die argmax and all-gathers four candidate values and indices
+    (``_WatcherCleanSampling1D._sample_argmax``), top-k/top-p takes a per-die
+    top-32 and all-gathers 32 values and indices. Which of the two is faster here
+    was measured, not assumed -- see ``sample_greedy_argmax``.
+
+``rotary`` (decode only)
+    ``ttnn.experimental.rotary_embedding_hf(is_decode_mode=True)`` reading a
+    per-user cos/sin pair **gathered on device** by ``ttnn.embedding`` from a
+    position tensor the trace advances with ``ttnn.plus_one``. The layer's
+    shipped spelling, ``ttnn.experimental.rotary_embedding``, takes the position
+    as a **Python int** compile-time argument and therefore cannot be replayed:
+    a captured trace would rotate every subsequent token at the position it was
+    captured at. Note this is the *HF* rotary, same ``rotate_half`` channel
+    convention -- so unlike stage 04's rejected ``rotary_embedding_llama`` lever
+    (README limitation 4) it needs no weight permutation, changes no KV-cache
+    channel convention and leaves prefill untouched.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import gc
+import json
+import math
+import os
+from collections.abc import Sequence
+from pathlib import Path
+
+import torch
+from safetensors import safe_open
+from transformers import AutoConfig
+
+import ttnn
+from models.common.modules.sampling.sampling_1d import Sampling1D, Sampling1DConfig
+
+from .functional_decoder import DecoderLayerConfig, KVCache
+from .multichip_decoder import (
+    MESH_SHAPE,
+    NUM_DEVICES,
+    TOPOLOGY,
+    MeshContext,
+    MeshDecoderConfig,
+    MultichipWeights,
+    _head_shard,
+    _norm_compute_config,
+    build_local_sparsity,
+    decode_residual_norm,
+    decoder_layer_decode_multichip,
+    decoder_layer_prefill_multichip,
+    fallback_audit,
+    mesh_context,
+    upload_multichip_weights,
+)
+from .precision import DEFAULT_PRECISION, PrecisionConfig, dtype_to_name
+
+HF_MODEL_ID = "Qwen/Qwen3-Coder-30B-A3B-Instruct"
+HF_REVISION = "b2cff646eb4bb1d68355c01b18ae02e7cf42d120"
+
+HIDDEN_SIZE = 2048
+VOCAB_SIZE = 151936
+NUM_LAYERS = 48
+HEAD_DIM = 128
+MAX_CONTEXT = 262144
+DEFAULT_PAGE_BLOCK_SIZE = 32
+DEFAULT_MAX_BATCH_SIZE = 1
+#: ``ttnn.sampling`` and ``nlp_create_qkv_heads_decode`` both address 32 fixed
+#: user slots; decode is always one 32-row tile regardless of the active batch.
+SAMPLING_SLOTS = 32
+#: Trace region per device. Two traces (model decode + sampling) over 48 layers.
+DEFAULT_TRACE_REGION_SIZE = 300_000_000
+#: RoPE table rows materialised at construction; grown on demand to the request
+#: horizon by ``ensure_rope_capacity`` so a short request never pays for 262144
+#: rows (which would be 64 MB/die of cos plus 64 MB of sin).
+DEFAULT_ROPE_CACHE_LEN = 8192
+
+#: ``lm_head`` weight dtype. bfloat8_b halves the 155 MB/die bf16 read that a
+#: decode step would otherwise make against a 2048x37984 weight.
+#:
+#: Since stage 07 this is an **alias** for ``DEFAULT_PRECISION.lm_head_dtype``,
+#: not the source of truth: a model built at a non-default ``PrecisionConfig``
+#: does not read it. See ``tt/precision.py``.
+LM_HEAD_WEIGHT_DTYPE = DEFAULT_PRECISION.lm_head_dtype
+#: The embedding table stays bf16: it is a gather, not a matmul, and bfloat8_b
+#: would quantise every token's hidden state at the very top of the stack.
+#: Alias for ``DEFAULT_PRECISION.embedding_dtype``, as above.
+EMBED_WEIGHT_DTYPE = DEFAULT_PRECISION.embedding_dtype
+
+
+#: ``_WatcherCleanSampling1D._sample_argmax``'s "not a winner" sentinel. Any value
+#: strictly greater than the vocabulary works; 2**20 is exact in int32 and leaves
+#: ``idx - BIG`` far from overflow.
+_DIST_ARGMAX_BIG = 1 << 20
+
+
+class _WatcherCleanSampling1D(Sampling1D):
+    """``Sampling1D`` with the force-argmax gather spelled the way this layer spells it.
+
+    Two overrides, for two different reasons.
+
+    ------------------------------------------------------------------------
+    ``_sample_argmax`` -- reduce first, gather second
+    ------------------------------------------------------------------------
+
+    ``Sampling1D._sample_argmax`` all-gathers the whole column-parallel logit
+    shard (37984 bf16 columns per die) up to the full 151936 on **every** die,
+    untilizes 151936 columns and runs one ``ttnn.argmax`` over them.
+    ``doc/full_model/tt_perf_report_full_model_decode.txt`` shows those two ops
+    at ``AllGatherAsync 889 us`` and ``ArgMax 859 us``. **Neither number is a
+    share of a token-out step, and the two must not be summed against one.**
+    That report is stage 05's **2-layer** window, which charges the terminal
+    path against two layers instead of 48 and so over-weights it by
+    construction -- the two rows are 27.5% and 26.5% of *that* window. And the
+    column is per-op device-kernel time summed over the op's own cores (2 for
+    the gather, 110 for the argmax), which is a different accounting from the
+    wall clock of a decode step. An earlier revision of this docstring set
+    ``889 + 859`` against "the 1.87 ms of non-layer work in a 22.079 ms
+    token-out step"; the near-agreement was a coincidence between two
+    incommensurable measurements and the claim is withdrawn.
+
+    The full 48-layer profile is the accounting that means something. On the
+    shipped tree ``doc/optimized_full_model/probes/profile_summary_decode.json``
+    puts the *whole* terminal block -- final norm, LM head, this sampler and the
+    token feedback -- at **366.5 us of an 18889.5 us decode iteration, 1.94%**,
+    of which this sampler is **126.2 us**. The baseline path was replaced before
+    that profile was taken and so has no 48-layer op row of its own; its
+    in-model price is a token-out delta and is quoted as one under
+    ``sample_greedy_argmax``.
+
+    The override computes the same token by reducing on each die first and
+    all-gathering only the four survivors::
+
+        rm         = untilize(local_shard)          # bf16   [1,1,32,37984]
+        rm         = rm[:, :, :B, :]                # bf16   [1,1,B,37984]
+        local_idx  = argmax(rm, -1, keepdim)        # uint32 [1,1,B,1]
+        local_max  = gather(rm, -1, local_idx)      # bf16   [1,1,B,1]
+        global_idx = local_idx + rank*37984         # int32, sharded constant
+        vals4      = all_gather(local_max)          # bf16   [1,1,B,4]
+        idx4       = all_gather(global_idx)         # int32  [1,1,B,4]
+        gmax       = max(vals4, -1, keepdim)
+        mask       = (vals4 == gmax)                # int32 0/1
+        token      = min(BIG + mask*(idx4 - BIG), -1)
+        token      = pad(token, to=32, value=0)     # uint32 [1,1,32]
+
+    ``doc/optimized_full_model/probes/distributed_argmax_probe.py`` measures the
+    two against each other at the shipped shape, trace-captured, median of 100:
+    **1.1432 ms baseline against 0.6275 ms, 1.82x**. Five things in that spelling
+    are load-bearing and were each established on the device, not assumed:
+
+    * **The local maximum must come from ``ttnn.gather``, not ``ttnn.max``.**
+      ``ttnn.max`` over the 37984-wide shard costs 0.494 ms -- more than the
+      ``ttnn.argmax`` over the same tensor (0.371 ms). ``ttnn.gather`` at the
+      index the argmax already produced costs 0.059 ms. That single substitution
+      is the difference between 1.05x and 1.82x.
+    * **Only the live user rows are reduced.** The logit tile is logically 32
+      rows because ``ttnn.sampling`` addresses 32 slots, but at batch ``B`` the
+      other ``32-B`` are zero-logit padding and reduce to token 0 by
+      construction. ``ttnn.argmax``'s kernel compares scalar-wise on a
+      data-movement RISC, so the cost is linear in rows: the whole reduction is
+      **631.6 us over 32 rows and 250.8 us over 1**, and the ``ttnn.pad(value=0)``
+      that restores the 32 slots writes back exactly the values the 32-row
+      reduction produced. ``argmax_outer_dim_probe.py`` checks that on the
+      device rather than asserting it.
+    * **Untilize before the argmax.** ``ttnn.argmax``'s multicore path needs
+      ROW_MAJOR; the TILE path is single-core and the whole leg becomes 23.25 ms.
+      The untilize itself is 0.075 ms.
+    * **Indices are INT32 end to end.** FLOAT32 elementwise rounds an index
+      through bf16 (36885 -> 36864), and ``ttnn.where`` on int32 operands returns
+      bit garbage -- hence the arithmetic select ``BIG + mask*(idx-BIG)`` rather
+      than a ``where``. ``ttnn.gather`` in turn demands a UINT32 index, which is
+      exactly what ``ttnn.argmax`` emits, so no cast happens on that edge.
+    * **The cross-die reduction is a ``min`` over masked indices, never a sum.**
+      On an exact tie both lanes survive the mask; ``sum(mask*idx)`` would add
+      the two indices together, ``min`` keeps the lower one. Because the dies own
+      contiguous ascending vocabulary ranges and ``ttnn.argmax`` returns the
+      first occurrence within a die, that is precisely ``torch.argmax``'s
+      first-maximal rule. The probe checks it with crafted cross-die, within-die
+      and triple ties, and checks the first-occurrence property of ``ttnn.argmax``
+      itself.
+
+    **Output contract.** The base writes the token into the caller's
+    ``tt_out_tok`` via ``ttnn.argmax(output_tensor=...)`` and returns it. The
+    traced decode loop feeds that same buffer back as the next token's input, so
+    returning a *new* tensor would silently break token feedback
+    (``models/common/sampling/generator.py::_validate_trace_inputs`` checks
+    identity, and this model's trace binds ``token`` as both sampler output and
+    model input). The override therefore ends in ``ttnn.copy`` into the caller's
+    buffer and returns that exact object -- same dtype (uint32), layout
+    (ROW_MAJOR), shape and buffer address.
+
+    **Fallback.** The fast path is only taken when the reduction it performs is
+    provably the same function as the base's. It falls back to
+    ``super()._sample_argmax`` whenever ``valid_vocab_size < vocab_size`` (a
+    padded vocabulary needs the invalid tail masked *before* the local argmax,
+    which ``_mask_invalid_vocab_logits`` /
+    ``_can_slice_valid_vocab_for_argmax`` do around the base's gather and this
+    path does not reproduce), whenever any invalid-vocab mask buffer is present,
+    on a single device, or when the logits do not arrive as an exact even shard.
+    For this model ``valid_vocab_size == vocab_size == 151936 == 4*37984``, so
+    the fast path is what runs -- but a token id >= the real vocabulary stays
+    impossible either way, because a padded vocabulary never reaches it.
+
+    ------------------------------------------------------------------------
+    ``_argmax_all_gather`` -- no ``Topology::Linear`` + ``num_workers_per_link=1``
+    ------------------------------------------------------------------------
+
+    Still overridden, still needed: the split top-k/top-p path is live for any
+    request with ``top_k > 1`` or ``top_p > 0`` (``sample_split``), and
+    ``_sample_argmax``'s fallback branch above uses it too.
+
+    ``ttnn.experimental.all_gather_async`` trips a
+    BRISC ``ASSERT`` in ``minimal_default_writer.cpp`` when it is given
+    ``topology=Topology::Linear`` **together with** ``num_workers_per_link=1``.
+    Neither alone does it; the pair does, at any width. The full A/B matrix is
+    ``doc/full_model/watcher_ab.log`` and the model-free reproducer is
+    ``doc/full_model/probes/ccl_watcher_ab.py --leg linear_workers1``.
+
+    ``Sampling1D._argmax_all_gather`` walks straight into that pair on any mesh
+    smaller than T3K. Its first branch -- Ring, no barrier -- is guarded by
+    ``default_topology(mesh) == Topology.Ring``, which is **False** on this 1x4
+    Blackhole mesh, so the branch is unreachable here. The fallback then runs
+    ``_get_argmax_all_gather_config``, which forces ``Topology.Linear`` for any
+    mesh under 8 devices, and the call below it hardcodes
+    ``num_workers_per_link=1``. Linear + 1 worker: exactly the tripping pair.
+
+    The decoder layer's own two all-reduces have been watcher-clean for four
+    stages, and the reason is visible in the same matrix: the layer never passes
+    ``num_workers_per_link`` at all, so the op picks its default. This override
+    does the same thing -- same op, same ``dim``, same semaphores, same
+    ``Topology.Ring`` the layer uses, and **no tuning knobs pinned**. The
+    matrix's ``sampler_shape_default_knobs`` leg is this exact call at this exact
+    shape, and it is clean.
+
+    This is a local workaround for an upstream bug, not a fix for it. Both
+    reports (the op, and ``sampling_1d.py``'s unreachable Ring branch) still
+    stand and should still be filed; this subclass just means stage 05 does not
+    ship an unchecked-but-violated device invariant while they are open. When
+    the op is fixed, delete this class and pass ``Sampling1D`` directly.
+
+    Subclassing is the seam because ``Sampling1D.from_config`` builds through
+    ``object.__new__(cls)`` and ``_bind_strategy`` binds
+    ``self._pre_argmax_gather = self._argmax_all_gather`` by attribute lookup on
+    the instance -- so the override is what gets bound. **No shared code is
+    edited.**
+    """
+
+    def _argmax_all_gather(self, logits):
+        cfg = self.config
+        return ttnn.experimental.all_gather_async(
+            logits,
+            persistent_output_buffer=None,
+            dim=3,
+            multi_device_global_semaphore=cfg.tt_ccl.get_and_cycle_ag_semaphore_handles(),
+            barrier_semaphore=cfg.tt_ccl.get_and_cycle_barrier_semaphore_handle(),
+            num_links=cfg.num_argmax_gather_links,
+            memory_config=logits.memory_config(),
+            topology=cfg.ag_topology,
+            # Deliberately NOT passing chunks_per_sync / num_workers_per_link /
+            # num_buffers_per_channel. Pinning num_workers_per_link=1 is the half
+            # of the tripping pair we control. See the class docstring.
+        )
+
+    # -- distributed argmax ---------------------------------------------------
+
+    def _distributed_argmax_local_vocab(self):
+        """Per-die vocabulary width if the distributed argmax applies, else ``None``.
+
+        Every condition here is a condition under which the reduction below is
+        *provably* the same function as ``Sampling1D._sample_argmax``'s. Anything
+        else falls back to the base implementation rather than being approximated.
+        """
+        cfg = self.config
+        if getattr(self, "_invalid_vocab_mask", None) is not None:
+            return None
+        if getattr(self, "_invalid_vocab_tail_mask", None) is not None:
+            return None
+        valid = cfg.valid_vocab_size if cfg.valid_vocab_size is not None else cfg.vocab_size
+        if valid != cfg.vocab_size:
+            # A padded vocabulary needs the invalid tail masked before the *local*
+            # argmax, which this path does not do. The base masks/slices around
+            # its full gather and stays correct; use it.
+            return None
+        num_devices = cfg.mesh_device.get_num_devices()
+        if num_devices < 2 or cfg.vocab_size % num_devices != 0:
+            return None
+        local = cfg.vocab_size // num_devices
+        if local % ttnn.TILE_SIZE != 0:
+            return None
+        return local
+
+    #: Live user rows in the sampler's 32-slot logit tile. ``None`` means "all 32"
+    #: and reproduces the pre-row-slicing behaviour exactly. The model sets it to
+    #: its own ``max_batch_size`` (1 by default) -- see ``_sample_argmax``, and
+    #: ``doc/optimized_full_model/probes/argmax_outer_dim_probe.py`` for why it is
+    #: worth 2.5x on the whole sampler.
+    _dist_active_rows = None
+
+    def _distributed_argmax_active_rows(self, slots: int) -> int:
+        rows = self._dist_active_rows
+        if rows is None:
+            return int(slots)
+        return max(1, min(int(rows), int(slots)))
+
+    # -- sampling penalties ---------------------------------------------------
+    #
+    # ``Sampling1D`` has no penalty stage at all, and the vLLM TT plugin does not
+    # route penalised requests to host sampling (``platform.py`` sends ``min_p``,
+    # ``bad_words``, ``logit_bias``, ``allowed_token_ids``, ``min_tokens``,
+    # ``prompt_logprobs`` and structured output to the host sampler -- penalties
+    # are deliberately *not* in that list). It packs all three into
+    # ``TTSamplingParams`` and hands the model the token history it needs
+    # (``model_runner.py``: ``prompt_tokens`` / ``output_tokens`` are populated
+    # "if penalties are needed (decode only)"), expecting the model's on-device
+    # sampler to apply them. This is that stage.
+    #
+    # ------------------------------------------------------------------------
+    # The shard-boundary problem, and why this spelling cannot get it wrong
+    # ------------------------------------------------------------------------
+    #
+    # Logits are column-parallel: die ``d`` holds vocabulary ids
+    # ``d*37984 .. d*37984+37983`` of the 151936, contiguous and ascending -- the
+    # same decomposition ``load_device_buffers`` above builds ``_dist_die_offset``
+    # from, and ``_dist_local_vocab`` is reused here rather than recomputed. A
+    # penalty is keyed by a **global** token id, so for id ``t`` only die
+    # ``t // 37984`` may touch column ``t % 37984``; penalising local index
+    # ``t % 37984`` on the *other three* dies would silently penalise three
+    # unrelated tokens and produce plausible-looking wrong output rather than an
+    # error.
+    #
+    # This stage never does that arithmetic in a kernel. The penalty operands are
+    # built on the host as **full-vocabulary** ``[1, 1, 32, 151936]`` tensors --
+    # indexed by global id, which is the only frame in which a penalty is
+    # defined -- and handed to the device through
+    # ``ttnn.ShardTensorToMesh(dim=-1)``, the *same* mapper and the same even
+    # 4-way split the logits themselves were produced under by the
+    # column-parallel LM head. Column ``t`` of the host tensor therefore lands on
+    # exactly the die and exactly the local column that holds logit ``t``, by
+    # construction rather than by a computed index. Every op below is
+    # elementwise between two tensors with identical per-die shapes, so no op
+    # ever needs to know a global id.
+    #
+    # The identity is *checked* rather than assumed:
+    # ``probes/penalty_shard_boundary_probe.py`` penalises one token in die 0's
+    # range and one in die 3's, and asserts both moved and that the same local
+    # index on the other dies did not.
+    #
+    # ------------------------------------------------------------------------
+    # The arithmetic
+    # ------------------------------------------------------------------------
+    #
+    # vLLM's ``model_executor/layers/utils.py::apply_penalties`` is the contract,
+    # and its order is load-bearing -- repetition first, on the raw logit:
+    #
+    #     repetition p (over prompt+output): x = x/p if x > 0 else x*p
+    #     frequency  f (over output):        x -= f * count(t in output)
+    #     presence   q (over output):        x -= q * (count(t in output) > 0)
+    #
+    # The repetition rule is sign-dependent, so it is *not* expressible as an
+    # additive delta. It is spelled as a per-column multiplicative factor whose
+    # two branches are both uploaded:
+    #
+    #     pos    = gtz(x)                       # 1.0 where x > 0, else 0.0
+    #     factor = rep_neg + pos * rep_dif      # rep_neg = p, rep_dif = 1/p - p
+    #     x      = x * factor
+    #     x      = x - add_delta                # f*count + q*presence, host-summed
+    #
+    # For a column no row penalises, the host writes ``rep_neg = 1.0``,
+    # ``rep_dif = 0.0``, ``add_delta = 0.0``: ``x * 1.0 - 0.0`` is **bit-exact**
+    # in bf16, so an unpenalised token is not merely close to unchanged, it is
+    # unchanged. That is what makes the cross-die non-perturbation claim a
+    # property of the arithmetic and not of a tolerance.
+    #
+    # Per-row isolation is likewise structural: the operands are ``[1,1,32,V]``
+    # and every op is elementwise, so row *i*'s columns are only ever combined
+    # with row *i*'s logits. Padding slots get the neutral row and are untouched.
+    #
+    # Baking the per-row scalars (p, 1/p, f, q) into the full-width tensors on
+    # the host, rather than broadcasting a ``[1,1,32,1]`` scalar column on
+    # device, costs one more upload but removes every H-broadcast from the traced
+    # graph -- and the host is rebuilding these rows anyway, because vLLM re-sends
+    # the whole token history each step.
+    #
+    # ------------------------------------------------------------------------
+    # Fast path
+    # ------------------------------------------------------------------------
+    #
+    # ``_penalty_mode`` is a *graph* property, not a value: 0 means the ops below
+    # are not in the captured trace at all, so an unpenalised request pays
+    # nothing -- no op, no buffer, no upload. Bit 0 is the repetition stage and
+    # bit 1 the additive stage, and they are independent, so a repetition-only
+    # request never pays for the additive tensor. The generator releases and
+    # re-captures the decode traces when the mode changes, exactly as it already
+    # does when ``_sampling_stochastic`` flips between the argmax and split
+    # strategies.
+
+    #: Bitmask: 1 = repetition stage in the graph, 2 = frequency/presence stage.
+    _penalty_mode = 0
+    _penalty_rep_neg = None
+    _penalty_add = None
+
+    def penalty_buffer_shape(self) -> tuple[int, int]:
+        """``(slots, vocab_size)`` the host-side penalty operands must have."""
+        cfg = self.config
+        return int(cfg.max_batch_size), int(cfg.vocab_size)
+
+    def penalty_shard_geometry(self) -> tuple[int, int]:
+        """``(num_devices, local_vocab)`` -- the split the operands must be staged in.
+
+        The **same** decomposition ``load_device_buffers`` builds
+        ``_dist_die_offset`` from, read off the same config rather than
+        recomputed, so the staging path and the distributed argmax cannot drift
+        apart.
+        """
+        cfg = self.config
+        devices = cfg.mesh_device.get_num_devices()
+        vocab = int(cfg.vocab_size)
+        if vocab % devices:
+            raise RuntimeError(f"penalties need an even column-parallel split; {vocab} % {devices} != 0")
+        return devices, vocab // devices
+
+    def allocate_penalty_buffers(self, mode: int) -> None:
+        """Allocate/free the per-stage operands for ``mode``.
+
+        Called by the generator **outside** any trace capture -- ``ttnn.from_torch``
+        inside ``begin_trace_capture`` raises and leaves the capture open (stage-04
+        ``work_log.md`` §6), which is the same reason ``load_device_buffers``
+        builds ``_dist_die_offset`` eagerly.
+        """
+        mode = int(mode)
+        if mode == self._penalty_mode:
+            return
+        cfg = self.config
+        slots, vocab = self.penalty_buffer_shape()
+        num_devices = cfg.mesh_device.get_num_devices()
+        if mode and vocab % num_devices != 0:
+            raise RuntimeError(f"penalties need an even column-parallel vocabulary split; {vocab} % {num_devices} != 0")
+
+        def _alloc(fill: float):
+            return ttnn.from_torch(
+                torch.full((1, 1, slots, vocab), fill, dtype=torch.bfloat16),
+                dtype=ttnn.bfloat16,
+                layout=ttnn.TILE_LAYOUT,
+                device=cfg.mesh_device,
+                memory_config=ttnn.DRAM_MEMORY_CONFIG,
+                mesh_mapper=ttnn.ShardTensorToMesh(cfg.mesh_device, dim=-1),
+            )
+
+        for want, names, fills in (
+            (mode & 1, ("_penalty_rep_neg",), (1.0,)),
+            (mode & 2, ("_penalty_add",), (0.0,)),
+        ):
+            for name, fill in zip(names, fills):
+                current = getattr(self, name, None)
+                if want and current is None:
+                    setattr(self, name, _alloc(fill))
+                elif not want and current is not None:
+                    ttnn.deallocate(current, True)
+                    setattr(self, name, None)
+        self._penalty_mode = mode
+
+    def penalty_device_buffers(self) -> dict:
+        """The live operands, keyed by name; the generator uploads into these."""
+        return {"rep_neg": self._penalty_rep_neg, "add": self._penalty_add}
+
+    def _apply_penalties(self, logits):
+        """``logits`` -> penalised logits, or ``logits`` itself when the mode is 0.
+
+        Returns ``(tensor, is_new)``; the caller deallocates when ``is_new``.
+        """
+        mode = self._penalty_mode
+        if not mode:
+            return logits, False
+        if int(logits.shape[-1]) != self.penalty_buffer_shape()[1] // self.config.mesh_device.get_num_devices():
+            # Already gathered, or some shape this stage was not built for. The
+            # penalty operands are per-die shards; refusing is the only safe
+            # answer, because applying them at the wrong width would penalise
+            # the wrong tokens.
+            raise RuntimeError(
+                f"penalty operands are per-die shards of width "
+                f"{self.penalty_buffer_shape()[1] // self.config.mesh_device.get_num_devices()}, "
+                f"got logits of width {int(logits.shape[-1])}"
+            )
+        out = logits
+        if mode & 1:
+            # ``rep_dif`` (= 1/p - p) is derived **on device** rather than
+            # uploaded. It used to be a second full-width operand, and staging one
+            # of those costs 2.049 ms of host time per decode step -- more than
+            # every device op in this stage put together. ``ttnn.reciprocal`` of
+            # the operand gives the same thing for free, because the operand is
+            # ``p`` at penalised columns and exactly ``1.0`` everywhere else.
+            #
+            # This is only allowed to be here because ``reciprocal(1.0)`` is
+            # **exactly** 1.0 on this device -- checked, not assumed
+            # (``penalty_shard_boundary_probe.py``'s reference and
+            # bit-identity legs both fail if it is not). That is what keeps the
+            # unpenalised column at ``x * 1.0 - 0.0``, i.e. bit-exact, which is
+            # the whole cross-die non-perturbation argument. At *penalised*
+            # columns the LLK reciprocal differs from a host-computed ``1/p`` by
+            # up to about one bf16 ulp (p=1.05: 0.95703 against 0.95313), which is
+            # inside the accuracy the bf16 operand already has.
+            inv = ttnn.reciprocal(self._penalty_rep_neg)
+            rep_dif = ttnn.subtract(inv, self._penalty_rep_neg)
+            ttnn.deallocate(inv)
+            # gtz, not a where: ttnn.where on this path is the op the argmax
+            # override already avoids, and gtz is a single unary.
+            pos = ttnn.gtz(out)
+            scaled = ttnn.multiply(pos, rep_dif)
+            ttnn.deallocate(pos)
+            ttnn.deallocate(rep_dif)
+            factor = ttnn.add(scaled, self._penalty_rep_neg)
+            ttnn.deallocate(scaled)
+            out = ttnn.multiply(out, factor)
+            ttnn.deallocate(factor)
+        if mode & 2:
+            penalised = ttnn.subtract(out, self._penalty_add)
+            if out is not logits:
+                ttnn.deallocate(out)
+            out = penalised
+        return out, True
+
+    def decode_forward(self, logits, **kwargs):
+        """Penalty stage, then ``Sampling1D``'s own routing -- unchanged.
+
+        Overriding here rather than in each strategy means both the argmax path
+        and the top-k/top-p split path get penalties from one place, applied
+        **before** any selection, which is the only order that is correct.
+        """
+        penalised, is_new = self._apply_penalties(logits)
+        try:
+            return super().decode_forward(penalised, **kwargs)
+        finally:
+            if is_new:
+                ttnn.deallocate(penalised)
+
+    def load_device_buffers(self):
+        """Base buffers, plus the per-die vocabulary offset the reduction adds.
+
+        Built here rather than lazily in ``_sample_argmax`` because the first
+        ``_sample_argmax`` may already be inside ``begin_trace_capture``, and
+        ``ttnn.from_torch`` inside a capture raises and leaves the capture open.
+        """
+        already_loaded = self._device_buffers_loaded
+        super().load_device_buffers()
+        if already_loaded and getattr(self, "_dist_die_offset", None) is not None:
+            return
+        local_vocab = self._distributed_argmax_local_vocab()
+        if local_vocab is None:
+            self._dist_die_offset = None
+            return
+        cfg = self.config
+        num_devices = cfg.mesh_device.get_num_devices()
+        # ``_dist_active_rows`` rows, not ``cfg.max_batch_size``: the reduction
+        # below runs over the live user rows only and the shapes must match
+        # exactly, because an H-broadcast here would silently re-expand the
+        # result back to 32 rows. See ``_sample_argmax``.
+        rows = self._distributed_argmax_active_rows(cfg.max_batch_size)
+        offsets = (
+            (
+                torch.arange(num_devices, dtype=torch.int64)
+                .reshape(1, 1, 1, num_devices)
+                .expand(1, 1, rows, num_devices)
+                * local_vocab
+            )
+            .contiguous()
+            .to(torch.int32)
+        )
+        # Sharded on the last dim: die d holds the single column ``d*local_vocab``.
+        self._dist_die_offset = ttnn.from_torch(
+            offsets,
+            dtype=ttnn.int32,
+            layout=ttnn.TILE_LAYOUT,
+            device=cfg.mesh_device,
+            memory_config=ttnn.DRAM_MEMORY_CONFIG,
+            mesh_mapper=ttnn.ShardTensorToMesh(cfg.mesh_device, dim=-1),
+        )
+        self._dist_local_vocab = local_vocab
+
+    def _sample_argmax(self, logits, tt_out_tok):
+        """Distributed argmax: reduce per die, all-gather 4 candidates, reduce again.
+
+        Honours ``Sampling1D._sample_argmax``'s contract exactly -- writes the
+        caller's ``tt_out_tok`` in place and returns ``(tt_out_tok, None)``; the
+        argmax path never emits logprobs. See the class docstring for why each
+        step is spelled the way it is and for the measured 1.82x.
+        """
+        self.load_device_buffers()
+        die_offset = getattr(self, "_dist_die_offset", None)
+        if die_offset is None or int(logits.shape[-1]) != self._dist_local_vocab:
+            # Not an even per-die shard (already gathered, padded vocab, 1x1
+            # mesh, ...): the base path is the one that is still correct.
+            return super()._sample_argmax(logits, tt_out_tok)
+
+        # -- per-die reduction, over this die's own columns only ---------------
+        # ROW_MAJOR because ttnn.argmax's TILE path is single-core (23 ms).
+        rm = ttnn.untilize(logits, use_multicore=True)
+        # **Reduce the live user rows, not the padding.** ``decode_terminal``
+        # hands the sampler a logically-32-row tile because ``ttnn.sampling``
+        # addresses 32 fixed slots, but at batch B only the first B rows carry a
+        # user: the rest are the zero rows ``ttnn.pad(..., value=0.0)`` put on the
+        # pre-head hidden, and ``lm_head`` has no bias, so their logits are exactly
+        # zero. ``ttnn.argmax``'s multicore kernel does the comparison as a scalar
+        # C++ ``>`` loop on the RISCV_1 *data-movement* core -- 32 x 37984 values
+        # over 110 cores is ~11k compares each, and that, not the 32-round
+        # semaphore barrier, is why the op costs 366 us and sits 75x off
+        # bandwidth. Dropping the padding rows drops the work proportionally.
+        #
+        # Measured standalone at the shipped shape
+        # (``doc/optimized_full_model/probes/argmax_outer_dim_probe.py``,
+        # trace-captured, median of 60; the harness floor is ~58 us):
+        #
+        #   argmax over 32 rows                        371.1 us
+        #   argmax over 32 rows, keepdim=False         309.3    (one barrier, not 32)
+        #   ROW_MAJOR slice to 1 row + argmax           58.0    (i.e. at the floor)
+        #   whole reduction, 32 rows                   631.6
+        #   whole reduction, 1 row                     250.8    **2.52x**
+        #
+        # ``keepdim=False`` is a real but small effect and is *not* taken: it buys
+        # 62 us on its own and nothing at all once the rows are sliced (251.0 vs
+        # 250.8), while costing a ``[1,1,B] -> [1,1,B,1]`` reshape.
+        #
+        # The substitution is exact, not an approximation. The probe's
+        # ``padding_rows_produce_token_zero`` leg checks on the device that a
+        # zero logit row reduces to token **0** on the shipped 32-row path -- all
+        # four dies tie at 0.0, so the masked ``min`` keeps global index 0 -- which
+        # is precisely the value the ``ttnn.pad`` below writes back.
+        slots = int(rm.shape[-2])
+        active = self._distributed_argmax_active_rows(slots)
+        if active < slots:
+            live = ttnn.slice(rm, [0, 0, 0, 0], [1, 1, active, self._dist_local_vocab])
+            ttnn.deallocate(rm)
+            rm = live
+        local_idx = ttnn.argmax(rm, dim=-1, keepdim=True)  # uint32 RM [1,1,B,1]
+        # ttnn.gather, NOT ttnn.max: 0.059 ms against 0.494 ms. This is the win.
+        local_max = ttnn.to_layout(ttnn.gather(rm, dim=-1, index=local_idx), ttnn.TILE_LAYOUT)
+        ttnn.deallocate(rm)
+        # INT32, not FLOAT32: fp32 elementwise rounds the index through bf16.
+        local_idx_i32 = ttnn.to_layout(ttnn.typecast(local_idx, ttnn.int32), ttnn.TILE_LAYOUT)
+        ttnn.deallocate(local_idx)
+        global_idx = ttnn.add(local_idx_i32, die_offset)
+        ttnn.deallocate(local_idx_i32)
+
+        # -- gather 4 candidates, not the whole vocabulary ---------------------
+        vals4 = self._argmax_all_gather(local_max)  # bf16  [1,1,B,4]
+        idx4 = self._argmax_all_gather(global_idx)  # int32 [1,1,B,4]
+        ttnn.deallocate(local_max)
+        ttnn.deallocate(global_idx)
+
+        # -- cross-die reduction ----------------------------------------------
+        gmax = ttnn.max(vals4, dim=-1, keepdim=True)
+        mask = ttnn.typecast(ttnn.eq(vals4, gmax), ttnn.int32)  # 0/1
+        # NOT sum(mask*idx): on a tie that adds the tied indices together.
+        # BIG + mask*(idx-BIG) sends losers to BIG and leaves every tied winner at
+        # its own global index, so min() keeps the lowest -- the first-maximal one,
+        # because die ranges ascend.
+        sel = ttnn.add(ttnn.multiply(mask, ttnn.subtract(idx4, _DIST_ARGMAX_BIG)), _DIST_ARGMAX_BIG)
+        token = ttnn.min(sel, dim=-1, keepdim=False)  # int32 TILE
+        for scratch in (vals4, idx4, gmax, mask, sel):
+            ttnn.deallocate(scratch)
+
+        # -- match ttnn.argmax's output contract: UINT32 / ROW_MAJOR ------------
+        token = ttnn.typecast(ttnn.to_layout(token, ttnn.ROW_MAJOR_LAYOUT), ttnn.uint32)
+        if active < slots:
+            # Restore the 32-slot vector ``tt_out_tok`` is. 0 is not a convenient
+            # filler, it is the token the shipped 32-row reduction *already*
+            # produces for a padding row (see the comment above the slice), so the
+            # buffer's contents are unchanged slot for slot. It also keeps every
+            # slot a valid id: ``embed_decode`` runs ``ttnn.embedding`` over all 32
+            # before slicing to ``batch``, and an out-of-vocabulary id there would
+            # be an out-of-bounds table read.
+            padded = ttnn.pad(token, [(0, 0), (0, 0), (0, slots - active)], value=0)
+            ttnn.deallocate(token)
+            token = padded
+        if tt_out_tok is None:
+            return token, None
+        # Write **into the caller's buffer**. The traced decode loop feeds this
+        # exact tensor back as the next token, so the object and its address must
+        # survive; returning a new tensor breaks token feedback silently.
+        ttnn.copy(ttnn.reshape(token, tt_out_tok.shape), tt_out_tok)
+        ttnn.deallocate(token)
+        return tt_out_tok, None
+
+
+def _resolve_precision(precision) -> PrecisionConfig:
+    """Accept a ``PrecisionConfig``, a dict, a path to JSON, or ``None``.
+
+    ``None`` is ``DEFAULT_PRECISION``, so every existing caller keeps the
+    shipped policy. The dict and path forms exist so a sweep runner -- and,
+    later, the vLLM construction path -- can pass the *artifact*
+    (``selected_precision_config.json``) rather than importing the dataclass,
+    which is what makes the artifact something the model consumes rather than
+    something written next to it.
+    """
+    if precision is None:
+        return DEFAULT_PRECISION
+    if isinstance(precision, PrecisionConfig):
+        return precision
+    if isinstance(precision, dict):
+        return PrecisionConfig.from_dict(precision)
+    if isinstance(precision, (str, Path)):
+        return PrecisionConfig.read_json(precision)
+    # A ``PrecisionConfig`` from a *duplicate copy* of ``tt.precision``, which is
+    # a real hazard in this tree and not a hypothetical one: ``tt/generator.py``
+    # imports ``tt.model`` by absolute path while tests and probes import it
+    # relatively, and under pytest's ``--import-mode=importlib`` (this repo's
+    # ``addopts``) with no ``models/__init__.py`` the two spellings produce two
+    # distinct module objects and therefore two distinct classes. ``isinstance``
+    # is then False for an object that is, by every meaning that matters, the
+    # right one. Rebuild it through the serialised form rather than refusing it.
+    if type(precision).__name__ == "PrecisionConfig" and hasattr(precision, "to_dict"):
+        return PrecisionConfig.from_dict(precision.to_dict())
+    raise TypeError(f"precision must be a PrecisionConfig, dict, path or None; got {type(precision).__name__}")
+
+
+def _lm_head_compute_config(device, precision: PrecisionConfig = DEFAULT_PRECISION):
+    return ttnn.init_device_compute_kernel_config(
+        device.arch(),
+        math_fidelity=precision.lm_head_fidelity,
+        math_approx_mode=False,
+        fp32_dest_acc_en=False,
+        packer_l1_acc=True,
+    )
+
+
+class ShardedCheckpoint:
+    """Read named tensors out of a sharded safetensors checkpoint on demand.
+
+    The full checkpoint is 30.5B parameters, ~61 GB in bf16. Materialising it as
+    one ``state_dict`` to build a model that uploads it layer by layer would
+    need that whole 61 GB of host RAM at once; this reads only the tensors asked
+    for, from only the shards that hold them, and holds nothing.
+    """
+
+    def __init__(self, path: str | Path):
+        self.path = Path(path)
+        index_path = self.path / "model.safetensors.index.json"
+        if not index_path.is_file():
+            raise FileNotFoundError(f"checkpoint index is missing: {index_path}")
+        self.weight_map: dict[str, str] = json.loads(index_path.read_text())["weight_map"]
+
+    def get(self, name: str) -> torch.Tensor:
+        shard = self.weight_map.get(name)
+        if shard is None:
+            raise KeyError(name)
+        with safe_open(self.path / shard, framework="pt") as f:
+            return f.get_tensor(name)
+
+    def layer(self, layer_idx: int) -> dict[str, torch.Tensor]:
+        """Every ``model.layers.<i>.*`` tensor, keyed layer-relative."""
+        prefix = f"model.layers.{layer_idx}."
+        by_shard: dict[str, list[str]] = {}
+        for name, shard in self.weight_map.items():
+            if name.startswith(prefix):
+                by_shard.setdefault(shard, []).append(name)
+        if not by_shard:
+            raise KeyError(f"no tensors for layer {layer_idx}")
+        out: dict[str, torch.Tensor] = {}
+        for shard, names in by_shard.items():
+            with safe_open(self.path / shard, framework="pt") as f:
+                for name in names:
+                    out[name[len(prefix) :]] = f.get_tensor(name)
+        return out
+
+
+def _validate_mesh(mesh_device) -> None:
+    shape = tuple(int(v) for v in mesh_device.shape)
+    if shape != MESH_SHAPE:
+        raise ValueError(f"Qwen3CoderModel requires mesh {MESH_SHAPE}, got {shape}")
+    if mesh_device.get_num_devices() != NUM_DEVICES:
+        raise ValueError(f"Qwen3CoderModel requires exactly {NUM_DEVICES} devices")
+
+
+def _rope_parameters(hf_config) -> dict:
+    """``rope_parameters`` on current transformers, ``rope_theta`` on older ones.
+
+    ``Qwen3MoeConfig`` no longer exposes a top-level ``rope_theta`` attribute --
+    reading it raises ``AttributeError`` rather than returning ``None`` -- so
+    the dict is the only spelling that works on both.
+    """
+    params = getattr(hf_config, "rope_parameters", None)
+    if params:
+        return dict(params)
+    return {"rope_theta": hf_config.rope_theta, "rope_type": "default"}
+
+
+def _rope_type(hf_config) -> str:
+    return str(_rope_parameters(hf_config).get("rope_type", "default"))
+
+
+def _rope_theta(hf_config) -> float:
+    return float(_rope_parameters(hf_config)["rope_theta"])
+
+
+def _rope_tables(hf_config, capacity: int) -> tuple[torch.Tensor, torch.Tensor]:
+    """The HF ``(cos, sin)`` tables for positions ``0..capacity-1``.
+
+    Built here rather than through ``Qwen3MoeRotaryEmbedding`` so that no
+    transformers model object is constructed at load time; the formula is the
+    default rope (``rope_scaling`` is null in this checkpoint, which
+    ``from_checkpoint`` asserts).
+    """
+    head_dim = int(getattr(hf_config, "head_dim", HEAD_DIM))
+    theta = _rope_theta(hf_config)
+    inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, dtype=torch.float64) / head_dim))
+    angles = torch.outer(torch.arange(capacity, dtype=torch.float64), inv_freq)
+    angles = torch.cat([angles, angles], dim=-1)
+    return angles.cos().float(), angles.sin().float()
+
+
+class Qwen3CoderModel:
+    """The 48-layer causal LM over the stage-04 multichip decoder layer."""
+
+    def __init__(
+        self,
+        *,
+        mesh_device,
+        hf_config,
+        checkpoint: ShardedCheckpoint,
+        max_batch_size: int = DEFAULT_MAX_BATCH_SIZE,
+        max_cache_len: int = MAX_CONTEXT,
+        num_layers: int = NUM_LAYERS,
+        page_block_size: int = DEFAULT_PAGE_BLOCK_SIZE,
+        rope_cache_len: int = DEFAULT_ROPE_CACHE_LEN,
+        precision: "PrecisionConfig | dict | str | Path | None" = None,
+    ) -> None:
+        _validate_mesh(mesh_device)
+        if not 1 <= int(max_batch_size) <= 32:
+            # nlp_create_qkv_heads_decode_device_operation.cpp:51 asserts
+            # num_users <= 32; a TTNN op limit, unchanged by TP.
+            raise ValueError(f"max_batch_size must be in [1,32], got {max_batch_size}")
+        if not 1 <= int(num_layers) <= int(hf_config.num_hidden_layers):
+            raise ValueError(f"num_layers must be in [1,{hf_config.num_hidden_layers}]")
+        if not 1 <= int(max_cache_len) <= int(hf_config.max_position_embeddings):
+            raise ValueError(f"max_cache_len must be in [1,{hf_config.max_position_embeddings}]")
+        if int(hf_config.hidden_size) != HIDDEN_SIZE or int(hf_config.vocab_size) != VOCAB_SIZE:
+            raise ValueError("HF config does not match the Qwen3-Coder-30B-A3B full-model contract")
+        if bool(hf_config.tie_word_embeddings):
+            raise ValueError("this checkpoint has an untied lm_head; tied weights would be a different contract")
+        if _rope_type(hf_config) != "default":
+            raise ValueError(f"rope_type {_rope_type(hf_config)!r} is not supported by this port's rotary tables")
+
+        # The precision policy, resolved once and then read by every builder and
+        # every forward below. ``None`` -> ``DEFAULT_PRECISION``, the shipped
+        # stage-06 policy, so a caller that says nothing gets exactly the model
+        # it got before this parameter existed. A ``dict`` or a path is accepted
+        # too, so a sweep runner can hand over a ``selected_precision_config.json``
+        # without importing the dataclass.
+        self.precision = _resolve_precision(precision)
+
+        self.mesh_device = mesh_device
+        self.hf_config = hf_config
+        self.max_batch_size = int(max_batch_size)
+        self.max_cache_len = int(max_cache_len)
+        self.num_layers = int(num_layers)
+        self.page_block_size = int(page_block_size)
+        self.hidden_size = HIDDEN_SIZE
+        self.vocab_size = VOCAB_SIZE
+        self.head_dim = int(getattr(hf_config, "head_dim", HEAD_DIM))
+        self.rms_norm_eps = float(hf_config.rms_norm_eps)
+        # Exact: 151936 = 4 * 37984 and 37984 = 32 * 1187.
+        assert self.vocab_size % (32 * NUM_DEVICES) == 0, self.vocab_size
+        self.local_vocab_size = self.vocab_size // NUM_DEVICES
+
+        #: Skip the expert work of decode rows that hold no live request. On by
+        #: default and a no-op at ``max_batch_size == 1``; see
+        #: ``_decode_active_mask``. ``QWEN3_DECODE_ACTIVE_ROW_GATING=0`` restores
+        #: the stage-08 graph exactly, which is what
+        #: ``doc/optimized_vllm/probes/inactive_row_gating_probe.py`` A/Bs
+        #: against for the token-equality leg.
+        self.active_row_gating = os.getenv("QWEN3_DECODE_ACTIVE_ROW_GATING", "1") not in ("0", "", "false", "no")
+
+        #: Width of the decode graph currently being built or captured.
+        #: Equal to ``max_batch_size`` everywhere except inside
+        #: ``decode_width_scope``, which the generator opens to capture a
+        #: **narrower** decode graph than the configured slot count -- see
+        #: ``doc/batch_scaling/README.md``. Every decode-path use of the row
+        #: count reads this, not ``max_batch_size``: the embedding slice, the
+        #: rotary gather and shard, and the active-row mask. Prefill and the
+        #: sampler are untouched -- ``decode_terminal`` pads to the 32 fixed
+        #: ``SAMPLING_SLOTS`` regardless, so the sampler never sees the width.
+        self.decode_width = self.max_batch_size
+
+        self.ctx: MeshContext = mesh_context(mesh_device)
+        self.config = MeshDecoderConfig.from_hf(hf_config)
+        self.global_config: DecoderLayerConfig = self.config.global_config
+
+        self.embed_tokens = self._build_embedding(checkpoint)
+        self.layers: list[MultichipWeights] = self._build_layers(checkpoint)
+        self.final_norm, self.final_norm_rm = self._build_final_norm(checkpoint)
+        self.lm_head = self._build_lm_head(checkpoint)
+
+        self.sparsity = build_local_sparsity(mesh_device, self.config.local_moe)
+        self.lm_head_compute_config = _lm_head_compute_config(mesh_device, self.precision)
+        self.norm_compute_config = _norm_compute_config(mesh_device, self.precision)
+        # Set by ``local_logits`` / the sampler-input path once a forward has
+        # run, so ``runtime_fallback_audit`` can report the dtypes the terminal
+        # path *produced* rather than the ones the config asked for.
+        self._observed_logits_dtype = None
+        self._observed_sampling_dtype = None
+
+        self.rope_cache_len = 0
+        self.cos_table = None
+        self.sin_table = None
+        self.ensure_rope_capacity(min(int(rope_cache_len), self.max_cache_len))
+
+        # ``_WatcherCleanSampling1D`` rather than ``Sampling1D``: same module,
+        # same strategies, the force-argmax gather spelled without the pinned
+        # ``num_workers_per_link`` that trips the watcher on this mesh. See the
+        # class docstring above and ``doc/full_model/watcher_ab.log``.
+        self.sampler = _WatcherCleanSampling1D.from_config(
+            Sampling1DConfig(
+                vocab_size=self.vocab_size,
+                valid_vocab_size=self.vocab_size,
+                mesh_device=mesh_device,
+                tt_ccl=self.ctx.ccl,
+                max_batch_size=32,
+                max_top_k=32,
+                num_gather_links=1,
+                sampling_memory_config=ttnn.DRAM_MEMORY_CONFIG,
+                allow_force_argmax=True,
+                num_argmax_gather_links=1,
+                ag_topology=TOPOLOGY,
+                # **False, and that is a measurement.** ``Sampling1D``'s comment
+                # calls the power-of-two pad a "big device-perf win for
+                # non-power-of-2 vocab on the multi-device path". For a per-die
+                # shard of 37984 it is the opposite: the pad is to 65536, a 1.73x
+                # blow-up of the tensor ``ttnn.topk`` then scans, and
+                # ``probes/sampler_probe.py`` measures the whole split path at
+                # **11.006 ms padded against 6.151 ms unpadded**, 1.79x, at the
+                # shipped logits shape with the sampled token unchanged.
+                pad_to_power_of_2=False,
+            )
+        )
+        # ``max_batch_size=32`` above is the *slot* count ``ttnn.sampling`` and
+        # ``decode_terminal`` address; this is how many of those slots carry a
+        # user. The distributed argmax reduces only those rows -- the rest are the
+        # zero-logit padding ``decode_terminal`` adds -- which is worth 2.52x on
+        # the whole sampler at batch 1. Set before ``load_device_buffers`` because
+        # the per-die offset constant is built to this row count.
+        self.sampler._dist_active_rows = self.max_batch_size
+        self.sampler.load_device_buffers()
+        self.kv_cache: list[KVCache] | None = None
+
+    # -- construction ---------------------------------------------------------
+
+    @classmethod
+    def from_checkpoint(
+        cls,
+        checkpoint_path: str | Path,
+        *,
+        mesh_device,
+        max_batch_size: int = DEFAULT_MAX_BATCH_SIZE,
+        max_cache_len: int = MAX_CONTEXT,
+        num_layers: int = NUM_LAYERS,
+        page_block_size: int = DEFAULT_PAGE_BLOCK_SIZE,
+        rope_cache_len: int = DEFAULT_ROPE_CACHE_LEN,
+        precision: "PrecisionConfig | dict | str | Path | None" = None,
+    ) -> "Qwen3CoderModel":
+        checkpoint_path = Path(checkpoint_path)
+        hf_config = AutoConfig.from_pretrained(checkpoint_path)
+        checkpoint = ShardedCheckpoint(checkpoint_path)
+        model = cls(
+            mesh_device=mesh_device,
+            hf_config=hf_config,
+            checkpoint=checkpoint,
+            max_batch_size=max_batch_size,
+            max_cache_len=max_cache_len,
+            num_layers=num_layers,
+            page_block_size=page_block_size,
+            rope_cache_len=rope_cache_len,
+            precision=precision,
+        )
+        gc.collect()
+        return model
+
+    def _build_embedding(self, checkpoint: ShardedCheckpoint) -> ttnn.Tensor:
+        host = checkpoint.get("model.embed_tokens.weight").float()
+        tensor = ttnn.from_torch(
+            host,
+            dtype=self.precision.embedding_dtype,
+            layout=ttnn.ROW_MAJOR_LAYOUT,
+            device=self.mesh_device,
+            memory_config=ttnn.DRAM_MEMORY_CONFIG,
+            mesh_mapper=ttnn.ReplicateTensorToMesh(self.mesh_device),
+        )
+        del host
+        gc.collect()
+        return tensor
+
+    def _build_layers(self, checkpoint: ShardedCheckpoint) -> list[MultichipWeights]:
+        from .weight_mapping import convert_layer_weights
+
+        layers = []
+        for layer_idx in range(self.num_layers):
+            sd = checkpoint.layer(layer_idx)
+            torch_weights = convert_layer_weights(sd, self.hf_config)
+            del sd
+            layers.append(
+                upload_multichip_weights(torch_weights, self.mesh_device, self.config, precision=self.precision)
+            )
+            del torch_weights
+            gc.collect()
+        return layers
+
+    def _build_final_norm(self, checkpoint: ShardedCheckpoint):
+        host = checkpoint.get("model.norm.weight").float().reshape(-1)
+        tiled = ttnn.from_torch(
+            host.reshape(1, 1, 1, -1),
+            dtype=self.precision.norm_weight_dtype,
+            layout=ttnn.TILE_LAYOUT,
+            device=self.mesh_device,
+            memory_config=ttnn.DRAM_MEMORY_CONFIG,
+            mesh_mapper=ttnn.ReplicateTensorToMesh(self.mesh_device),
+        )
+        # The layout the sharded rms_norm program factory reads; see
+        # ``multichip_decoder.upload_multichip_weights.norm_row_major``.
+        row_major = ttnn.from_torch(
+            host.reshape(1, 1, host.numel() // 32, 32).contiguous(),
+            dtype=self.precision.norm_weight_dtype,
+            layout=ttnn.ROW_MAJOR_LAYOUT,
+            device=self.mesh_device,
+            memory_config=ttnn.DRAM_MEMORY_CONFIG,
+            mesh_mapper=ttnn.ReplicateTensorToMesh(self.mesh_device),
+        )
+        return tiled, row_major
+
+    def _build_lm_head(self, checkpoint: ShardedCheckpoint) -> ttnn.Tensor:
+        host = checkpoint.get("lm_head.weight").float().transpose(-2, -1).contiguous()
+        assert tuple(host.shape) == (self.hidden_size, self.vocab_size), tuple(host.shape)
+        tensor = ttnn.from_torch(
+            host.reshape(1, 1, self.hidden_size, self.vocab_size),
+            dtype=self.precision.lm_head_dtype,
+            layout=ttnn.TILE_LAYOUT,
+            device=self.mesh_device,
+            memory_config=ttnn.DRAM_MEMORY_CONFIG,
+            mesh_mapper=ttnn.ShardTensorToMesh(self.mesh_device, dim=-1),
+        )
+        del host
+        gc.collect()
+        return tensor
+
+    # -- rotary ---------------------------------------------------------------
+
+    def ensure_rope_capacity(self, required_len: int) -> bool:
+        """Grow the device cos/sin tables to cover ``required_len`` positions."""
+        required_len = int(required_len)
+        if required_len <= self.rope_cache_len:
+            return False
+        if required_len > self.max_cache_len:
+            raise ValueError(f"rotary capacity {required_len} exceeds context {self.max_cache_len}")
+        capacity = min(self.max_cache_len, max(32, 1 << (required_len - 1).bit_length()))
+        cos, sin = _rope_tables(self.hf_config, capacity)
+        new = []
+        for host in (cos, sin):
+            new.append(
+                ttnn.from_torch(
+                    host.reshape(1, 1, capacity, self.head_dim),
+                    dtype=ttnn.bfloat16,
+                    layout=ttnn.TILE_LAYOUT,
+                    device=self.mesh_device,
+                    memory_config=ttnn.DRAM_MEMORY_CONFIG,
+                    mesh_mapper=ttnn.ReplicateTensorToMesh(self.mesh_device),
+                )
+            )
+        old = (self.cos_table, self.sin_table)
+        self.cos_table, self.sin_table = new
+        for tensor in old:
+            if tensor is not None:
+                ttnn.deallocate(tensor, True)
+        self.rope_cache_len = capacity
+        return True
+
+    def rope_decode_tables(self, rotary_position: ttnn.Tensor):
+        """Per-user ``(cos, sin)`` for one decode step, gathered **on device**.
+
+        ``rotary_position`` is a ``[1, batch]`` uint32 device tensor. The gather
+        is ``ttnn.embedding`` against the replicated cos/sin tables, so the
+        position never leaves the device and the whole thing is capturable; the
+        trace advances ``rotary_position`` itself with ``ttnn.plus_one``.
+
+        Returns the height-sharded ``[1, batch, 1, head_dim]`` pair that
+        ``rotary_embedding_hf(is_decode_mode=True)`` requires -- one core per
+        user, the same ``_head_shard`` layout ``nlp_create_qkv_heads_decode``
+        emits for Q and K.
+        """
+        batch = self.decode_width
+        shard = _head_shard(32, self.head_dim, batch)
+        out = []
+        for table in (self.cos_table, self.sin_table):
+            # [1, batch] -> [1, batch, head_dim] -> [1, 1, batch, head_dim]
+            # -> [1, batch, 1, head_dim], the layout rotary_embedding_hf's decode
+            # factory reads. Same sequence as ``RotarySetup1D.decode_forward``.
+            gathered = ttnn.unsqueeze_to_4D(ttnn.embedding(rotary_position, table, layout=ttnn.TILE_LAYOUT))
+            transposed = ttnn.transpose(gathered, 1, 2)
+            if int(transposed.shape[1]) != batch:
+                trimmed = ttnn.slice(transposed, [0, 0, 0, 0], [1, batch, 1, self.head_dim])
+                ttnn.deallocate(transposed, True)
+                transposed = trimmed
+            out.append(ttnn.interleaved_to_sharded(transposed, shard))
+            ttnn.deallocate(transposed, True)
+        return out[0], out[1]
+
+    def _rope_decode(self, tensor: ttnn.Tensor, cos_sharded, sin_sharded, _token_index):
+        """The ``rope=`` seam handed to ``decoder_layer_decode_multichip``.
+
+        ``_token_index`` is accepted and ignored: the position lives in the
+        cos/sin pair, which is what makes this spelling replayable where the
+        layer's default one is not.
+        """
+        shard = _head_shard(32, self.head_dim, self.decode_width)
+        staged = ttnn.to_memory_config(tensor, shard)
+        rotated = ttnn.experimental.rotary_embedding_hf(staged, cos_sharded, sin_sharded, is_decode_mode=True)
+        ttnn.deallocate(staged, True)
+        out = ttnn.to_memory_config(rotated, ttnn.DRAM_MEMORY_CONFIG)
+        ttnn.deallocate(rotated, True)
+        return out
+
+    # -- KV cache -------------------------------------------------------------
+
+    def allocate_kv_cache(
+        self,
+        *,
+        max_cache_len: int | None = None,
+        num_blocks: int | None = None,
+        page_table: ttnn.Tensor | None = None,
+    ) -> list[KVCache]:
+        """One paged ``KVCache`` per layer, 1 local KV head per die.
+
+        512 B per token per layer per die -- a quarter of the single-die 2048 --
+        which is what makes the advertised 262144 context fit; see
+        ``config/context_contract.json``.
+        """
+        cache_len = self.max_cache_len if max_cache_len is None else int(max_cache_len)
+        blocks_per_seq = math.ceil(cache_len / self.page_block_size)
+        total_blocks = self.max_batch_size * blocks_per_seq if num_blocks is None else int(num_blocks)
+        local = self.config.local_attention
+        caches = []
+        for _ in range(self.num_layers):
+            k, v = (
+                ttnn.from_torch(
+                    torch.zeros(total_blocks, local.num_key_value_heads, self.page_block_size, local.head_dim),
+                    dtype=self.precision.kv_cache_dtype,
+                    layout=ttnn.TILE_LAYOUT,
+                    device=self.mesh_device,
+                    memory_config=ttnn.DRAM_MEMORY_CONFIG,
+                    mesh_mapper=ttnn.ReplicateTensorToMesh(self.mesh_device),
+                )
+                for _ in range(2)
+            )
+            caches.append(KVCache(k=k, v=v, page_table=page_table, block_size=self.page_block_size))
+        return caches
+
+    def ensure_internal_kv_cache(self, page_table: ttnn.Tensor | None = None) -> list[KVCache]:
+        if self.kv_cache is None:
+            self.kv_cache = self.allocate_kv_cache(page_table=page_table)
+        return self.kv_cache
+
+    @staticmethod
+    def bind_page_table(kv_cache: Sequence[KVCache], page_table: ttnn.Tensor | None) -> list[KVCache]:
+        """Point every layer's cache at ``page_table`` in place.
+
+        The page table is a *persistent device tensor* owned by the caller (the
+        generator, or vLLM later). Rebinding mutates the ``KVCache`` records
+        rather than reallocating, so the tensor identity a captured trace
+        recorded is preserved and an unchanged page table costs nothing.
+        """
+        for cache in kv_cache:
+            cache.page_table = page_table
+        return list(kv_cache)
+
+    def reset_kv_cache(self, kv_cache: Sequence[KVCache] | None = None) -> None:
+        selected = self.ensure_internal_kv_cache() if kv_cache is None else kv_cache
+        for cache in selected:
+            ttnn.fill(cache.k, 0.0, memory_config=cache.k.memory_config(), output_tensor=cache.k)
+            ttnn.fill(cache.v, 0.0, memory_config=cache.v.memory_config(), output_tensor=cache.v)
+
+    # -- forward: prefill -----------------------------------------------------
+
+    def embed_prefill(self, tokens: ttnn.Tensor) -> ttnn.Tensor:
+        """``[1, S]`` uint32 -> replicated ``[1, 1, S, 2048]``, no collective."""
+        hidden = ttnn.embedding(
+            tokens,
+            self.embed_tokens,
+            layout=ttnn.TILE_LAYOUT,
+            memory_config=ttnn.DRAM_MEMORY_CONFIG,
+            dtype=self.precision.activation_dtype,
+        )
+        hidden = ttnn.unsqueeze_to_4D(hidden)
+        return ttnn.reshape(hidden, (1, 1, int(hidden.shape[-2]), self.hidden_size))
+
+    def prefill_hidden(
+        self,
+        tokens: ttnn.Tensor,
+        *,
+        kv_cache: Sequence[KVCache] | None = None,
+        user_id: int = 0,
+        start_pos: int = 0,
+        chunk_page_table=None,
+        fill_page_table=None,
+        fill_len: int | None = None,
+    ) -> ttnn.Tensor:
+        """Run the whole stack over one user's prompt. ``S`` is arbitrary.
+
+        Nothing here constrains ``S``: the collectives scatter on dim 3 (hidden,
+        2048, fixed), ``attention_prefill`` slices RoPE's tile padding back, and
+        ``moe_prefill_optimized`` pads to its chunk internally and slices back.
+        """
+        caches = self.ensure_internal_kv_cache() if kv_cache is None else kv_cache
+        if len(caches) != self.num_layers:
+            raise ValueError(f"kv_cache has {len(caches)} layers, expected {self.num_layers}")
+        hidden = self.embed_prefill(tokens)
+        seq_len = int(hidden.shape[-2])
+        # A split prefill's suffix occupies absolute positions
+        # [start_pos, start_pos + seq_len), so the tables must cover the END of
+        # the range, not its length.
+        self.ensure_rope_capacity(start_pos + seq_len)
+        # Exactly ``seq_len`` rows, including non-tile-aligned lengths -- the
+        # same shape the single-layer prefill gates pass at S = 33/100/257.
+        # RoPE is applied at ABSOLUTE positions: the suffix of a split prefill
+        # must rotate at [start_pos, start_pos + seq_len), not from 0, or its keys
+        # disagree with the ones already in the cache. Identical to the shipped
+        # slice when start_pos == 0.
+        # When the requested window is the WHOLE table, ``ttnn.slice`` hands
+        # back a view of its input rather than a copy -- as a different Python
+        # object, so an ``is`` guard does not catch it. The deallocate at the
+        # end of this method would then free ``self.cos_table``'s own DRAM, and
+        # the next prefill dies on "Input Tensor is not allocated" rather than
+        # anywhere near here. Identical hazard, and identical fix, to the
+        # one-token case in ``select_prefill_rows``.
+        #
+        # It is reachable whenever ``start_pos == 0`` and the length equals the
+        # rope capacity, which ``ensure_rope_capacity`` rounds to a power of
+        # two -- so a 8192- or 16384-token prompt, or any prefill bucket rung
+        # that is a power of two. Found by the bucket ladder's 8192 rung.
+        whole_table = start_pos == 0 and seq_len >= self.rope_cache_len
+        if whole_table:
+            cos = ttnn.clone(self.cos_table, memory_config=ttnn.DRAM_MEMORY_CONFIG)
+            sin = ttnn.clone(self.sin_table, memory_config=ttnn.DRAM_MEMORY_CONFIG)
+        else:
+            cos = ttnn.slice(self.cos_table, [0, 0, start_pos, 0], [1, 1, start_pos + seq_len, self.head_dim])
+            sin = ttnn.slice(self.sin_table, [0, 0, start_pos, 0], [1, 1, start_pos + seq_len, self.head_dim])
+        for layer_idx in range(self.num_layers):
+            hidden = decoder_layer_prefill_multichip(
+                hidden,
+                self.layers[layer_idx],
+                self.config,
+                self.ctx,
+                cos,
+                sin,
+                self.sparsity,
+                kv_cache=caches[layer_idx],
+                user_id=user_id,
+                precision=self.precision,
+                start_pos=start_pos,
+                chunk_page_table=chunk_page_table,
+                fill_page_table=fill_page_table,
+                fill_len=fill_len,
+            )
+        ttnn.deallocate(cos, True)
+        ttnn.deallocate(sin, True)
+        return hidden
+
+    def prefill_norm(self, hidden: ttnn.Tensor) -> ttnn.Tensor:
+        return ttnn.rms_norm(
+            hidden,
+            weight=self.final_norm,
+            epsilon=self.rms_norm_eps,
+            compute_kernel_config=self.norm_compute_config,
+            memory_config=ttnn.DRAM_MEMORY_CONFIG,
+        )
+
+    def select_prefill_rows(self, hidden: ttnn.Tensor, rows: Sequence[int]) -> ttnn.Tensor:
+        """Keep only ``rows`` of a ``[1, 1, S, H]`` prefill result."""
+        seq_len = int(hidden.shape[-2])
+        pieces = []
+        for row in rows:
+            if not 0 <= int(row) < seq_len:
+                raise ValueError(f"prefill row {row} is outside [0,{seq_len})")
+            if seq_len == 1:
+                # At a **one-token prompt** the requested slice covers the whole
+                # tensor, and ``ttnn.slice`` then hands back a view of its input
+                # rather than a copy -- as a *different* Python object, so an
+                # ``is`` guard does not catch it. The caller deallocates
+                # ``hidden`` immediately afterwards, leaving the retained row
+                # pointing at freed DRAM; that does not raise, it **segfaults**
+                # in whatever reads it next (the final norm here). Copy instead.
+                # `probes/prompt_len_1_repro.py` is the four-line reproduction.
+                pieces.append(ttnn.clone(hidden, memory_config=ttnn.DRAM_MEMORY_CONFIG))
+                continue
+            pieces.append(
+                ttnn.slice(
+                    hidden,
+                    [0, 0, int(row), 0],
+                    [1, 1, int(row) + 1, self.hidden_size],
+                    memory_config=ttnn.DRAM_MEMORY_CONFIG,
+                )
+            )
+        if len(pieces) == 1:
+            return pieces[0]
+        out = ttnn.concat(pieces, dim=2, memory_config=ttnn.DRAM_MEMORY_CONFIG)
+        for piece in pieces:
+            ttnn.deallocate(piece, True)
+        return out
+
+    def local_logits(self, normed: ttnn.Tensor) -> ttnn.Tensor:
+        """``[1, 1, rows, 2048]`` -> this die's ``[1, 1, rows, 37984]`` logits."""
+        out = ttnn.linear(
+            normed,
+            self.lm_head,
+            compute_kernel_config=self.lm_head_compute_config,
+            memory_config=ttnn.DRAM_MEMORY_CONFIG,
+            dtype=self.precision.logits_dtype,
+        )
+        # Observed, not asserted: the dtype the produced tensor actually carries.
+        # ``runtime_fallback_audit`` reports it so ``logits_dtype`` is verified
+        # off a real tensor rather than echoed back out of the config -- see the
+        # ``*_observed`` entries there.
+        self._observed_logits_dtype = out.dtype
+        return out
+
+    def gather_logits_to_torch(self, local_logits: ttnn.Tensor, *, valid_rows: int | None = None) -> torch.Tensor:
+        """Host-side full-vocabulary logits. **Not** on the token-out path.
+
+        Used by ``return_all_logits`` prefill checks and the host-sampling
+        compatibility mode only; the measured decode path never calls this.
+        """
+        gathered = ttnn.all_gather(
+            local_logits,
+            dim=3,
+            num_links=1,
+            memory_config=ttnn.DRAM_MEMORY_CONFIG,
+            topology=TOPOLOGY,
+        )
+        host = ttnn.to_torch(ttnn.get_device_tensors(gathered)[0]).float()
+        ttnn.deallocate(gathered, True)
+        if valid_rows is not None:
+            host = host[..., : int(valid_rows), :]
+        return host[..., : self.vocab_size]
+
+    # -- forward: decode ------------------------------------------------------
+
+    def embed_decode(self, tokens: ttnn.Tensor) -> ttnn.Tensor:
+        """``[1, 1, 1, 32]`` uint32 -> replicated ``[1, 1, batch, 2048]``."""
+        hidden = ttnn.embedding(
+            tokens,
+            self.embed_tokens,
+            layout=ttnn.TILE_LAYOUT,
+            memory_config=ttnn.DRAM_MEMORY_CONFIG,
+            dtype=self.precision.activation_dtype,
+        )
+        hidden = ttnn.unsqueeze_to_4D(hidden)
+        flat = ttnn.reshape(hidden, (1, 1, int(hidden.shape[-2]), self.hidden_size))
+        if int(flat.shape[-2]) == self.decode_width:
+            return flat
+        sliced = ttnn.slice(
+            flat, [0, 0, 0, 0], [1, 1, self.decode_width, self.hidden_size], memory_config=ttnn.DRAM_MEMORY_CONFIG
+        )
+        ttnn.deallocate(flat, True)
+        return sliced
+
+    @contextlib.contextmanager
+    def decode_width_scope(self, width: int):
+        """Build the decode graph ``width`` rows wide instead of ``max_batch_size``.
+
+        The narrow graph is legal because **nothing in decode binds a user to a
+        slot index except the three per-row inputs** -- ``current_pos``, the
+        rotary position and the page-table row. The KV cache is fully paged:
+        ``paged_update_cache`` and ``paged_scaled_dot_product_attention_decode``
+        both reach the cache only through ``page_table_tensor`` rows and
+        ``cur_pos_tensor`` entries, and neither takes a ``batch_offset``. So a
+        request can be decoded in *any* row provided its page-table row, its
+        position and its token travel with it; no cache page moves.
+
+        What the width actually changes is the amount of work: the expert
+        ``ttnn.sparse_matmul`` visits ``width x local_experts`` slots per layer
+        with ``nnz=None``, the expert tail is dense over the same product, the
+        router runs one ``topk`` per row, and paged SDPA reads one window per
+        row. That cost is paid per row *configured*, not per row live -- see
+        ``doc/optimized_vllm/README.md``'s control curve -- so narrowing the
+        graph is the only lever that removes it.
+
+        Sampling is deliberately outside the scope: ``decode_terminal`` pads to
+        the 32 fixed ``SAMPLING_SLOTS`` whatever the width is, so ``tt_out_tok``
+        keeps its ``[1,1,1,32]`` shape and the sampler's per-slot parameters
+        keep their meaning. Row *i* of a narrow graph is sampling slot *i*.
+        """
+        width = int(width)
+        if not 1 <= width <= self.max_batch_size:
+            raise ValueError(f"decode width must be in [1,{self.max_batch_size}], got {width}")
+        previous = self.decode_width
+        self.decode_width = width
+        try:
+            yield width
+        finally:
+            self.decode_width = previous
+
+    def _decode_active_mask(self, current_pos: ttnn.Tensor):
+        """``[1, 1, batch, 1]`` of 1.0 for live slots and 0.0 for inactive ones.
+
+        A serving decode batch is always the configured ``max_num_seqs`` rows --
+        vLLM pads it so the trace shape is constant -- with inactive slots
+        carrying ``current_pos = -1``. Those rows still embed a token, still run
+        attention, and, critically, still route to a full top-8 of experts, so
+        their ``(row, expert)`` pairs land in ``sparse_matmul``'s sparsity and
+        cost real expert weight reads and real math. Multiplying the routing
+        vector by this mask takes them out of the sparsity instead
+        (``decoder_layer_decode_multichip``).
+
+        **Why it is derived on device rather than passed in.** ``current_pos`` is
+        already a persistent trace input, and the traced graph advances it with
+        ``ttnn.plus_one(..., skip_negative_entries=True)`` -- an inactive row
+        stays at ``-1`` through any number of replays, and a slot only becomes
+        active through a host reinstall of ``current_pos``. So a mask computed
+        from it inside the same graph is correct by construction on every replay,
+        with no extra trace input to refresh and no way for it to go stale. A
+        host-supplied mask would be one more thing that has to be right.
+
+        Returns ``None`` at ``max_batch_size == 1``, where there is no inactive
+        row to skip: the graph is then byte-for-byte the one stage 08 shipped and
+        the single-user headline cannot be perturbed by this change.
+        """
+        if self.decode_width <= 1 or not self.active_row_gating:
+            return None
+        row = ttnn.to_layout(ttnn.reshape(current_pos, (1, 1, 1, self.decode_width)), ttnn.TILE_LAYOUT)
+        # bf16 cannot represent every position exactly at 262144, but it
+        # represents every position's *sign* exactly, and ``gez`` only reads the
+        # sign. -1 -> 0.0, everything >= 0 -> 1.0.
+        as_float = ttnn.typecast(row, ttnn.bfloat16)
+        ttnn.deallocate(row, True)
+        live_row = ttnn.gez(as_float)
+        ttnn.deallocate(as_float, True)
+        mask = ttnn.transpose(live_row, -2, -1)
+        ttnn.deallocate(live_row, True)
+        return mask
+
+    def decode_hidden(
+        self,
+        tokens: ttnn.Tensor,
+        *,
+        current_pos: ttnn.Tensor,
+        rotary_position: ttnn.Tensor,
+        kv_cache: Sequence[KVCache] | None = None,
+    ) -> ttnn.Tensor:
+        caches = self.ensure_internal_kv_cache() if kv_cache is None else kv_cache
+        if len(caches) != self.num_layers:
+            raise ValueError(f"kv_cache has {len(caches)} layers, expected {self.num_layers}")
+        hidden = self.embed_decode(tokens)
+        cos, sin = self.rope_decode_tables(rotary_position)
+        # Computed once per decode step and shared by all 48 layers.
+        active_mask = self._decode_active_mask(current_pos)
+        for layer_idx in range(self.num_layers):
+            hidden = decoder_layer_decode_multichip(
+                hidden,
+                self.layers[layer_idx],
+                self.config,
+                self.ctx,
+                cos,
+                sin,
+                caches[layer_idx],
+                current_pos,
+                0,  # token_index: unused by the rope seam below, see _rope_decode
+                rope=self._rope_decode,
+                precision=self.precision,
+                active_mask=active_mask,
+            )
+        ttnn.deallocate(cos, True)
+        ttnn.deallocate(sin, True)
+        if active_mask is not None:
+            ttnn.deallocate(active_mask, True)
+        return hidden
+
+    def decode_terminal(self, hidden: ttnn.Tensor) -> ttnn.Tensor:
+        """Final norm + column-parallel ``lm_head``, sampler-ready local logits.
+
+        The norm is the layer's own width-sharded decode kernel, and the shard
+        it emits is exactly the width-sharded L1 config the projections read, so
+        crossing into the head costs one sharded-to-interleaved.
+        """
+        normed_sharded = decode_residual_norm(hidden, self.final_norm_rm, self.rms_norm_eps, self.precision)
+        normed = ttnn.sharded_to_interleaved(normed_sharded, ttnn.DRAM_MEMORY_CONFIG)
+        ttnn.deallocate(normed_sharded, True)
+        # ``ttnn.sampling`` addresses 32 fixed user slots, and it compares the
+        # *logical* shapes of its values and indices, so the logits handed to it
+        # must be logically 32 rows and not ``batch`` rows padded to a tile.
+        # The rows are already physically there -- ``batch <= 32`` and decode is
+        # one 32-row tile -- so this only rewrites the logical shape.
+        rows = int(normed.shape[-2])
+        if rows < SAMPLING_SLOTS:
+            padded = ttnn.pad(
+                normed,
+                [(0, 0), (0, 0), (0, SAMPLING_SLOTS - rows), (0, 0)],
+                value=0.0,
+                memory_config=ttnn.DRAM_MEMORY_CONFIG,
+            )
+            ttnn.deallocate(normed, True)
+            normed = padded
+        logits = self.local_logits(normed)
+        ttnn.deallocate(normed, True)
+        if self.precision.sampling_dtype != logits.dtype:
+            # Equal on the shipped path, so this is dead code at the default and
+            # the traced decode graph is byte-for-byte what stage 06 captured.
+            cast = ttnn.typecast(logits, self.precision.sampling_dtype)
+            ttnn.deallocate(logits, True)
+            logits = cast
+        self._observed_sampling_dtype = logits.dtype
+        return logits
+
+    def decode_forward_from_ttnn_inputs(
+        self,
+        tokens: ttnn.Tensor,
+        current_pos: ttnn.Tensor,
+        *,
+        rotary_position: ttnn.Tensor,
+        kv_cache: Sequence[KVCache] | None = None,
+        advance_position: bool = True,
+    ) -> ttnn.Tensor:
+        """Token in -> sampler-ready local logits out, entirely on device.
+
+        With ``advance_position`` the two position tensors are incremented
+        **inside** this graph, so a captured trace steps its own positions on
+        replay and the host never refreshes them per token.
+        """
+        hidden = self.decode_hidden(
+            tokens,
+            current_pos=current_pos,
+            rotary_position=rotary_position,
+            kv_cache=kv_cache,
+        )
+        logits = self.decode_terminal(hidden)
+        if advance_position:
+            ttnn.plus_one(current_pos, skip_negative_entries=True)
+            ttnn.plus_one(rotary_position)
+        return logits
+
+    # -- sampling -------------------------------------------------------------
+
+    def sample_split(self, logits, *, k, p, temp, seeds=None, tt_out_tok=None):
+        """Canonical split sampling: local top-32 -> all-gather -> ``ttnn.sampling``.
+
+        ``k=1, p=0, temp=1`` is **semantically greedy**: the global argmax is by
+        construction inside some die's local top-32, and the all-gather makes
+        all four dies' candidates visible before the top-1 is taken.
+        """
+        return self.sampler.decode_forward(
+            logits, k=k, p=p, temp=temp, seeds=seeds, tt_out_tok=tt_out_tok, enable_log_probs=False
+        )[0]
+
+    def sample_greedy_argmax(self, logits, *, tt_out_tok=None):
+        """``Sampling1D``'s force-argmax path, on this model's distributed override.
+
+        Still the common module, still on device, still traced, still writes the
+        sampled token straight into ``tt_out_tok`` -- it is a different strategy
+        inside the same implementation, not a custom sampler. The strategy body
+        is ``_WatcherCleanSampling1D._sample_argmax``: reduce on each die, then
+        all-gather the four survivors instead of all-gathering the vocabulary.
+
+        **This is what greedy uses**, because at this vocabulary it is 6.6x
+        faster than the top-k/top-p split path (0.928 ms against 6.155 ms in the
+        48-layer model, both rows of
+        ``doc/optimized_full_model/probes/perf_full_model.csv``, which is the
+        **shipped** measurement) and produces the same token -- both rows sample
+        token 16 on that run. Whole-model ``token_out`` on that same run is
+        **19.693 ms, 50.78 t/s/u**.
+
+        Stage 05 shipped the same choice at 1.125 ms against 6.155 ms, and two
+        changes inside this override moved the greedy row since, each with its
+        own token-out delta at a like-for-like context:
+
+        * the **distributed reduction** above -- 22.079 ms to 21.461 ms
+          (45.29 -> 46.60 t/s/u), both at ``context`` 4096,
+          ``../full_model/probes/perf_full_model.json`` against
+          ``doc/optimized_full_model/probes/perf_full_model_part1_preadoption.json``;
+        * the **live-row slice** (reduce ``max_batch_size`` rows, not 32) --
+          20.146 ms to 19.693 ms at ``context`` 8192,
+          ``doc/optimized_full_model/probes/perf_full_model_p128_after.json``
+          against
+          ``doc/optimized_full_model/probes/perf_full_model_p128_argmaxrows.json``.
+
+        The remaining step between them is the paged SDPA program config in
+        ``tt/multichip_decoder.py`` and is not this sampler's. The moment any
+        slot asks for ``top_k > 1`` or ``top_p > 0`` the generator switches back
+        to ``sample_split``.
+        """
+        return self.sampler.decode_forward(logits, tt_out_tok=tt_out_tok, enable_log_probs=False)[0]
+
+    # -- audit ----------------------------------------------------------------
+
+    def runtime_fallback_audit(self, batch: int | None = None) -> dict:
+        """The layer audit, plus the boundaries this wrapper owns."""
+        batch = self.max_batch_size if batch is None else int(batch)
+        audit = fallback_audit(self.layers[0], self.config, batch, self.precision)
+        audit.update(
+            {
+                "num_layers": self.num_layers,
+                "embedding": "replicated_bf16_no_collective",
+                "residual_contract": "replicated [1,1,B,2048] bf16 TILE DRAM, no inter-layer collective",
+                "final_norm": "replicated, width-sharded decode kernel",
+                "lm_head_parallelism": "column_parallel_over_vocab",
+                "lm_head_local_vocab": self.local_vocab_size,
+                "lm_head_weight_dtype": str(self.lm_head.dtype),
+                "embedding_weight_dtype": str(self.embed_tokens.dtype),
+                "precision": self.precision.to_dict(),
+                "vocab_padding": 0,
+                "decode_rope": "rotary_embedding_hf(is_decode_mode=True), device position gather",
+                "decode_rope_position_source": "device tensor advanced by ttnn.plus_one inside the trace",
+                "sampling_greedy": (
+                    "Sampling1D force-argmax, distributed: per-die untilize/argmax/gather -> "
+                    "all-gather 4 candidates -> masked-min, traced, writes tt_out_tok"
+                ),
+                "sampling_topk_topp": "Sampling1D split (local topk -> all-gather 32 candidates -> ttnn.sampling)",
+                "sampling_pad_to_power_of_2": False,
+                "host_logit_readback_on_token_out_path": False,
+                "host_argmax_on_token_out_path": False,
+                # Read off the allocated cache when one exists, so a swept
+                # ``kv_cache_dtype`` is *observed* rather than asserted. This
+                # was a hard-coded "bfloat16" until stage 07's sweep, which
+                # would have silently mislabelled every non-default KV row.
+                # Falls back to the configured value before allocation.
+                # Emitted as the PLAIN name ("bfloat16"), not ``str(dtype)``
+                # ("DataType.BFLOAT16"), because that is the existing contract:
+                # doc/optimized_full_model's committed runtime_fallback_audit.json
+                # and check_published_figures.py both pin the plain spelling, and
+                # they are stage evidence that must keep passing. The sibling
+                # ``device_*`` fields use str(dtype) and are left alone.
+                "kv_cache_dtype": dtype_to_name(
+                    self.kv_cache[0].k.dtype if self.kv_cache else self.precision.kv_cache_dtype
+                ),
+                "kv_cache_dtype_source": "device_readback" if self.kv_cache else "config_not_yet_allocated",
+                # -- the four fields stage 07's selection proof could not check --
+                #
+                # Before the stage-07 review these were the only swept fields
+                # with no audit entry at all, so ``R03_lmhead_lofi``,
+                # ``R21_norm_hifi2`` and ``R22_logits_sampling_bfp8`` produced
+                # ``device_audit`` blocks byte-identical to the baseline's and
+                # "this lever does nothing" was indistinguishable from "this
+                # lever is not wired up". For ``norm_fidelity`` it was the
+                # second: ``decode_residual_norm`` built its compute config from
+                # the module default and never saw ``self.precision``.
+                #
+                # The two fidelities are read off the ``compute_kernel_config``
+                # objects the ops are actually handed (built here, passed at the
+                # call site), so they verify the config -> compute-config
+                # threading. The two dtypes are read off the **produced
+                # tensors** and are ``None`` until a forward has run.
+                "lm_head_math_fidelity": str(self.lm_head_compute_config.math_fidelity),
+                "norm_math_fidelity": str(self.norm_compute_config.math_fidelity),
+                "logits_dtype_observed": (
+                    None if self._observed_logits_dtype is None else dtype_to_name(self._observed_logits_dtype)
+                ),
+                "sampling_dtype_observed": (
+                    None if self._observed_sampling_dtype is None else dtype_to_name(self._observed_sampling_dtype)
+                ),
+                "terminal_dtype_source": ("device_readback" if self._observed_logits_dtype else "no_forward_yet"),
+                "kv_cache_paged": True,
+                "page_block_size": self.page_block_size,
+                "collective_topology": str(TOPOLOGY),
+                "prefill_num_links": self.ctx.num_links,
+                "decode_num_links": self.ctx.decode_num_links,
+            }
+        )
+        return audit
+
+    def teardown(self) -> None:
+        if self.kv_cache is not None:
+            for cache in self.kv_cache:
+                ttnn.deallocate(cache.k, True)
+                ttnn.deallocate(cache.v, True)
+            self.kv_cache = None
+
+
+__all__ = [
+    "DEFAULT_MAX_BATCH_SIZE",
+    "DEFAULT_PAGE_BLOCK_SIZE",
+    "DEFAULT_ROPE_CACHE_LEN",
+    "DEFAULT_TRACE_REGION_SIZE",
+    "HF_MODEL_ID",
+    "HF_REVISION",
+    "MAX_CONTEXT",
+    "NUM_LAYERS",
+    "Qwen3CoderModel",
+    "ShardedCheckpoint",
+]

code/models/demos/blackhole/qwen3_coder_30b_a3b/tt/multichip_decoder.py

@@ -0,0 +1,1982 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+"""Multichip TTNN decoder layer for Qwen3-Coder-30B-A3B-Instruct, 4 Blackhole dies.
+
+Stages 03 and 04. Stage 04 optimized this file **in place**; the parallelisation
+below is stage 03's and is unchanged, and what stage 04 changed is where the
+activations live inside a layer:
+
+* **Both residual RMSNorms are width-sharded over 8 cores** rather than running
+  on one (``decode_residual_norm``). 19.82 -> 4.92 us each, and *more* accurate
+  than the call they replace. The shard spec is deliberately
+  ``_width_sharded_l1(2048)``, so the first norm's output feeds the qkv
+  projection with no conversion at all.
+* **The router projection reads that L1 shard** instead of DRAM-interleaved.
+  24.62 -> **5.85 us** at the shipped 8-core norm shard (the same sweep's 4-core
+  leg reads 4.30, but 4 cores is not what ships -- the norm shards over
+  ``_NORM_SHARD_CORES = 8``), output **bit-identical**, which is what keeps the
+  four dies agreeing on the top-8. In the layer it is row 182, 6.241 us.
+* **The two collectives use caller-owned persistent buffers**
+  (``_decode_ccl_buffers``), so nothing in the forward path allocates inside the
+  trace.
+* **Decode collectives use one ethernet link, not two** (``NUM_LINKS_DECODE``).
+  Stage 03 measured this at 0.6% and kept 2 for a single code path; against the
+  stage-04 layer it is **1.22%**, over six passes with the leg order alternating
+  so that a position effect cannot be read as a link effect. Prefill keeps both.
+
+Decode layer device time 414.661 -> 362.828 us on device 0 (1.143x); traced
+decode at ctx 128, 0.4767 -> 0.4286 ms (1.112x), and 0.4700 -> 0.4282 measured
+before and after in one process by ``probes/layer_levers.py``. The inter-layer contract is untouched: a layer
+takes and returns a replicated ``[1, 1, B, 2048]`` bf16 TILE DRAM tensor with no
+collective, gather or reshard between layers. Everything is in
+``doc/optimized_multichip_decoder/``.
+
+The single-chip baseline is ``optimized_decoder.py`` -- every program
+config, dtype and fidelity constant it measured is imported rather than
+re-derived, and the multichip path is the *same graph* with three changes:
+
+1. **Attention is tensor-parallel by 4.** Each die owns 8 Q heads, 1 K head and
+   1 V head, so ``wqkv`` is ``[2048, 1280]`` per die and ``wo`` is
+   ``[1024, 2048]``. Both still satisfy ``_dram_sharded_ok`` (1280 = 5x256,
+   1024 = 4x256), so stage 02's DRAM-sharded decode projections survive intact.
+2. **Experts are expert-parallel by 4.** Each die owns 32 whole experts. M, N
+   and K of both ``sparse_matmul`` calls are unchanged, which is the entire
+   reason EP was chosen over splitting ``moe_intermediate`` -- see
+   ``doc/multichip_decoder/mesh_plan.md`` section 2.
+3. **Two all-reduces per layer**, one after ``wo`` and one after the expert
+   reduce, so the residual stream stays a replicated ``[1, 1, B, 2048]``. That
+   makes the layer's input contract identical to its output contract and lets 48
+   of them stack with no boundary conversion.
+
+Router, both RMSNorms and the residual are replicated. Routing needs a global
+view of 128 logits for top-8, and a 128-wide ``topk`` occupies one core, so
+there is nothing to fracture; the price is that 25.18% of the single-die decode
+layer is replicated work -- 129.09 us of 512.65, the two residual RMSNorms
+(40.23) plus the router block (88.86) -- which caps decode at 3.97x even at
+infinite dies.
+
+Mesh and fabric
+---------------
+1x4, ``FabricConfig.FABRIC_1D_RING`` before mesh open, ``Topology.Ring`` and
+``num_links=2`` on every collective. All three are deliberate:
+``tt_ccl.default_topology()`` returns ``Topology.Linear`` for a 4-device mesh
+(it only special-cases 8-device T3K/Galaxy), and the cluster descriptor for this
+host -- ``ClusterType.P300_X2``, two p300 boards -- shows a genuine closed
+4-ring with two ethernet links on every hop. Measured cost of taking the default
+instead: 1.21x at decode size, 1.79x at 2 MB
+(``doc/multichip_decoder/mesh_plan.md`` section 5).
+
+Both all-reduces, in both modes, are reduce-scatter followed by all-gather. The
+design phase expected decode to want AG-of-partials instead, on a standalone
+sweep that measured 19.96 us against RS+AG's 23.69 at ``[1,1,32,2048]``; the
+shipped decode tensor has **one** logical row rather than 32, which makes
+``ttnn.sum`` pull a ``FillPad``, and measured on the real layer the order
+reverses. See ``all_reduce`` for the profile rows and the A/B.
+
+The ``nnz`` contract, which is a device hang if you get it wrong
+--------------------------------------------------------------
+``ttnn.sparse_matmul`` bakes ``nnz`` into the kernel as a compile-time arg and
+requires ``count_nonzero(sparsity) == nnz`` exactly;
+``sparse_matmul_device_operation.cpp:205-211`` says a mismatch *deadlocks the
+device* (tt-metal #45943), silently unless the watcher is on.
+
+A TTNN mesh op is SPMD -- one program, one ``nnz``, four dies. Under EP the
+number of locally-live experts is the number of the global top-8 that landed in
+this die's 32-expert window: data-dependent, different on every die, anywhere in
+0..8. There is no single correct value, so **decode must pass ``nnz=None``**,
+which switches the sender to reading the sparsity page at runtime. That costs
+0.79 us per slot per matmul, measured at the shipped shapes: the pair costs
+158.01 us dynamic against 107.73 with an exact nnz, **1.47x**
+(``probes/nnz_cost_probe.py``). That is 50 us over 32 slots -- affordable only
+because EP already cut E from 128 to 32, where the same rate would cost 200. The
+two decisions are coupled.
+
+**Prefill keeps an exact ``nnz``.** Its sparsity is per 32-token tile and with
+32 tokens x top-8 = 256 selections over 128 experts essentially every expert is
+live, so the shipped path uses an all-ones mask; under EP that means all 32
+local experts are live on every die, deterministically, and
+``nnz = 32 * group_size`` is exact and identical across dies.
+
+Rejected alternatives, with the measurement, are in
+``doc/multichip_decoder/mesh_plan.md`` section 6.
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass, field, replace
+
+import torch
+
+import ttnn
+from models.common.modules.tt_ccl import TT_CCL
+
+from .functional_decoder import (
+    AttentionConfig,
+    AttentionWeights,
+    DecoderLayerConfig,
+    KVCache,
+    MoEConfig,
+    apply_rope_llama,
+    attention_prefill,
+    rope_transformation_matrix,
+)
+
+# The four precision constants this module used to import from here are gone:
+# every one of them is now read off the ``PrecisionConfig`` threaded through the
+# functions below, so importing the import-time default would have been the bug
+# this stage exists to remove. ``tt/precision.py`` holds the values.
+from .optimized_decoder import (
+    _DRAM_BANKS,
+    OptimizedWeights,
+    _attention_compute_kernel_config,
+    _bank_row,
+    _dram_sharded_ok,
+    _expert_compute_kernel_config,
+    _ones_column,
+    _tuned_sparse_matmul_config,
+    _width_sharded_l1,
+    attention_decode_optimized,
+    moe_prefill_optimized,
+)
+from .precision import DEFAULT_PRECISION, PrecisionConfig  # noqa: F401  (re-exported)
+from .weight_mapping import hf_to_meta_channels, permute_head_vector_to_meta, permute_wqkv_to_meta
+
+# The target mesh. This module deliberately supports exactly one shape: the
+# goal is the best use of *this* machine, and every constant below -- the head
+# split, the 32-expert window, the ring topology, the two links -- is chosen
+# against it. A 1x2 or 1x8 mesh would want different answers, not a scaled
+# version of these ones.
+MESH_SHAPE = (1, 4)
+NUM_DEVICES = 4
+
+# Ring, not Linear. See the module docstring; this must be passed explicitly
+# because tt_ccl.default_topology() returns Linear for a 4-device mesh.
+TOPOLOGY = ttnn.Topology.Ring
+
+# Both ethernet links on every hop are used **in prefill**, where the payload is
+# large enough to be bandwidth-bound: 1.84x at 2 MB.
+NUM_LINKS = 2
+
+# Decode uses **one**. Stage 03 measured 1 link at 0.4738 ms against 2 links'
+# 0.4766, called 0.6% noise-level, and kept 2 for a single code path. Against
+# the stage-04 layer -- where the collectives are a larger share because
+# everything around them got smaller, and where they no longer allocate -- the
+# gap is 1.22% and output is bit-identical (``probes/links_probe.py``, six
+# passes with the **leg order alternating**, so that a position effect cannot be
+# read as a link effect):
+#
+#     posA  2 links 0.4342  0.4341  0.4340     1 link 0.4290  0.4288  0.4286
+#     posB  2 links 0.4341  0.4337  0.4339     1 link 0.4291  0.4283  0.4287
+#
+#     mean  2 links 0.43400      1 link 0.42875      1.22%
+#
+# Each configuration reads the same at both positions, which is what rules out
+# the alternative explanation -- that the leg running first in a pass is simply
+# slower. That control was added because review found ``_links`` had stopped
+# honouring an explicit ``num_links=2``, leaving the probe unable to tell its
+# own legs apart; the figure survived the repair, its reproducibility did not
+# and now does. 5.25 us on the layer against a leg-against-itself spread of
+# 0.5-0.8 us. A decode
+# collective moves 128 KB per die and is latency-bound, so the second link buys
+# no bandwidth and costs the split and merge. ``all_reduce`` branches on the
+# same ``S <= 32`` test ``_decode_ccl_buffers`` uses, so prefill keeps both.
+NUM_LINKS_DECODE = 1
+
+
+def _links(x: ttnn.Tensor, ctx: "MeshContext") -> int:
+    """``ctx.num_links`` for prefill, ``ctx.decode_num_links`` for decode.
+
+    The two counts are **separate fields** rather than one field plus a
+    "differs from the default means override" test. That test was the first
+    spelling here and it is not expressible: a caller asking explicitly for
+    ``num_links=2`` at decode passes ``ctx.num_links == NUM_LINKS``, which the
+    test read as "no override" and silently gave 1 link. ``links_probe.py``
+    builds its two-link leg exactly that way, so the probe that established
+    ``NUM_LINKS_DECODE`` could not have been re-run against it -- review caught
+    this. Two fields make each mode's count independently settable and the
+    probe's legs actually different.
+    """
+    return ctx.decode_num_links if int(x.shape[-2]) <= 32 else ctx.num_links
+
+
+# Decode's two expert intermediates under EP:
+#
+#     batch * 32 experts * 32 padded rows * (2*768 + 2048) cols * 2 B
+#         = batch * 7,340,032 B = batch * 7.34 MB
+#
+# a quarter of the single-die figure, because EP fractures the expert dimension
+# these tensors are indexed by. Stage 02's 40 MB threshold sat between its batch
+# 1 (29.4 MB) and batch 2 (58.8 MB) and its own comment calls it "asserted, not
+# measured"; inherited here it would have admitted batch 5 by accident. Swept
+# instead (``probes/l1_budget_probe.py``, eager, ms per decode step):
+#
+#     batch   intermediates      L1        DRAM
+#         1          7.34 MB   1.7336    1.6924
+#         2         14.68      1.7816    1.7930
+#         4         29.36      1.8887    1.9753
+#         8         58.72      2.4969    2.9082
+#        16        117.44      3.3040    4.0386
+#        32        234.88   allocator refuses  (bank_manager.cpp:462)
+#
+# L1 wins from batch 2 to 16 and stops being allocatable at 32, so the threshold
+# goes between 117.44 and 234.88 MB. The batch-1 row of that sweep reads the
+# other way, but it is an *eager* measurement where host dispatch is most of the
+# 1.7 ms; the warmed traced A/B that decides the shipped configuration says L1,
+# clearly -- **0.4766 ms against DRAM's 0.5128, 7.6%**
+# (``probes/decode_levers.py``). Batch 1 is also the latency target.
+_DECODE_EXPERT_L1_BUDGET_BYTES = 128 * 1024 * 1024
+
+
+def _decode_expert_memory_config(batch: int, local_moe: MoEConfig) -> ttnn.MemoryConfig:
+    padded_rows = batch * local_moe.num_experts * 32
+    nbytes = padded_rows * (2 * local_moe.moe_intermediate_size + local_moe.hidden_size) * 2
+    return ttnn.L1_MEMORY_CONFIG if nbytes <= _DECODE_EXPERT_L1_BUDGET_BYTES else ttnn.DRAM_MEMORY_CONFIG
+
+
+# --- Meta-ordered rotary for decode (stage 04) --------------------------------
+
+
+def _meta_rope(ctx: MeshContext, cos_cache: ttnn.Tensor, sin_cache: ttnn.Tensor, head_dim: int):
+    """Return a ``rope(t, cos, sin, token_index)`` callable using the llama op.
+
+    **Measured, and not adopted.** Kept runnable rather than deleted, on the
+    same principle as ``router_forward_threshold``: the finding is the useful
+    part, and a future stage that changes prefill should not have to rediscover
+    it. Nothing on the shipped path calls this --
+    ``decoder_layer_decode_multichip`` ships the HF op, and
+    ``upload_multichip_weights`` builds the Meta weights only under
+    ``meta_rope=True``, so the shipped upload pays no DRAM for it.
+
+    ``rotary_embedding_llama`` costs **1.26 us** against the shipped HF op's
+    3.84 at the per-die decode shape, with ``max|diff|`` exactly 0.0 and PCC
+    1.0000000 (``probes/rope_probe.py``). Both run on one core: the llama
+    decode factory shards over *batch*, not heads, so at batch 1 none of the
+    3.05x is parallelism -- it is the activation living in L1 and a kernel that
+    multiplies by a resident 32x32 matrix instead of gathering a cos/sin row out
+    of a DRAM cache. Same lever as the router projection, different op.
+
+    Two things are hoisted out of the forward path, both on the first (eager)
+    call and cached on ``ctx``:
+
+    * the **Meta cos/sin** for this ``token_index``, read off the HF device
+      cache once, permuted on the host and uploaded already sharded. This hoist
+      is what makes *this* wiring unreplayable, and it is a property of the
+      wiring rather than of the op: ``rotary_embedding_llama`` takes cos/sin as
+      tensors and no position argument at all
+      (``rotary_embedding_llama_nanobind.cpp:38-44``), so it can be driven from
+      a position tensor inside a trace. The shipped
+      ``rotary_embedding(..., token_index)`` genuinely cannot, which is why
+      stage 05 moved decode to ``rotary_embedding_hf``. An earlier revision of
+      this docstring claimed neither spelling could; that was wrong.
+    * the **transformation matrix**, which is position-independent.
+
+    The Meta *channel order* is not established here at all: it is a property of
+    ``ctx``-independent weights, applied once by
+    ``weight_mapping.permute_wqkv_to_meta`` at upload.
+
+    **Why it is not adopted.** RoPE runs *before* K is written to the cache, so
+    the cache inherits the rotary's channel convention. Prefill is untouched by
+    this lever and writes HF-ordered keys; a Meta-ordered decode Q then scores
+    against them, and the dot products are meaningless.
+    ``probes/rope_layer_probe.py``:
+
+        fresh KV cache          PCC 0.9999697    the rotary itself is right
+        prefill-primed cache    PCC 0.1932974    the cache convention is not
+
+    The op-level probe looked clean precisely because its cache was fresh. So
+    the lever is not decode-local: adopting it means adopting the llama rotary
+    in **prefill** as well, permuting the interleaved ``wqkv`` prefill copy, and
+    changing the KV cache's channel convention -- which
+    ``test_per_die_kv_heads_stitched`` compares against a single-chip cache and
+    which ``config/context_contract.json`` describes. That is a whole-layer change,
+    not the in-place decode optimization this stage is.
+
+    A second cost, smaller and independent: the qkv weight dtype
+    (``PrecisionConfig.attention_qkv_dtype``) is ``bfloat8_b`` by default, and
+    bfloat8_b's 16-element blocks share an exponent, so permuting
+    channels **regroups the blocks** and requantizes. The two paths therefore
+    are not bit-identical in the layer even where the ops are -- attention out
+    ``max|diff|`` 1.221e-04 on a fresh cache, and the K cache differs by
+    3.125e-01 after permuting back. "Bit-identical" is a property of the op at
+    fixed input, not of the layer at permuted weights.
+    """
+    st = ctx.rope_meta
+
+    def trans_mat(batch: int):
+        # One 32x32 copy **per batch core**, because the decode factory shards
+        # over batch: at batch 1 that is a single tile on a single core, at
+        # batch 32 it is 32 of them. Keyed by batch for that reason.
+        t = st.get(("tm", batch))
+        if t is None:
+            t = st[("tm", batch)] = ttnn.from_torch(
+                rope_transformation_matrix().repeat(1, 1, batch, 1),
+                device=ctx.mesh,
+                layout=ttnn.TILE_LAYOUT,
+                dtype=ttnn.bfloat16,
+                memory_config=_head_shard(32, 32, batch),
+                mesh_mapper=ttnn.ReplicateTensorToMesh(ctx.mesh),
+            )
+        return t
+
+    def rope(t: ttnn.Tensor, _cos, _sin, token_index: int) -> ttnn.Tensor:
+        batch = int(t.shape[1])
+        key = (int(token_index), batch, head_dim)
+        pair = st.get(key)
+        if pair is None:
+            if "host" not in st:
+                # Read the HF cos/sin cache back once and permute on the host.
+                # Replicated, so die 0's copy is the whole tensor.
+                comp = ttnn.ConcatMeshToTensor(ctx.mesh, dim=0)
+                st["host"] = (
+                    ttnn.to_torch(cos_cache, mesh_composer=comp)[:1].float(),
+                    ttnn.to_torch(sin_cache, mesh_composer=comp)[:1].float(),
+                )
+            perm = hf_to_meta_channels(head_dim)
+            mem = _head_shard(32, head_dim, batch)
+            up = []
+            for c in st["host"]:
+                row = c[:, :, token_index : token_index + 1, :]
+                row = row.expand(1, 1, 32, head_dim).contiguous()[..., perm]
+                up.append(
+                    ttnn.from_torch(
+                        row.expand(1, batch, 32, head_dim).contiguous(),
+                        device=ctx.mesh,
+                        layout=ttnn.TILE_LAYOUT,
+                        dtype=ttnn.bfloat16,
+                        memory_config=mem,
+                        mesh_mapper=ttnn.ReplicateTensorToMesh(ctx.mesh),
+                    )
+                )
+            pair = st[key] = tuple(up)
+        sharded = ttnn.to_memory_config(t, _head_shard(32, head_dim, int(t.shape[1])))
+        out = apply_rope_llama(sharded, pair[0], pair[1], trans_mat(batch))
+        ttnn.deallocate(sharded)
+        return out
+
+    return rope
+
+
+def _head_shard(rows: int, cols: int, batch: int) -> ttnn.MemoryConfig:
+    """The height-sharded L1 config ``nlp_create_qkv_heads_decode`` emits and
+    ``rotary_embedding_llama``'s decode factory requires: one core per user,
+    each holding that user's whole ``[32 padded heads, head_dim]`` block."""
+    gx = min(batch, 8)
+    while batch % gx:
+        gx -= 1
+    gy = batch // gx
+    return ttnn.MemoryConfig(
+        ttnn.TensorMemoryLayout.HEIGHT_SHARDED,
+        ttnn.BufferType.L1,
+        ttnn.ShardSpec(
+            ttnn.CoreRangeSet({ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(gx - 1, gy - 1))}),
+            [rows, cols],
+            ttnn.ShardOrientation.ROW_MAJOR,
+        ),
+    )
+
+
+# --- mesh plumbing ------------------------------------------------------------
+
+
+@dataclass
+class MeshContext:
+    """Mesh, CCL semaphores and the collective parameters, owned explicitly.
+
+    ``TT_CCL`` is instantiated directly rather than through
+    ``tt_ccl.get_tt_ccl()``: that helper caches by ``mesh_device.id()`` in a
+    module-global dict, and the pytest ``mesh_device`` fixture is
+    function-scoped, so a later mesh can be handed a recycled id and inherit
+    semaphores belonging to a closed device. Ownership here is per-caller and
+    dies with the caller.
+
+    Global semaphores are hardware resources allocated at construction time,
+    which is also what makes this trace-safe: nothing in the forward path
+    allocates one.
+    """
+
+    mesh: ttnn.MeshDevice
+    ccl: TT_CCL
+    num_devices: int = NUM_DEVICES
+    num_links: int = NUM_LINKS
+    topology: ttnn.Topology = TOPOLOGY
+    # Links for a *decode* collective, separately settable. See ``_links``.
+    decode_num_links: int = NUM_LINKS_DECODE
+    # Stage 04. Meta-ordered rotary state for the decode path: the 32x32
+    # transformation matrix, the host-side Meta cos/sin caches, and the
+    # per-position sharded cos/sin pair. Keyed by ``token_index``, which is a
+    # Python int here exactly as it is for the shipped HF op -- the rotary
+    # position is baked into a traced program either way, so the gather is
+    # hoisted out of the forward path rather than run per token. Allocated on a
+    # miss, so the first call at each position must be eager, which is the same
+    # discipline ``ccl_buffers`` below already imposes. See ``_meta_rope``.
+    rope_meta: dict = field(default_factory=dict)
+    # Stage 04. Persistent collective buffers, keyed by (logical shape, padded
+    # shape, dtype), so that neither the reduce-scatter nor the all-gather
+    # allocates inside the trace. See ``_decode_ccl_buffers``. Owned by the
+    # context and therefore by the caller, exactly like the semaphores above.
+    ccl_buffers: dict = field(default_factory=dict)
+
+
+def mesh_context(mesh_device) -> MeshContext:
+    """Build the CCL context for the 4-die mesh, asserting the shape."""
+    n = mesh_device.get_num_devices()
+    assert n == NUM_DEVICES, (
+        f"multichip_decoder targets exactly {NUM_DEVICES} dies (the full P300_X2 mesh); got {n}. "
+        "Smaller meshes are out of scope by design -- the head split, expert window and ring "
+        "topology are all chosen against the 4-die shape."
+    )
+    return MeshContext(mesh=mesh_device, ccl=TT_CCL(mesh_device))
+
+
+def all_reduce(x: ttnn.Tensor, ctx: MeshContext, precision: PrecisionConfig = DEFAULT_PRECISION) -> ttnn.Tensor:
+    """All-reduce a ``[1, 1, ., H]`` partial as reduce-scatter then all-gather.
+
+    **One spelling for both modes, and that is a change from the plan.**
+    ``mesh_plan.md`` §5 chose AG-of-partials-plus-local-sum for decode on a
+    standalone sweep that measured 19.96 us against RS+AG's 23.69 at
+    ``[1,1,32,2048]``. That sweep used the wrong shape. The shipped decode tensor
+    is ``[1,1,1,2048]`` -- **one** logical row padded to a tile -- and
+    ``ttnn.sum`` over a tensor whose last two dims are not both tile-aligned
+    drags a ``FillPad`` behind it (``fill_pad.cpp:17-24``), which is precisely
+    the hazard stage 02 removed from the router. Read off
+    ``ops_perf_multichip_decode_agsum.csv`` -- a profile of this layer with the
+    plan's spelling, kept precisely because the shipped path no longer produces
+    those rows (``probes/profile_layer.py decode-agsum``) -- the local sum is not
+    one op but four:
+
+        AllGatherAsync 22.31 us + FillPad 5.89 + FastReduceNC 2.44 + Slice 1.32
+            = 31.96 us   (attention all-reduce)
+        AllGatherAsync 18.65 + FillPad 5.64 + FastReduceNC 2.43 + Slice 1.31
+            = 28.04 us   (expert all-reduce)
+
+    against the 19.96 the probe promised for each. Measured on the whole traced
+    layer at ctx 128, median of 100 (``probes/allreduce_ab.py``):
+
+        AG(dim 0) + ttnn.sum   0.4801 ms
+        reduce-scatter + all-gather   **0.4760 ms**   <- adopted
+
+    0.9%, which is small -- but it is also three fewer ops, one code path
+    instead of two, and it is the direction the standalone probe got backwards.
+    A third leg that kept the single collective and reshaped the logical shape
+    up to the padded 32 rows to dodge the ``FillPad`` **did not run**:
+    ``reshape_common.cpp:50`` rejects it, ``new_volume == old_volume``.
+
+    Prefill was RS+AG already and for the reason that still holds: 76.85 us
+    against AG-of-partials' 121.72 at ``[1,1,512,2048]``, because past ~128 KB
+    per device the collective is bandwidth-bound and RS+AG moves a quarter of
+    the bytes on each of its two hops.
+
+    The scatter axis is dim 3 (hidden, 2048), which is independent of the
+    sequence length -- that is what keeps non-aligned S working through the
+    collective without any padding of its own.
+
+    ``precision.ccl_dtype`` is ``None`` on the shipped path, which means "run
+    the collective at whatever dtype the partial arrives in" -- no cast, no
+    extra op, the behaviour every stage-02..06 number was measured at. A named
+    dtype casts in before the reduce-scatter and back out after the all-gather,
+    so a sweep can price a narrower wire without touching the arithmetic that
+    feeds it. The cast is deliberately *outside* the buffer cache key's reach
+    only in the sense that the cache keys on ``x.dtype`` already -- casting
+    first means the cached buffers are allocated at the wire dtype, which is the
+    point.
+    """
+    # The cast allocates a *new* tensor and leaves ``x`` alone: every caller
+    # deallocates the partial it passed in, so freeing it here would be a double
+    # free the moment ``ccl_dtype`` was set.
+    wire_dtype = precision.ccl_dtype
+    restore_dtype = None
+    cast_in = None
+    if wire_dtype is not None and x.dtype != wire_dtype:
+        restore_dtype = x.dtype
+        cast_in = ttnn.typecast(x, wire_dtype)
+        x = cast_in
+    bufs = _decode_ccl_buffers(x, ctx)
+    num_links = _links(x, ctx)
+    scattered = ttnn.experimental.reduce_scatter_minimal_async(
+        x,
+        persistent_output_buffers=None if bufs is None else bufs[0],
+        dim=3,
+        multi_device_global_semaphore=ctx.ccl.get_and_cycle_rs_semaphore_handles(),
+        barrier_semaphore=ctx.ccl.get_and_cycle_barrier_semaphore_handle(),
+        num_links=num_links,
+        memory_config=ttnn.DRAM_MEMORY_CONFIG,
+        intermediate_memory_config=ttnn.DRAM_MEMORY_CONFIG,
+        topology=ctx.topology,
+    )
+    gathered = ttnn.experimental.all_gather_async(
+        scattered,
+        persistent_output_buffer=None if bufs is None else bufs[1],
+        dim=3,
+        multi_device_global_semaphore=ctx.ccl.get_and_cycle_ag_semaphore_handles(),
+        barrier_semaphore=ctx.ccl.get_and_cycle_barrier_semaphore_handle(),
+        num_links=num_links,
+        memory_config=ttnn.DRAM_MEMORY_CONFIG,
+        topology=ctx.topology,
+    )
+    if bufs is None:
+        ttnn.deallocate(scattered)
+        out = gathered
+    else:
+        # ``gathered`` *is* the persistent buffer, which the caller is about to
+        # deallocate. Hand back a copy so the buffer survives the next token.
+        out = ttnn.clone(gathered, memory_config=ttnn.DRAM_MEMORY_CONFIG)
+    if restore_dtype is not None:
+        cast_out = ttnn.typecast(out, restore_dtype)
+        ttnn.deallocate(out)
+        out = cast_out
+    if cast_in is not None:
+        ttnn.deallocate(cast_in)
+    return out
+
+
+def _decode_ccl_buffers(x: ttnn.Tensor, ctx: MeshContext):
+    """``([rs intermediate, rs output, penult], ag output)`` for a decode-shaped
+    ``x``, allocated once per (logical shape, padded shape, dtype) and cached on
+    the context.
+
+    ``None`` for anything taller than one 32-row tile: prefill runs at a
+    different ``S`` on every call, so caching there would allocate a set per
+    sequence length for a lever worth 0.2% at decode and nothing measurable at
+    prefill (the prefill collective is bandwidth-bound, not allocation-bound).
+
+    **Allocated on the first call at each shape, so that call must be eager.**
+    ``ttnn.from_torch`` inside ``begin_trace_capture`` raises "Writes are not
+    supported during trace capture" and leaves the trace open, which is a hung
+    mesh and a ``tt-smi -r``. Every harness here runs the layer once before
+    capturing, which is also what the semaphores in ``MeshContext`` already
+    require; the constraint is not new, only wider.
+
+    All 48 layers of the stacked model share the cache, and so does every token.
+    That is safe because the trace serialises the collectives and each result is
+    cloned out before the next one starts -- but it is exactly the property a
+    future change has to preserve, so it is exercised by
+    ``test_multichip_decode_20_steps_deterministic`` running 20 tokens through
+    the same buffers and by ``test_two_layers_stacked``.
+
+    The layer's *two* all-reduces do **not** share a set; see the key below.
+
+    Measured: 0.4343 / 0.4337 ms against the allocating path's 0.4348 / 0.4346,
+    over two interleaved passes (``probes/layer_levers3.py``), and 0.4335 /
+    0.4333 against 0.4348 / 0.4346 in ``probes/layer_levers2.py``.
+    """
+    if int(x.shape[-2]) > 32:
+        return None
+    # The key must carry the **logical** shape, not just the padded one.
+    #
+    # The layer's two all-reduces are both [1,1,batch,2048] and *do* share one
+    # set, correctly: the attention partial is ``batch`` rows because
+    # ``_concat_heads_decode`` slices the padded tile back before ``wo``, and the
+    # expert partial is ``batch`` by construction. What collides is the priming
+    # prefill: at ``S <= 32`` it takes this branch too, and a 32-token prefill
+    # and a decode at ``batch < 32`` have the same *padded* shape, one 32-row
+    # tile. A persistent output buffer imposes *its* logical shape on the op's
+    # result, so keyed on the padded shape alone the decode layer inherited the
+    # prefill's 32 rows and silently returned a 32-row tensor. Not hypothetical
+    # -- six decode tests caught it (``work_log.md`` section 5).
+    key = (tuple(int(v) for v in x.shape), tuple(int(v) for v in x.padded_shape), str(x.dtype))
+    entry = ctx.ccl_buffers.get(key)
+    if entry is None:
+        interm, penult = ttnn.experimental.reduce_scatter_minimal_async_create_intermediate_buffer(
+            x, dim=3, topology=ctx.topology, cluster_axis=None
+        )
+        shape = list(x.shape)
+        shape[3] //= ctx.num_devices
+
+        def zeros(s):
+            return ttnn.from_torch(
+                torch.zeros(s),
+                device=ctx.mesh,
+                layout=ttnn.TILE_LAYOUT,
+                dtype=x.dtype,
+                memory_config=ttnn.DRAM_MEMORY_CONFIG,
+                mesh_mapper=ttnn.ReplicateTensorToMesh(ctx.mesh),
+            )
+
+        rs_bufs = [interm, zeros(shape)] + ([penult] if penult is not None else [])
+        entry = (rs_bufs, zeros(list(x.shape)))
+        ctx.ccl_buffers[key] = entry
+    return entry
+
+
+# Kept as names so callers read as prefill/decode; both are the same collective.
+all_reduce_prefill = all_reduce
+all_reduce_decode = all_reduce
+
+
+# --- configuration ------------------------------------------------------------
+
+
+@dataclass(frozen=True)
+class MeshDecoderConfig:
+    """The global layer config plus the per-die views the kernels actually see.
+
+    ``local_attention`` carries 8 Q heads and 1 KV head; ``local_moe`` carries
+    32 experts. Every op below is handed a *local* config, which is what makes
+    the multichip layer literally the single-chip code at a quarter of the
+    shape rather than a reimplementation of it.
+    """
+
+    global_config: DecoderLayerConfig
+    local_attention: AttentionConfig
+    local_moe: MoEConfig
+    num_devices: int = NUM_DEVICES
+
+    @classmethod
+    def from_hf(cls, hf_config, num_devices: int = NUM_DEVICES) -> "MeshDecoderConfig":
+        return cls.from_global(DecoderLayerConfig.from_hf(hf_config), num_devices)
+
+    @classmethod
+    def from_global(cls, config: DecoderLayerConfig, num_devices: int = NUM_DEVICES) -> "MeshDecoderConfig":
+        a, m = config.attention, config.moe
+        assert a.num_attention_heads % num_devices == 0, f"{a.num_attention_heads} Q heads / {num_devices}"
+        assert a.num_key_value_heads % num_devices == 0, (
+            f"{a.num_key_value_heads} KV heads / {num_devices} -- this is the hard cap on the TP factor; "
+            "TP=8 would need KV-head replication and would also take wqkv's N to 640, which is not a "
+            "multiple of 8 banks x 32 = 256, so the DRAM-sharded attention path would silently vanish"
+        )
+        assert m.num_experts % num_devices == 0, f"{m.num_experts} experts / {num_devices}"
+        local_attention = AttentionConfig(
+            hidden_size=a.hidden_size,
+            num_attention_heads=a.num_attention_heads // num_devices,
+            num_key_value_heads=a.num_key_value_heads // num_devices,
+            head_dim=a.head_dim,
+            rms_norm_eps=a.rms_norm_eps,
+        )
+        local_moe = MoEConfig(
+            hidden_size=m.hidden_size,
+            num_experts=m.num_experts // num_devices,
+            num_experts_per_tok=m.num_experts_per_tok,
+            moe_intermediate_size=m.moe_intermediate_size,
+            norm_topk_prob=m.norm_topk_prob,
+        )
+        return cls(
+            global_config=config,
+            local_attention=local_attention,
+            local_moe=local_moe,
+            num_devices=num_devices,
+        )
+
+
+# --- weights ------------------------------------------------------------------
+
+
+def head_interleaved_wqkv(wqkv: torch.Tensor, config: AttentionConfig, num_devices: int) -> torch.Tensor:
+    """Permute the fused QKV columns so a contiguous 4-way split is the TP split.
+
+    **This is the one weight transform that a naive port gets wrong.** The
+    checkpoint's fused weight is ``[Wq(4096) | Wk(512) | Wv(512)]``, so a plain
+    ``ShardTensorToMesh(dim=-1)`` hands die 0 nothing but Q heads and die 3
+    nothing but K and V. Die *d* must instead own Q heads ``8d..8d+7``, K head
+    ``d`` and V head ``d``, laid out as ``[Q_local | K_local | V_local]`` --
+    which is what ``nlp_create_qkv_heads_decode(num_heads=8, num_kv_heads=1)``
+    reads on the other side.
+
+    Rebuilding the tensor in that order here means the runtime split stays a
+    plain contiguous shard, and the failure mode -- which produces no shape
+    error, only a wrong answer -- cannot come back through a different mapper.
+
+    ``wqkv`` is ``[..., hidden, (n_heads + 2*n_kv) * head_dim]``; the return has
+    the same shape with its last dim permuted.
+    """
+    n_heads, n_kv, hd = config.num_attention_heads, config.num_key_value_heads, config.head_dim
+    q_per, kv_per = n_heads // num_devices, n_kv // num_devices
+    q_end, k_end = n_heads * hd, n_heads * hd + n_kv * hd
+
+    cols = []
+    for d in range(num_devices):
+        cols.append(wqkv[..., d * q_per * hd : (d + 1) * q_per * hd])
+        cols.append(wqkv[..., q_end + d * kv_per * hd : q_end + (d + 1) * kv_per * hd])
+        cols.append(wqkv[..., k_end + d * kv_per * hd : k_end + (d + 1) * kv_per * hd])
+    out = torch.cat(cols, dim=-1)
+    assert out.shape == wqkv.shape
+    return out
+
+
+@dataclass
+class MultichipWeights:
+    """Everything one multichip decoder layer reads.
+
+    ``experts`` is an ``OptimizedWeights`` whose tensors are mesh-sharded: the
+    two expert weights on the expert dimension, ``wqkv``/``wo`` on the
+    head-interleaved column and Q-head row split respectively. The dataclass is
+    reused unchanged so ``attention_decode_optimized`` and
+    ``moe_prefill_optimized`` can be called directly.
+
+    ``expert_window`` is the only genuinely *device-varying* constant in the
+    layer: a one-hot ``[1, 1, 128, 32]`` matrix, different on every die, that
+    slices this die's 32-expert column window out of the replicated dense
+    routing vector. See ``router_forward_multichip``.
+    """
+
+    input_layernorm: ttnn.Tensor
+    post_attention_layernorm: ttnn.Tensor
+    router: ttnn.Tensor
+    expert_window: ttnn.Tensor
+    experts: OptimizedWeights
+    # Stage 04. ``ttnn.rms_norm``'s sharded program factory reads its weight as a
+    # ROW_MAJOR ``[1, 1, dim/32, 32]`` tensor rather than the tiled ``[1,1,1,dim]``
+    # the interleaved factory takes, so decode carries a second copy of each of
+    # the two residual norm vectors. 4 KB each against the layer's 95.5 MB.
+    input_layernorm_rm: ttnn.Tensor | None = None
+    post_attention_layernorm_rm: ttnn.Tensor | None = None
+    # Stage 04. The same ``OptimizedWeights`` with the Q and K channels of
+    # ``wqkv_decode`` -- and of ``q_norm``/``k_norm``, which Qwen3 applies
+    # between the head split and RoPE -- reordered to the Meta convention
+    # ``rotary_embedding_llama`` requires. Decode-only: ``wqkv`` (prefill's
+    # interleaved copy), ``wo`` and the expert weights are the *same objects*,
+    # not copies, so this costs one extra DRAM-sharded qkv (11.14 MB/4 per die)
+    # and two 128-element vectors, and prefill cannot reach it.
+    experts_meta: OptimizedWeights | None = None
+
+
+# SDPA-decode's tree reduction is capped at 6 rounds, i.e. 64 cores per KV head
+# (``sdpa_decode_program_factory.cpp:245``). With no program config the op sets
+# ``max_cores_per_head = num_cores_available``, so at TP=4 -- one KV head per die
+# -- batch 1 asks for all 110 worker cores on that single head and the op raises
+#
+#     Tree reduction max 6 rounds (64 cores/head), got 110 cores/head
+#
+# This is a *new* failure created by the head split: at the single-chip 4 KV
+# heads the same arithmetic gives 27 cores/head. It only bites the contiguous
+# cache path -- the paged one runs at the default -- and only at small batch,
+# because ``num_cores_per_head`` divides by the batch. Capping the per-head core
+# budget at the op's own limit fixes it without giving up any parallelism the op
+# would have been allowed to use.
+_SDPA_MAX_CORES_PER_HEAD = 64
+
+# The **paged** path -- the one the full model actually runs -- had no program
+# config at all until stage 06, because the cap above was added to clear a
+# ``TT_FATAL`` the paged path never raised. Running at the op default is not
+# free: with no config the op picks its own ``k_chunk_size`` and core split, and
+# the result is a decode cost that is **linear in ``cur_pos``** rather than
+# flat. Measured at the shipped per-die decode shapes -- 8 Q heads, 1 KV head,
+# head_dim 128, page 32, batch 1, **bfloat16** cache -- with PCC taken against a
+# float32 reference built from the same cache the kernel reads, not against the
+# default leg (``probes/sdpa_sweep_confirm.py``, median of 5 blocks of 50):
+#
+#     cur_pos      default   k256/c16   speedup   default PCC   k256/c16 PCC
+#         127     23.72 us   19.00 us     1.25x      0.999734       0.999707
+#        1023    120.51      22.02        5.47x      0.999714       0.999703
+#        4095    451.85      30.75       14.69x      0.999519       0.999655
+#        8191    893.14      38.15       23.41x      0.999024       0.999590
+#       16383   1777.60      49.83       35.67x      0.993199       0.999577
+#       32767   3545.05      74.13       47.82x      0.989703       0.999692
+#
+# Two things in that table, not one. The speed column is the expected one. The
+# **PCC columns are the surprise**: the default's accuracy *decays with depth* --
+# 0.9932 at 16k and 0.9897 at 32k, through this project's 0.995 layer bar --
+# while the configured path holds 0.9996-0.9997 flat from 127 to 32767. So this
+# is not a speed-for-accuracy trade. At the context this model advertises the
+# config is strictly better on both axes, and the shipped default was the *less*
+# accurate of the two.
+#
+# **The cache dtype is why this took two passes, and it is the lesson.** The
+# stage-06 lever analysis recommended ``k_chunk_size=512`` on the strength of
+# probes that allocated the cache as ``bfloat8_b``; ``create_mesh_kv_cache``
+# allocates ``ttnn.bfloat16`` (see below, ~line 1167). Re-run at the real dtype,
+# 512 loses its edge -- and, far worse, **512 is wrong in-model**:
+# ``test_multichip_decode_batch`` (128-position paged cache, cur_pos 32) returns
+# PCC **-0.04 to -0.17** against HF with it, nondeterministically in 2-3 of its 4
+# batch sizes, across nine runs. Sweeping that real test pins the boundary
+# exactly -- ``k_chunk`` in {32, 64, 128, 256} passes 4/4 at every
+# ``max_cores_per_head_batch`` in {16, 32, 64}; only 512 fails -- so
+# ``max_cores`` is innocent and ``k_chunk`` is the whole effect.
+#
+# No standalone construction reproduces it. ``probes/sdpa_kchunk_rule_probe.py``
+# re-runs the op at bfloat16, at the failing 128-deep cache, with an 8-user paged
+# page table laid out exactly as ``create_mesh_kv_cache`` lays it out, and reads
+# PCC 0.9997 at k512 at every depth from 128 to 4096;
+# ``probes/sdpa_shallow_cache_probe.py`` finds nothing either. The leading
+# explanation is **L1 pressure**: standalone the op owns the whole of L1, while
+# in-model it is co-resident with the layer's sharded activations, expert
+# weights and CCL buffers, and a 512-deep bf16 K chunk is exactly the size that
+# stops fitting. That the boundary is dtype-linked is independently visible --
+# ``k1024/c64`` fails to *build* at bfloat16 (``program.cpp:1722``) and builds
+# fine at bfloat8_b. It is recorded as unexplained-in-detail rather than argued;
+# what is measured is that 512 is unsafe in-model and 256 is not.
+#
+# So: the sweep was redone at bfloat16, 6 x 4 points at five positions
+# (``probes/sdpa_sweep_probe.py``), finalists re-timed over nine positions, and
+# the choice restricted to the in-model-safe ``k_chunk <= 256``. **256/16 is the
+# uniform winner** -- fastest of the safe configs at cur_pos 4095 and above,
+# within 0.6% at 511-2047, and its worst point is +6.4% at cur_pos 127 (19.00 vs
+# 17.86 us for 256/8, i.e. 0.05 ms on a 20 ms iteration). There is no
+# context-dependence worth a runtime switch, so it is **fixed**; a traced decode
+# could not vary it per step anyway. ``q_chunk_size`` stays 32: decode has one
+# query row and 32 is the tile height.
+_SDPA_PAGED_K_CHUNK = 256
+_SDPA_PAGED_MAX_CORES_PER_HEAD = 16
+
+
+def _paged_cache_depth(kv_cache) -> int:
+    """Positions allocated **per user** in a paged cache.
+
+    ``page_table`` is ``[max_batch, blocks_per_seq]`` and every block holds
+    ``block_size`` positions, so this is the length of the logical sequence the
+    cache can hold for one user -- which is the quantity ``k_chunk_size`` has to
+    respect. See ``_sdpa_k_chunk``.
+    """
+    return int(kv_cache.page_table.shape[-1]) * int(kv_cache.block_size)
+
+
+def _sdpa_k_chunk(kv_cache) -> int:
+    """``_SDPA_PAGED_K_CHUNK``, clamped to what the cache can actually supply.
+
+    **``k_chunk_size`` must not exceed the cache's per-user allocated depth**, and
+    exceeding it does not raise -- it silently returns garbage. This is the whole
+    reason the first adoption of this lever failed its gates, and it is worth
+    stating precisely because nothing in the op signature hints at it.
+
+    How it presents: ``test_multichip_decode_batch`` allocates a **128**-position
+    paged cache. At ``k_chunk_size=256`` it returns PCC **-0.10 to +0.06** against
+    HF -- noise, not a degraded answer -- but *only when another test has run
+    before it in the same process*; run alone it passes. Run the same test after
+    ``test_router_windows_partition_global_routing`` and it fails 4/4, at every
+    ``max_cores_per_head_batch``. Sweeping ``k_chunk`` through that reproducer
+    puts the boundary exactly at the cache depth:
+
+        k_chunk  32   64   128  ->  7 passed, at max_cores in {8, 16, 32, 64}
+        k_chunk 256          ->  4 failed
+
+    That order-dependence is the tell, and it is what makes the bug so easy to
+    miss: the op reads a full ``k_chunk`` past the end of the cache buffer, and
+    whether that hurts depends on what the allocator last left there. On a fresh
+    device it is zeros and the softmax mask hides it; after another test has
+    allocated and freed tensors it is live garbage. **Every standalone probe
+    misses this by construction** -- ``probes/sdpa_shallow_cache_probe.py`` and
+    ``probes/sdpa_kchunk_rule_probe.py`` both reproduce the shapes, the dtype, the
+    128-deep cache and the multi-user page table exactly, and both read PCC 0.9997
+    at k512, because in a probe the cache is the only thing allocated. This is the
+    same shape of miss as the stage-04 ``rotary_embedding_llama`` rejection: a
+    probe that structurally cannot see the state interaction.
+
+    So the clamp is not defensive coding, it is the operating range. At the
+    shipped ``max_context_len`` (4096 and up, contract 262144) it never binds and
+    the config is the tuned 256; at the tests' 128-deep caches it drops to 128,
+    which the sweep prices at +6% on the op at cur_pos 127 and 0% past 511.
+    """
+    depth = _paged_cache_depth(kv_cache)
+    # The ``max(32, ...)`` floor exists because SDPA will not take a chunk below
+    # one tile. It is the one input that could make this function *violate* the
+    # invariant in its own first line, and only when ``block_size < 32``, which
+    # this model never configures (the block size is 32 and the page table is at
+    # least one block per user). Assert it rather than leave a silent hole: a
+    # shallower cache than one tile would return a chunk deeper than the cache.
+    assert depth >= 32, (
+        f"paged cache depth {depth} is below one tile, so the 32-row floor below would return a "
+        "k_chunk_size deeper than the per-user allocated depth -- which SDPA reads past without "
+        "raising. Raise block_size (currently "
+        f"{int(kv_cache.block_size)}) or the page table width ({int(kv_cache.page_table.shape[-1])})."
+    )
+    chunk = min(_SDPA_PAGED_K_CHUNK, max(32, depth))
+    # SDPA wants a power-of-two chunk; take the largest one that still fits.
+    return 1 << (chunk.bit_length() - 1)
+
+
+#: Program configs are immutable and there are at most a handful of distinct
+#: ones, but they are built **per layer per call** -- 48 times a token on the
+#: decode path and 48 times a prefill chunk. Each build calls
+#: ``device.compute_with_storage_grid_size()``, which is a device query, not a
+#: Python attribute. On the traced decode path that is capture-only and free; on
+#: the *untraced* paths (``run_teacher_forcing``, ``run_prefill_check``, eager
+#: decode) it is 96 device queries per token of pure host time. Memoised on the
+#: grid size rather than the device handle so the cache survives device reopen.
+_SDPA_CONFIG_CACHE: dict = {}
+
+
+def _cached_sdpa_config(grid, q_chunk, k_chunk, max_cores=None):
+    key = (grid.x, grid.y, q_chunk, k_chunk, max_cores)
+    cfg = _SDPA_CONFIG_CACHE.get(key)
+    if cfg is None:
+        kwargs = {} if max_cores is None else {"max_cores_per_head_batch": max_cores}
+        cfg = _SDPA_CONFIG_CACHE[key] = ttnn.SDPAProgramConfig(
+            compute_with_storage_grid_size=grid, q_chunk_size=q_chunk, k_chunk_size=k_chunk, **kwargs
+        )
+    return cfg
+
+
+def _sdpa_program_config(device, kv_cache=None):
+    """Program config for SDPA-decode; ``kv_cache`` selects the tuned paged form.
+
+    Both spellings are the same op family and the same maths; they differ only
+    in chunking and core budget. Neither touches dtype, fidelity, the KV cache
+    layout, or any collective -- this is a program config on a call the model
+    already makes.
+    """
+    paged = kv_cache is not None and kv_cache.is_paged
+    return _cached_sdpa_config(
+        device.compute_with_storage_grid_size(),
+        32,
+        _sdpa_k_chunk(kv_cache) if paged else 32,
+        _SDPA_PAGED_MAX_CORES_PER_HEAD if paged else _SDPA_MAX_CORES_PER_HEAD,
+    )
+
+
+# Prefill has the *same* gap and it is larger in absolute terms:
+# ``attention_prefill`` also called SDPA with no program config, and the op
+# default is quadratic-with-a-bad-constant in S. Same shapes, bfloat16
+# (``doc/optimized_full_model/probes/sdpa_prefill_confirm.py``):
+#
+#       S    default   q128/k128   q256/k256
+#     128    23.92 us    25.72 us    32.68 us
+#     512    58.96       54.14       87.18
+#    1024   230.58       88.36      127.54
+#    2048   741.08      216.03      207.28
+#    4096  2850.25      882.61      451.04
+#    8192 10938.43     2907.58     1956.15
+#   16384 44456.67    11364.22     6527.48
+#
+# The winner *is* length-dependent here, unlike decode, and the two legs cross
+# at S ~= 2048. Prefill could pick at call time -- it is eager, not traced, so
+# the branch is a Python ``if`` and there is no captured trace to invalidate.
+#
+# **It is nevertheless NOT adopted.** ``decoder_layer_prefill_multichip`` passes
+# ``sdpa_program_config=None`` and prefill runs at the op default. Two measured
+# reasons, in this order:
+#
+# 1. **It costs accuracy on the one gate that can see it.** With this config
+#    wired, ``run_teacher_forcing`` reads top-1 **0.980** against a baseline of
+#    **0.990** on the same tree (top-5 and top-100 stay 1.000). Bisected: the
+#    *decode* config alone holds 0.990, the *prefill* config alone drops it to
+#    0.980, so the flip is this and not the lever above
+#    (``logs/run_teacher_forcing_leg_prefill.log`` /
+#    ``logs/run_teacher_forcing_leg_decode.log``). One greedy token in a hundred
+#    is small, but the stage bar is "do not spend accuracy for speed", and here
+#    there is no speed to buy it with, which is reason 2.
+# 2. **At the length actually being served it is a loss, not a win.** The
+#    readiness reference prompt is **158 tokens**. The table above says the
+#    config is *behind* the default below S ~= 384, and the measured TTFTs agree:
+#    3448.79 ms baseline against 3445.31 ms with prefill configured -- noise. The
+#    6.8x is real but it lives at S >= 4096, which nothing in the current gate
+#    set exercises.
+#
+# So this is a **verified-fast, accuracy-ungated** lever, left wired and
+# documented rather than taken. What it needs before adoption is a readiness
+# reference with a multi-thousand-token prompt, so the regime where it pays
+# (S >= 4096, 6.3-6.8x on the SDPA op, and 48 of them per prefill) is the same
+# regime the accuracy gate covers. The seam in ``attention_prefill`` exists for
+# exactly that -- same pattern as ``_meta_rope``, which is also built, measured
+# and not adopted.
+#
+# **Arbitrary S keeps working**, so nothing here is blocked on alignment. This
+# was checked and not assumed: S in
+# {1, 3, 31, 33, 100, 129, 255, 257, 1000, 1023, 1025, 2049, 4095, 4097, 5000}
+# all build and run under both chunkings with PCC identical to the default's to
+# five decimals. Prefill is not chunked in this model -- ``prefill_forward``
+# feeds each user's whole logical length in -- so that property is load-bearing
+# for the stage contract, not a nicety.
+#
+# ``q512/k512`` is **rejected** outright: it fails to build at *every* length
+# including 128 (``program.cpp:1722``), so it is a resource limit, not an
+# alignment rule.
+_SDPA_PREFILL_CROSSOVER = 2048
+
+
+def _sdpa_prefill_program_config(device, seq_len: int):
+    """Built and measured; **not wired in**. See the note above before adopting."""
+    q_chunk = 256 if seq_len >= _SDPA_PREFILL_CROSSOVER else 128
+    return _cached_sdpa_config(device.compute_with_storage_grid_size(), q_chunk, q_chunk)
+
+
+# --- decode residual RMSNorm, width-sharded (stage 04) ------------------------
+#
+# The stage-03 decode layer spent 40.21 us -- 9.7% -- in its two residual
+# RMSNorms, and the profile says why: both run on **one core**
+# (``../doc/multichip_decoder/ops_perf_multichip_decode.csv.gz``, device 0, rows
+# 134 and 159, 20.081 and 20.127 us, ``CORE COUNT`` 1). A 2048-wide bf16 norm
+# over one 32-row tile is 128 KB in 20 us, i.e. 6.5 GB/s, which is a single
+# core's share of L1 bandwidth and nothing else.
+#
+# ``ttnn.rms_norm`` has a sharded program factory that splits the row across a
+# core grid. Feeding it the same L1 width-shard the DRAM-sharded qkv projection
+# already wants -- 8 cores, one per DRAM bank, ``[32, 256]`` -- gives
+# (``probes/norm_accuracy_probe.py``, trace slope, median of 30):
+#
+#     interleaved, no compute config (shipped)   19.82 us   max|err vs fp64| 6.711e-02
+#     sharded  4 cores, HiFi4 fp32acc             7.53          1.439e-02
+#     sharded  8 cores, default                   4.26          3.586e-02
+#     sharded  8 cores, HiFi4 fp32acc             4.92          1.686e-02
+#                                        i2s 0.51 us, s2i 0.53 us
+#
+# 8 cores at HiFi4 with fp32 accumulation is **4.0x faster and 4.0x more
+# accurate** than the shipped call, which passes no compute config at all and so
+# accumulates the sum of squares in bf16. The reference is torch fp64 over the
+# bf16-rounded inputs the device actually sees, so "more accurate" is against
+# the mathematical answer rather than against the other kernel.
+#
+# 16 cores and beyond do not pay: the norm itself stops improving (a 2048-wide
+# row is 64 tiles, so 8 cores already hold 8 tiles each) while the resharding at
+# both ends grows with the core count.
+_NORM_SHARD_CORES = _DRAM_BANKS
+
+
+def _norm_shard_config(dim: int) -> ttnn.MemoryConfig:
+    """The L1 width-shard the sharded norm reads and writes.
+
+    Deliberately ``_width_sharded_l1(dim)``'s spec: at ``dim == hidden_size``
+    this is bit-for-bit the memory config ``attention_decode_optimized`` reshards
+    its input into, so the first norm's output feeds the qkv projection with no
+    conversion at all.
+    """
+    return ttnn.MemoryConfig(
+        ttnn.TensorMemoryLayout.WIDTH_SHARDED,
+        ttnn.BufferType.L1,
+        ttnn.ShardSpec(_bank_row(_NORM_SHARD_CORES), [32, dim // _NORM_SHARD_CORES], ttnn.ShardOrientation.ROW_MAJOR),
+    )
+
+
+def _norm_program_config(dim: int):
+    block_w = dim // _NORM_SHARD_CORES // 32
+    subblock_w = next(w for w in (4, 3, 2, 1) if block_w % w == 0)
+    return ttnn.LayerNormShardedMultiCoreProgramConfig(
+        compute_with_storage_grid_size=[_NORM_SHARD_CORES, 1],
+        subblock_w=subblock_w,
+        block_h=1,  # decode's padded M is exactly one 32-row tile; batch is capped at 32
+        block_w=block_w,
+        inplace=False,
+    )
+
+
+def _norm_compute_config(device, precision: PrecisionConfig = DEFAULT_PRECISION):
+    return ttnn.init_device_compute_kernel_config(
+        device.arch(),
+        math_fidelity=precision.norm_fidelity,
+        math_approx_mode=False,
+        fp32_dest_acc_en=True,
+        packer_l1_acc=True,
+    )
+
+
+def decode_residual_norm(
+    x: ttnn.Tensor, weight_rm: ttnn.Tensor, eps: float, precision: PrecisionConfig = DEFAULT_PRECISION
+) -> ttnn.Tensor:
+    """One residual-stream RMSNorm at decode shape, width-sharded across 8 cores.
+
+    Takes a DRAM-interleaved ``[1, 1, B, H]`` (B <= 32, padded to one tile) and
+    returns an **L1 width-sharded** tensor in ``_norm_shard_config(H)``. Callers
+    that need it interleaved say so; the first norm's consumer does not.
+    """
+    dim = int(x.shape[-1])
+    assert int(x.shape[-2]) <= 32, (
+        f"decode_residual_norm shards a single 32-row tile; got {int(x.shape[-2])} rows. "
+        "Prefill uses the interleaved rms_norm."
+    )
+    mc = _norm_shard_config(dim)
+    xs = ttnn.to_memory_config(x, mc)
+    out = ttnn.rms_norm(
+        xs,
+        weight=weight_rm,
+        epsilon=eps,
+        program_config=_norm_program_config(dim),
+        memory_config=mc,
+        # ``precision`` rather than the default: this is the only site
+        # ``norm_fidelity`` reaches. It was called with the module default until
+        # the stage-07 review, which meant the field was a documented knob with
+        # no effect and ``R21_norm_hifi2`` measured nothing. The prefill norms
+        # (``decoder_layer_prefill_multichip``) pass no compute config at all and
+        # still take the op default -- ``norm_fidelity`` is a decode-path field,
+        # which is the path the stage ranks on.
+        compute_kernel_config=_norm_compute_config(x.device(), precision),
+    )
+    ttnn.deallocate(xs)
+    return out
+
+
+def _exact_matmul_config(device, precision: PrecisionConfig = DEFAULT_PRECISION):
+    """HiFi4, so the one-hot window matmul is a copy rather than an approximation.
+
+    The matmul default is LoFi, which keeps ~5 mantissa bits, and that is fine
+    for everything else in this layer -- but here the operand is 0/1 and the
+    intent is to *select* a routing weight, not to compute with it. Measured, the
+    LoFi spelling moved the stitched windows 9.77e-4 away from the single-chip
+    dense routing (one bf16 ulp at these magnitudes) where HiFi4 reproduces them
+    bit-for-bit. The tensor is 4 tiles by 1, so exactness is free.
+    """
+    return ttnn.init_device_compute_kernel_config(
+        device.arch(),
+        math_fidelity=precision.router_window_fidelity,
+        math_approx_mode=False,
+        fp32_dest_acc_en=False,
+        packer_l1_acc=False,
+    )
+
+
+def _expert_window_matrix(mesh_device, num_experts: int, num_devices: int) -> ttnn.Tensor:
+    """Per-die one-hot selector ``[1, 1, E, E/num_devices]``.
+
+    A TTNN mesh op is SPMD: one program on four dies, so ``ttnn.slice`` cannot
+    take a different start offset per die and there is no way to ask for
+    "columns 32d..32d+31" directly. The device-varying constant is built the
+    only way a mesh tensor can vary by device -- a leading dim of ``num_devices``
+    sharded on dim 0 -- and applied as a matmul.
+
+    The matmul is exact, not approximate: the operand is 0/1, the accumulator is
+    fp32 and the output is bf16, so each selected weight is copied bit-for-bit.
+    K = 128 is 4 tiles and N = 32 is 1, so it is the cheapest op in the router
+    block, and it *replaces* work rather than adding it -- the divide that
+    follows now runs over 32 columns instead of 128.
+    """
+    local = num_experts // num_devices
+    sel = torch.zeros(num_devices, 1, num_experts, local)
+    for d in range(num_devices):
+        for j in range(local):
+            sel[d, 0, d * local + j, j] = 1.0
+    return ttnn.from_torch(
+        sel,
+        dtype=ttnn.bfloat16,
+        layout=ttnn.TILE_LAYOUT,
+        device=mesh_device,
+        memory_config=ttnn.DRAM_MEMORY_CONFIG,
+        mesh_mapper=ttnn.ShardTensorToMesh(mesh_device, dim=0),
+    )
+
+
+def upload_multichip_weights(
+    torch_weights: dict[str, torch.Tensor],
+    mesh_device,
+    config: MeshDecoderConfig,
+    expert_dtype=None,
+    meta_rope: bool = False,
+    precision: PrecisionConfig = DEFAULT_PRECISION,
+) -> MultichipWeights:
+    """Shard and upload one layer's weights across the mesh.
+
+    Per-die footprint at the shipped dtypes, which is the table
+    ``doc/multichip_decoder/mesh_plan.md`` section 2 computes:
+
+        gate_up  [1, 32, 2048, 1536] bfloat4_b   56.623 MB
+        down     [1, 32,  768, 2048] bfloat4_b   28.312 MB
+        wqkv     [2048, 1280] bfloat8_b x2 copies 5.570 MB
+        wo       [1024, 2048] bfloat8_b x2 copies 4.456 MB
+        router   [2048, 128]  bf16               0.524 MB
+        norms + qk-norms                         0.009 MB
+                                          total 95.49 MB / layer / die
+
+    Every division is exact -- 2048/4, 32/4, 4/4, 128/4 -- so this scheme needs
+    **zero load-time padding**, and the contract's allowance for it goes unused.
+    """
+    a = config.global_config.attention
+    n = config.num_devices
+    # ``expert_dtype`` is the stage-04 spelling and still wins when given (the
+    # multichip tests sweep it); otherwise every dtype below comes from
+    # ``precision``, whose defaults are the values this docstring's table was
+    # measured at.
+    gate_up_dtype = expert_dtype if expert_dtype is not None else precision.experts_gate_up_dtype
+    down_dtype = expert_dtype if expert_dtype is not None else precision.experts_down_dtype
+
+    def replicate(t: torch.Tensor, tensor_dtype, memory_config=ttnn.DRAM_MEMORY_CONFIG) -> ttnn.Tensor:
+        return ttnn.from_torch(
+            t.contiguous().float(),
+            dtype=tensor_dtype,
+            layout=ttnn.TILE_LAYOUT,
+            device=mesh_device,
+            memory_config=memory_config,
+            mesh_mapper=ttnn.ReplicateTensorToMesh(mesh_device),
+        )
+
+    def shard(t: torch.Tensor, dim: int, tensor_dtype, memory_config=ttnn.DRAM_MEMORY_CONFIG) -> ttnn.Tensor:
+        return ttnn.from_torch(
+            t.contiguous().float(),
+            dtype=tensor_dtype,
+            layout=ttnn.TILE_LAYOUT,
+            device=mesh_device,
+            memory_config=memory_config,
+            mesh_mapper=ttnn.ShardTensorToMesh(mesh_device, dim=dim),
+        )
+
+    def as_4d(t: torch.Tensor, pad_to_4d: bool = False) -> torch.Tensor:
+        if pad_to_4d:
+            t = t.reshape(1, 1, 1, -1)
+        while t.dim() < 4:
+            t = t.unsqueeze(0)
+        return t
+
+    wqkv = head_interleaved_wqkv(as_4d(torch_weights["wqkv"]), a, n)
+    wo = as_4d(torch_weights["wo"])
+
+    # Per-die shapes, used by both the shard spec and the assertions below.
+    k_qkv, n_qkv = int(wqkv.shape[-2]), int(wqkv.shape[-1]) // n
+    k_o, n_o = int(wo.shape[-2]) // n, int(wo.shape[-1])
+    assert _dram_sharded_ok(k_qkv, n_qkv), (
+        f"per-die wqkv [{k_qkv}, {n_qkv}] is not bank-divisible; the DRAM-sharded decode "
+        "projections would silently fall back to interleaved and give back stage 02's 1.11x"
+    )
+    assert _dram_sharded_ok(k_o, n_o), f"per-die wo [{k_o}, {n_o}] is not bank-divisible"
+
+    def dram_sharded(t: torch.Tensor, dim: int, k: int, n_local: int, tensor_dtype) -> ttnn.Tensor:
+        """Width-shard the per-die weight one shard per DRAM bank, then mesh-shard it.
+
+        Two independent shardings compose here and it is worth being explicit
+        about which is which: ``mesh_mapper`` fractures the tensor *across dies*
+        (TP), while ``memory_config`` fractures each die's piece across that
+        die's 8 DRAM banks (stage 02's decode projection layout). The shard spec
+        is therefore written in per-die elements, not global ones.
+        """
+        return shard(
+            t,
+            dim,
+            tensor_dtype,
+            ttnn.MemoryConfig(
+                ttnn.TensorMemoryLayout.WIDTH_SHARDED,
+                ttnn.BufferType.DRAM,
+                ttnn.ShardSpec(_bank_row(_DRAM_BANKS), [k, n_local // _DRAM_BANKS], ttnn.ShardOrientation.ROW_MAJOR),
+            ),
+        )
+
+    experts = OptimizedWeights(
+        # [E, 2I, H] -> [1, E, H, 2I], sharded on the expert dim.
+        gate_up_proj=shard(torch_weights["experts_gate_up"].transpose(-2, -1).unsqueeze(0), 1, gate_up_dtype),
+        # [E, H, I] -> [1, E, I, H], sharded on the expert dim.
+        down_proj=shard(torch_weights["experts_down"].transpose(-2, -1).unsqueeze(0), 1, down_dtype),
+        attention=AttentionWeights(
+            # Column split (head-interleaved, see head_interleaved_wqkv).
+            wqkv=shard(wqkv, -1, precision.attention_qkv_dtype),
+            # Row split by Q head. Contiguous *because* the Q head assignment
+            # above is contiguous per die: die d owns rows 1024d..1024d+1023.
+            wo=shard(wo, -2, precision.attention_wo_dtype),
+            q_norm=replicate(as_4d(torch_weights["q_norm"], pad_to_4d=True), precision.norm_weight_dtype),
+            k_norm=replicate(as_4d(torch_weights["k_norm"], pad_to_4d=True), precision.norm_weight_dtype),
+        ),
+        wqkv_decode=dram_sharded(wqkv, -1, k_qkv, n_qkv, precision.attention_qkv_dtype),
+        wo_decode=dram_sharded(wo, -2, k_o, n_o, precision.attention_wo_dtype),
+    )
+
+    # Stage 04. The Meta-ordered decode twin, built **only when asked**. It is
+    # not the shipped path (see ``_meta_rope``); it exists so
+    # ``probes/rope_layer_probe.py`` can re-measure the rejection rather than
+    # cite it. Off by default, so the shipped upload pays no extra DRAM.
+    experts_meta = None
+    if meta_rope:
+        # The channel permutation is applied to the *pre-interleave* wqkv, which
+        # is safe because it reorders channels **within** a head and
+        # ``head_interleaved_wqkv`` only reorders whole heads -- the two commute.
+        # V is untouched, and so are ``wo`` and every expert weight, which are
+        # shared objects here rather than copies.
+        wqkv_meta = head_interleaved_wqkv(
+            permute_wqkv_to_meta(
+                as_4d(torch_weights["wqkv"]),
+                n_heads=a.num_attention_heads,
+                n_kv_heads=a.num_key_value_heads,
+                head_dim=a.head_dim,
+            ),
+            a,
+            n,
+        )
+        experts_meta = replace(
+            experts,
+            attention=replace(
+                experts.attention,
+                q_norm=replicate(
+                    as_4d(permute_head_vector_to_meta(torch_weights["q_norm"], head_dim=a.head_dim), pad_to_4d=True),
+                    precision.norm_weight_dtype,
+                ),
+                k_norm=replicate(
+                    as_4d(permute_head_vector_to_meta(torch_weights["k_norm"], head_dim=a.head_dim), pad_to_4d=True),
+                    precision.norm_weight_dtype,
+                ),
+            ),
+            wqkv_decode=dram_sharded(wqkv_meta, -1, k_qkv, n_qkv, precision.attention_qkv_dtype),
+        )
+
+    router = torch_weights["router"]
+
+    def norm_row_major(t: torch.Tensor) -> ttnn.Tensor:
+        """The same vector the tiled copy holds, in the layout the sharded
+        ``rms_norm`` program factory reads: ROW_MAJOR ``[1, 1, dim/32, 32]``."""
+        flat = t.reshape(-1)
+        return ttnn.from_torch(
+            flat.reshape(1, 1, flat.numel() // 32, 32).contiguous().float(),
+            dtype=precision.norm_weight_dtype,
+            layout=ttnn.ROW_MAJOR_LAYOUT,
+            device=mesh_device,
+            memory_config=ttnn.DRAM_MEMORY_CONFIG,
+            mesh_mapper=ttnn.ReplicateTensorToMesh(mesh_device),
+        )
+
+    return MultichipWeights(
+        input_layernorm=replicate(torch_weights["input_layernorm"].reshape(1, 1, 1, -1), precision.norm_weight_dtype),
+        post_attention_layernorm=replicate(
+            torch_weights["post_attention_layernorm"].reshape(1, 1, 1, -1), precision.norm_weight_dtype
+        ),
+        router=replicate(router.T.contiguous().reshape(1, 1, router.shape[1], router.shape[0]), precision.router_dtype),
+        expert_window=_expert_window_matrix(mesh_device, config.global_config.moe.num_experts, n),
+        experts=experts,
+        experts_meta=experts_meta,
+        input_layernorm_rm=norm_row_major(torch_weights["input_layernorm"]),
+        post_attention_layernorm_rm=norm_row_major(torch_weights["post_attention_layernorm"]),
+    )
+
+
+def create_mesh_kv_cache(
+    mesh_device,
+    config: MeshDecoderConfig,
+    max_batch: int,
+    max_seq_len: int,
+    block_size: int | None = None,
+    precision: PrecisionConfig = DEFAULT_PRECISION,
+) -> KVCache:
+    """Allocate the *local* KV cache: 1 KV head per die, not 4.
+
+    This is where the TP factor buys capacity rather than speed. Per die the
+    cache is ``[.., 1, .., 128]`` instead of ``[.., 4, .., 128]``, i.e. 512 B
+    per token per layer instead of 2048 -- 6.44 GB at the advertised 262144
+    context over 48 layers, against 25.77 GB on one die. One die cannot hold
+    this model at full context; four can, with room to spare. See
+    ``config/context_contract.json``.
+
+    The buffers are *replicated at allocation* because they are zeros, and
+    diverge the moment the first token is written -- each die holds a different
+    KV head. The page table is genuinely identical on every die: paging is a
+    logical-to-physical block mapping and does not depend on which head lives
+    where.
+    """
+    local = config.local_attention
+    if block_size is None:
+        shape = (max_batch, local.num_key_value_heads, max_seq_len, local.head_dim)
+        page_table = None
+    else:
+        blocks_per_seq = math.ceil(max_seq_len / block_size)
+        shape = (max_batch * blocks_per_seq, local.num_key_value_heads, block_size, local.head_dim)
+        page_table = ttnn.from_torch(
+            torch.arange(max_batch * blocks_per_seq, dtype=torch.int32).reshape(max_batch, blocks_per_seq),
+            dtype=ttnn.int32,
+            layout=ttnn.ROW_MAJOR_LAYOUT,
+            device=mesh_device,
+            mesh_mapper=ttnn.ReplicateTensorToMesh(mesh_device),
+        )
+
+    k, v = (
+        ttnn.from_torch(
+            torch.zeros(shape),
+            dtype=precision.kv_cache_dtype,
+            layout=ttnn.TILE_LAYOUT,
+            device=mesh_device,
+            memory_config=ttnn.DRAM_MEMORY_CONFIG,
+            mesh_mapper=ttnn.ReplicateTensorToMesh(mesh_device),
+        )
+        for _ in range(2)
+    )
+    return KVCache(k=k, v=v, page_table=page_table, block_size=block_size or 0)
+
+
+def build_local_sparsity(mesh_device, local_moe: MoEConfig) -> ttnn.Tensor:
+    """All-ones prefill sparsity over this die's 32 experts, replicated."""
+    return ttnn.from_torch(
+        torch.ones(1, 1, 1, local_moe.num_experts, dtype=torch.bfloat16),
+        layout=ttnn.ROW_MAJOR_LAYOUT,
+        dtype=ttnn.bfloat16,
+        device=mesh_device,
+        mesh_mapper=ttnn.ReplicateTensorToMesh(mesh_device),
+    )
+
+
+# --- router -------------------------------------------------------------------
+
+
+def router_forward_threshold(
+    x: ttnn.Tensor,
+    w_router: ttnn.Tensor,
+    window: ttnn.Tensor,
+    config: MoEConfig,
+    local_moe: MoEConfig,
+) -> ttnn.Tensor:
+    """``router_forward_multichip`` with the dense vector built by a threshold
+    comparison instead of a scatter, so nothing leaves TILE layout.
+
+    Stage 03 inherited stage 02's routing tail, whose shape is
+    ``topk -> untilize(zeros) / untilize(indices) / untilize(values) ->
+    scatter -> tilize``: ``ttnn.scatter`` only accepts ROW_MAJOR, and every
+    consumer of the dense vector (both matmuls, the divide) needs TILE. Stage 02
+    recorded that round trip as *not removable* on those grounds. It is
+    removable -- by not scattering.
+
+    ``topk(sorted=True)`` already returns the 8th-largest logit in column 7, and
+    the top-8 set is exactly ``{j : logit_j >= that}``. So the same dense vector
+    is
+
+        dense = exp(logits - top_max) * (logits >= top_logits[..., 7])
+
+    computed over all 128 columns, entirely in TILE. The surviving values are
+    ``ttnn.exp`` of the same fp32 inputs the scatter path fed it, so the result
+    is bit-identical -- **unless two logits tie exactly at rank 8**, in which
+    case this selects both and the scatter path selects one. With fp32 logits
+    accumulated over K=2048 that does not happen on real weights, and
+    ``test_router_windows_partition_global_routing`` asserts the equality at
+    ``max |diff| = 0.0``, so a tie would fail loudly rather than drift.
+
+    **Measured, and rejected.** It removes rows 190-197 of the stage-04 decode
+    profile -- ``zeros_like`` 1.210, two ``typecast`` 2.537, three ``untilize``
+    4.654, ``scatter`` 3.030, ``tilize`` 5.576 = **17.007 us** -- and is
+    nonetheless **0.8% slower on the layer**: 0.4382 / 0.4382 ms against the
+    shipped 0.4348 / 0.4346 over two interleaved passes
+    (``doc/optimized_multichip_decoder/probes/layer_levers2.py``). Widening the
+    softmax's ``sub`` and ``exp`` from 8 columns to 128, plus the ``ge`` and the
+    ``mul``, costs more than the layout conversions save. The output is
+    bit-identical on all four dies (``max|diff| 0.000e+00``), which is also the
+    evidence that no two logits tie at rank 8.
+
+    Kept rather than deleted because the arithmetic is the useful part: stage
+    02 recorded this round trip as *not removable*, and it is.
+    """
+    assert config.norm_topk_prob
+    e = config.num_experts
+
+    logits = ttnn.linear(x, w_router, dtype=ttnn.float32, memory_config=ttnn.DRAM_MEMORY_CONFIG)
+    top_logits, _idx = ttnn.topk(logits, k=config.num_experts_per_tok, dim=-1, largest=True, sorted=True)
+
+    rows = top_logits.shape[2]
+    top_max = ttnn.slice(top_logits, [0, 0, 0, 0], [1, 1, rows, 1])
+    cutoff = ttnn.slice(top_logits, [0, 0, 0, config.num_experts_per_tok - 1], [1, 1, rows, config.num_experts_per_tok])
+
+    # exp(l - max) over the whole row; every entry is in (0, 1], so nothing can
+    # overflow and the losers underflow towards zero before the mask even runs.
+    weights = ttnn.exp(ttnn.sub(logits, top_max))
+    dense = ttnn.typecast(ttnn.mul(weights, ttnn.ge(logits, cutoff)), ttnn.bfloat16)
+
+    total = ttnn.matmul(dense, _ones_column(x.device(), e), dtype=ttnn.bfloat16)
+    local = ttnn.matmul(dense, window, dtype=ttnn.bfloat16, compute_kernel_config=_exact_matmul_config(x.device()))
+    guarded = ttnn.maximum(total, 1e-30)
+    normalised = ttnn.div(local, guarded)
+    assert int(normalised.shape[-1]) == local_moe.num_experts
+    for t in (logits, top_logits, _idx, top_max, cutoff, weights, dense, total, local, guarded):
+        ttnn.deallocate(t)
+    return normalised
+
+
+def router_forward_multichip(
+    x: ttnn.Tensor,
+    w_router: ttnn.Tensor,
+    window: ttnn.Tensor,
+    config: MoEConfig,
+    local_moe: MoEConfig,
+    precision: PrecisionConfig = DEFAULT_PRECISION,
+) -> ttnn.Tensor:
+    """Replicated global routing, returning this die's ``[1, 1, S, 32]`` window.
+
+    Identical arithmetic to ``optimized_decoder.router_forward_optimized`` --
+    selection on raw fp32 logits, softmax over the 8 survivors, neither keepdim
+    reduction spelled as a ttnn reduction -- with one op added and one op made
+    four times narrower.
+
+    **Why the whole router is replicated.** Top-8 of 128 needs the global logit
+    vector, and there is nothing worth fracturing anyway: the router matmul has
+    N = 128 = 4 tiles so it can use 4 cores, and ``ttnn.topk`` over a single
+    128-wide row occupies exactly 1. Splitting N four ways would give each die
+    one tile and one core, and would additionally need a collective *inside* the
+    routing path to reassemble the logits before the top-k. So each die computes
+    the full 128-way routing on the bit-identical replicated activation and
+    takes its own window out of the result -- no collective, at the price of
+    88.9 us of decode device time that four dies pay in full.
+
+    **The correctness assumption this makes, stated plainly.** The four windows
+    are a partition of the global top-8 only if all four dies agree on which 8
+    experts won. The inputs are bit-identical and the program is the same, so
+    ``ttnn.topk`` should return identical indices -- but that is a *tie-breaking
+    determinism* claim, and if it ever failed the layer would be silently wrong
+    with no shape error and only a PCC drift to show for it. It is asserted
+    directly by ``test_topk_is_identical_across_dies`` rather than argued.
+    """
+    assert config.norm_topk_prob, (
+        "router selects on raw logits, which relies on the softmax denominator "
+        "cancelling during top-k renormalisation; that only holds when norm_topk_prob is True"
+    )
+
+    logits = ttnn.linear(x, w_router, dtype=ttnn.float32, memory_config=ttnn.DRAM_MEMORY_CONFIG)
+    return _router_tail(logits, window, config, local_moe, x.device(), precision)
+
+
+def _router_tail(
+    logits,
+    window,
+    config: MoEConfig,
+    local_moe: MoEConfig,
+    device,
+    precision: PrecisionConfig = DEFAULT_PRECISION,
+) -> ttnn.Tensor:
+    """Top-8, softmax over the survivors, this die's 32-expert window.
+
+    Split out of ``router_forward_multichip`` so a probe can vary where the
+    logits are produced without duplicating the tail.
+    """
+    top_logits, top_indices = ttnn.topk(logits, k=config.num_experts_per_tok, dim=-1, largest=True, sorted=True)
+
+    top_max = ttnn.slice(top_logits, [0, 0, 0, 0], [1, 1, top_logits.shape[2], 1])
+    exp_logits = ttnn.exp(ttnn.sub(top_logits, top_max))
+
+    zeros = ttnn.typecast(ttnn.zeros_like(logits), ttnn.bfloat16)
+    dense = ttnn.scatter(zeros, dim=-1, index=top_indices, src=ttnn.typecast(exp_logits, ttnn.bfloat16))
+
+    # The denominator is the sum over all 128 -- which is the sum over the 8
+    # survivors, since the scatter fills a field of exact zeros -- and must stay
+    # global: normalising within a window would renormalise each die's share to
+    # 1 and the four contributions would sum to 4.
+    total = ttnn.matmul(dense, _ones_column(device, config.num_experts), dtype=ttnn.bfloat16)
+    local = ttnn.matmul(
+        dense, window, dtype=ttnn.bfloat16, compute_kernel_config=_exact_matmul_config(device, precision)
+    )
+    # Same clamp as the single-chip router, and for the same reason: after the
+    # scatter the divide runs over whole tiles, and the tile row-padding has a
+    # zero numerator *and* a zero denominator, which unguarded ttnn.div returns
+    # as +inf. Every real row's denominator is >= 1 (sorted=True makes column 0
+    # of exp_logits exactly exp(0) = 1), so the clamp cannot touch one.
+    guarded = ttnn.maximum(total, 1e-30)
+    normalised = ttnn.div(local, guarded)
+    assert int(normalised.shape[-1]) == local_moe.num_experts
+    for t in (logits, top_logits, top_indices, top_max, exp_logits, dense, total, local, guarded):
+        ttnn.deallocate(t)
+    return normalised
+
+
+# --- experts ------------------------------------------------------------------
+
+
+def moe_decode_multichip(
+    x: ttnn.Tensor,
+    routing: ttnn.Tensor,
+    weights: OptimizedWeights,
+    local_moe: MoEConfig,
+    precision: PrecisionConfig = DEFAULT_PRECISION,
+) -> ttnn.Tensor:
+    """Decode expert pass over this die's 32 experts. Returns a *partial* sum.
+
+    Structurally ``optimized_decoder.moe_decode_optimized`` with the local
+    expert count, and one difference that is not cosmetic: **``nnz`` is
+    ``None``**.
+
+    Stage 02 passes ``nnz = top_k * batch``, exact because every one of the
+    global top-8 is computed on the single die. Under EP the count of live
+    experts in *this* die's window is data-dependent -- 0 to 8, mean 2, and
+    different on each die -- while a mesh op is SPMD and compiles one kernel for
+    all four. Passing any host-computed value would deadlock the board the first
+    time the routing was unbalanced (``sparse_matmul_device_operation.cpp``
+    205-211, tt-metal #45943), silently unless the watcher is on. ``nnz=None``
+    switches the in0 sender to reading the sparsity page at runtime and
+    multicasting a per-slot valid flag; the loop still visits all 32 slots but
+    only reads weights and does math for the live ones.
+
+    Measured cost of dynamic mode, decode M=1, bfp4/LoFi, trace-slope:
+
+        E=128 nnz=8 (single-die baseline)   139.45 + 125.20 = 264.65 us
+        E=32  nnz=None (this path)           60.67 +  63.29 = 123.96 us  2.13x
+        E=32  nnz=8  (exact, illegal here)   82.72 +  58.21 = 140.93 us
+        E=128 nnz=None (dynamic at full E)  243.08 + 249.03 = 492.11 us  0.54x
+
+    **That 2.13x did not survive.** Re-measured at the shipped shapes -- E=32,
+    M=1, bfloat4_b, LoFi, L1 output, stage 02's tuned block widths -- dynamic
+    mode costs 158.01 us against an exact ``nnz``'s 107.73, **1.47x**, and the
+    multichip decode profile reads 82.65 us for the pair against the single
+    chip's 92.06, i.e. **1.11x, not 2.13x**
+    (``probes/nnz_cost_probe.py``, ``doc/multichip_decoder/work_log.md`` section
+    8). The sweep above was a DRAM-out, random-weight microbenchmark whose E=128
+    baseline read 264.65 where the profiled layer reads 92.06; only ratios were
+    taken from it, and the ratio was still wrong, because the overhead it hid is
+    additive rather than proportional.
+
+    The E=32/nnz=8 row is also the answer to "why not capacity padding": the
+    only capacity that can never be exceeded is 8, and building a fixed-count
+    sparsity on device needs a second ``topk`` over the local 32 -- **26.32 us**
+    on one core in the decode profile, more than the ~26 us it would save.
+    Any smaller capacity drops experts, which changes the model output.
+    """
+    batch = x.shape[2]
+    n_experts = local_moe.num_experts
+    hidden_size = local_moe.hidden_size
+    inter = local_moe.moe_intermediate_size
+
+    sparsity = ttnn.to_layout(routing, ttnn.ROW_MAJOR_LAYOUT)
+    expert_memory_config = _decode_expert_memory_config(batch, local_moe)
+    output_tile = ttnn.Tile([32, 32])
+    compute_config = _expert_compute_kernel_config(x.device(), precision)
+    gate_up_config = _tuned_sparse_matmul_config(1, 2 * inter, hidden_size, precision.experts_gate_up_in0_block_w)
+    down_config = _tuned_sparse_matmul_config(1, hidden_size, inter, precision.experts_down_in0_block_w)
+
+    x_batched = ttnn.reshape(x, (1, batch, 1, hidden_size))
+    fused = ttnn.sparse_matmul(
+        x_batched,
+        weights.gate_up_proj,
+        sparsity=sparsity,
+        nnz=None,  # see docstring -- a host-computed nnz deadlocks the board here
+        memory_config=expert_memory_config,
+        output_tile=output_tile,
+        program_config=gate_up_config,
+        compute_kernel_config=compute_config,
+        dtype=precision.activation_dtype,
+    )
+    packed_width = fused.shape[-1]
+    fused = ttnn.reshape(fused, (batch, n_experts, packed_width))
+
+    half = packed_width // 2
+    gate = ttnn.slice(fused, [0, 0, 0], [batch, n_experts, half])
+    up = ttnn.slice(fused, [0, 0, half], [batch, n_experts, packed_width])
+    ttnn.deallocate(fused)
+
+    down_input = ttnn.reshape(ttnn.mul(ttnn.silu(gate), up), (batch, n_experts, 1, half))
+    ttnn.deallocate(gate)
+    ttnn.deallocate(up)
+
+    down = ttnn.sparse_matmul(
+        down_input,
+        weights.down_proj,
+        sparsity=sparsity,
+        nnz=None,
+        memory_config=expert_memory_config,
+        output_tile=output_tile,
+        program_config=down_config,
+        is_input_a_sparse=True,
+        is_input_b_sparse=False,  # selects batch_length_A = B * E; see the single-chip docstring
+        compute_kernel_config=compute_config,
+        dtype=precision.activation_dtype,
+    )
+    ttnn.deallocate(down_input)
+
+    # The multiply by the routing weight is what makes a skipped slot harmless:
+    # a die holding none of the global top-8 multiplies 32 untouched output
+    # slots by exact zero and contributes an exact zero to the all-reduce.
+    # test_expert_window_can_be_empty pins that, because "untouched" would not
+    # be enough if the op left a NaN there.
+    states = ttnn.reshape(down, (batch, n_experts, hidden_size))
+    states = ttnn.mul(states, ttnn.reshape(routing, (batch, n_experts, 1)))
+    states = ttnn.unsqueeze_to_4D(ttnn.sum(states, dim=1))
+    return ttnn.reshape(states, (1, 1, batch, hidden_size), (1, 1, max(32, batch), hidden_size))
+
+
+# --- the layer ----------------------------------------------------------------
+
+
+def decoder_layer_prefill_multichip(
+    x: ttnn.Tensor,
+    weights: MultichipWeights,
+    config: MeshDecoderConfig,
+    ctx: MeshContext,
+    cos_cache: ttnn.Tensor,
+    sin_cache: ttnn.Tensor,
+    sparsity: ttnn.Tensor,
+    kv_cache: KVCache | None = None,
+    user_id: int = 0,
+    precision: PrecisionConfig = DEFAULT_PRECISION,
+    start_pos: int = 0,
+    chunk_page_table=None,
+    fill_page_table=None,
+    fill_len: int | None = None,
+) -> ttnn.Tensor:
+    """Prefill one layer on the mesh. ``x`` / return replicated ``[1, 1, S, 2048]``.
+
+    ``S`` is arbitrary. Nothing in the multichip path adds an alignment
+    constraint: the collectives scatter on dim 3 (hidden, 2048, fixed), and the
+    only padding in play is ``moe_prefill_optimized``'s internal chunk padding,
+    which is the single-chip behaviour and is sliced back inside that function.
+    """
+    eps = config.global_config.rms_norm_eps
+
+    normed = ttnn.rms_norm(x, weight=weights.input_layernorm, epsilon=eps)
+    attn_partial = attention_prefill(
+        normed,
+        weights.experts.attention,
+        config.local_attention,
+        cos_cache,
+        sin_cache,
+        kv_cache,
+        user_id,
+        # ``None`` at the shipped precision, which is the op default and what
+        # every prefill number was measured at; see
+        # ``optimized_decoder._attention_compute_kernel_config``.
+        compute_kernel_config=_attention_compute_kernel_config(x.device(), precision),
+        activation_dtype=precision.activation_dtype,
+        # NOT adopted -- see _sdpa_prefill_program_config. The seam is wired and
+        # the config is built and measured; passing it costs a top-1 point on
+        # run_teacher_forcing, so prefill stays at the op default.
+        sdpa_program_config=None,
+        start_pos=start_pos,
+        chunk_page_table=chunk_page_table,
+        fill_page_table=fill_page_table,
+        fill_len=fill_len,
+    )
+    ttnn.deallocate(normed)
+    attn_out = all_reduce_prefill(attn_partial, ctx, precision)
+    ttnn.deallocate(attn_partial)
+    hidden = ttnn.add(x, attn_out)
+    ttnn.deallocate(attn_out)
+
+    normed = ttnn.rms_norm(hidden, weight=weights.post_attention_layernorm, epsilon=eps)
+    routing = router_forward_multichip(
+        normed, weights.router, weights.expert_window, config.global_config.moe, config.local_moe, precision
+    )
+    moe_partial = moe_prefill_optimized(normed, routing, weights.experts, config.local_moe, sparsity, precision)
+    ttnn.deallocate(normed)
+    ttnn.deallocate(routing)
+    moe_out = all_reduce_prefill(moe_partial, ctx, precision)
+    ttnn.deallocate(moe_partial)
+
+    out = ttnn.add(hidden, moe_out)
+    ttnn.deallocate(hidden)
+    ttnn.deallocate(moe_out)
+    return out
+
+
+def decoder_layer_decode_multichip(
+    x: ttnn.Tensor,
+    weights: MultichipWeights,
+    config: MeshDecoderConfig,
+    ctx: MeshContext,
+    cos_cache: ttnn.Tensor,
+    sin_cache: ttnn.Tensor,
+    kv_cache: KVCache,
+    current_pos: ttnn.Tensor,
+    token_index: int,
+    rope=None,
+    precision: PrecisionConfig = DEFAULT_PRECISION,
+    active_mask: ttnn.Tensor | None = None,
+) -> ttnn.Tensor:
+    """Decode one token per user on the mesh. ``x`` / return ``[1, 1, B, 2048]``.
+
+    Input and output layouts are the same replicated tensor, which is the point:
+    48 of these stack with no boundary conversion, and the stacked model pays
+    the two all-reduces per layer and nothing else. **That is the inter-layer
+    residual layout contract**, and stage 04 keeps it unchanged while moving
+    every *intra*-layer boundary it can into L1 shards -- see
+    ``doc/optimized_multichip_decoder/README.md``.
+
+    Both residual norms are width-sharded (``decode_residual_norm``). The first
+    one's output is already in ``attention_decode_optimized``'s qkv input shard,
+    so it crosses into attention with no conversion; the second one's output
+    feeds the router projection sharded and the expert path interleaved.
+    """
+    eps = config.global_config.rms_norm_eps
+
+    normed = decode_residual_norm(x, weights.input_layernorm_rm, eps, precision)
+    # ``rope`` is the stage-05 seam. It defaults to ``None`` and therefore to
+    # ``_apply_rope`` -- ``ttnn.experimental.rotary_embedding`` with a **Python
+    # int** ``token_index``, which is what every stage-03/04 number was measured
+    # at and what the single-layer tests still exercise. That spelling cannot be
+    # replayed: the position is a compile-time argument, so a captured trace
+    # rotates every later token at the position it was captured at. The full
+    # model therefore passes ``model._rope_decode``, which is
+    # ``ttnn.experimental.rotary_embedding_hf(is_decode_mode=True)`` reading a
+    # **per-user cos/sin pair gathered on device** from a position tensor the
+    # trace itself advances. Same HF ``rotate_half`` channel convention, so the
+    # KV cache convention, prefill, and every weight are untouched -- which is
+    # exactly what stage 04's rejected ``rotary_embedding_llama`` lever could not
+    # offer (README limitation 4).
+    # The rotary stays the **HF** op. ``rotary_embedding_llama`` is 3.05x faster
+    # standalone and bit-identical there, but it cannot be adopted for decode
+    # alone: it needs Meta channel order, and the KV cache prefill already wrote
+    # is in HF order, so SDPA would score a Meta-ordered Q against HF-ordered
+    # keys. Measured, not argued -- PCC 0.193 against a prefill-primed cache
+    # where a fresh cache reads 0.99997 (``probes/rope_layer_probe.py``). See
+    # ``_meta_rope`` and ``README.md`` limitation 4.
+    attn_partial = attention_decode_optimized(
+        normed,
+        weights.experts,
+        config.local_attention,
+        cos_cache,
+        sin_cache,
+        kv_cache,
+        current_pos,
+        token_index,
+        # Both paths are configured now. The contiguous one needs the 64-core cap
+        # to clear a TT_FATAL; the paged one -- what the full model runs -- takes
+        # the swept k256/c16 config (k clamped to the cache depth), which is flat
+        # in cur_pos where the op default is linear in it. See
+        # _sdpa_program_config and _sdpa_k_chunk.
+        sdpa_program_config=_sdpa_program_config(x.device(), kv_cache),
+        rope=rope,
+        precision=precision,
+    )
+    ttnn.deallocate(normed)
+    attn_out = all_reduce_decode(attn_partial, ctx, precision)
+    ttnn.deallocate(attn_partial)
+    hidden = ttnn.add(x, attn_out)
+    ttnn.deallocate(attn_out)
+
+    normed_sharded = decode_residual_norm(hidden, weights.post_attention_layernorm_rm, eps, precision)
+    # The router projection reads the shard directly -- N = 128 is 4 tiles, so
+    # the matmul uses 4 cores either way, but a width-sharded L1 in0 turns a
+    # 24.62 us DRAM-interleaved read into 5.85 us of L1 with bit-identical
+    # output (``probes/norm_router_probe.py``, max|diff| exactly 0.0).
+    routing = router_forward_multichip(
+        normed_sharded, weights.router, weights.expert_window, config.global_config.moe, config.local_moe, precision
+    )
+    if active_mask is not None:
+        # Zero the routing weights of every slot that holds no live request.
+        #
+        # ``routing`` *is* ``sparse_matmul``'s sparsity tensor, and its nonzero
+        # count is the amount of expert math the op does: with ``nnz=None`` the
+        # kernel reads the sparsity page at runtime and only fetches weights and
+        # multiplies for the live ``(row, expert)`` pairs. A serving decode batch
+        # is padded to the configured ``max_num_seqs`` with inactive rows, and an
+        # inactive row's garbage hidden state still routes to a full top-8 -- so
+        # without this a 32-slot server does 32 rows of expert work no matter how
+        # many users are actually connected. See
+        # ``doc/optimized_vllm/probes/batch_decode_control.py``.
+        #
+        # ``active_mask`` is derived **on device** from ``current_pos`` inside the
+        # same traced graph (``Qwen3CoderModel._decode_active_mask``), so it can
+        # never be stale: ``ttnn.plus_one(..., skip_negative_entries=True)`` leaves
+        # an inactive row at ``-1`` forever, and a row that changes hands only does
+        # so through a host reinstall of ``current_pos``.
+        gated = ttnn.mul(routing, active_mask)
+        ttnn.deallocate(routing)
+        routing = gated
+    # ``sparse_matmul``'s in0 is DRAM-interleaved, so the expert path pays one
+    # sharded-to-interleaved (0.53 us) rather than the norm paying 15.
+    normed = ttnn.sharded_to_interleaved(normed_sharded, ttnn.DRAM_MEMORY_CONFIG)
+    ttnn.deallocate(normed_sharded)
+    moe_partial = moe_decode_multichip(normed, routing, weights.experts, config.local_moe, precision)
+    ttnn.deallocate(normed)
+    ttnn.deallocate(routing)
+    moe_out = all_reduce_decode(moe_partial, ctx, precision)
+    ttnn.deallocate(moe_partial)
+
+    out = ttnn.add(hidden, moe_out)
+    ttnn.deallocate(hidden)
+    ttnn.deallocate(moe_out)
+    return out
+
+
+# Bytes one 32x32 tile occupies, per dtype. Spelled out because
+# ``Tensor.element_size()`` raises for the block-float types -- their storage is
+# a byte (or nibble) per element *plus* a shared exponent per 16-element face
+# row, i.e. 1024 + 64 for bfloat8_b and 512 + 64 for bfloat4_b -- and the
+# expert weights, which are the whole point of measuring this, are block-float.
+_TILE_BYTES = {
+    str(ttnn.float32): 4096,
+    str(ttnn.bfloat16): 2048,
+    str(ttnn.bfloat8_b): 1088,
+    str(ttnn.bfloat4_b): 576,
+}
+
+
+def _tensor_bytes(t: ttnn.Tensor) -> int | None:
+    """Device bytes one mesh-sharded tensor occupies **per die**, or ``None``.
+
+    A mesh tensor's shape is already the *local* (per-die) shape, so this is the
+    allocation a dtype change actually moves -- which is the observable
+    ``tests/test_precision_config.py`` asserts on. ``None`` for a dtype with no
+    entry above rather than a wrong number.
+    """
+    tile_bytes = _TILE_BYTES.get(str(t.dtype))
+    if tile_bytes is None:
+        return None
+    shape = [int(v) for v in t.padded_shape]
+    tiles = math.prod(shape[:-2]) * math.ceil(shape[-2] / 32) * math.ceil(shape[-1] / 32)
+    return tiles * tile_bytes
+
+
+def fallback_audit(
+    weights: MultichipWeights,
+    config: MeshDecoderConfig,
+    batch: int,
+    precision: PrecisionConfig = DEFAULT_PRECISION,
+) -> dict:
+    """Every runtime fallback the imported single-chip code can still take.
+
+    Three of stage 02's helpers choose a slower path silently rather than
+    raising, and all three have different inputs under TP/EP than they were
+    tuned against, so "it still passes PCC" would not notice any of them:
+
+    * ``_dram_sharded_usable`` -- falls back to the interleaved ``attention_decode``
+      if the weight dims were not bank-divisible at upload or the batch exceeds
+      32. Per-die N is now 1280 rather than 5120 and per-die K 1024 rather than
+      4096, and 1280 = 5x256 is only one factor of two away from failing.
+    * ``_tuned_sparse_matmul_config`` -- silently lowers ``in0_block_w`` to the
+      largest divisor of K in tiles. EP leaves K alone (2048 and 768), so the
+      tuned 16 and 12 must survive; if they did not, this would be scheme A's
+      regression arriving by the back door.
+    * ``_decode_expert_memory_config`` -- moves the expert intermediates from L1
+      to DRAM past a byte budget, which EP shrank 4x.
+
+    Since stage 07 it also reports what the *precision config actually put on
+    the device*: the dtypes read back off the uploaded tensors (not the config's
+    own fields -- those would only prove the dataclass round-trips), the block
+    widths the program configs resolved to, and the fidelities the compute
+    configs carry. That is what ``tests/test_precision_config.py`` asserts
+    against when it constructs at a non-default value.
+
+    Returned as data so a test can assert on it and the work log can quote it.
+    """
+    a = config.local_attention
+    m = config.local_moe
+    k_qkv = int(weights.experts.wqkv_decode.shape[-2]) if weights.experts.wqkv_decode is not None else None
+    n_qkv = int(weights.experts.wqkv_decode.shape[-1]) if weights.experts.wqkv_decode is not None else None
+    k_o = int(weights.experts.wo_decode.shape[-2]) if weights.experts.wo_decode is not None else None
+    n_o = int(weights.experts.wo_decode.shape[-1]) if weights.experts.wo_decode is not None else None
+    gate_up = _tuned_sparse_matmul_config(
+        1, 2 * m.moe_intermediate_size, m.hidden_size, precision.experts_gate_up_in0_block_w
+    )
+    down = _tuned_sparse_matmul_config(1, m.hidden_size, m.moe_intermediate_size, precision.experts_down_in0_block_w)
+    return {
+        "batch": batch,
+        "dram_sharded_qkv": (k_qkv, n_qkv),
+        "dram_sharded_wo": (k_o, n_o),
+        "dram_sharded_taken": weights.experts.wqkv_decode is not None
+        and weights.experts.wo_decode is not None
+        and batch <= 32,
+        "gate_up_in0_block_w": gate_up.in0_block_w,
+        "down_in0_block_w": down.in0_block_w,
+        "expert_intermediate_buffer": "L1"
+        if _decode_expert_memory_config(batch, m) == ttnn.L1_MEMORY_CONFIG
+        else "DRAM",
+        "local_heads": (a.num_attention_heads, a.num_key_value_heads),
+        "local_experts": m.num_experts,
+        # Stage 04. Not a fallback in the "silently slower path" sense -- a
+        # mismatch here raises rather than degrades -- but it is the same class
+        # of risk, so it is reported as data: if the sharded norm's output shard
+        # ever stops being *exactly* the one the DRAM-sharded qkv projection
+        # wants, TTNN inserts a reshard between them and the layer gets slower
+        # with no error at all. That single equality is what removed stage-03
+        # row 135 from the profile.
+        "norm_shard_cores": _NORM_SHARD_CORES,
+        "norm_shard_feeds_qkv_directly": _norm_shard_config(m.hidden_size) == _width_sharded_l1(m.hidden_size),
+        "decode_ccl_buffers_persistent": True,
+        # -- what the precision config actually produced on device -------------
+        # Read off the uploaded tensors, so these differ from
+        # ``precision.<field>`` if any of the threading above is broken.
+        "device_experts_gate_up_dtype": str(weights.experts.gate_up_proj.dtype),
+        "device_experts_down_dtype": str(weights.experts.down_proj.dtype),
+        "device_attention_qkv_dtype": str(weights.experts.attention.wqkv.dtype),
+        "device_attention_wo_dtype": str(weights.experts.attention.wo.dtype),
+        "device_attention_qkv_decode_dtype": (
+            None if weights.experts.wqkv_decode is None else str(weights.experts.wqkv_decode.dtype)
+        ),
+        "device_router_dtype": str(weights.router.dtype),
+        "device_norm_weight_dtype": str(weights.input_layernorm.dtype),
+        # Bytes one layer's expert weights occupy per die -- the allocation-size
+        # consequence of the two expert dtypes, in a form a sweep can diff.
+        "device_expert_bytes_per_die": (
+            _tensor_bytes(weights.experts.gate_up_proj) + _tensor_bytes(weights.experts.down_proj)
+        ),
+        "expert_math_fidelity": str(precision.experts_fidelity),
+        "attention_math_fidelity": None if precision.attention_fidelity is None else str(precision.attention_fidelity),
+        "router_window_math_fidelity": str(precision.router_window_fidelity),
+        "ccl_dtype": str(precision.effective_ccl_dtype),
+        "activation_dtype": str(precision.activation_dtype),
+    }
+
+
+__all__ = [
+    "MESH_SHAPE",
+    "NUM_DEVICES",
+    "NUM_LINKS",
+    "NUM_LINKS_DECODE",
+    "TOPOLOGY",
+    "MeshContext",
+    "MeshDecoderConfig",
+    "MultichipWeights",
+    "all_reduce",
+    "all_reduce_decode",
+    "all_reduce_prefill",
+    "build_local_sparsity",
+    "create_mesh_kv_cache",
+    "decoder_layer_decode_multichip",
+    "decoder_layer_prefill_multichip",
+    "fallback_audit",
+    "head_interleaved_wqkv",
+    "mesh_context",
+    "decode_residual_norm",
+    "moe_decode_multichip",
+    "router_forward_multichip",
+    "router_forward_threshold",
+    "upload_multichip_weights",
+]

code/models/demos/blackhole/qwen3_coder_30b_a3b/tt/optimized_decoder.py

@@ -0,0 +1,1093 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+"""Optimized TTNN decoder layer for Qwen3-Coder-30B-A3B-Instruct.
+
+Same semantics as ``functional_decoder`` -- prefill/decode contract, paged KV
+cache, non-aligned sequence lengths, determinism -- with the measured path
+retuned. Every number below is from real checkpoint weights on a Blackhole
+p300c, 1x1 mesh; the sweeps behind them are in ``doc/optimized_decoder/``.
+
+    prefill  536.54 -> 69.12 us/token           at S=512   (7.76x)
+    decode     1.5655 -> 0.5634 ms/token traced at ctx128   (2.78x)
+
+Both lines are cells of ``doc/{functional,optimized}_decoder/perf_prefill.csv``
+and ``perf_decode.csv``, which every run of ``tests/test_perf.py`` rewrites;
+the third significant figure moves between runs.
+
+What changed, in order of how much it mattered
+----------------------------------------------
+**1. ``in0_block_w`` (3.0x on prefill).** Stage 01 inherited ``in0_block_w=1``
+from the exemplar's config helper. With K = 2048 (64 tiles) that feeds the
+kernel one tile of the inner dimension at a time, which is what held the expert
+matmuls at ~5.4% of peak FLOPs -- not the core count, and not precision.
+
+**2. bfloat4_b expert weights (2.2x on prefill).** Only visible *after* the
+block-width fix: at ``in0_block_w=1`` the kernel is latency-bound, so weight
+dtype cannot matter. The two knobs interact and sweeping either alone finds the
+wrong optimum -- see ``EXPERT_IN0_BLOCK_W_GATE_UP``.
+
+**3. DRAM-sharded decode attention projections (1.11x on decode).** Once the
+experts were fast, ``o_proj`` and ``qkv`` were 21% of decode device time and
+the stage-01 "attention is 0.08% of prefill, no action" call went stale.
+``MatmulMultiCoreReuseMultiCastDRAMShardedProgramConfig`` with the weight
+width-sharded across the 8 DRAM banks takes qkv 68.3 -> 46.8 us and wo
+96.0 -> 41.7 us at the op level, and the whole traced layer 0.6508 -> 0.5863 ms
+at ctx128 measured like for like -- both legs otherwise at the configuration
+shipped at the time, on the same bfloat8_b weights, and both before lever 7
+below, which is why the fast leg reads 0.5863 rather than today's 0.5634. A
+``0.697 -> 0.587`` pair (1.19x)
+that this file and the docs used to carry is withdrawn; see
+``attention_decode_optimized`` and ``doc/optimized_decoder/work_log.md`` §5.
+
+**4. Packing gate and up (1.09x).** ``_sparse_matmul_config`` parallelises only
+over N, so N tiles cap the usable cores: 768 -> 24, 1536 -> 48, 2048 -> 64.
+Packing doubles gate/up's cores. Measured against a *properly tuned separate*
+candidate it is only 1.09x (2 x 1.476 = 2.952 ms -> 2.699 ms); against the
+untuned stage-01 candidate it looks like 1.66x, but most of that belongs to the
+block-width fix. Re-confirmed at the end of the stage on the whole layer:
+separate 0.735 ms vs packed 0.673 ms traced decode.
+
+**5. LoFi on the expert matmuls (1.05x on prefill).** bfp4 weights carry 4
+mantissa bits, so HiFi4's extra passes have nothing to work on. Prefill S=512
+72.46 -> 69.13 us/token, decode 0.6746 -> 0.6638 ms, and layer PCC is 0.99910 at
+LoFi vs 0.99909 at HiFi4 -- i.e. very slightly *better*.
+
+**6. bfloat8_b attention projections (1.02x on decode).** 0.6726 -> 0.6605 ms at
+PCC 0.99906 vs 0.99909. bfloat4_b is 0.6595 ms but drops layer PCC to 0.9928,
+below the 0.995 bar, so it is rejected.
+
+**7. The router's two keepdim reductions (1.045x on decode).** The router and
+its routing prep were 111.6 us of decode device time -- 20.9% of it, more than
+either matmul family -- and had no audit finding of their own until the fourth
+review. Two thirds of the removable part was not arithmetic at all: ``ttnn.max``
+and ``ttnn.sum`` each pull a ``FillPad`` behind them on a tensor whose last two
+dims are not tile-aligned, 10.42 and 10.41 us against 1.43 and 1.41 us of actual
+reduction. ``router_forward_optimized`` deletes both -- the max is column 0 of
+the sorted top-k, and the sum moves after the scatter where the reduction length
+is a whole number of tiles -- for 0.5866 -> 0.5615 ms traced at ctx128 (both
+legs in one process), with the routing itself unchanged
+(``test_optimized_router_matches_functional`` asserts identical expert
+selection). Moving the sum past the scatter puts the divide over whole tiles,
+whose row-padding then divides 0 by 0; the divisor is clamped so that padding
+stays exactly zero, which costs +1.6 us and is why ``perf_decode.csv`` reads
+0.5634 rather than 0.5615. See that function and ``work_log.md`` §7.
+
+
+Rejected, with measurements
+---------------------------
+**Per-token sparsity in prefill.** Prefill hands ``sparse_matmul`` a sparsity
+tensor of shape ``[1, 1, group_size, E]`` whose rows are 32-token *tiles*, so an
+expert counts as active if any of the tile's 32 tokens picked it. With 256
+selections landing across 128 slots essentially every expert is hit -- hence
+``active=128/128``. Making it per-token requires tokens to be *batch* indices
+(``sparse_matmul`` indexes sparsity by batch dims, not by M), i.e. a
+``[1, T, 1, H]`` layout. Measured rather than assumed: it runs, cuts nnz 16x
+from 4096 to 256, and is **2.1x slower** (14.35 ms vs 6.70 ms), because M
+collapses to 1 and the op pads M to a full 32-row tile. Decode keeps real
+per-token sparsity, which is free there because M is genuinely 1.
+
+**1x32 output tiles on the decode sparse matmuls.** The M padding above is the
+single largest remaining inefficiency: at decode M=1 the gate/up matmul writes
+12 MB and ``down`` writes 16 MB where 0.4/0.5 MB is real, and the reshapes that
+compact it away cost 31 + 33 + 46 us. ``output_tile=ttnn.Tile([1, 32])``
+removes the padding at the source and is 1.07x faster end to end -- but no
+downstream op consumes the result correctly. Measured, in this order:
+``slice`` rejects it (``slice_device_operation.cpp:165`` hardcodes
+``TILE_HEIGHT``), ``ttnn.sum`` and ``ttnn.reshape`` raise
+``MeshBuffer must be large enough``, ``untilize`` returns wrong data without
+erroring, and ``fast_reduce_nc`` returns all zeros. Only eltwise ops read it
+correctly, and they immediately re-pad to 32 rows. Blocked on TTNN support for
+non-32 tile heights outside matmul, not on this model.
+
+**Folding the routing weight in before ``down``.** ``down`` is linear, so
+scaling its *input* by the routing probability is equivalent to scaling its
+output, and the input is the compact ``[B, E, I]`` tensor rather than the
+32x-row-padded ``[B, E, 1, H]`` one. It looked like it should collapse the
+whole tail into one reduce. Measured at ctx128 in one run, before the §7 router
+change (so its shipped-tail leg is the 0.5862 ms configuration of the time, not
+today's 0.5634): shipped tail 0.5862 ms, folded with a compact sum 0.5852 ms,
+folded with ``fast_reduce_nc`` straight off the padded tensor 0.6316 ms. A
+second run of the same three legs read 0.5870 / 0.5853 / 0.6326 and was quoted
+in parallel with this one; this triple is the one whose shipped leg matches the
+``perf_decode.csv`` ctx128 cell of that day, and it is now the only one quoted
+anywhere. The first two are a tie and the third is 8% *worse* --
+``fast_reduce_nc`` also promotes ``down``'s tile padding into the logical shape,
+so recovering ``[1,1,B,H]`` needs a permute plus a slice that together cost more
+than the ops they replaced. The shipped tail stays. (The intermediate version of
+this that used a plain reshape instead of the permute was faster still, 0.550 ms
+-- and silently wrong for every user but the first, which is how the permute
+came to be needed. ``test_optimized_decode_batch`` caught it.)
+
+**Keeping the expert path rank-6.** The obvious reading of the profile is that
+the rank-changing reshapes are pure overhead. They are not: they compact the 32x
+M padding away, so the elementwise ops that follow touch 192 tiles instead of
+6144. Staying rank-6 and dropping all three reshapes measured **0.713 ms vs
+0.673** -- 6% slower.
+
+**Everything else tt-perf-report suggested**, each measured on the traced layer
+against the tuned baseline: in0 in L1 on the sparse rows 1.001x, on the
+attention rows 0.998x, ``out_subblock_w=2`` 1.001x, ``=4`` 1.001x. All noise;
+none adopted. HiFi2 on the sparse rows is covered by lever 5 -- LoFi is both
+faster and no less accurate.
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+
+import torch
+
+import ttnn
+
+from .functional_decoder import (  # noqa: F401  (re-exported for callers)
+    AttentionConfig,
+    AttentionWeights,
+    DecoderLayerConfig,
+    DecoderLayerWeights,
+    KVCache,
+    MoEConfig,
+    _apply_rope,
+    _concat_heads_decode,
+    _per_head_rms_norm,
+    _sparse_matmul_config,
+    attention_decode,
+    attention_prefill,
+    build_expert_sparsity,
+    build_rope_cache,
+    create_kv_cache,
+    upload_layer_weights,
+    upload_router_weight,
+)
+from .precision import DEFAULT_PRECISION, PrecisionConfig  # noqa: F401  (re-exported)
+
+# Tokens per expert-path chunk. Kept at one tile: sparse_matmul folds the group
+# dimension into M, so a larger chunk grows num_blocks_y and can overflow the
+# core grid. Chunking at 32 keeps all blocking in N.
+EXPERT_CHUNK_SIZE = 32
+
+# Expert matmul precision and inner-block width.
+#
+# These two knobs INTERACT, and sweeping either alone finds the wrong optimum.
+# At in0_block_w=1 the kernel is latency-bound, so weight dtype makes no
+# measurable difference -- which is exactly the null result stage 02 first
+# recorded, and wrongly concluded from. Widening the block makes the matmul
+# bandwidth-bound, at which point precision becomes the dominant lever.
+#
+# Joint sweep, real weights, M=32, ms (grid 8x6 for gate/up, 8x8 for down):
+#
+#   packed gate+up (K=2048, 64 tiles)      down (K=768, 24 tiles)
+#   blk   bf16   bfp8   bfp4  bfp4/LoFi    blk   bf16   bfp8   bfp4
+#     4  2.697  2.371  2.367      2.366      4  1.489  1.062  1.058
+#     8  2.726  1.669  1.445      1.431      6  1.507  1.002  0.785
+#    16  2.932  1.806  1.259      1.149      8  1.507  0.993  0.712
+#    32  2.910  1.734  1.158      1.153     12  1.513  0.941  0.654
+#    64  3.180  1.789  1.372      1.204     24  1.622  0.982  0.733
+#
+# bf16's best is 2.697 + 1.489 = 4.186 ms; bfp4's is 1.259 + 0.654 = 1.913 ms.
+# Block width must divide K in tiles, and the two matmuls have different K, so
+# the widths are per-role rather than one shared constant.
+#
+# The table above is a matmul microbenchmark; the widths were re-confirmed on
+# the whole layer, where the interaction with fidelity reverses the gate/up
+# choice (prefill S=512 us/token, real weights):
+#
+#   blk       8      16      32      64
+#   HiFi4  78.27   72.46   69.96   75.56
+#   LoFi   78.12   69.13   69.27   70.88
+#
+# 16 at LoFi is the minimum, so 16 stays.
+#
+# **These five names are now aliases, not the source of truth.** The values
+# themselves live in ``precision.PrecisionConfig``, whose defaults are exactly
+# what was written here before stage 07; the names survive because probes under
+# ``doc/`` and the stage-02/04 tests import them, and because a reader arriving
+# at the sweep comments above should find the value they describe next to them.
+# Anything that needs to *vary* the policy must take a ``PrecisionConfig``
+# instead -- these are bound at import and cannot follow a non-default model.
+EXPERT_WEIGHT_DTYPE = DEFAULT_PRECISION.experts_gate_up_dtype
+EXPERT_IN0_BLOCK_W_GATE_UP = DEFAULT_PRECISION.experts_gate_up_in0_block_w  # divides 2048/32 = 64
+EXPERT_IN0_BLOCK_W_DOWN = DEFAULT_PRECISION.experts_down_in0_block_w  # divides 768/32 = 24
+
+# bfp4 weights carry 4 mantissa bits, so HiFi4's extra passes have nothing left
+# to resolve. LoFi is 4.6% faster on prefill and 1.6% on decode at PCC 0.99910
+# vs HiFi4's 0.99909. This also answers tt-perf-report's "HiFi2 may also work"
+# on the sparse rows: HiFi2 measured 69.78 us/token, between the two.
+EXPERT_MATH_FIDELITY = DEFAULT_PRECISION.experts_fidelity
+
+# Attention projections. bf16 -> bfloat8_b costs 0.00003 PCC and buys 1.8% of
+# decode; bfloat4_b buys another 0.1% but drops layer PCC to 0.9928, under the
+# 0.995 bar, so it is rejected. q_norm/k_norm stay bf16 -- they are norms, not
+# projections, and weigh 4 KB.
+ATTENTION_WEIGHT_DTYPE = DEFAULT_PRECISION.attention_qkv_dtype
+
+# Blackhole p300c has 8 DRAM banks. The DRAM-sharded matmul wants the weight
+# width-sharded one shard per bank, and both the activation and the output
+# width-sharded in L1 over the matching core row.
+_DRAM_BANKS = 8
+
+
+def _expert_compute_kernel_config(device, precision: PrecisionConfig = DEFAULT_PRECISION):
+    """LoFi, and ``fp32_dest_acc_en`` deliberately OFF.
+
+    ``fp32_dest_acc_en`` looks like the natural next lever but must not be used
+    here: it halves the matmul dest from 8 tiles to 4, which corrupts expert
+    output on Blackhole (tt-metal #49068, hit on BH-QB-2). It is therefore
+    **not** a ``PrecisionConfig`` field -- a sweep must not be able to turn it
+    on.
+    """
+    return ttnn.init_device_compute_kernel_config(
+        device.arch(),
+        math_fidelity=precision.experts_fidelity,
+        math_approx_mode=False,
+        fp32_dest_acc_en=False,
+        packer_l1_acc=False,
+    )
+
+
+def _attention_compute_kernel_config(device, precision: PrecisionConfig = DEFAULT_PRECISION):
+    """``None`` at the default, which is what the projections have always passed.
+
+    ``attention_fidelity=None`` means "leave the op at its own default", so this
+    returns ``None`` and the ``compute_kernel_config=`` argument is a no-op. Any
+    named fidelity builds a real config; the remaining flags mirror
+    ``_expert_compute_kernel_config``'s, which is the closest measured
+    neighbour.
+    """
+    if precision.attention_fidelity is None:
+        return None
+    return ttnn.init_device_compute_kernel_config(
+        device.arch(),
+        math_fidelity=precision.attention_fidelity,
+        math_approx_mode=False,
+        fp32_dest_acc_en=False,
+        packer_l1_acc=False,
+    )
+
+
+def _tuned_sparse_matmul_config(m: int, n: int, k: int, target_blk: int):
+    """``_sparse_matmul_config`` with a tuned inner block width.
+
+    ``k`` is the inner dimension in elements; the block width must divide it in
+    tiles, so this falls back to the largest legal divisor at or below the
+    target rather than failing.
+    """
+    k_tiles = max(1, k // 32)
+    blk = min(target_blk, k_tiles)
+    while blk > 1 and k_tiles % blk:
+        blk -= 1
+    return _sparse_matmul_config(m, n, in0_block_w=blk)
+
+
+# Column of ones used to sum the dense routing row (see ``router_forward_optimized``).
+# Cached per (device, length) because it is a constant, and because allocating a
+# tensor inside a trace capture is illegal -- every caller runs the layer eagerly
+# once to compile before capturing, which is what populates this.
+#
+# The key carries ``id(device)`` but the *value* carries the device object itself.
+# ``mesh_device`` is function-scoped in ``conftest.py`` and is closed and deleted
+# after each test, and CPython reuses freed addresses, so a later device could be
+# handed the same id and collide with an entry bound to a destroyed one. Holding
+# the object in the value makes the pin explicit: the address cannot be recycled
+# while the entry lives, so equal ids imply the same live device. The identity
+# check below is then a belt-and-braces assertion, not a hope.
+_ONES_COLUMN: dict[tuple[int, int], tuple[object, ttnn.Tensor]] = {}
+
+
+def _ones_column(device, n: int) -> ttnn.Tensor:
+    key = (id(device), n)
+    entry = _ONES_COLUMN.get(key)
+    if entry is not None and entry[0] is device:
+        return entry[1]
+    cached = ttnn.from_torch(
+        torch.ones(1, 1, n, 1),
+        dtype=ttnn.bfloat16,
+        layout=ttnn.TILE_LAYOUT,
+        device=device,
+        memory_config=ttnn.DRAM_MEMORY_CONFIG,
+    )
+    _ONES_COLUMN[key] = (device, cached)
+    return cached
+
+
+def router_forward_optimized(x: ttnn.Tensor, w_router: ttnn.Tensor, config: MoEConfig) -> ttnn.Tensor:
+    """``router_forward`` with both keepdim reductions removed. Same result.
+
+    The routing maths is unchanged from ``functional_decoder.router_forward``,
+    including the part that is load-bearing for correctness: selection happens
+    on the **raw fp32 logits** and the softmax is taken over the 8 survivors
+    only. A 128-wide bf16 softmax misroutes 83/128 tokens and is not an option
+    here; see that function's docstring for the algebra.
+
+    What changes is how the two reductions are spelled. ``ttnn.max``/``ttnn.sum``
+    call ``fill_implicit_tile_padding`` whenever *either* of the last two dims is
+    unaligned (``fill_pad.cpp:17-24``); the top-k tensor is 8 wide, and decode's
+    is 1 row tall, so each keepdim reduction dragged a ``FillPad`` behind it --
+    **10.421 µs and 10.413 µs** in the archived stage-01 decode profile
+    (``doc/functional_decoder/ops_perf_decode_paged32.csv`` rows 73 and 77),
+    against **1.432 and 1.407 µs** for the reductions themselves (rows 74 and
+    78). Both are avoided rather than tuned:
+
+    * the **max** is column 0 of the top-k output, which ``sorted=True``
+      guarantees is the largest, so one 0.87 µs ``slice`` replaces
+      ``FillPad + Reduce``;
+    * the **sum** moves *after* the scatter and becomes a matmul against a
+      column of ones. Over the dense row the reduction length is
+      ``num_experts`` = 128 — a whole number of tiles — so no padding lane can
+      enter the sum. That is also why this is preferred to the same matmul over
+      the 8-wide tensor, whose K padding would carry whatever ``topk`` left
+      behind. Normalising after the scatter is legal because the scatter is a
+      permutation of the 8 survivors into a field of exact zeros, so the sum
+      over 128 *is* the sum over the 8.
+
+    **The padding cost of moving the sum.** Dividing after the scatter divides
+    over whole tiles, and the tile *row* padding -- rows S..ceil(S/32)*32 -- has
+    ``dense`` = 0 and therefore ``total`` = 0 too. Unguarded, ``ttnn.div``
+    returns **+inf** there (not NaN; measured at S = 33 and 100, where every one
+    of the 31 and 28 padding rows came back +inf), where the functional router,
+    which divided before the scatter, returned exact zeros. Nothing observable
+    leaked -- ``ttnn.to_torch`` returns the logical shape, the sparsity path
+    drops the padding in ``to_layout(ROW_MAJOR)``, and the scale multiply,
+    ``rms_norm`` and ``fast_reduce_nc`` all reduce along axes that are either
+    tile-aligned or not the padded one -- but it is a hazard the functional path
+    did not have, so the divisor is clamped (``ttnn.maximum(total, 1e-30)``).
+    Decode is the one case that was already clean: at S = 1 the padding came
+    back exact zero unguarded. The clamp is one extra op -- 1.12 µs in the
+    decode profile, +1.6 µs on the traced layer, 0.28% -- and
+    ``test_optimized_router_padding_is_zero`` stops it being optimized back out.
+    See ``work_log.md`` §7 for the padding table and the rejected free version.
+
+    Measured on the whole traced layer at ctx128, real weights, median of 100:
+    **0.5866 -> 0.5615 ms** (``perf_decode.csv`` reads 0.5634, which is that
+    configuration in its own run), and the router block **111.6 -> 87.8 µs** of
+    decode device time -- rows 68-88 of
+    ``doc/functional_decoder/ops_perf_decode_paged32.csv``, which still holds the
+    pre-fix block, against the same block in the pre-guard optimized profile.
+    With the divisor guard the block is 88.9 µs, rows 69-88 of
+    ``doc/optimized_decoder/ops_perf_optimized_decode.csv``. Layer PCC is
+    0.99901 either way (prefill S=128 vs HF: 0.9990057 after, 0.9990050
+    before). ``doc/optimized_decoder/work_log.md`` §7
+    carries the rejected variants, including ``ttnn.softmax`` over the 8
+    survivors — faster still, and wrong: it reduces over the whole 32-wide tile,
+    so the weights sum to 0.9736 instead of 1.
+    """
+    assert config.norm_topk_prob, (
+        "router selects on raw logits, which relies on the softmax denominator "
+        "cancelling during top-k renormalisation; that only holds when "
+        "norm_topk_prob is True"
+    )
+
+    logits = ttnn.linear(x, w_router, dtype=ttnn.float32, memory_config=ttnn.DRAM_MEMORY_CONFIG)
+    top_logits, top_indices = ttnn.topk(logits, k=config.num_experts_per_tok, dim=-1, largest=True, sorted=True)
+
+    # Subtracting the max is for exp() range only; any shared shift cancels in
+    # the division. sorted=True means column 0 already is that max.
+    top_max = ttnn.slice(top_logits, [0, 0, 0, 0], [1, 1, top_logits.shape[2], 1])
+    exp_logits = ttnn.exp(ttnn.sub(top_logits, top_max))
+
+    zeros = ttnn.typecast(ttnn.zeros_like(logits), ttnn.bfloat16)
+    dense = ttnn.scatter(
+        zeros,
+        dim=-1,
+        index=top_indices,
+        src=ttnn.typecast(exp_logits, ttnn.bfloat16),
+    )
+    # Sum over the dense row == sum over the 8 survivors; see the docstring.
+    total = ttnn.matmul(dense, _ones_column(x.device(), config.num_experts), dtype=ttnn.bfloat16)
+    # Guard the divisor's tile row-padding, which is 0 where ``dense`` is also 0
+    # (see the padding note in the docstring). Every real row's denominator is
+    # >= 1, because sorted=True makes column 0 of ``exp_logits`` exactly
+    # exp(0) = 1, so clamping at 1e-30 cannot touch one: measured bit-identical
+    # on the real rows at S = 1, 33, 100. In the padding it turns 0/0 into
+    # 0/1e-30 = 0, which is what the functional router returned there.
+    guarded = ttnn.maximum(total, 1e-30)
+    normalised = ttnn.div(dense, guarded)
+    for t in (logits, top_logits, top_indices, top_max, exp_logits, dense, total, guarded):
+        ttnn.deallocate(t)
+    return normalised
+
+
+def _bank_row(n: int) -> ttnn.CoreRangeSet:
+    return ttnn.CoreRangeSet({ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(n - 1, 0))})
+
+
+# The L1 shard height below, and ``per_core_M=1`` in the program config, are
+# both decode's padded M of one 32-row tile. That is what caps this path at
+# batch 32; see ``_dram_sharded_usable``.
+_DRAM_SHARDED_MAX_BATCH = 32
+
+
+def _width_sharded_l1(width: int) -> ttnn.MemoryConfig:
+    """L1 width-sharded over one core per DRAM bank, 32 rows (decode's padded M)."""
+    return ttnn.MemoryConfig(
+        ttnn.TensorMemoryLayout.WIDTH_SHARDED,
+        ttnn.BufferType.L1,
+        ttnn.ShardSpec(_bank_row(_DRAM_BANKS), [32, width // _DRAM_BANKS], ttnn.ShardOrientation.ROW_MAJOR),
+    )
+
+
+def _dram_sharded_ok(k: int, n: int) -> bool:
+    """Both dims must split evenly into whole tiles across the banks."""
+    return k % (_DRAM_BANKS * 32) == 0 and n % (_DRAM_BANKS * 32) == 0
+
+
+# Decode's two expert intermediates are 97% M padding -- ``sparse_matmul`` pads
+# M=1 back to a 32-row tile -- so they are large in absolute terms:
+#
+#     batch * 128 experts * 32 rows * (1536 + 2048) cols * 2 B = batch * 29.4 MB
+#
+# Blackhole offers ~160 MB of allocatable L1 (110 banks x 1.46 MB, as the
+# allocator reports it on this p300c). Holding both in L1 is therefore a
+# batch-1 affordance, not a general one: at batch 8 the allocator rejects
+# ``down``'s 134 MB output outright. Past the budget the pair goes to DRAM,
+# which is what prefill already does at every length.
+#
+# The 40 MB threshold itself is **asserted, not measured**: it is one comfortable
+# step above batch 1's 29.4 MB and below batch 2's 58.8 MB, so it separates the
+# only two cases that exist here, and no sweep was run to find where L1 actually
+# stops paying. What is measured is the pair of endpoints -- B=1 in L1 is the
+# shipped, profiled configuration, and B=8 in L1 does not allocate at all.
+_DECODE_EXPERT_L1_BUDGET_BYTES = 40 * 1024 * 1024
+
+
+def _decode_expert_memory_config(batch: int, config: MoEConfig) -> ttnn.MemoryConfig:
+    """L1 for the intermediates while they fit the budget above, else DRAM."""
+    padded_rows = batch * config.num_experts * 32
+    nbytes = padded_rows * (2 * config.moe_intermediate_size + config.hidden_size) * 2
+    return ttnn.L1_MEMORY_CONFIG if nbytes <= _DECODE_EXPERT_L1_BUDGET_BYTES else ttnn.DRAM_MEMORY_CONFIG
+
+
+def _dram_sharded_usable(weights: "OptimizedWeights", batch: int) -> bool:
+    """Whether decode may take the DRAM-sharded projections at this batch.
+
+    Two independent conditions:
+
+    * the weight dims divided evenly across the banks at upload time, so a
+      sharded copy exists at all (``_dram_sharded_ok``);
+    * the batch still fits decode's single 32-row M tile. ``_width_sharded_l1``
+      hardcodes a 32-row shard and ``_dram_sharded_program_config`` sets
+      ``per_core_M=1``, so at ``batch > 32`` the activation no longer matches
+      its shard spec. Without this check that surfaces as a shard-shape
+      mismatch deep in the matmul rather than as a fallback.
+    """
+    if weights.wqkv_decode is None or weights.wo_decode is None:
+        return False
+    return batch <= _DRAM_SHARDED_MAX_BATCH
+
+
+def _dram_sharded_program_config(k: int, n: int):
+    return ttnn.MatmulMultiCoreReuseMultiCastDRAMShardedProgramConfig(
+        in0_block_w=k // _DRAM_BANKS // 32,
+        per_core_M=1,
+        per_core_N=n // _DRAM_BANKS // 32,
+        fused_activation=None,
+    )
+
+
+@dataclass
+class OptimizedWeights:
+    """Device weights for the optimized layer.
+
+    Two things live here that ``upload_layer_weights`` does not provide:
+
+    * experts with gate and up kept as one weight. ``weight_mapping`` already
+      produces the checkpoint's fused ``[E, 2I, H]`` tensor; stage 01 split it
+      apart at upload time to mirror the exemplars, so packing is *undoing*
+      that split rather than inventing a layout.
+    * two copies of the attention projections. Decode uses a DRAM
+      width-sharded copy for ``MatmulMultiCoreReuseMultiCastDRAMShardedProgramConfig``;
+      prefill cannot -- a plain ``ttnn.linear`` on a DRAM-sharded weight throws
+      ``Only L1 buffers can have an associated circular buffer`` -- so an
+      interleaved copy is kept for it. At bfloat8_b (1.0625 B/elem, because
+      each 16-element block carries its own exponent byte) wqkv is 11.14 MB and
+      wo 8.91 MB, so the duplicate copy costs **20.05 MB** and the pair 40.11 MB,
+      against 24 GB available. An 18.9 MB figure that this file used to carry
+      came from rounding bfloat8_b to 1 B/elem and is withdrawn.
+      ``config/context_contract.json`` now carries 20.05 MB too; an earlier
+      revision of it made the same rounding error and called the pair a wash
+      against stage 01's single bf16 copy, which its ``optimized_note`` records.
+    """
+
+    gate_up_proj: ttnn.Tensor  # [1, num_experts, hidden, 2 * intermediate]
+    down_proj: ttnn.Tensor  # [1, num_experts, intermediate, hidden]
+    attention: AttentionWeights  # interleaved, for prefill
+    wqkv_decode: ttnn.Tensor | None  # DRAM width-sharded, for decode
+    wo_decode: ttnn.Tensor | None
+
+
+# Stage-01 name, kept so existing callers and docs still resolve.
+PackedExpertWeights = OptimizedWeights
+
+
+def upload_optimized_weights(
+    torch_weights,
+    device,
+    config: MoEConfig,
+    dtype=None,
+    *,
+    precision: PrecisionConfig = DEFAULT_PRECISION,
+) -> OptimizedWeights:
+    """Upload experts packed along the output dim, plus both attention copies.
+
+    ``precision`` supplies every weight dtype. ``dtype``, the stage-02 spelling,
+    still overrides **both** expert dtypes when given -- several stage-02 tests
+    sweep it directly -- but new callers should pass a ``PrecisionConfig``,
+    which can also move gate/up and down apart.
+
+    The expert dtype is a parameter at all because it and ``in0_block_w`` are
+    **not** independent: precision only pays once the block width is wide enough
+    for the matmul to become bandwidth-bound. Sweeping either alone finds the
+    wrong optimum, which is why both live in the same config object.
+    """
+    fused = torch_weights["experts_gate_up"]  # [E, 2I, H], gate first
+    gate_up_dtype = dtype if dtype is not None else precision.experts_gate_up_dtype
+    down_dtype = dtype if dtype is not None else precision.experts_down_dtype
+
+    def up(t: torch.Tensor, tensor_dtype, memory_config=ttnn.DRAM_MEMORY_CONFIG) -> ttnn.Tensor:
+        return ttnn.from_torch(
+            t.contiguous().float(),
+            dtype=tensor_dtype,
+            layout=ttnn.TILE_LAYOUT,
+            device=device,
+            memory_config=memory_config,
+        )
+
+    def as_4d(t: torch.Tensor, pad_to_4d: bool = False) -> torch.Tensor:
+        if pad_to_4d:
+            t = t.reshape(1, 1, 1, -1)
+        while t.dim() < 4:
+            t = t.unsqueeze(0)
+        return t
+
+    wqkv, wo = as_4d(torch_weights["wqkv"]), as_4d(torch_weights["wo"])
+
+    def dram_sharded(t: torch.Tensor, tensor_dtype) -> ttnn.Tensor | None:
+        k, n = int(t.shape[-2]), int(t.shape[-1])
+        if not _dram_sharded_ok(k, n):
+            return None
+        return up(
+            t,
+            tensor_dtype,
+            ttnn.MemoryConfig(
+                ttnn.TensorMemoryLayout.WIDTH_SHARDED,
+                ttnn.BufferType.DRAM,
+                ttnn.ShardSpec(_bank_row(_DRAM_BANKS), [k, n // _DRAM_BANKS], ttnn.ShardOrientation.ROW_MAJOR),
+            ),
+        )
+
+    return OptimizedWeights(
+        gate_up_proj=up(fused.transpose(-2, -1).unsqueeze(0), gate_up_dtype),
+        down_proj=up(torch_weights["experts_down"].transpose(-2, -1).unsqueeze(0), down_dtype),
+        attention=AttentionWeights(
+            wqkv=up(wqkv, precision.attention_qkv_dtype),
+            wo=up(wo, precision.attention_wo_dtype),
+            q_norm=up(as_4d(torch_weights["q_norm"], pad_to_4d=True), precision.norm_weight_dtype),
+            k_norm=up(as_4d(torch_weights["k_norm"], pad_to_4d=True), precision.norm_weight_dtype),
+        ),
+        wqkv_decode=dram_sharded(wqkv, precision.attention_qkv_dtype),
+        wo_decode=dram_sharded(wo, precision.attention_wo_dtype),
+    )
+
+
+# Stage-01 name, kept so existing callers still resolve.
+upload_packed_expert_weights = upload_optimized_weights
+
+
+def attention_decode_optimized(
+    x: ttnn.Tensor,
+    weights: OptimizedWeights,
+    config: AttentionConfig,
+    cos_cache: ttnn.Tensor,
+    sin_cache: ttnn.Tensor,
+    kv_cache: KVCache,
+    current_pos: ttnn.Tensor,
+    token_index: int,
+    sdpa_program_config=None,
+    rope=None,
+    precision: PrecisionConfig = DEFAULT_PRECISION,
+) -> ttnn.Tensor:
+    """``attention_decode`` with the two projections run DRAM-sharded.
+
+    ``sdpa_program_config`` is passed straight through to the SDPA-decode op and
+    defaults to ``None``, which is what every single-chip caller uses and what
+    every number in this file was measured at. It exists for the multichip path:
+    at one KV head the op's default core assignment tries to put all 110 worker
+    cores on the single head and its tree reduction refuses past 64
+    (``sdpa_decode_program_factory.cpp:245``). See
+    ``multichip_decoder._sdpa_program_config``.
+
+    At decode M=1 both projections are pure weight reads, so what limits them is
+    how well the read spreads over the DRAM banks. The interleaved layout gave
+    383 GB/s on qkv and 235 GB/s on wo; sharding the weight one shard per bank
+    and keeping the activation and output width-sharded in L1 measures
+
+        qkv (K=2048, N=5120)   68.3 -> 46.8 us   1.46x
+        wo  (K=4096, N=2048)   96.0 -> 41.7 us   2.30x
+
+    at the op level, and 0.6508 -> 0.5863 ms (1.11x) on the whole traced layer
+    (both legs measured before the §7 router change, hence 0.5863 rather than
+    today's 0.5634)
+    at ctx128 -- both legs otherwise at the shipped configuration and on the
+    same bfloat8_b weights, so only the program config and shard layout differ.
+    Core count was swept: 8 (one per bank) beats 16, 32 and 64 on both matmuls,
+    because past one shard per bank the extra cores only add mcast traffic. The
+    8 is the tuned quantity; the profiler reports ``CORE COUNT`` 80 for these
+    rows and ``tt-perf-report`` prints 12, and neither of those was chosen.
+
+    In the archived profiles the two projections go 57.06 -> 27.33 us (qkv) and
+    72.80 -> 21.91 us (wo). Only ``wo`` was SLOW-classified interleaved; qkv was
+    already DRAM-classified, so its gain is duration, not a change of class.
+
+    Batch is capped at 32 here -- see ``_dram_sharded_usable`` -- which is where
+    ``nlp_create_qkv_heads_decode`` caps it anyway, on either path.
+
+    Everything between the two projections is identical to ``attention_decode``,
+    including the Blackhole staging workarounds, so the two stay diffable.
+    """
+    if not _dram_sharded_usable(weights, int(x.shape[-2])):
+        return attention_decode(x, weights.attention, config, cos_cache, sin_cache, kv_cache, current_pos, token_index)
+
+    k_cache, v_cache, page_table = kv_cache.k, kv_cache.v, kv_cache.page_table
+    k_qkv, n_qkv = int(weights.wqkv_decode.shape[-2]), int(weights.wqkv_decode.shape[-1])
+    k_o, n_o = int(weights.wo_decode.shape[-2]), int(weights.wo_decode.shape[-1])
+
+    attn_compute_config = _attention_compute_kernel_config(x.device(), precision)
+    x_sharded = ttnn.to_memory_config(x, _width_sharded_l1(k_qkv))
+    xqkv = ttnn.linear(
+        x_sharded,
+        weights.wqkv_decode,
+        program_config=_dram_sharded_program_config(k_qkv, n_qkv),
+        memory_config=_width_sharded_l1(n_qkv),
+        dtype=precision.activation_dtype,
+        compute_kernel_config=attn_compute_config,
+    )
+    ttnn.deallocate(x_sharded)
+
+    # nlp_create_qkv_heads_decode wants interleaved L1. (It also must not be
+    # handed a DRAM tensor at all on Blackhole -- tt-metal #16667 zeroes
+    # odd-indexed Q rows via a NoC DRAM-read alignment violation.)
+    xqkv = ttnn.to_memory_config(xqkv, ttnn.L1_MEMORY_CONFIG)
+    q, k, v = ttnn.experimental.nlp_create_qkv_heads_decode(
+        xqkv,
+        num_heads=config.num_attention_heads,
+        num_kv_heads=config.num_key_value_heads,
+        memory_config=ttnn.L1_HEIGHT_SHARDED_MEMORY_CONFIG,
+    )
+    ttnn.deallocate(xqkv)
+
+    # rms_norm wants interleaved DRAM while paged_update_cache requires a
+    # *sharded* update tensor, so remember the split's layout and restore it.
+    kv_sharded_mem = k.memory_config()
+    q = _per_head_rms_norm(
+        ttnn.to_memory_config(q, ttnn.DRAM_MEMORY_CONFIG), weights.attention.q_norm, config.rms_norm_eps
+    )
+    k = _per_head_rms_norm(
+        ttnn.to_memory_config(k, ttnn.DRAM_MEMORY_CONFIG), weights.attention.k_norm, config.rms_norm_eps
+    )
+    # ``rope`` defaults to ``None`` and therefore to ``_apply_rope``, which is
+    # what every caller uses -- including the shipped multichip decode path --
+    # and what every number in this file was measured at. It is a seam, not a
+    # switch: stage 04 used it to build and measure a Meta-ordered
+    # ``rotary_embedding_llama`` alternative (3.05x faster standalone and
+    # bit-identical) without disturbing the 1x1 baseline the multichip documents
+    # compare against. That alternative is **rejected** -- the KV cache carries
+    # the rotary's channel convention and prefill writes HF-ordered keys, so it
+    # is not a decode-local change. See ``multichip_decoder._meta_rope`` and
+    # ``doc/optimized_multichip_decoder/README.md`` limitation 4.
+    _rope = _apply_rope if rope is None else rope
+    q = _rope(q, cos_cache, sin_cache, token_index)
+    k = ttnn.to_memory_config(_rope(k, cos_cache, sin_cache, token_index), kv_sharded_mem)
+
+    # Deliberately NOT cast to the cache dtype, unlike the prefill fill writers.
+    # ``paged_update_cache`` requires a FLOAT32/BFLOAT16 update and converts into
+    # the cache itself (measured: bfp8 cache + bf16 update round-trips at PCC
+    # 0.999969, bfp8 update is rejected at
+    # ``paged_update_cache_device_operation.cpp:296``). See
+    # ``functional_decoder.match_cache_dtype`` for the full table.
+    ttnn.experimental.paged_update_cache(k_cache, k, update_idxs_tensor=current_pos, page_table=page_table)
+    ttnn.experimental.paged_update_cache(v_cache, v, update_idxs_tensor=current_pos, page_table=page_table)
+    ttnn.deallocate(k)
+    ttnn.deallocate(v)
+
+    if kv_cache.is_paged:
+        attn = ttnn.transformer.paged_scaled_dot_product_attention_decode(
+            q,
+            k_cache,
+            v_cache,
+            page_table_tensor=page_table,
+            cur_pos_tensor=current_pos,
+            scale=config.head_dim**-0.5,
+            memory_config=ttnn.DRAM_MEMORY_CONFIG,
+            program_config=sdpa_program_config,
+        )
+    else:
+        attn = ttnn.transformer.scaled_dot_product_attention_decode(
+            q,
+            k_cache,
+            v_cache,
+            cur_pos_tensor=current_pos,
+            scale=config.head_dim**-0.5,
+            memory_config=ttnn.DRAM_MEMORY_CONFIG,
+            program_config=sdpa_program_config,
+        )
+    ttnn.deallocate(q)
+
+    attn = ttnn.to_memory_config(_concat_heads_decode(attn, config), _width_sharded_l1(k_o))
+    out = ttnn.linear(
+        attn,
+        weights.wo_decode,
+        program_config=_dram_sharded_program_config(k_o, n_o),
+        memory_config=_width_sharded_l1(n_o),
+        dtype=precision.activation_dtype,
+        compute_kernel_config=attn_compute_config,
+    )
+    ttnn.deallocate(attn)
+    return ttnn.to_memory_config(out, ttnn.DRAM_MEMORY_CONFIG)
+
+
+def _experts_chunk_packed(
+    hidden: ttnn.Tensor,
+    routing: ttnn.Tensor,
+    weights: OptimizedWeights,
+    config: MoEConfig,
+    sparsity_base: ttnn.Tensor,
+    precision: PrecisionConfig = DEFAULT_PRECISION,
+) -> ttnn.Tensor:
+    """One 32-token chunk with gate and up computed in a single matmul.
+
+    The win is core occupancy, not the saved kernel launch.
+    ``_sparse_matmul_config`` parallelises only over N, so the usable core count
+    is capped by the number of N tiles:
+
+        gate or up alone   N = 768  -> 24 tiles -> 24 cores
+        gate+up packed     N = 1536 -> 48 tiles -> 48 cores
+        down               N = 2048 -> 64 tiles -> 64 cores
+
+    which is also why the stage-01 profile showed down running at 127 GB/s
+    while gate/up sat at 64 GB/s. Worth 1.09x against a *tuned* separate
+    candidate (2 x 1.476 = 2.952 ms -> 2.699 ms); the larger figure it shows
+    against an untuned one belongs to the block-width fix, not to packing.
+    """
+    chunk_len = hidden.shape[2]
+    n_experts = config.num_experts
+    hidden_size = config.hidden_size
+    inter = config.moe_intermediate_size
+    group_size = chunk_len // EXPERT_CHUNK_SIZE
+
+    device = hidden.device()
+    compute_config = _expert_compute_kernel_config(device, precision)
+    output_tile = ttnn.Tile([32, 32])
+    # PREFILL blocking. Separate from the decode fields (see precision.py): these
+    # matmuls run at M = EXPERT_CHUNK_SIZE here and M = 1 in decode, and the
+    # shipped values were tuned only for the latter. Defaults are identical, so
+    # this is a no-op until the prefill fields are changed.
+    gate_up_config = _tuned_sparse_matmul_config(
+        EXPERT_CHUNK_SIZE, 2 * inter, hidden_size, precision.prefill_experts_gate_up_in0_block_w
+    )
+    down_config = _tuned_sparse_matmul_config(
+        EXPERT_CHUNK_SIZE, hidden_size, inter, precision.prefill_experts_down_in0_block_w
+    )
+
+    hidden_grouped = ttnn.reshape(hidden, (1, group_size, EXPERT_CHUNK_SIZE, hidden_size))
+    sparsity = ttnn.repeat(sparsity_base, (1, 1, group_size, 1))
+    nnz = n_experts * group_size
+
+    fused = ttnn.sparse_matmul(
+        hidden_grouped,
+        weights.gate_up_proj,
+        sparsity=sparsity,
+        nnz=nnz,
+        memory_config=ttnn.DRAM_MEMORY_CONFIG,
+        output_tile=output_tile,
+        program_config=gate_up_config,
+        compute_kernel_config=compute_config,
+        dtype=precision.activation_dtype,
+    )
+    ttnn.deallocate(hidden_grouped)
+    packed_width = fused.shape[-1]
+    fused = ttnn.reshape(ttnn.transpose(fused, 1, 3), (1, n_experts, chunk_len, packed_width))
+
+    # gate is the first half -- matches Qwen3MoeExperts.forward's chunk(2, dim=-1)
+    half = packed_width // 2
+    gate = ttnn.slice(fused, [0, 0, 0, 0], [1, n_experts, chunk_len, half])
+    up = ttnn.slice(fused, [0, 0, 0, half], [1, n_experts, chunk_len, packed_width])
+    ttnn.deallocate(fused)
+
+    down_input = ttnn.reshape(ttnn.mul(ttnn.silu(gate), up), (1, n_experts, chunk_len, half))
+    ttnn.deallocate(gate)
+    ttnn.deallocate(up)
+
+    down = ttnn.sparse_matmul(
+        down_input,
+        weights.down_proj,
+        sparsity=sparsity_base,
+        nnz=n_experts,
+        memory_config=ttnn.DRAM_MEMORY_CONFIG,
+        output_tile=output_tile,
+        program_config=down_config,
+        is_input_a_sparse=True,
+        compute_kernel_config=compute_config,
+        dtype=precision.activation_dtype,
+    )
+    ttnn.deallocate(down_input)
+
+    states = ttnn.reshape(down, (1, n_experts, chunk_len, hidden_size))
+    states = ttnn.mul(states, ttnn.permute(routing, (0, 3, 2, 1)))
+    states = ttnn.unsqueeze_to_4D(ttnn.experimental.fast_reduce_nc(states, dims=[1]))
+    return ttnn.reshape(states, (1, 1, chunk_len, hidden_size))
+
+
+def moe_prefill_optimized(
+    x: ttnn.Tensor,
+    routing: ttnn.Tensor,
+    weights: OptimizedWeights,
+    config: MoEConfig,
+    sparsity_base: ttnn.Tensor,
+    precision: PrecisionConfig = DEFAULT_PRECISION,
+) -> ttnn.Tensor:
+    """Expert pass over a sequence. ``x`` ``[1, 1, S, H]``, any S.
+
+    Non-aligned lengths are zero-padded to a chunk boundary and sliced back;
+    padded rows carry an all-zero routing vector so they contribute nothing.
+    """
+    seq_len = x.shape[2]
+    padded_len = math.ceil(seq_len / EXPERT_CHUNK_SIZE) * EXPERT_CHUNK_SIZE
+
+    if padded_len != seq_len:
+        pad = [(0, 0), (0, 0), (0, padded_len - seq_len), (0, 0)]
+        x = ttnn.pad(x, pad, value=0.0)
+        routing = ttnn.pad(routing, pad, value=0.0)
+
+    outputs = []
+    for start in range(0, padded_len, EXPERT_CHUNK_SIZE):
+        end = start + EXPERT_CHUNK_SIZE
+        outputs.append(
+            _experts_chunk_packed(
+                ttnn.slice(x, [0, 0, start, 0], [1, 1, end, config.hidden_size]),
+                ttnn.slice(routing, [0, 0, start, 0], [1, 1, end, config.num_experts]),
+                weights,
+                config,
+                sparsity_base,
+                precision,
+            )
+        )
+    out = outputs[0] if len(outputs) == 1 else ttnn.concat(outputs, dim=2)
+    if padded_len != seq_len:
+        out = ttnn.slice(out, [0, 0, 0, 0], [1, 1, seq_len, config.hidden_size])
+    return out
+
+
+def moe_decode_optimized(
+    x: ttnn.Tensor,
+    routing: ttnn.Tensor,
+    weights: OptimizedWeights,
+    config: MoEConfig,
+    precision: PrecisionConfig = DEFAULT_PRECISION,
+) -> ttnn.Tensor:
+    """Decode MoE with gate/up packed and per-token sparsity. ``x`` ``[1, 1, batch, H]``.
+
+    Tokens are carried as *batch* indices (``[1, B, 1, H]``) rather than along M.
+    ``sparse_matmul`` indexes its sparsity tensor by batch dims, so this is what
+    makes the pattern per-token -- and unlike prefill it costs nothing here,
+    because decode's M is genuinely 1.
+
+    The two matmuls need different sparsity flags, which is not obvious and is
+    what previously limited this path to a single user. From
+    ``sparse_matmul_device_operation.cpp``::
+
+        a_sparse && b_sparse -> batch_length = batch_length_B
+        a_sparse             -> batch_length = batch_length_A
+        neither              -> batch_length = batch_length_A * batch_length_B
+
+    ``is_input_b_sparse`` defaults to true. gate/up takes the third branch and
+    gets ``B * E``, which matches a ``[1, 1, B, E]`` sparsity tensor. The down
+    projection has a sparse activation, so it would take the *first* branch and
+    get just ``E`` -- ignoring the batch entirely and rejecting any B > 1.
+    Passing ``is_input_b_sparse=False`` selects the second branch instead, so
+    down sees ``batch_length_A = B * E`` and matches.
+
+    The rank juggling is deliberate. ``sparse_matmul`` returns
+    ``[1, 1, B, E, 1, N]`` and pads that M=1 to a full 32-row tile, so the
+    result is 97% padding; reshaping to a compact ``[B, E, N]`` before the
+    elementwise work makes those ops touch 192 tiles instead of 6144. Dropping
+    the reshapes and staying rank-6 measured 6% *slower*.
+
+    Applying the routing weight to ``down``'s *input* instead -- equivalent,
+    since ``down`` is linear, and the input is the compact tensor -- was
+    measured and is a tie (0.5852 vs 0.5862 ms, one run, at the pre-§7
+    configuration); see the module docstring for why the variant that looked
+    much better than that was not.
+    """
+    batch = x.shape[2]
+    n_experts = config.num_experts
+    hidden_size = config.hidden_size
+    inter = config.moe_intermediate_size
+    nnz = config.num_experts_per_tok * batch
+
+    sparsity = ttnn.to_layout(routing, ttnn.ROW_MAJOR_LAYOUT)
+    expert_memory_config = _decode_expert_memory_config(batch, config)
+    output_tile = ttnn.Tile([32, 32])
+    compute_config = _expert_compute_kernel_config(x.device(), precision)
+    gate_up_config = _tuned_sparse_matmul_config(1, 2 * inter, hidden_size, precision.experts_gate_up_in0_block_w)
+    down_config = _tuned_sparse_matmul_config(1, hidden_size, inter, precision.experts_down_in0_block_w)
+
+    x_batched = ttnn.reshape(x, (1, batch, 1, hidden_size))
+    fused = ttnn.sparse_matmul(
+        x_batched,
+        weights.gate_up_proj,
+        sparsity=sparsity,
+        nnz=nnz,
+        memory_config=expert_memory_config,
+        output_tile=output_tile,
+        program_config=gate_up_config,
+        compute_kernel_config=compute_config,
+        dtype=precision.activation_dtype,
+    )
+    packed_width = fused.shape[-1]
+    fused = ttnn.reshape(fused, (batch, n_experts, packed_width))
+
+    # gate is the first half -- matches Qwen3MoeExperts.forward's chunk(2, dim=-1)
+    half = packed_width // 2
+    gate = ttnn.slice(fused, [0, 0, 0], [batch, n_experts, half])
+    up = ttnn.slice(fused, [0, 0, half], [batch, n_experts, packed_width])
+    ttnn.deallocate(fused)
+
+    down_input = ttnn.reshape(ttnn.mul(ttnn.silu(gate), up), (batch, n_experts, 1, half))
+    ttnn.deallocate(gate)
+    ttnn.deallocate(up)
+
+    down = ttnn.sparse_matmul(
+        down_input,
+        weights.down_proj,
+        sparsity=sparsity,
+        nnz=nnz,
+        memory_config=expert_memory_config,
+        output_tile=output_tile,
+        program_config=down_config,
+        is_input_a_sparse=True,
+        is_input_b_sparse=False,  # see docstring: selects batch_length_A = B * E
+        compute_kernel_config=compute_config,
+        dtype=precision.activation_dtype,
+    )
+    ttnn.deallocate(down_input)
+
+    states = ttnn.reshape(down, (batch, n_experts, hidden_size))
+    states = ttnn.mul(states, ttnn.reshape(routing, (batch, n_experts, 1)))
+    states = ttnn.unsqueeze_to_4D(ttnn.sum(states, dim=1))
+    return ttnn.reshape(states, (1, 1, batch, hidden_size), (1, 1, max(32, batch), hidden_size))
+
+
+def decoder_layer_prefill_optimized(
+    x: ttnn.Tensor,
+    weights: DecoderLayerWeights,
+    config: DecoderLayerConfig,
+    cos_cache: ttnn.Tensor,
+    sin_cache: ttnn.Tensor,
+    sparsity: ttnn.Tensor,
+    packed_experts: OptimizedWeights,
+    kv_cache: KVCache | None = None,
+    user_id: int = 0,
+    precision: PrecisionConfig = DEFAULT_PRECISION,
+) -> ttnn.Tensor:
+    """Optimized prefill. Same contract as ``decoder_layer_prefill``.
+
+    Attention runs the interleaved bfloat8_b copy: the DRAM-sharded program
+    config is decode-only (``per_core_M=1``), and a plain ``ttnn.linear`` cannot
+    read a DRAM-sharded weight at all.
+    """
+    eps = config.rms_norm_eps
+
+    normed = ttnn.rms_norm(x, weight=weights.input_layernorm, epsilon=eps)
+    attn_out = attention_prefill(
+        normed, packed_experts.attention, config.attention, cos_cache, sin_cache, kv_cache, user_id
+    )
+    ttnn.deallocate(normed)
+    hidden = ttnn.add(x, attn_out)
+    ttnn.deallocate(attn_out)
+
+    normed = ttnn.rms_norm(hidden, weight=weights.post_attention_layernorm, epsilon=eps)
+    routing = router_forward_optimized(normed, weights.router, config.moe)
+    moe_out = moe_prefill_optimized(normed, routing, packed_experts, config.moe, sparsity, precision)
+    ttnn.deallocate(normed)
+    ttnn.deallocate(routing)
+
+    out = ttnn.add(hidden, moe_out)
+    ttnn.deallocate(hidden)
+    ttnn.deallocate(moe_out)
+    return out
+
+
+def decoder_layer_decode_optimized(
+    x: ttnn.Tensor,
+    weights: DecoderLayerWeights,
+    config: DecoderLayerConfig,
+    cos_cache: ttnn.Tensor,
+    sin_cache: ttnn.Tensor,
+    kv_cache: KVCache,
+    current_pos: ttnn.Tensor,
+    token_index: int,
+    *,
+    packed_experts: OptimizedWeights,
+    precision: PrecisionConfig = DEFAULT_PRECISION,
+) -> ttnn.Tensor:
+    """Optimized decode. Decode already used per-token sparsity in stage 01."""
+    eps = config.rms_norm_eps
+
+    normed = ttnn.rms_norm(x, weight=weights.input_layernorm, epsilon=eps)
+    attn_out = attention_decode_optimized(
+        normed,
+        packed_experts,
+        config.attention,
+        cos_cache,
+        sin_cache,
+        kv_cache,
+        current_pos,
+        token_index,
+        precision=precision,
+    )
+    ttnn.deallocate(normed)
+    hidden = ttnn.add(x, attn_out)
+    ttnn.deallocate(attn_out)
+
+    normed = ttnn.rms_norm(hidden, weight=weights.post_attention_layernorm, epsilon=eps)
+    routing = router_forward_optimized(normed, weights.router, config.moe)
+    moe_out = moe_decode_optimized(normed, routing, packed_experts, config.moe, precision)
+    ttnn.deallocate(normed)
+    ttnn.deallocate(routing)
+
+    out = ttnn.add(hidden, moe_out)
+    ttnn.deallocate(hidden)
+    ttnn.deallocate(moe_out)
+    return out
+
+
+__all__ = [
+    "EXPERT_CHUNK_SIZE",
+    "router_forward_optimized",
+    "EXPERT_WEIGHT_DTYPE",
+    "EXPERT_MATH_FIDELITY",
+    "ATTENTION_WEIGHT_DTYPE",
+    "PrecisionConfig",
+    "DEFAULT_PRECISION",
+    "moe_prefill_optimized",
+    "moe_decode_optimized",
+    "attention_decode_optimized",
+    "OptimizedWeights",
+    "PackedExpertWeights",
+    "upload_optimized_weights",
+    "upload_packed_expert_weights",
+    "decoder_layer_prefill_optimized",
+    "decoder_layer_decode_optimized",
+    "build_rope_cache",
+    "build_expert_sparsity",
+    "create_kv_cache",
+    "upload_layer_weights",
+    "upload_router_weight",
+    "DecoderLayerConfig",
+    "DecoderLayerWeights",
+    "KVCache",
+    "MoEConfig",
+    "AttentionConfig",
+    "AttentionWeights",
+]

code/models/demos/blackhole/qwen3_coder_30b_a3b/tt/precision.py

@@ -0,0 +1,378 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+"""The model's precision policy, as one value that construction consumes.
+
+Before this module the policy was module-level constants in
+``optimized_decoder.py`` (``EXPERT_WEIGHT_DTYPE``, ``EXPERT_MATH_FIDELITY``,
+``ATTENTION_WEIGHT_DTYPE``, the two ``EXPERT_IN0_BLOCK_W_*``) plus a handful of
+literal ``dtype=ttnn.bfloat16`` arguments scattered through ``model.py`` and
+``multichip_decoder.py``. That is enough to *ship* a policy but not enough to
+*sweep* one: varying any of it meant editing source between runs, and a JSON
+file written next to unchanged source would be a claim rather than a
+configuration.
+
+``PrecisionConfig`` is the value those constants became. It is frozen, it
+round-trips through JSON losslessly (:meth:`to_json` / :meth:`from_json`), and
+every field below is read at model-construction or forward time by code that
+would behave differently if the field changed. The module constants still exist
+-- probes and stage-02/04 tests import them by name -- but they are now
+*derived* from ``DEFAULT_PRECISION`` rather than being the source of truth, so
+there is exactly one place a shipped value is written down.
+
+**The default is the shipped policy, and stage 07 moved it.** When this module
+was introduced its default reproduced stages 02-06 exactly. It no longer does:
+the stage-07 sweep selected new values for the two expert block widths
+(``experts_gate_up_in0_block_w`` 16 -> 64, ``experts_down_in0_block_w`` 12 ->
+24) and ``DEFAULT_PRECISION`` carries them, because the goal requires the
+selection to be what the default construction path consumes. **Every other
+field still reads back exactly what stages 02-06 measured**, and the two that
+moved are scheduling choices rather than numerical ones -- the graph is the same
+graph and the tokens are the same tokens; only the expert matmuls' inner
+blocking differs. See the block-width fields below and
+``doc/datatype_sweep/README.md``.
+
+Three fields are ``None`` by default and that is deliberate rather than an
+omission:
+
+``attention_fidelity``
+    The attention projections pass no ``compute_kernel_config`` today, so they
+    take the op default. ``None`` reproduces that exactly; any other value
+    builds a config and passes it. Encoding "op default" as an explicit
+    ``MathFidelity`` would be a guess about what the op picks.
+``ccl_dtype``
+    The collectives run at whatever dtype the activation arrives in. ``None``
+    means "inherit", which is today's behaviour and costs no ops; a named dtype
+    casts into and out of the collective.
+``experts_gate_up_fidelity`` has no ``None`` counterpart -- the expert matmuls
+    have always passed an explicit config -- so it is a plain value.
+
+The block widths live here too. They are not dtypes, but they were tuned
+*against* the dtype, so a sweep that varies expert dtype and cannot vary the
+block width alongside it would be measuring a mis-tuned point.
+
+That mattered more than expected. Stage 07's sweep found the two block widths
+to be **the only fields worth moving in the entire twenty-field config**: taking
+each to its full-K ceiling (gate_up 16 -> 64, down 12 -> 24) bought +2.83%
+traced decode at bit-identical accuracy, while every dtype lever the sweep tried
+either regressed, landed inside the run-to-run band, or hit a TTNN blocker. The
+old values were inherited from single-chip stage-02 tuning and expert
+parallelism had since cut per-die N four-fold, changing which blocking the
+matmul wants. See ``doc/datatype_sweep/README.md``.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import asdict, dataclass, fields, replace
+from pathlib import Path
+
+import ttnn
+
+__all__ = [
+    "PrecisionConfig",
+    "DEFAULT_PRECISION",
+    "dtype_from_name",
+    "dtype_to_name",
+    "fidelity_from_name",
+    "fidelity_to_name",
+]
+
+
+# --- name <-> object tables ---------------------------------------------------
+#
+# Spelled out rather than derived from ``str(dtype)`` so the JSON is a stable
+# contract: a rename inside ttnn's binding would silently invalidate every
+# archived config if the names were scraped.
+
+_DTYPES: dict[str, "ttnn.DataType"] = {
+    "bfloat16": ttnn.bfloat16,
+    "bfloat8_b": ttnn.bfloat8_b,
+    "bfloat4_b": ttnn.bfloat4_b,
+    "float32": ttnn.float32,
+    "uint8": ttnn.uint8,
+    "uint16": ttnn.uint16,
+    "int32": ttnn.int32,
+    "uint32": ttnn.uint32,
+}
+_DTYPE_NAMES: dict["ttnn.DataType", str] = {v: k for k, v in _DTYPES.items()}
+
+_FIDELITIES: dict[str, "ttnn.MathFidelity"] = {
+    "LoFi": ttnn.MathFidelity.LoFi,
+    "HiFi2": ttnn.MathFidelity.HiFi2,
+    "HiFi3": ttnn.MathFidelity.HiFi3,
+    "HiFi4": ttnn.MathFidelity.HiFi4,
+}
+_FIDELITY_NAMES: dict["ttnn.MathFidelity", str] = {v: k for k, v in _FIDELITIES.items()}
+
+
+def dtype_from_name(name):
+    """``"bfloat4_b"`` -> ``ttnn.bfloat4_b``. ``None`` and ttnn dtypes pass through."""
+    if name is None or isinstance(name, ttnn.DataType):
+        return name
+    try:
+        return _DTYPES[str(name)]
+    except KeyError:
+        raise ValueError(f"unknown dtype {name!r}; known: {sorted(_DTYPES)}") from None
+
+
+def dtype_to_name(dtype):
+    if dtype is None:
+        return None
+    try:
+        return _DTYPE_NAMES[dtype]
+    except KeyError:
+        raise ValueError(f"dtype {dtype!r} has no serialised name; add it to precision._DTYPES") from None
+
+
+def fidelity_from_name(name):
+    """``"LoFi"`` -> ``ttnn.MathFidelity.LoFi``. ``None`` and fidelities pass through."""
+    if name is None or isinstance(name, ttnn.MathFidelity):
+        return name
+    try:
+        return _FIDELITIES[str(name)]
+    except KeyError:
+        raise ValueError(f"unknown math fidelity {name!r}; known: {sorted(_FIDELITIES)}") from None
+
+
+def fidelity_to_name(fidelity):
+    if fidelity is None:
+        return None
+    try:
+        return _FIDELITY_NAMES[fidelity]
+    except KeyError:
+        raise ValueError(f"fidelity {fidelity!r} has no serialised name") from None
+
+
+# Which coercion each field takes on the way in from JSON. Every field of
+# ``PrecisionConfig`` must appear here or in ``_INT_FIELDS``; ``__post_init__``
+# asserts that, so a field added without a serialisation rule fails at import
+# rather than producing a JSON file that silently drops it.
+_DTYPE_FIELDS = frozenset(
+    {
+        "experts_gate_up_dtype",
+        "experts_down_dtype",
+        "attention_qkv_dtype",
+        "attention_wo_dtype",
+        "lm_head_dtype",
+        "router_dtype",
+        "embedding_dtype",
+        "norm_weight_dtype",
+        "activation_dtype",
+        "ccl_dtype",
+        "kv_cache_dtype",
+        "logits_dtype",
+        "sampling_dtype",
+    }
+)
+_FIDELITY_FIELDS = frozenset(
+    {
+        "experts_fidelity",
+        "attention_fidelity",
+        "router_window_fidelity",
+        "lm_head_fidelity",
+        "norm_fidelity",
+    }
+)
+_INT_FIELDS = frozenset(
+    {
+        "experts_gate_up_in0_block_w",
+        "experts_down_in0_block_w",
+        "prefill_experts_gate_up_in0_block_w",
+        "prefill_experts_down_in0_block_w",
+    }
+)
+
+
+@dataclass(frozen=True)
+class PrecisionConfig:
+    """One model's dtype / fidelity policy.
+
+    Constructed with no arguments this is ``DEFAULT_PRECISION``: the shipped
+    stage-06 policy. Vary a field with :meth:`with_overrides` (or plain
+    ``dataclasses.replace``) to get a different model out of the same source.
+    """
+
+    # -- weight dtypes, per group ---------------------------------------------
+    # The two expert weights are separate fields even though they ship at the
+    # same dtype: they have different K, different tuned block widths, and
+    # different sensitivity (down feeds the residual directly), so a sweep that
+    # could only move them together would be unable to price them apart.
+    experts_gate_up_dtype: "ttnn.DataType" = ttnn.bfloat4_b
+    experts_down_dtype: "ttnn.DataType" = ttnn.bfloat4_b
+    # Likewise qkv and wo: both bfloat8_b today, but wo is the one whose output
+    # goes straight into the attention all-reduce.
+    attention_qkv_dtype: "ttnn.DataType" = ttnn.bfloat8_b
+    attention_wo_dtype: "ttnn.DataType" = ttnn.bfloat8_b
+    lm_head_dtype: "ttnn.DataType" = ttnn.bfloat8_b
+    router_dtype: "ttnn.DataType" = ttnn.bfloat16
+    embedding_dtype: "ttnn.DataType" = ttnn.bfloat16
+    # RMSNorm weights and the per-head q_norm/k_norm vectors. 4 KB each; here
+    # for completeness of the picture rather than because it is a lever.
+    norm_weight_dtype: "ttnn.DataType" = ttnn.bfloat16
+
+    # -- per-group compute fidelity -------------------------------------------
+    experts_fidelity: "ttnn.MathFidelity" = ttnn.MathFidelity.LoFi
+    # ``None`` == the op default, which is what the projections take today.
+    attention_fidelity: "ttnn.MathFidelity | None" = None
+    # HiFi4 so the one-hot expert-window matmul selects rather than approximates
+    # -- see ``multichip_decoder._exact_matmul_config``. Lowering this is a
+    # correctness change, not a speed/accuracy trade; it is configurable so a
+    # sweep can *demonstrate* that rather than assert it.
+    router_window_fidelity: "ttnn.MathFidelity" = ttnn.MathFidelity.HiFi4
+    lm_head_fidelity: "ttnn.MathFidelity" = ttnn.MathFidelity.HiFi2
+    norm_fidelity: "ttnn.MathFidelity" = ttnn.MathFidelity.HiFi4
+
+    # -- expert matmul inner block widths -------------------------------------
+    # Tuned against ``experts_fidelity``; see the module docstring.
+    #
+    # **Stage 07 moved these, and they are the only fields the sweep moved.**
+    # They were 16 and 12, inherited from the single-chip stage-02 tuning. The
+    # 48-layer sweep measured both brackets end to end and found each monotonic
+    # upward to its full-K ceiling, at *identical* accuracy:
+    #
+    #   gate_up (K = hidden_size 2048 = 64 tiles):  8 -> 41.33, 16 -> 42.34,
+    #                                              32 -> 42.94, 64 -> 43.23 t/s/u
+    #   down    (K = moe_intermediate_size 768 = 24 tiles):
+    #                                               6 -> 41.67, 12 -> 42.34,
+    #                                              24 -> 42.99 t/s/u
+    #
+    # and the combination at both ceilings measured 43.54 t/s/u -- +2.83% over
+    # the shipped default, top-1 0.990 / top-5 1.000 / top-100 1.000, i.e. no
+    # accuracy cost at all, because a block width is a *scheduling* choice and
+    # not a numerical one. The stage-02 comment that 16 wins at LoFi predates
+    # expert parallelism, which cut per-die N four-fold and changed which
+    # blocking the matmul wants.
+    #
+    # Both values are exact divisors of K in tiles, so
+    # ``_tuned_sparse_matmul_config`` does not clamp them; ``fallback_audit``
+    # reports the resolved widths and stage 07 asserts on those, not on these.
+    # See ``doc/datatype_sweep/README.md``.
+    experts_gate_up_in0_block_w: int = 64
+    experts_down_in0_block_w: int = 24
+
+    # PREFILL-only inner block widths. Default to the decode values above, so
+    # the shipped graph is byte-identical until a sweep says otherwise.
+    #
+    # Why they exist separately: the two widths above were selected by stage 07,
+    # the DATATYPE sweep, optimised against single-token DECODE (52.05 t/s/u),
+    # where the expert matmuls run at M = 1. Prefill runs the same matmuls at
+    # M = EXPERT_CHUNK_SIZE = 32 and simply inherited decode's tuning; an
+    # optimum for a latency-bound M=1 matmul has no reason to be the optimum at
+    # M=32. Consumed at optimized_decoder.py:782-784 (prefill) while 926-927
+    # (decode, M=1) keeps reading the fields above.
+    #
+    # Scheduling only, not numerics: the graph and the tokens are the same, only
+    # the matmuls' inner blocking differs.
+    #
+    # SHIPPED 16/12, which are the PRE-stage-07 values. Stage 07 moved the decode
+    # fields 16 -> 64 and 12 -> 24 and, because prefill shared them, regressed
+    # prefill by ~4 % without measuring it. Sweeping prefill on its own lands
+    # exactly back on the old pair:
+    #
+    #   64/24 (stage 07)  3.741 s      16/24  3.664 s
+    #   32/12             3.604 s      16/12  3.596 s   <- 1.0404x
+    #
+    # at 4,096 tokens, 15 configurations, every one PCC 1.0000000000000058 with
+    # identical greedy tokens -- scheduling only, no numerical dimension. The 4 %
+    # is modest but free, and decode keeps its own 64/24 above.
+    # See doc/batch_scaling/README.md, "Where prefill time is NOT going".
+    prefill_experts_gate_up_in0_block_w: int = 16
+    prefill_experts_down_in0_block_w: int = 12
+
+    # -- activations ----------------------------------------------------------
+    # The dtype of every hidden state, including the inter-layer residual. The
+    # residual *layout* (replicated ``[1, 1, rows, 2048]``, DRAM interleaved) is
+    # a contract and is not configurable here.
+    activation_dtype: "ttnn.DataType" = ttnn.bfloat16
+    # ``None`` == run the collective at the activation dtype, no cast. A named
+    # dtype casts in and out around the reduce-scatter/all-gather pair.
+    ccl_dtype: "ttnn.DataType | None" = None
+
+    # -- kv cache -------------------------------------------------------------
+    kv_cache_dtype: "ttnn.DataType" = ttnn.bfloat16
+
+    # -- terminal path --------------------------------------------------------
+    logits_dtype: "ttnn.DataType" = ttnn.bfloat16
+    # What the sampler is handed. Equal to ``logits_dtype`` by default, so the
+    # shipped path casts nothing.
+    sampling_dtype: "ttnn.DataType" = ttnn.bfloat16
+
+    def __post_init__(self) -> None:
+        names = {f.name for f in fields(self)}
+        unclassified = names - _DTYPE_FIELDS - _FIDELITY_FIELDS - _INT_FIELDS
+        assert not unclassified, f"PrecisionConfig fields with no serialisation rule: {sorted(unclassified)}"
+        for name in _DTYPE_FIELDS:
+            object.__setattr__(self, name, dtype_from_name(getattr(self, name)))
+        for name in _FIDELITY_FIELDS:
+            object.__setattr__(self, name, fidelity_from_name(getattr(self, name)))
+        for name in _INT_FIELDS:
+            value = int(getattr(self, name))
+            if value < 1:
+                raise ValueError(f"{name} must be >= 1, got {value}")
+            object.__setattr__(self, name, value)
+        # ``None`` is legal only where the docstring says it is.
+        for name in _DTYPE_FIELDS | _FIDELITY_FIELDS:
+            if getattr(self, name) is None and name not in ("ccl_dtype", "attention_fidelity"):
+                raise ValueError(f"{name} may not be None")
+
+    # -- convenience ----------------------------------------------------------
+
+    def with_overrides(self, **overrides) -> "PrecisionConfig":
+        """A copy with ``overrides`` applied; values may be names or objects."""
+        unknown = set(overrides) - {f.name for f in fields(self)}
+        if unknown:
+            raise ValueError(f"unknown precision fields: {sorted(unknown)}")
+        return replace(self, **overrides)
+
+    @property
+    def effective_ccl_dtype(self):
+        """The dtype the collectives actually run at, resolving ``None``."""
+        return self.activation_dtype if self.ccl_dtype is None else self.ccl_dtype
+
+    # -- serialisation --------------------------------------------------------
+
+    def to_dict(self) -> dict:
+        """JSON-ready ``{field: name}``. Every field of the dataclass appears."""
+        out: dict = {}
+        for name, value in asdict(self).items():
+            if name in _DTYPE_FIELDS:
+                out[name] = dtype_to_name(value)
+            elif name in _FIDELITY_FIELDS:
+                out[name] = fidelity_to_name(value)
+            else:
+                out[name] = int(value)
+        return out
+
+    @classmethod
+    def from_dict(cls, data: dict) -> "PrecisionConfig":
+        known = {f.name for f in fields(cls)}
+        unknown = set(data) - known
+        if unknown:
+            raise ValueError(f"unknown precision fields in config: {sorted(unknown)}")
+        return cls(**{k: v for k, v in data.items() if k in known})
+
+    def to_json(self, indent: int = 2) -> str:
+        return json.dumps(self.to_dict(), indent=indent, sort_keys=True) + "\n"
+
+    @classmethod
+    def from_json(cls, text: str) -> "PrecisionConfig":
+        return cls.from_dict(json.loads(text))
+
+    def write_json(self, path: str | Path) -> Path:
+        path = Path(path)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        path.write_text(self.to_json())
+        return path
+
+    @classmethod
+    def read_json(cls, path: str | Path) -> "PrecisionConfig":
+        return cls.from_json(Path(path).read_text())
+
+
+#: The shipped policy, as stage 07 selected it.
+#:
+#: Every stage-02..06 number was measured at this config **except for the two
+#: expert block widths**, which stage 07 moved from 16/12 to 64/24; those stages
+#: ran at 16/12. Nothing else here has changed since stage 02.
+DEFAULT_PRECISION = PrecisionConfig()

code/models/demos/blackhole/qwen3_coder_30b_a3b/tt/weight_mapping.py

@@ -0,0 +1,184 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+"""HuggingFace -> TTNN weight mapping for Qwen3-Coder-30B-A3B-Instruct.
+
+RoPE convention (the decision this file encodes)
+------------------------------------------------
+TTNN offers two rotary embedding ops, and each demands a different weight
+layout. They are interchangeable only as a *pair*; crossing them runs fine and
+silently produces garbage.
+
+  Meta-style  ``rotary_embedding_llama``     head channels interleaved
+              -> q/k rows must be reordered (``reverse_permute``), and because
+                 that reorders channels *within* a head, Qwen3's per-head
+                 QK-norm weights must be reordered to match as well.
+
+  HF-style    ``ttnn.experimental.rotary_embedding``   HF's native layout
+              -> no weight transformation at all.
+
+This port uses **HF-style**, matching models/demos/gemma4 (the Blackhole
+exemplar this attention block is written against, and the only other supported
+model with Qwen3-shaped per-head QK-norm). Keeping the checkpoint layout
+untouched removes both permutation steps, so the QK-norm weights are copied
+verbatim. If the RoPE op in functional_decoder.py is ever swapped for the llama
+variant, both permutations have to come back together.
+
+Expert fusion (MoE)
+-------------------
+The checkpoint stores 3 tensors per expert. TTNN wants them batched, with gate
+and up fused as ``[gate ; up]`` along the output dim, matching
+``Qwen3MoeExperts.forward``'s ``chunk(2, dim=-1)``.
+
+The fused QKV layout follows models/tt_transformers/tt/attention.py: transpose
+each projection to ``[in, out]``, then concatenate ``[q, k, v]`` along the
+output dim.
+"""
+
+from __future__ import annotations
+
+import torch
+
+
+def convert_attention_weights(
+    sd: dict[str, torch.Tensor],
+    *,
+    n_heads: int,
+    n_kv_heads: int,
+    head_dim: int,
+) -> dict[str, torch.Tensor]:
+    """Return the fused ``wqkv`` plus ``wo`` and the QK-norm weights.
+
+    Channel order is left exactly as HuggingFace stores it -- see the RoPE
+    convention note in the module docstring. The only structural change is
+    fusing Q, K and V into one matmul.
+
+    ``sd`` uses layer-relative keys (``self_attn.q_proj.weight``, ...).
+    Single-device layout; tensor-parallel chunking is a later stage.
+    """
+    q = sd["self_attn.q_proj.weight"].float()
+    k = sd["self_attn.k_proj.weight"].float()
+    v = sd["self_attn.v_proj.weight"].float()
+    o = sd["self_attn.o_proj.weight"].float()
+
+    assert q.shape[0] == n_heads * head_dim, f"q_proj {tuple(q.shape)} != {n_heads}x{head_dim}"
+    assert k.shape[0] == n_kv_heads * head_dim, f"k_proj {tuple(k.shape)} != {n_kv_heads}x{head_dim}"
+
+    # torch stores nn.Linear as [out, in]; TTNN matmuls right-multiply, so
+    # transpose to [in, out] before concatenating along the output dim.
+    wqkv = torch.cat([q.T, k.T, v.T], dim=-1).unsqueeze(0).unsqueeze(0)
+
+    out = {
+        "wqkv": wqkv,
+        "wo": o.T.contiguous(),
+    }
+
+    for name, key in (("q_norm", "self_attn.q_norm.weight"), ("k_norm", "self_attn.k_norm.weight")):
+        if key in sd:
+            out[name] = sd[key].float()
+
+    return out
+
+
+def hf_to_meta_channels(head_dim: int) -> torch.Tensor:
+    """Index vector mapping an HF-ordered head to Meta (llama) channel order.
+
+    HF's ``rotate_half`` pairs channel ``i`` with ``i + head_dim/2``; the Meta
+    convention that ``rotary_embedding_llama``'s transformation matrix encodes
+    pairs ``2i`` with ``2i + 1``. So ``meta[2i] = hf[i]`` and
+    ``meta[2i+1] = hf[i + head_dim/2]``, i.e. an interleave of the two halves.
+
+    The *same* vector converts a cos/sin row, because HF stores
+    ``[c0 .. c_{d/2-1}, c0 .. c_{d/2-1}]`` and Meta stores ``[c0, c0, c1, c1, ...]``.
+
+    Applied per head, so it commutes with any reordering of whole heads --
+    which is what lets ``multichip_decoder`` apply it before
+    ``head_interleaved_wqkv`` splits the heads across dies.
+    """
+    assert head_dim % 2 == 0, head_dim
+    half = head_dim // 2
+    return torch.stack([torch.arange(half), torch.arange(half) + half], dim=1).reshape(-1)
+
+
+def permute_wqkv_to_meta(wqkv: torch.Tensor, *, n_heads: int, n_kv_heads: int, head_dim: int) -> torch.Tensor:
+    """Reorder the Q and K channels of a fused ``[..., in, out]`` wqkv to Meta order.
+
+    **V is deliberately untouched** -- RoPE is applied to Q and K only, so V's
+    channels keep their HF meaning, and so does ``wo``, which consumes the
+    attention output in V's space.
+
+    This is the ``reverse_permute`` step the module docstring says comes back
+    the moment the llama rotary op is used. It is a *checkpoint layout* change
+    and is applied once at upload, never per token.
+    """
+    out = wqkv.clone()
+    perm = hf_to_meta_channels(head_dim)
+    q_end = n_heads * head_dim
+    k_end = q_end + n_kv_heads * head_dim
+    for start, count in ((0, n_heads), (q_end, n_kv_heads)):
+        for h in range(count):
+            lo = start + h * head_dim
+            out[..., lo : lo + head_dim] = wqkv[..., lo : lo + head_dim][..., perm]
+    assert torch.equal(out[..., k_end:], wqkv[..., k_end:]), "V must not be permuted"
+    return out
+
+
+def permute_head_vector_to_meta(vec: torch.Tensor, *, head_dim: int) -> torch.Tensor:
+    """Reorder a per-head vector (Qwen3's ``q_norm`` / ``k_norm``) to Meta order.
+
+    Qwen3 applies these **between the head split and RoPE**, so they index the
+    same channels the rotary op does. Permuting Q/K without permuting these
+    scales the wrong channel and is silent -- it does not change any shape and
+    it does not raise. ``test_meta_rope_weights_match_hf`` is the assertion that
+    catches it.
+    """
+    flat = vec.reshape(-1)
+    assert flat.numel() == head_dim, (flat.shape, head_dim)
+    return flat[hf_to_meta_channels(head_dim)].reshape(vec.shape)
+
+
+def convert_moe_weights(sd: dict[str, torch.Tensor], *, n_experts: int) -> dict[str, torch.Tensor]:
+    """Batch the per-expert checkpoint tensors and fuse gate/up.
+
+    gate first, then up -- ``Qwen3MoeExperts.forward`` chunks the matmul output
+    in half and treats the FIRST half as gate.
+    """
+    gate_up = torch.stack(
+        [
+            torch.cat(
+                [
+                    sd[f"mlp.experts.{e}.gate_proj.weight"].float(),
+                    sd[f"mlp.experts.{e}.up_proj.weight"].float(),
+                ],
+                dim=0,
+            )
+            for e in range(n_experts)
+        ]
+    )
+    down = torch.stack([sd[f"mlp.experts.{e}.down_proj.weight"].float() for e in range(n_experts)])
+    return {
+        "router": sd["mlp.gate.weight"].float(),
+        "experts_gate_up": gate_up,
+        "experts_down": down,
+    }
+
+
+def convert_norm_weights(sd: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+    """Decoder-layer RMSNorm weights. Plain RMSNorm -- no zero-centering."""
+    return {
+        "input_layernorm": sd["input_layernorm.weight"].float(),
+        "post_attention_layernorm": sd["post_attention_layernorm.weight"].float(),
+    }
+
+
+def convert_layer_weights(sd: dict[str, torch.Tensor], config) -> dict[str, torch.Tensor]:
+    """Full layer conversion: attention + MoE + norms."""
+    weights = convert_attention_weights(
+        sd,
+        n_heads=config.num_attention_heads,
+        n_kv_heads=config.num_key_value_heads,
+        head_dim=config.head_dim,
+    )
+    weights.update(convert_moe_weights(sd, n_experts=config.num_experts))
+    weights.update(convert_norm_weights(sd))
+    return weights

code/models/demos/blackhole/qwen3_coder_30b_a3b/vllm_bundle/qwen3_coder_30b_a3b_instruct/tt_qwen3_coder_30b_a3b_instruct.py

@@ -0,0 +1,35 @@
+# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
+# SPDX-License-Identifier: Apache-2.0
+
+"""Bundle entry point for ``EXTRA_MODELS_DIR``.
+
+``vllm_tt_plugin.platform.register_tt_models()`` runs
+``_register_models_from_extra_dir(ModelRegistry)`` as its **first** action --
+"so a distributed bundle can supply a model without touching this file". That
+hook appends this folder to ``sys.path`` and lazily registers
+``vllm_metadata.json``'s ``main_class`` under the plugin's ``TT``-prefixed
+convention, so the model ends up registered *by* ``register_tt_models()`` with
+no edit to the plugin checkout.
+
+Registration is lazy: vLLM resolves the ``"module:Class"`` string later, in the
+API-server process and again in each EngineCore worker. This module therefore
+has to be importable on its own, which means it cannot assume the tt-metal
+checkout is already on ``sys.path`` -- an EngineCore worker's working directory
+is not guaranteed. It appends the repository root (never ``insert(0)``, matching
+the hook's own rule that an installed package of the same name must still win)
+and re-exports the real adapter, which lives with the model it adapts.
+"""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+# .../models/demos/blackhole/<model>/vllm_bundle/<bundle>/this_file.py
+_REPO_ROOT = Path(__file__).resolve().parents[6]
+if str(_REPO_ROOT) not in sys.path:
+    sys.path.append(str(_REPO_ROOT))
+
+from models.demos.blackhole.qwen3_coder_30b_a3b.tt.generator_vllm import Qwen3CoderForCausalLM  # noqa: E402
+
+__all__ = ["Qwen3CoderForCausalLM"]

code/models/demos/blackhole/qwen3_coder_30b_a3b/vllm_bundle/qwen3_coder_30b_a3b_instruct/vllm_metadata.json

@@ -0,0 +1,8 @@
+{
+  "arch": "Qwen3MoeForCausalLM",
+  "main_class": "tt_qwen3_coder_30b_a3b_instruct:Qwen3CoderForCausalLM",
+  "hf_model": "Qwen/Qwen3-Coder-30B-A3B-Instruct",
+  "adapter": "models/demos/blackhole/qwen3_coder_30b_a3b/tt/generator_vllm.py",
+  "registered_as": "TTQwen3MoeForCausalLM",
+  "notes": "Self-contained bundle for vllm_tt_plugin.platform.register_tt_models() -> _register_models_from_extra_dir(). Point EXTRA_MODELS_DIR at the parent vllm_bundle/ directory. 'arch' is this checkout's config.json architecture; the plugin prefixes it with TT. 'main_class' resolves through the sibling shim module, which is importable because the plugin appends this folder to sys.path."
+}

image/blobs/sha256/0926a8eb0e608a5c6888d1cd5594184bdf3ed3aa311dba5b42a547caefdc6f2e

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0926a8eb0e608a5c6888d1cd5594184bdf3ed3aa311dba5b42a547caefdc6f2e
+size 29752807

image/blobs/sha256/24cba7375920bef8d4cc4f0ce4294f8f70c65b9a37c4ed6f6c2d63405c76ba3c

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:24cba7375920bef8d4cc4f0ce4294f8f70c65b9a37c4ed6f6c2d63405c76ba3c
+size 1385020

image/blobs/sha256/3de1f5eb93e54b4561afa733cf844d1914a9ba260f49e86fe600f344c5dd025c

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3de1f5eb93e54b4561afa733cf844d1914a9ba260f49e86fe600f344c5dd025c
+size 35141596

image/blobs/sha256/530b0e35f44c6f963e06fdaacdbfecb2021d4c55d49fe4c0019e161d94c18de3

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:530b0e35f44c6f963e06fdaacdbfecb2021d4c55d49fe4c0019e161d94c18de3
+size 10547489

image/blobs/sha256/540cf00275e913a9bccc49fe7beba58037696661bde6ae083a3ac84c5b160e67

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:540cf00275e913a9bccc49fe7beba58037696661bde6ae083a3ac84c5b160e67
+size 9565684

image/blobs/sha256/8753e0cfbd424e962ccaf50aaaf02fd06ff2efb8677219657e751a53922efa9f

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8753e0cfbd424e962ccaf50aaaf02fd06ff2efb8677219657e751a53922efa9f
+size 1076048

image/blobs/sha256/b6df468b82a4b2f9ee3ca3a79a6bbe99b5cda02ddc63b7e3c89e7bd08ef41706

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b6df468b82a4b2f9ee3ca3a79a6bbe99b5cda02ddc63b7e3c89e7bd08ef41706
+size 48137518

image/blobs/sha256/c00314a02c644cf2aea5a8ae4e3ee5ee383f4072d3e35248fe20dea57a90c46c

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c00314a02c644cf2aea5a8ae4e3ee5ee383f4072d3e35248fe20dea57a90c46c
+size 1727127131

image/blobs/sha256/c18d0f3c8022bcd5a8059f67fc0e5cfd53ef36a44997335d6cd0aa6b19db140d

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c18d0f3c8022bcd5a8059f67fc0e5cfd53ef36a44997335d6cd0aa6b19db140d
+size 35141895

image/blobs/sha256/c3bc7b373b4523cabcdd9c64ab10d32510ef61bace60873a279ccc4902738989

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c3bc7b373b4523cabcdd9c64ab10d32510ef61bace60873a279ccc4902738989
+size 139055657

image/blobs/sha256/ca0b072b65f8c21199f96e7498f2ccbc252490ce39060800647332751f857287

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ca0b072b65f8c21199f96e7498f2ccbc252490ce39060800647332751f857287
+size 35988690

image/blobs/sha256/cbb77a738c7df827819a8b8bf87682eeca8bdb434d41411a9682c38717f2f187

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cbb77a738c7df827819a8b8bf87682eeca8bdb434d41411a9682c38717f2f187
+size 82239359

image/blobs/sha256/e19b6f1fb65dd2888d9003ef9513a21d129041a328bb8a9a4164d29ef0382b16

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e19b6f1fb65dd2888d9003ef9513a21d129041a328bb8a9a4164d29ef0382b16
+size 138938826

image/blobs/sha256/fdc1ed79ffd24d66f8be3754ec4dc80b1ab0fcc8c8165a6007748c464c94f897

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fdc1ed79ffd24d66f8be3754ec4dc80b1ab0fcc8c8165a6007748c464c94f897
+size 415353