Optimized build (2026-10-03): model call 87 -> 28 ms, pose path 38 ms

code/ updated to the independently verified optimized build (12x10 compute grid). README Demo & Performances, GPU_COMPARISON.md update, OPT_REPORT/OPT_BASELINE, VERIFICATION_2026-10-03.md and patches/tt-metal-eth-dispatch.patch added. tt-model.yaml and SERVING.md are unchanged; they still describe the container image.

GPU_COMPARISON.md

@@ -184,3 +184,46 @@ HF_HUB_OFFLINE=1 /home/deepgadget/experiments/tt-models/.venv-gpu/main/bin/pytho
 ```
 
 GPU released after the run: `nvidia-smi --query-compute-apps=pid --format=csv,noheader` -> empty.
+
+## Update 2026-10-03: optimized build and RTX 5090
+
+The GPU numbers above are unchanged. They come from the 2026-09-14 run. This section compares them with the optimized build at commit `ad39150` (see `OPT_REPORT.md`).
+
+### New Blackhole numbers (independent verification, 2026-10-03)
+
+The verifier measured these numbers with `bench_breakdown.py` (warm, batch 1, one 512x512 pair, 30 iterations per run, median / min). The hardware is a Blackhole chip with a 12x10 compute grid. Details: `VERIFICATION_2026-10-03.md`.
+
+| Blackhole metric | Configuration | Value |
+|---|---|---:|
+| Model call `model(img1, img2)` (H2D + trace + D2H of both maps) | served path config: worker dispatch, 2 CQ, uint8 input | 28.23 / 27.55 and 27.99 / 27.68 ms (2 runs) |
+| Device trace (host wall, synchronized) | served path config | 25.39 / 24.82 and 25.23 / 24.52 ms |
+| Model call | ETH dispatch, 1 CQ, float input | 29.99 / 29.84 and 29.54 / 29.27 ms |
+| Device trace | ETH dispatch, 1 CQ, float input | 24.72 / 24.47 and 25.03 / 24.41 ms |
+| Device span (profiler cycles at nominal 1.35 GHz, 3 replays) | ETH dispatch, 1 CQ | 23.484 / 23.473 / 23.485 ms |
+| `return_pose` forward: `sym` graph / two-pass | served path config, real pair | 38.00 / 55.82 ms |
+
+For the ratios, this section uses 28.1 ms for the model call and 25.3 ms for the device trace. Each value is the mean of the two run medians (served path config). The model call matches the GPU `incl_h2d` definition. The device trace matches the GPU `excl_h2d` definition. The previous Blackhole value for the model call was 73.9 ms.
+
+### Recomputed ratios (Blackhole ms / GPU ms; > 1 means the GPU is faster)
+
+| GPU variant / precision | GPU incl_h2d | ratio vs 28.1 ms | GPU excl_h2d | ratio vs 25.3 ms |
+|---|---:|---:|---:|---:|
+| ref, fp32 strict | 103.409 | 0.27 (Blackhole 3.68x faster) | 102.278 | 0.25 (Blackhole 4.04x faster) |
+| ref, tf32 | 70.753 | 0.40 (Blackhole 2.52x faster) | 69.983 | 0.36 (Blackhole 2.77x faster) |
+| ref, bf16 autocast | 63.753 | 0.44 (Blackhole 2.27x faster) | 62.613 | 0.40 (Blackhole 2.47x faster) |
+| ref, fp16 autocast | 56.552 | 0.50 (Blackhole 2.01x faster) | 55.505 | 0.46 (Blackhole 2.19x faster) |
+| bf16_weights, eager | 42.329 | 0.66 (Blackhole 1.51x faster) | 41.277 | 0.61 (Blackhole 1.63x faster) |
+| fp16_weights, eager | 39.459 | 0.71 (Blackhole 1.40x faster) | 38.308 | 0.66 (Blackhole 1.51x faster) |
+| SDPA, bf16_weights, eager | 35.327 | 0.80 (Blackhole 1.26x faster) | 34.285 | 0.74 (Blackhole 1.36x faster) |
+| fp16 autocast + compile reduce-overhead | 23.214 | 1.21 (GPU 1.21x faster) | 22.149 | 1.14 (GPU 1.14x faster) |
+| bf16_weights + compile reduce-overhead | 21.106 | 1.33 (GPU 1.33x faster) | 20.109 | 1.26 (GPU 1.26x faster) |
+
+### Reading
+
+- The optimized build is faster than every eager GPU variant, including bf16 weights with SDPA attention (1.26x on the model call).
+- `torch.compile` with CUDA graphs is still faster on the GPU: 1.21-1.33x on the model call and 1.14-1.26x on the device forward.
+- On the measurement host, the Blackhole chip has an x1 PCIe link. The overlapped D2H readback stalls one device kernel and costs about 1.2 ms per pair (`OPT_REPORT.md`, round 9).
+- Under load, the chip clock drops to about 1.26-1.30 GHz. Thus, the trace wall time is about 1.3-1.7 ms longer than the device span.
+- With `return_pose`, the `sym` graph takes 38.0 ms. Two GPU forwards with bf16 weights and SDPA take 2 x 35.3 = 70.6 ms. This comparison is specific to the pose mode.
+- The served `/predict` total stays host-bound on both sides. The npz encode takes about 160-172 ms, and the build did not change the encode code. The verifier did not measure a served `/predict` total for this build.
+- The verifier did not measure p150a power, so this section makes no efficiency comparison.

OPT_BASELINE.md

@@ -0,0 +1,228 @@
+# mast3r-p150 — optimization baseline (Galaxy BH chip 9)
+
+Date: 2026-10-02. tt-metal `8b98410e730` + ETH-dispatch patch (shared tree, unmodified by me).
+Model code is the baseline commit `20bdff0`, unmodified. Weights: `naver/DUSt3R_ViTLarge_BaseDecoder_512_dpt` @ `61c57447d7b0` (shared HF cache).
+Default path = `TT_FUSED=1` fused graph (llama RoPE, `dit_minimal_matmul_addcmul_fused` residual linears, SDPA 128/256 chunks, cached conv weights) captured into one metal trace.
+
+## Chip 9 vs p150
+
+| | p150a (authors) | Galaxy BH chip 9 |
+|---|---|---|
+| Tensix grid, default dispatch | 11x10 | 12x10 (13 Tensix columns, 1 used for dispatch) |
+| Tensix grid, ETH dispatch | 12x10 | 13x10 (policy: capped at 12x10, chipenv sets `TT_METAL_CORE_GRID_OVERRIDE_TODEPRECATE=11,9`) |
+| PCIe | x16 | x1 Gen5. H2D and D2H are much slower, and readback is dominated by host untilize. |
+| CQs with ETH dispatch | — | 1 only. This model uses 1 CQ, so ETH dispatch is usable. |
+
+I reproduced the p150 11x10 grid with `TT_METAL_CORE_GRID_OVERRIDE_TODEPRECATE=10,9` and Tensix dispatch (`bench_breakdown.py --mode worker11`).
+
+## How to run
+
+```bash
+ROOT=/home/ttuser/experiments/tt-models; M=$ROOT/models/mast3r-p150
+# one time: $ROOT/tools/mkvenv.sh $M   (done; numpy/cv2/hf_hub come from the shared env, nothing extra installed)
+cd $M/code
+source $ROOT/tools/chipenv.sh 9 $M/.venv      # forces the 12x10 cap
+export TT_WEIGHTS_REVISION=61c57447d7b0adc8a1a30b2b0adec7a8935aa2a3
+python test_mast3r.py --layer end_to_end --runs 25        # authors' gate: synthetic pair, PCC >= 0.99, best-of-N latency (traced)
+python test_mast3r.py --layer full_encoder|full_decoder|dpt_head_device|dpt_head_device_2 --runs 5   # eager per-stage gates
+python tools_prof/real_pair_acc.py                       # real pair media/source_{1,2}.png vs fp32 torch (pts3d/conf PCC)
+python bench_breakdown.py --mode worker|worker11|eth12   # per-phase timing: host prep / H2D / trace / D2H / e2e
+MAST3R_TRACE=0 python -m tracy -r -p -v -o $TT_METAL_PROFILER_DIR/mast3r_eager --op-support-count 6000 tools_prof/prof_eager.py   # device op profile
+python tools_prof/d2h_probe.py                           # readback-format micro-benchmark
+```
+The first run on a new grid JIT-compiles about 460 kernels, which takes about 80 s. After that the kernel cache in `chipstate/chip9/cache` is warm and the first call takes about 4.5 s.
+
+## Accuracy (baseline, all PASS)
+
+| Test | Result |
+|---|---|
+| `end_to_end` synthetic pair, 11x10 Tensix | PCC 0.9987 (head1 0.9983 / head2 0.9985). README: 0.9987 |
+| `end_to_end` synthetic pair, 12x10 Tensix | PCC 0.9984 (0.9983 / 0.9982) |
+| `bench_breakdown` 12x10 ETH | head1 0.99836 / head2 0.99895 |
+| `full_encoder` / `full_decoder` / `dpt_head_device` / `dpt_head_device_2` (eager) | 0.9997 / 1.0001 (b1 0.9999, b2 1.0000) / 1.0000 / 0.9999 |
+| Real pair (kitchen source_1/2), 12x10 | pts3d PCC head1 0.98985 / head2 0.99078; conf PCC 0.9924 / 0.9944; raw PCC 0.9919 / 0.9765; median rel depth err 1.3%. README range: pts3d 0.987-0.9995, conf ~0.99 |
+
+Proposed accuracy gates for optimization work, so we stay at or above the baseline: synthetic e2e PCC ≥ 0.998; real pair pts3d PCC ≥ 0.989 for each head; conf PCC ≥ 0.99.
+
+## Performance (warm, synthetic 512x512 pair, chip 9)
+
+`bench_breakdown.py`, 30 iterations, median (min):
+
+| Phase | 11x10 Tensix (p150-equiv) | 12x10 Tensix | 12x10 ETH |
+|---|---|---|---|
+| host prep (cat, bf16, host tensor) | 0.85 | 0.86 | 0.85 |
+| H2D input `[2,3,512,512]` bf16 RM (3 MB) | 1.07 | 1.06 | 1.05 |
+| **trace (device forward, execute + sync)** | **62.62** (62.49) | **61.82** (61.73) | **62.19** (62.12) |
+| D2H, 2 heads (`to_torch` + reshape/permute/float) | 24.89 | 24.01 | 23.19 |
+| **e2e model call** | 89.16 (84.39) | 87.20 (84.00) | 86.39 (83.94) |
+
+`test_mast3r.py --layer end_to_end --runs 25` best-of-25: **84.91 ms** on 11x10 and **83.64 ms** on 12x10. The README reports 73.1 ms on p150; the 11 ms gap is the x1 PCIe readback.
+
+- Moving from 11x10 to 12x10 gains only 0.8 ms (1.3%) on the device forward, because most ops do not scale to the full grid (see grid usage below).
+- ETH dispatch vs Tensix dispatch at the same 12x10 grid makes no difference (62.2 vs 61.8 ms). The trace is kernel-bound: the sum of device kernel durations is 60.97 ms against 62.2 ms of trace, so dispatch gaps total only about 1.2 ms over 1252 ops (about 1 µs per op).
+
+## Host vs device; every host↔device transfer per forward
+
+Everything from the patch-embed im2col to the final 1x1 head conv runs on device inside one trace. No host round-trips happen mid-graph on the fused path: RoPE LUTs, `trans_mat`, the `ones` vectors and all weights (including prepared conv weights) are uploaded during the eager warm-up, before capture.
+
+On the host, per call:
+1. `torch.cat` + bf16 cast + `ttnn.from_torch` RM: about 0.85 ms.
+2. Post-processing of the outputs: reshape `(1,512,512,4)`, permute to NCHW, `.float()`. About 0.4 ms per head. This is included in D2H.
+
+Transfers per forward (traced path):
+
+| # | Direction | Tensor | Bytes on the wire | Time |
+|---|---|---|---|---|
+| 1 | H2D | input `[2,3,512,512]` bf16 ROW_MAJOR (`copy_host_to_device_tensor`) | 3.1 MB | 1.0 ms |
+| 2 | D2H | head1 output `[1,1,262144,4]` bf16 **TILE, padded to 32 channels** | **16.8 MB** (8x the useful 2 MB) | ~12 ms |
+| 3 | D2H | head2 output, same | 16.8 MB | ~12 ms |
+
+The readback probe (`tools_prof/d2h_probe.py`, one head, so these are per-head times):
+- `to_torch` of the padded TILE output takes 28-30 ms in isolation (12 ms in the model loop). Of that, raw `from_device` DMA is only 5.4 ms; the rest is host untilize of 16 MB.
+- A 2 MB ROW_MAJOR tensor with 8-byte pages (`[262144,4]`) also costs 5.7 ms, because the cost is per page.
+- A 2 MB ROW_MAJOR tensor with wide pages (`[1024,1024]`) costs **0.93 ms**, and a 2 MB TILE `[1024,1024]` takes 0.73 ms raw or 1.43 ms with `to_torch`.
+
+**Conclusion:** writing the output in a compact, wide-page layout on device inside the trace would cut D2H from about 24 ms to about 2 ms. Examples: a channel-4 ROW_MAJOR tensor reshaped to `[512, 2048]`, or both heads packed into one `[2,512,2048]` RM tensor.
+
+## Grid / core usage
+
+- **There is no hard-coded 11x10 or core grid in the port.** Every explicit program config takes `device.compute_with_storage_grid_size()` (`ttnn_dust3r.py:_fused_consts`, lines 108-135): `MinimalMatmulConfig` (8,4,4 / 4,4,4, 2x2 subblocks) and `SDPAProgramConfig(q=128, k=256)`. Everything else uses ttnn auto configs: `ttnn.linear` with no program config, interleaved `layer_norm`, NLP create/concat heads, conv2d auto sharding, and upsample. So the code follows whatever grid the device exposes. Its own docstrings note that the block sizes were tuned on the 11x10 p150.
+- Actual core use per op, from the device profiler on 12x10 (120 cores), for one forward:
+
+| Cores used | Ops | Kernel time |
+|---|---|---|
+| 120 (full) | 483 | 29.2 ms (48%) |
+| 96-119 | 288 | 18.0 ms (30%). Auto `ttnn.linear` (MatmulMultiCoreReuseMultiCast on 96 cores, `in0_block_w=1`), conv2d 108/114, tilize 119 |
+| 64-95 | 263 | 9.5 ms (16%). Encoder LayerNorm and NlpCreateHeads (64 cores = tile rows), conv2d at 80 cores |
+| < 64 | 218 | 4.3 ms (7%). Decoder LayerNorm, create/concat heads (32 cores = 1024/32 tile rows) |
+
+  About half the device time runs on ops that do not use the full grid, which is why the 11→12 column change gave so little. Getting real value from 12x10 means sharded or explicit program configs sized to a 12-column grid (for example M=2048 → 64 tile rows, and N splits that divide 12).
+
+## What is fused / traced already; op count
+
+- Whole graph traced: one `execute_trace` per pair. Persistent input buffer; outputs live in trace-owned buffers.
+- **1252 device programs per forward**, by phase: patch-embed 5 (0.78 ms), encoder 289 (25.9 ms), decoder 534 (16.4 ms), DPT x2 424 (17.9 ms). Python-level ttnn calls are about 946; conv2d expands into halo, move, I2S/S2I and DRAM-slice ops.
+- Already fused:
+  - RoPE: 1 `rotary_embedding_llama` per q/k, with the channel permutation folded into the weights.
+  - The 5 residual linears per block run as `dit_minimal_matmul_addcmul_fused` (bias + residual add fused).
+  - SDPA flash attention.
+  - Fused QKV / fused cross-attention KV plus `split_query_key_value_and_split_heads`.
+  - `concatenate_heads`.
+  - Relu fused into the conv1 / head2 convs.
+  - TILE reshape of the DPT taps.
+  - The encoder runs both views at B=2.
+- **Not fused:** GELU. `ttnn.linear(..., activation="gelu")` with bias and an auto program config produces `fused_activation=nullopt` plus a **separate GELU UnaryDeviceOperation**: 48 ops, 4.97 ms (encoder 24 x 149 µs, decoder 24 x 57 µs). Also unfused:
+  - LayerNorm, separate from the following matmul.
+  - Create-heads separate from QKV.
+  - Pre-activation relu in ResConvUnit: 14 ops, 0.25 ms.
+  - Residual adds in the DPT (BinaryNg: 20 ops, 0.5 ms).
+  - Layout churn around the upsamples (tilize, untilize, S2I, I2S: about 4 ms in total).
+- Weights: all bf16 (bf8 was rejected by the authors for accuracy), matmul HiFi2, conv HiFi4 + fp32 accumulation, LN HiFi4.
+
+## Biggest device-time sinks (kernel time per forward, 12x10)
+
+| Op | Count | ms | % |
+|---|---|---|---|
+| `ttnn.linear` (MatmulDeviceOperation: qkv, fc1, cq, ckv, decoder_embed, patch, DPT 1x1) | 165 | 13.44 | 22.0 |
+| MinimalMatmul (dit fused residual: proj, fc2, cproj) | 120 | 7.76 | 12.7 |
+| Conv2d (DPT) | 60 | 6.96 | 11.4 |
+| SDPA | 72 | 5.75 | 9.4 |
+| Unary (48 GELU + 14 relu) | 62 | 5.21 | 8.5 |
+| RotaryEmbeddingLlama | 144 | 4.07 | 6.7 |
+| LayerNorm | 147 | 3.83 | 6.3 |
+| NlpCreateHeads | 72 | 3.08 | 5.1 |
+| Upsample (bilinear) | 10 | 1.87 | 3.1 |
+| Tilize | 11 | 1.80 | 3.0 |
+| Sharded→Interleaved / Interleaved→Sharded / Halo / Move / slices (conv plumbing) | ~270 | ~4.6 | 7.5 |
+
+Top individual items:
+- encoder fc1 matmul: 24 x 172 µs = 4.1 ms, 96 cores
+- encoder fc2 dit: 24 x 153 µs = 3.7 ms
+- encoder GELU: 3.6 ms
+- encoder qkv: 24 x 148 µs = 3.5 ms
+- encoder SDPA: 24 x 139 µs = 3.3 ms
+- DPT refinenet-1 resconv 3x3 256→256 at 128x128: 8 x 355 µs = 2.8 ms, 80 cores
+- decoder SDPA: 48 x 51 µs = 2.4 ms
+- RoPE: 4.1 ms. The encoder RoPE costs 44 µs per call, about a third of an SDPA.
+- head tilize `[262144,128]`: 2 x 628 µs
+- head upsample 256→512: 2 x 588 µs
+- final 128→4 1x1 matmul: 2 x 288 µs, a wasted N=32 pad
+
+For scale: the encoder qkv matmul does 12.9 GFLOP in 148 µs, about 87 TFLOP/s. That is well below Blackhole bf16 peak, so matmul tuning has headroom.
+
+## Ranked optimization opportunities
+
+Estimates are for the 12x10 grid. "e2e" is the model call including transfers.
+
+1. **Compact output readback** (−20 ms e2e, about 23%; device +≤0.5 ms). Inside the trace:
+   - Drop the 32-channel TILE padding (untilize to RM, or a wide-page reshape such as `[512, 2048]`).
+   - Optionally pack both heads into one buffer and read them with a single `to_torch`.
+   - Also stop running the 128→4 head conv as an N=32-padded TILE matmul.
+
+   This is the single biggest win on Galaxy's x1 PCIe, and it helps p150 as well. It does not change numerics.
+2. **Fuse GELU into fc1** (−4.5 to −5 ms device). Give an explicit `MatmulMultiCoreReuseMultiCastProgramConfig` with `fused_activation=(GELU, approx=False)`, or a `minimal_matmul` that fuses GELU, and check accuracy. Done together with item 3.
+3. **Explicit 12x10 program configs for the qkv / fc1 / ckv / cq / embed linears** (−4 to −6 ms). These currently run the auto MultiCast config on 96 cores with `in0_block_w=1`. Options: `minimal_matmul` / explicit block sizes for M=64 (encoder) and M=32 (decoder) tile rows on 12 columns, or 2D-sharded activations. A larger `in0_block_w` alone should help a lot. The authors' `MAST3R_FUSED_MM=minimal` was slower on p150 at (8,4,4); retune it for 12x10.
+4. **Sharded LayerNorm + create-heads / concat-heads on the full grid** (−3 to −4 ms). LN currently uses 64/32 cores at 42/18 µs, create-heads 59/35 µs, concat 25/13 µs. Width/block-sharded LN fused with the following matmul input (or `ttnn.rms/layer_norm` with a sharded program config) and `nlp_create_qkv_heads` from sharded input should be about 3x faster. Longer term: an L1-resident block pipeline (activations sharded in L1 across the whole block instead of DRAM-interleaved round trips), which is the "megakernel-lite" step.
+5. **DPT head restructuring** (−5 to −8 ms of the 17.9 ms DPT):
+   - Refinenet-1 3x3 convs at 128x128x256 run at 80 cores (2.8 ms). Force a 120-core height-sharded config and enable `enable_act_double_buffer` / `enable_weights_double_buffer`.
+   - Remove the tilize→upsample→S2I/I2S churn: about 4 ms of tilize, untilize, sharded↔interleaved, halo and slice ops.
+   - The final 512x512 head2 conv is DRAM-sliced into several conv passes. Doing the upsample+conv on an L1-sharded layout, or fusing bilinear ×2 and a 3x3 conv as transposed conv/sub-pixel, avoids the slicing.
+   - Run the two heads' independent work interleaved or batched where the weights allow; they cannot share weights, but layout plumbing can overlap.
+6. **RoPE** (−2 to −3 ms). 144 calls take 4.1 ms. Options:
+   - Fold RoPE into create-heads, using the `nlp_create_qkv_heads` + rotary fused variant (`rotary_embedding_llama_fused_qk` takes q and k in one call; check prefill support).
+   - At minimum, apply q and k in one op: 72 instead of 144 programs.
+7. **SDPA tuning** (−1 to −2 ms). Encoder 139 µs, decoder 51 µs. Re-sweep q/k chunk sizes for 12x10 (for example 256/256 or 128/512) and try `exp_approx_mode`, with an accuracy check.
+8. **bf8 weights / LoFi for the MLP only** (−3 to −5 ms, accuracy risk). The authors saw head1 pts3d PCC drop to 0.98 when all linears were bf8. Trying bfp8 only for fc1/fc2 of the encoder (or the decoder), with a real-pair gate, may be acceptable. This is lower priority because of the accuracy constraints.
+9. **ETH dispatch** (about 0 now). Dispatch overhead is only about 1.2 ms per trace, so moving dispatch to ETH matters only for the extra column. On Galaxy, Tensix dispatch already gives 12x10, and the 13x10 grid is forbidden by policy. ETH dispatch costs about 0.9 µs per op (dispatch_s off), so after the op count drops, compare both again. Keep ETH + 12x10 as the p150-representative configuration.
+10. **Host side** (−1 ms). Cast/cat into a pinned preallocated host tensor (0.85 ms). Upload the input as bf16 in a single `[2,3,512,512]` page-friendly layout; H2D is already 1 ms. Trace region: 512 MiB is reserved; measure the real size and shrink it to free DRAM (no speed gain).
+
+Device-only potential: about 62 → 40-45 ms with items 2-7. E2E: about 84 → 45-50 ms on Galaxy, dominated by item 1.
+
+## Artifacts
+
+- `code/bench_breakdown.py`: per-phase timing (host prep / H2D / trace / D2H / e2e) plus PCC, on any dispatch mode.
+- `code/tools_prof/prof_eager.py`: device-profiler driver (eager fused graph, 2nd iteration = steady state).
+- `code/tools_prof/real_pair_acc.py`: real-image accuracy gate.
+- `code/tools_prof/d2h_probe.py`: readback-format micro-benchmark.
+- Profiler CSV: `/home/ttuser/experiments/tt-models/chipstate/chip9/profiler/mast3r_eager/reports/2026_10_02_10_14_18/ops_perf_results_2026_10_02_10_14_18.csv` (2 iterations, 1252 ops each).
+
+## Post-crash re-verification (2026-10-02 11:18 UTC)
+
+After the unexpected host shutdown, chip 9 opened cleanly with no reset needed and the kernel cache was 100% warm (456/456 hits). The baseline was re-run with the later optimization pass disabled (`MAST3R_OPT=none python test_mast3r.py --layer end_to_end --runs 25`, 12x10):
+- PCC: 0.9984, PASS.
+- Best-of-25 latency: 85.43 ms. Before the crash it was 83.64 ms; the difference is host-load noise.
+
+The baseline numbers above still hold. Later commits (`0e84074` compact output, WIP `8f82803`) are optimization work beyond this baseline and are gated by `MAST3R_OPT`.
+
+## Post-crash #2 re-verification (2026-10-02 16:23 UTC)
+
+After the second host crash, chip 9 opened cleanly with no reset, and the kernel cache was 456/456 warm. I re-ran the baseline (`MAST3R_OPT=none`, 12x10) and it reproduces:
+- `test_mast3r.py --layer end_to_end --runs 25`: PCC 0.9984 (head1 0.9983 / head2 0.9982), PASS, best-of-25 85.53 ms.
+- `bench_breakdown.py --mode eth12`, median:
+  - host_prep 0.81 ms, H2D 1.06 ms, trace 62.20 ms, D2H 23.52 ms, e2e 84.85 ms.
+  - PCC head1 0.99836 / head2 0.99895.
+- `real_pair_acc.py`: pts3d PCC 0.98988 / 0.99076, conf 0.99237 / 0.99439. Identical to the first baseline.
+
+I also checked the current default path (`MAST3R_OPT` unset, giving out, mm, dpt, l1, dptf). It includes the WIP `dptf` change from snapshot `13363fc`, which was unverified at the time of the crash. It now passes every proposed gate:
+- e2e PCC 0.9988, PASS, best-of-25 53.66 ms.
+- bench eth12, median: trace 49.30 ms, D2H 2.47 ms, e2e 53.78 ms.
+- Real pair: pts3d 0.99032 / 0.99081, conf 0.99232 / 0.99466. Against the half-pixel-upsample reference, pts3d is 0.99988 / 0.99997.
+
+So `dptf` is verified: the trace drops from 50.3 to 49.3 ms with no loss of accuracy. Compared with the baseline, the optimization knobs so far bring the device trace from 62.2 to 49.3 ms (-21%) and e2e from 84.9 to 53.8 ms (-37%), on the same chip, the same 12x10 grid and the same timing method.
+
+## Post-crash #3 re-verification (2026-10-03 00:46 UTC)
+
+After the reboot, chip 9 (BDF 0000:42:00.0, /dev/tenstorrent/25) opened cleanly with no reset, and the JIT cache was 456/456 warm. Results:
+
+- Baseline (`MAST3R_OPT=none`, 12x10):
+  - `test_mast3r.py --layer end_to_end --runs 25`: PCC 0.9984, PASS, best-of-25 **84.49 ms**.
+  - `bench_breakdown.py --mode eth12`, median (min):
+    - host_prep 0.82 ms, H2D 1.05 ms
+    - trace **62.21 ms** (62.14)
+    - D2H 23.57 ms (20.43)
+    - e2e 92.37 ms (84.76). The e2e median is noisier than before because of D2H spikes up to 37 ms; the device trace is unchanged.
+    - PCC head1 0.99836 / head2 0.99895.
+  - `real_pair_acc.py`: pts3d 0.98988 / 0.99076, conf 0.99237 / 0.99439. This is bit-identical to the earlier runs.
+- Current default path (all OPT knobs): bench eth12 gives PCC 0.99852 / 0.99896, trace 49.28 ms, D2H 3.30 ms, e2e 54.44 ms (53.86).
+
+The baseline numbers and the ranked opportunity list above still hold. The WIP snapshot `8db75c8` (`tools_prof/fid_probe.py`, `gelu_probe.py`) contains only probe scripts and does not affect the model path.

README.md

@@ -23,7 +23,7 @@ base_model:
 
 # mast3r-p150
 
-DUSt3R two-view 3D reconstruction (naver/DUSt3R_ViTLarge_BaseDecoder_512_dpt: ViT-L/16 encoder, dual-branch decoder, two DPT heads; the backbone of MASt3R) running entirely on one Tenstorrent Blackhole p150a via tt-nn: an image pair in, two dense 512×512 pointmaps with confidence (plus optional PairViewer pose) out.
+NAVER DUSt3R ViT-L two-view 3D reconstruction (the MASt3R backbone) port on one Tenstorrent Blackhole p150a.
 Weights: [naver/DUSt3R_ViTLarge_BaseDecoder_512_dpt](https://huggingface.co/naver/DUSt3R_ViTLarge_BaseDecoder_512_dpt) · Paper: [arXiv:2312.14132](https://arxiv.org/abs/2312.14132) · Upstream code: [naver/mast3r](https://github.com/naver/mast3r) · Port: [changh95/tt-mast3r](https://github.com/changh95/tt-mast3r)
 
 Runs on **p150** (mesh `P150`).
@@ -37,7 +37,7 @@ tt-model pull  changh95/mast3r-p150 --with-weights
 tt-model serve changh95/mast3r-p150
 ```
 
-- Weights [`naver/DUSt3R_ViTLarge_BaseDecoder_512_dpt`](https://huggingface.co/naver/DUSt3R_ViTLarge_BaseDecoder_512_dpt) at `61c57447d7b0` go to your HF cache; the image does not contain them.
+- Weights [`naver/DUSt3R_ViTLarge_BaseDecoder_512_dpt`](https://huggingface.co/naver/DUSt3R_ViTLarge_BaseDecoder_512_dpt) go to your HF cache; the image does not contain them.
 - Serves on port 20000 (or the next free port); ready when the log says `Application startup complete`.
 
 ### Run with tt-cli
@@ -68,6 +68,7 @@ tt model stop changh95/mast3r-p150
 
 - `npz_b64` is a base64 `.npz` holding `pts3d1`/`pts3d2` float32 (512,512,3) and `conf1`/`conf2` float16 (512,512); both pointmaps are in the camera-1 frame, depth is `pts3d[..., 2]` in DUSt3R's own scale (not metres), `conf = 1 + exp(c) >= 1`. `output_format: "png"` returns per-view 16-bit depth + 8-bit confidence PNGs instead.
 - `preprocess` maps original pixels to the 512 grid, `(u, v) = ((x + offx) * scale, (y + offy) * scale)` (`mode` = `pad`, the default gray pad-to-square, or `crop` = centre-crop to square via the `MAST3R_PREPROC` env; offsets are negative for `crop`). With `return_pose: true`, `pose` holds `R`, `t` (view 2 in the view-1 frame, `X_cam2 = R @ X_cam1 + t`), `focal_1`, `focal_2` (512-px units) and `used_known_intrinsics`; `null` when PnP fails.
+- The `timing_ms` values above come from the container image. With the optimized `code/`, `forward` is about 28 ms (see below). The host `encode` code did not change.
 
 ### Demo
 
@@ -75,28 +76,49 @@ tt model stop changh95/mast3r-p150
 
 Kitchen frames 00 / 03 (VGGT example scene) → both predicted pointmaps in the camera-1 frame, coloured by the source pixels (served `npz` output, `code/make_demo.py`).
 
-### Accuracy and speed
+### Demo & Performances
 
-| Metric | Value |
+Warm, batch 1, one 512×512 pair, fused traced graph (`MAST3R_OPT=all`), `bench_breakdown.py` with 30 iterations per run (median / min). The served path config is worker dispatch, 2 CQ and uint8 input (the `tt-model.yaml` serve env).
+
+| Metric | Performance |
 |---|---:|
-| Torch fp32 reference vs upstream DUSt3R (`AsymmetricCroCo3DStereo`, same inputs), pts3d / conf PCC | 1.00000 / 1.00000 (max abs diff 7e-7) -- the reference is the upstream network since the 2026-09-14 decoder-tap fix |
-| End-to-end pointmap PCC vs fp32 torch reference (synthetic pair, `test_mast3r.py`) | 0.9987 (legacy graph `TT_FUSED=0`: 0.9989) |
-| Activated pts3d PCC vs reference, 4 real stereo pairs (CO3Dv2 apple, kitchen 00/03 and 00/08, LLFF fern 000/002), head 1 / head 2 range | 0.987-0.9995 / 0.991-0.9996 (2026-09-14, `logs/pointmap/mast3r/after_*_metrics.md`) |
-| Two-head coherence: median NN distance between the view-1 and view-2 pointmaps / scene scale, device (torch reference) | apple 0.012 (0.008) · kitchen 0.004 (0.003) · kitchen08 0.004 (0.0035); before the fix 0.39-0.65 (two separate sheets) |
-| PairViewer rotation vs the torch reference on the same pairs (est. focal) | 0.07-0.31° |
-| Inference, served over HTTP (warm, 30 requests, one 512×512 pair, npz; median / min / max; 2026-09-13, graph +2 `layer_norm` since) | 73.6 / 73.2 / 79.9 ms device forward · 241 / 240 / 280 ms end-to-end (160 ms of it is the host npz encode; ~4.1 pairs/s); 1.4 s with `return_pose` |
-| `test_mast3r.py --layer end_to_end`, best-of-25 (2026-09-14) | 73.1 ms per pair (legacy graph: 238.9 ms) |
-| Same forward on an RTX 5090 (same host, port's torch reference, eager PyTorch, batch 1, incl. H2D/D2H; bf16 / fp16 autocast) | 63.8 / 56.6 ms → GPU 1.2× / 1.3× faster than the p150a's 73.9 ms (both views); fp32-strict GPU 103.4 ms (p150a 1.4× faster); fp16 weights resident 39.5 ms (1.9×); best `torch.compile` 21.1 ms (3.5×) |
+| Model call `model(img1, img2)` (H2D + trace + D2H of both maps), served path config | **28.1 ms median** (27.99–28.23, 2 runs) · 27.55 ms min |
+| Device trace (ViT-L encoder ×2, decoder, two DPT heads), served path config | **25.2–25.4 ms median** · 24.52 ms min |
+| Model call, ETH dispatch, 1 CQ, float input | 29.5–30.0 ms median · 29.27 ms min |
+| Device trace, ETH dispatch, 1 CQ | **24.7–25.0 ms median** · 24.41 ms min |
+| Device span (profiler cycles at 1.35 GHz, ETH dispatch, 1 CQ, 3 replays) | 23.47–23.49 ms · 688 programs |
+| Host prep · H2D, served path config (uint8 pixels) | 0.52 ms · 0.60 ms |
+| D2H of both maps, ETH dispatch, 1 CQ | 2.95 ms |
+| `return_pose` forward on the real pair, served path config: symmetric graph (`sym`) · two-pass | **38.0 ms median** · 55.8 ms median |
+
+The measurement hardware is a Blackhole chip with a 12×10 compute grid and dispatch on ETH cores. The "served path config" rows use the same 12×10 compute grid with worker dispatch and 2 CQ. Accuracy against the torch fp32 reference on the real kitchen pair (served path config): pts3d PCC 0.99021 / 0.99086 (gate 0.989), conf PCC 0.99232 / 0.99497 (gate 0.99); synthetic pair `test_mast3r.py` end_to_end PCC 0.9985. Details: [`VERIFICATION_2026-10-03.md`](VERIFICATION_2026-10-03.md).
+
+RTX 5090 reference measurements (2026-09-14) are unchanged. They use the port's torch reference in PyTorch 2.11, batch 1, on the kitchen demo pair. The "incl. H2D/D2H" column compares with our model call (28.1 ms). The "forward only" column compares with our device trace (25.3 ms). Full table: [`GPU_COMPARISON.md`](GPU_COMPARISON.md).
+
+| RTX 5090 precision | GPU incl. H2D/D2H | vs current build (28.1 ms) | GPU forward only vs ours (25.3 ms) |
+|---|---:|---:|---:|
+| fp32 strict, eager | 103.4 ms | **Blackhole 3.68× faster** | 102.3 ms: Blackhole 4.04× faster |
+| tf32, eager | 70.8 ms | **Blackhole 2.52× faster** | 70.0 ms: Blackhole 2.77× faster |
+| bf16 autocast, eager | 63.8 ms | **Blackhole 2.27× faster** | 62.6 ms: Blackhole 2.47× faster |
+| fp16 autocast, eager | 56.6 ms | **Blackhole 2.01× faster** | 55.5 ms: Blackhole 2.19× faster |
+| bf16 weights, eager | 42.3 ms | **Blackhole 1.51× faster** | 41.3 ms: Blackhole 1.63× faster |
+| bf16 weights + SDPA, eager | 35.3 ms | **Blackhole 1.26× faster** | 34.3 ms: Blackhole 1.36× faster |
+| bf16 weights + `torch.compile` (CUDA graphs) | 21.1 ms | GPU 1.33× faster | 20.1 ms: GPU 1.26× faster |
+
+The Blackhole lead comes from one metal trace that runs fused, block-sharded bf16 kernels on all 120 compute cores. The eager GPU rows pay for kernel launches and explicit softmax attention. The autocast rows also cast the fp32 weights again on each forward. `torch.compile` with CUDA graphs removes these costs, and then the GPU is about 1.3× faster. The served `/predict` total stays host-bound on both sides, because the npz encode takes about 160 ms.
 
 ### Caveats
 
-- Fixed 512×512 input, exactly one image pair per request (batch 1); non-square images are gray-padded to square (`MAST3R_PREPROC=crop` centre-crops instead). The estimated focal ran ~5-8 % high on a CO3Dv2 pair with known K — pass `intrinsics` when you know them.
-- Point-map quality fix 2026-09-14: earlier revisions tapped the wrong decoder blocks for the DPT heads (two non-overlapping pointmaps, 0.98 PCC vs upstream); numbers published before that date (CO3Dv2 12-pair PCC / AUC@30, ETH3D) describe the old network and are not repeated here.
-- DUSt3R backbone only: no MASt3R matcher / descriptor head and no N-view global alignment; pose is single-pair PairViewer (estimated-focal PnP-RANSAC on one pair, no CO3Dv2 AUC re-measured on this host).
-- bf16 on device: activated pointmaps 0.987-0.9996 PCC vs fp32 per head on real pairs, confidence 0.99 PCC (raw confidence channel 0.85-0.96) — loosen thresholds tuned on the reference. Dense outputs come base64-encoded (`.npz`, or 16-bit/8-bit PNG).
+- Does not scale to multiple p150a in a mesh configuration. The current build uses a 12x10 compute grid of Tensix cores. To get this grid, the dispatch functions move from 10 Tensix cores to ETH cores (`patches/tt-metal-eth-dispatch.patch`). Thus, this build assumes that you do not need chip-to-chip ethernet communication.
+- The serve config sets `MAST3R_CQS=2`. Two CQs need worker (Tensix) dispatch. With ETH dispatch, set `MAST3R_CQS=1`. The outputs are bit-identical.
+- Fixed 512×512 input, exactly one image pair per request (batch 1). Non-square images are gray-padded to square (`MAST3R_PREPROC=crop` centre-crops instead). The estimated focal can be 5-8 % high, so pass `intrinsics` when you know them.
+- Numbers published before 2026-09-14 describe an older network with wrong decoder taps. This card does not repeat them.
+- DUSt3R backbone only: no MASt3R matcher / descriptor head and no N-view global alignment. The pose is single-pair PairViewer (estimated-focal PnP-RANSAC). With `return_pose`, one symmetric graph (`sym`) computes the pose maps. Its outputs are bit-identical to the two-pass path.
+- bf16 on device: on the kitchen pair, pts3d PCC vs fp32 is 0.990 / 0.991 per head and conf PCC is 0.992 / 0.995. Loosen thresholds that you tuned on the reference. Some optimizations change numerics (SDPA chunks, GELU polynomial, HiFi3 head convs, upsample kernel). [`OPT_REPORT.md`](OPT_REPORT.md) lists each change.
+- Dense outputs come base64-encoded (`.npz`, or 16-bit/8-bit PNG).
 - Weights are NAVER's DUSt3R checkpoint under CC-BY-NC-SA-4.0: non-commercial use only, share-alike.
-- Not an OpenAI-compatible API; `GET /v1/models` is a stub so the tt-model ready card does not 404. Validated on tt-metal `v0.78.0-dev20260820` (main `8b98410e730`), single p150a only.
-- GPU comparison: GPU fp32-strict eager is 1.4× slower than the p150a; GPU bf16/fp16 autocast 1.2–1.3× faster, fp16 weights resident 1.9×, compiled 3.5×; the served `npz` request is host-bound on both sides (~160 ms compression). RTX 5090 rows (2026-09-14): same host, the port's own torch reference (same weights) run eagerly in PyTorch 2.11 cu128 with fp32 weights + autocast unless stated, no TensorRT; medians of 50 iterations after warm-up, H2D/D2H included; the p150a rows are the served bf16 fused path incl. upload/readback. p150a power was not measured, so no efficiency comparison is made. Full table: [`GPU_COMPARISON.md`](GPU_COMPARISON.md).
+- Not an OpenAI-compatible API; `GET /v1/models` is a stub so the tt-model ready card does not 404.
+- p150a power was not measured, so no efficiency comparison is made.
 
 ### Licensing
 
@@ -105,11 +127,11 @@ Kitchen frames 00 / 03 (VGGT example scene) → both predicted pointmaps in the
 
 ## Provenance
 
-The exact sources the image was built from — `code/` in this repo is byte-identical to the model code inside the image:
+These are the exact sources the container image was built from. **`code/` has since been updated (2026-10-03 optimized build; see [`OPT_REPORT.md`](OPT_REPORT.md)) and is newer than the image.** `tt-model serve` runs the image's code until the image is rebuilt. `tt-model.yaml` and `SERVING.md` still describe the image:
 
 | component | built from |
 | --- | --- |
 | tt-metal | [`8b98410e730bb504fea43a88609756e34821d91d`](https://github.com/tenstorrent/tt-metal/commit/8b98410e730bb504fea43a88609756e34821d91d) |
-| `code/` digest | `b45444934e1d76bf` (sha256, first 16 hex digits) |
+| `code/` digest (image) | `b45444934e1d76bf` (sha256, first 16 hex digits; the current `code/` differs) |
 | built | 2026-09-14T06:35:13+00:00 by tt-model 0.1.0 |
 

VERIFICATION_2026-10-03.md

@@ -0,0 +1,85 @@
+# mast3r-p150: independent verification, 2026-10-03
+
+## Verdict
+
+**PASS.** An independent verifier measured commit `ad39150` again. Every claimed speed number matches within 1-2 %. Every accuracy number matches the claim to the printed digit. All accuracy gates pass. The verifier found no cached outputs, no skipped compute and no loosened gates in the diff `5b918db..ad39150`. No chip faulted.
+
+## Setup
+
+| item | value |
+|---|---|
+| Hardware | a Blackhole chip with a 12×10 compute grid (x1 PCIe link) |
+| Dispatch configurations | ETH dispatch with 1 CQ ("eth12"), and worker (Tensix) dispatch with 2 CQ (the served path config). Both use the 12×10 compute grid. |
+| tt-metal | `8b98410e730` with `patches/tt-metal-eth-dispatch.patch` |
+| Verified commit | `ad39150` (the `code/` in this repo) |
+| Baseline commit | `81ff339`. Its model code is identical to the previous release `20bdff0`. It is the earliest commit with `bench_breakdown.py`. |
+| Workload | one 512×512 image pair, batch 1, warm, traced fused graph (`MAST3R_OPT=all`). The real-pair checks use the kitchen pair `media/source_1.png` + `media/source_2.png`. |
+| Method | The verifier checked out baseline and final in separate git worktrees. The verifier ran the same `bench_breakdown.py` method on both. One workload ran at a time. |
+| Scripts | `code/bench_breakdown.py`, `code/tools_prof/r3_validate.sh`, `code/tools_prof/real_pair_acc.py`, `code/tools_prof/sym_check.py`, `code/test_mast3r.py`, `code/tools_prof/trace_span.py` (device span under tracy) |
+
+## Measured baseline vs final
+
+Median / min in ms. Each bench run has 30 iterations. Where the verifier ran a configuration twice, the table shows both runs.
+
+| measurement | baseline `81ff339` | final `ad39150` | change |
+|---|---:|---:|---:|
+| device trace, ETH dispatch, 1 CQ, float | 62.21 / 62.09 and 62.23 / 62.17 | **24.72 / 24.47** and **25.03 / 24.41** | 2.5× faster |
+| model call (H2D + trace + D2H), ETH dispatch, 1 CQ, float | 97.02 / 96.55 and 97.19 / 91.57 | **29.99 / 29.84** and **29.54 / 29.27** | about 3.3× faster |
+| D2H of both maps, ETH dispatch, 1 CQ | 29.5 | **2.95** | |
+| device trace, worker dispatch | 61.84 / 61.77 (1 CQ) | 24.95 / 24.48 (2 CQ, float) | |
+| model call, worker dispatch | 97.77 / 97.49 (1 CQ) | 28.42 / 27.83 (2 CQ, float) | |
+| device trace, served path config (worker, 2 CQ, uint8 input) | | **25.39 / 24.82** and **25.23 / 24.52** | |
+| model call, served path config | | **28.23 / 27.55** and **27.99 / 27.68** | about 3.5× faster than the baseline (about 98 ms) |
+| host prep / H2D, served path config | | 0.52 / 0.60 | |
+| `test_mast3r.py --layer end_to_end` | | 27.71 | |
+| real pair, served path config: single pair / `sym` pose / two-pass pose | | 27.90 / 38.00 / 55.82 (min 27.43 / 37.84 / 54.55) | |
+| device span (tracy, ETH dispatch, 1 CQ, 3 replays) | | **23.484 / 23.473 / 23.485** (688 programs, kernel sum 22.85, FW gaps 0.244) | |
+
+`MAST3R_OPT=none` at the final commit gives a device trace of 62.19 / 62.13 ms and a model call of 91.76 / 86.78 ms. This confirms the baseline.
+
+The claimed device span is 23.491 / 23.479 / 23.490 ms. The verified span matches it.
+
+## Accuracy gate results
+
+All gates pass on `ad39150`. The gate code did not change.
+
+| metric | gate | final `ad39150` | baseline `81ff339` |
+|---|---:|---:|---:|
+| synthetic randn pair, eth12 float, PCC head 1 / head 2 / all | ≥ 0.998 | **0.99848 / 0.99855 / 0.99875** | 0.99836 / 0.99895 / 0.99909 |
+| `test_mast3r.py --layer end_to_end` PCC | ≥ 0.998 | **0.9985** (PASS) | |
+| real pair, float: pts3d PCC head 1 / head 2 | ≥ 0.989 | **0.99035 / 0.99081** | 0.98985 / 0.99078 |
+| real pair, float: conf PCC head 1 / head 2 | ≥ 0.99 | **0.99262 / 0.99511** | 0.99237 / 0.99438 |
+| real pair, served uint8: pts3d PCC head 1 / head 2 | ≥ 0.989 | **0.99021 / 0.99086** | |
+| real pair, served uint8: conf PCC head 1 / head 2 | ≥ 0.99 | **0.99232 / 0.99497** | |
+| real pair, raw PCC vs the half-pixel reference, float / uint8 | | 0.999882 / 0.999814 · 0.999886 / 0.999782 | |
+| synthetic uint8 pair (not gated) | | 0.99733 / 0.99941 | |
+
+Other checks:
+
+- The e2e outputs are equal to the synchronized readback and to the first call.
+- `sym_check`: `out_ii`, `out_ji` and `out_jj` of the `sym` graph are bit-identical to the two-pass path.
+- On the real pair, the final build is equal to or better than the baseline on every gated metric.
+
+## Disclosed numerics changes
+
+The optimization rounds changed some numerics. [`OPT_REPORT.md`](OPT_REPORT.md) lists each change with its measured effect. The main changes are these:
+
+- `sdpa2`: larger SDPA chunks (encoder q160 / k512, decoder q224 / k512).
+- `mc2d` and `mm32`: the plain and residual linears use 2-D multicast matmuls with fp32 dest accumulation.
+- `gpoly`: the fc1 GELU uses a model-local erf-GELU polynomial (max abs error 6.2e-5 vs exact erf GELU).
+- `hf3`: HiFi3 instead of HiFi4 for the two output-resolution head convs.
+- `tups`: a model-local bilinear upsample kernel. It is closer to the exact bilinear result than the stock upsample.
+- `hostcol`: the served path uploads uint8 pixels. The device applies the normalization.
+- `resmm` and `resmmd`: the residual adds move into the matmul epilogues.
+
+The synthetic head 2 PCC is 0.99855, below the baseline 0.99895. It stays above the 0.998 gate. The three steps of the last round (`esplit`, `demb`, `hrqk`) are bit-identical. `gelua` and `lofi` stay off because of accuracy.
+
+## Review findings
+
+- The verifier reviewed the diff `5b918db..ad39150`. The gate and bench scripts did not change (`real_pair_acc`, `test_mast3r`, `bench_breakdown`, `sym_check`, `r3_validate`).
+- `hrqk` runs the same per-tile op sequence for q and k. `esplit` changes only the output base address and start tile of each writer. `demb` uses the existing dual-matmul path.
+- New diagnostic defines (`D_NOCOMP`, `D_NOREAD`, `D_NOWRITE`) and a reader override switch on only with `MAST3R_HROPE_DIAG` / `MAST3R_LN_READER*`. These variables are empty by default. `tt-model.yaml` and the server do not set them.
+- Minor: the verifier did not measure the round-start device span (23.531 ms) again. The last round gained only 0.044 ms (0.19 %). Host wall time cannot show this gain. The final absolute span reproduces.
+- Minor: `trace_span.py` converts cycles to time at the nominal 1.35 GHz. Under load, the clock drops to about 1.26-1.30 GHz. Thus, the true device time is about 4 % above the printed span.
+- Context: the RTX 5090 with `torch.compile` and CUDA graphs (21.1 ms with H2D/D2H, 20.1 ms forward only) is still faster than this build (about 28 ms model call, about 25 ms trace). This build is faster than every eager GPU variant, including bf16 weights with SDPA (35.3 ms). See [`GPU_COMPARISON.md`](GPU_COMPARISON.md).
+- The verifier did not run the server smoke test. The developer reports PASS 4 / 4 with the `tt-model.yaml` serve env (worker dispatch, 2 CQ).

code/bench_breakdown.py

@@ -0,0 +1,147 @@
+#!/usr/bin/env python3
+"""Per-phase timing of the traced fused DUSt3R forward (baseline profiling helper).
+
+Phases per call (same work as TtDust3r.__call__ / test_mast3r.py end_to_end):
+  host_prep  : torch.cat + bf16 cast + host ttnn tensor (ROW_MAJOR)
+  h2d        : copy_host_to_device_tensor into the persistent trace input (blocking)
+  trace      : execute_trace + synchronize_device (pure device time of the whole graph)
+  d2h        : two to_torch readbacks of the [1,1,H*W,4] head outputs (+ host untilize/reshape)
+  e2e        : the public model call (host_prep + h2d + trace + d2h, non-blocking trace)
+
+Usage (after chipenv.sh):
+  python bench_breakdown.py --mode worker11   # p150-equivalent: Tensix dispatch, grid capped 11x10
+  python bench_breakdown.py --mode worker     # Galaxy default Tensix dispatch (12x10, Galaxy-only A/B)
+  python bench_breakdown.py --mode eth12      # ETH dispatch, capped 12x10
+  MAST3R_CQS=2 python bench_breakdown.py --mode worker   # Tensix dispatch 12x10, 2 CQs (split trace, overlapped head-1 read)
+  add --u8 for uint8 pixel views (HWC-backed, as the server's PIL preprocess produces them)
+
+With MAST3R_CQS=2 "trace" is both traces back to back (A + B, synchronized), and "d2h" is the two reads
+after that sync (no overlap); only "e2e" shows the overlap.
+"""
+from __future__ import annotations
+
+import argparse
+import os
+import statistics
+import sys
+import time
+
+_CODE_ROOT = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, _CODE_ROOT)
+sys.path.insert(0, "/home/ttuser/experiments/tt-models/tools")
+
+ap = argparse.ArgumentParser()
+ap.add_argument("--mode", default="worker11", choices=["worker11", "worker", "eth12"])
+ap.add_argument("--iters", type=int, default=30)
+ap.add_argument("--no-ref", action="store_true")
+ap.add_argument("--u8", action="store_true", help="uint8 pixel views (exact reference input (v/255-0.5)/0.5)")
+ap.add_argument("--dump", default="", help="save the first-call outputs (torch.save) for bit-exact A/B")
+ap.add_argument("--cmp", default="", help="compare the first-call outputs with a --dump file")
+ap.add_argument("--u8-float", action="store_true", help="the --u8 images, fed to the model as float (A/B)")
+args = ap.parse_args()
+
+# grid cap must be set before ttnn creates its runtime options
+if args.mode == "worker11":
+    os.environ["TT_METAL_CORE_GRID_OVERRIDE_TODEPRECATE"] = "10,9"
+
+import torch  # noqa: E402
+import ttnn  # noqa: E402
+from models.demos.mast3r.reference.torch_dust3r import load_checkpoint, load_dust3r  # noqa: E402
+from models.demos.mast3r.tt.ttnn_dust3r import fused_config, get_model, release_device_caches  # noqa: E402
+
+
+def pcc(a, b):
+    a = a.float().flatten(); b = b.float().flatten()
+    a = a - a.mean(); b = b - b.mean()
+    return float((a @ b) / (a.norm() * b.norm()))
+
+
+cfg = fused_config()
+kw = cfg.open_device_kwargs(l1_small_size=32 * 1024)
+print(f"# MAST3R_CQS={cfg.num_cqs} opts={','.join(cfg.opts)}")
+if args.mode == "eth12":
+    from open_device_12x10 import open_device
+    device = open_device(grid="12x10", **kw)
+else:
+    device = ttnn.open_device(device_id=0, **kw)
+g = device.compute_with_storage_grid_size()
+print(f"# mode={args.mode} grid={g.x}x{g.y}")
+
+torch.manual_seed(0)
+state = load_checkpoint()
+img1 = torch.randn(1, 3, 512, 512)
+img2 = torch.randn(1, 3, 512, 512)
+rin1, rin2 = img1, img2
+if args.u8 or args.u8_float:
+    # smooth-ish random uint8 images (HWC arrays viewed as (1, 3, H, W), like the server's preprocess)
+    def u8img():
+        lo = torch.rand(1, 3, 32, 32) * 255
+        im = torch.nn.functional.interpolate(lo, size=(512, 512), mode="bicubic", align_corners=False)
+        im = (im + torch.randn_like(im) * 20).clamp(0, 255).round().to(torch.uint8)
+        return im[0].permute(1, 2, 0).contiguous().permute(2, 0, 1).unsqueeze(0)   # HWC storage
+    img1, img2 = u8img(), u8img()
+    rin1, rin2 = ((img1.float() / 255 - 0.5) / 0.5), ((img2.float() / 255 - 0.5) / 0.5)
+    if args.u8_float:
+        img1, img2 = rin1, rin2
+ref = None
+if not args.no_ref:
+    t0 = time.perf_counter()
+    with torch.no_grad():
+        ref = load_dust3r(state)(rin1, rin2)
+    print(f"# torch fp32 reference: {(time.perf_counter() - t0):.1f} s")
+
+try:
+    model = get_model(state, device)
+    t0 = time.perf_counter()
+    out = model(img1, img2)          # eager warm + capture + first replay
+    print(f"# first call (compile + warm + capture): {(time.perf_counter() - t0):.1f} s")
+    if ref is not None:
+        p1, p2 = pcc(ref[0], out[0]), pcc(ref[1], out[1])
+        pall = pcc(torch.stack(ref), torch.stack(out))
+        print(f"# PCC head1={p1:.5f} head2={p2:.5f} all={pall:.5f}")
+
+    if args.dump:
+        torch.save(out, args.dump)
+    if args.cmp:
+        prev = torch.load(args.cmp)
+        for h in range(2):
+            d = (prev[h] - out[h]).abs()
+            print(f"# vs {args.cmp}: head{h+1} bit-identical={torch.equal(prev[h], out[h])} n_diff={int((d > 0).sum())} "
+                  f"max_abs={float(d.max()):.3g} PCC(prev)={pcc(prev[h], out[h]):.6f}")
+    B, _, H, W = img1.shape
+    ph = {k: [] for k in ("host_prep", "h2d", "trace", "d2h", "e2e")}
+    kind = model.input_kind(img1)
+    kd = model._kinds[kind]
+    for _ in range(args.iters):
+        t0 = time.perf_counter()
+        host_in = model._host_input(img1, img2)
+        t1 = time.perf_counter()
+        ttnn.copy_host_to_device_tensor(host_in, kd["in"], cq_id=0)
+        ttnn.synchronize_device(device)
+        t2 = time.perf_counter()
+        for tid in kd["tids"]:
+            ttnn.execute_trace(device, tid, cq_id=0, blocking=False)
+        ttnn.synchronize_device(device)
+        t3 = time.perf_counter()
+        _ = model.read_outputs(kind, B, H, W)      # synchronized (no events pending)
+        t4 = time.perf_counter()
+        ph["host_prep"].append((t1 - t0) * 1e3); ph["h2d"].append((t2 - t1) * 1e3)
+        ph["trace"].append((t3 - t2) * 1e3); ph["d2h"].append((t4 - t3) * 1e3)
+    for _ in range(args.iters):
+        t0 = time.perf_counter()
+        _ = model(img1, img2)
+        ph["e2e"].append((time.perf_counter() - t0) * 1e3)
+    for k, v in ph.items():
+        print(f"{k:10s} median {statistics.median(v):8.2f} ms   min {min(v):8.2f}   max {max(v):8.2f}")
+    o = kd["outs"][0]
+    print(f"# output tensor: shape={tuple(o.shape)} padded={tuple(o.padded_shape)} layout={o.layout} dtype={o.dtype}")
+    print(f"# input tensor: shape={tuple(kd['in'].shape)} layout={kd['in'].layout} dtype={kd['in'].dtype} traces={len(kd['tids'])}")
+    # bit-exactness of the e2e path (overlapped reads with MAST3R_CQS=2) vs the synchronized reads above
+    a = model(img1, img2)
+    ttnn.synchronize_device(device)
+    b = model.read_outputs(kind, B, H, W)
+    print(f"# e2e outputs == synchronized readback: {all(torch.equal(x, y) for x, y in zip(a, b))}; "
+          f"== first call: {all(torch.equal(x, y) for x, y in zip(a, out))}")
+finally:
+    release_device_caches()
+    ttnn.close_device(device)

code/models/demos/mast3r/postprocess.py

@@ -58,7 +58,7 @@ def preprocess_mode(env: Optional[Mapping[str, str]] = None) -> str:
     return v
 
 
-def preprocess_image(img: Image.Image, size: int = IMG_SIZE, mode: str = DEFAULT_PREPROC):
+def preprocess_image(img: Image.Image, size: int = IMG_SIZE, mode: str = DEFAULT_PREPROC, uint8: bool = False):
     """Square ``size x size`` network input in [-1, 1] plus a preprocess record.
 
     ``mode="pad"`` (default): pad to a square (gray 128) then bicubic-resize -- preserves
@@ -69,6 +69,10 @@ def preprocess_image(img: Image.Image, size: int = IMG_SIZE, mode: str = DEFAULT
     with ``scale = out_size / pad_side``; ``pad_side`` is the side of the square canvas
     (``max(W, H)`` for pad, ``min(W, H)`` for crop) and the offsets are >= 0 for pad,
     <= 0 for crop. :func:`intrinsics_to_canonical` applies the same map to ``K``.
+
+    ``uint8=True``: return the resized pixels themselves as a ``(1, 3, size, size)`` torch.uint8
+    view of the HWC array (no copy); the network input is exactly ``(v/255 - 0.5)/0.5`` of it,
+    which the fused device graph applies itself (``TtDust3r`` uint8 path).
     """
     if mode not in PREPROC_MODES:
         raise ValueError(f"preprocess mode must be one of {PREPROC_MODES}, got {mode!r}")
@@ -85,6 +89,9 @@ def preprocess_image(img: Image.Image, size: int = IMG_SIZE, mode: str = DEFAULT
         canvas = img.crop((x0, y0, x0 + s, y0 + s))
         offx, offy = -x0, -y0
     canvas = canvas.resize((size, size), Image.BICUBIC)
+    rec = {"orig_W": W, "orig_H": H, "offx": offx, "offy": offy, "pad_side": s, "out_size": size, "mode": mode}
+    if uint8:
+        return torch.from_numpy(np.array(canvas, dtype=np.uint8)).permute(2, 0, 1).unsqueeze(0), rec
     arr = np.asarray(canvas, dtype=np.float32) / 255.0
     arr = (arr - 0.5) / 0.5
     tensor = torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0).contiguous()
@@ -92,10 +99,10 @@ def preprocess_image(img: Image.Image, size: int = IMG_SIZE, mode: str = DEFAULT
                     "pad_side": s, "out_size": size, "mode": mode}
 
 
-def load_image_for_dust3r(path, size: int = IMG_SIZE, mode: str = DEFAULT_PREPROC):
+def load_image_for_dust3r(path, size: int = IMG_SIZE, mode: str = DEFAULT_PREPROC, uint8: bool = False):
     """File-path convenience wrapper around :func:`preprocess_image` (eval scripts)."""
     with Image.open(Path(path)) as im:
-        return preprocess_image(im, size, mode)
+        return preprocess_image(im, size, mode, uint8=uint8)
 
 
 def intrinsics_to_canonical(K: np.ndarray, preproc: dict) -> np.ndarray:

code/models/demos/mast3r/tt/fattn.py

@@ -0,0 +1,130 @@
+"""MAST3R_OPT "fattn": model-local SDPA program (``ttnn.generic_op``; tt-metal untouched).
+
+ttnn's own streaming SDPA compute kernel (``sdpa.cpp`` -> ``sdpa_standard_v2``) and its own writer, with
+  * a single K chunk = the whole key sequence (exact softmax: no online max / sum rescaling between K chunks),
+  * a model-local reader (``kernels/fattn_reader.cpp``) that keeps each head's K^T / V RESIDENT in L1 (loaded once
+    per core and head instead of once per q chunk, no inter-core KV chains),
+  * the flat B*H*q-chunk work split to the tile row (enc 9 / 8 rows per core instead of 2 x 5 = 10, no padded rows).
+The compile-time configuration mirrors ttnn's SDPAProgramFactory for this case (non-causal, no mask, bf16, default
+compute config HiFi2 / bf16 dest / approx, exp_approx_mode False).
+"""
+from __future__ import annotations
+
+import math
+import os
+import struct
+
+import ttnn
+
+_KDIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "kernels")
+_SDPA_K = "ttnn/cpp/ttnn/operations/transformer/sdpa/device/kernels/"
+_INACTIVE = 0xFFFFFFFF
+
+
+def _largest_subblock(bh, bw, dst, max_h=1 << 30, max_w=1 << 30):
+    for h, w in ((2, 4), (4, 2), (1, 8), (8, 1), (1, 7), (7, 1), (2, 3), (3, 2), (1, 6), (6, 1),
+                 (1, 5), (5, 1), (2, 2), (1, 4), (4, 1), (1, 3), (3, 1), (1, 2), (2, 1), (1, 1)):
+        if h * w > dst or h > max_h or w > max_w or bh % h or bw % w:
+            continue
+        return h, w
+    return 1, 1
+
+
+def _gran(n, mx):
+    g = min(n, mx)
+    while g > 1 and n % g:
+        g -= 1
+    return g
+
+
+def fattn(q, k, v, device, q_chunk_t=None, out_mc=None, ncores=None):
+    """softmax(q k^T / sqrt(dh)) v for q [B, H, Sq, dh], k / v [B, H, Sk, dh] (TILE bf16, interleaved)."""
+    out_mc = out_mc or ttnn.L1_MEMORY_CONFIG
+    B, NQH, Sq, DH = (int(d) for d in q.shape)
+    NKH, Sk = int(k.shape[1]), int(k.shape[2])
+    assert Sq % 32 == 0 and Sk % 32 == 0 and DH % 32 == 0 and NQH % NKH == 0
+    Sqt, Skt, DHt = Sq // 32, Sk // 32, DH // 32
+    vDHt = DHt
+    Sq_chunk_t = q_chunk_t or int(os.environ.get("MAST3R_FATTN_Q", "1"))
+    assert Sqt % Sq_chunk_t == 0
+    q_num_chunks = Sqt // Sq_chunk_t
+    Sk_chunk_t, k_num_chunks = Skt, 1
+    grid = device.compute_with_storage_grid_size()
+    gx, gy = grid.x, grid.y
+    n = ncores or gx * gy
+    total = B * NQH * q_num_chunks
+    base, extra = divmod(total, n)
+    max_chunks = base + (1 if extra else 0)
+    dst = 8
+    qk_h, qk_w = _largest_subblock(Sq_chunk_t, Sk_chunk_t, dst)
+    o_h, o_w = _largest_subblock(Sq_chunk_t, vDHt, dst, max_h=int(os.environ.get("MAST3R_FATTN_OH", "2")))
+    qktv_h = 2 if (o_h == 1 and 2 * o_w <= dst and Sq_chunk_t >= 2) else o_h
+    out0_t = 2 * qktv_h * vDHt
+    scale_bits = struct.unpack("<I", struct.pack("<f", 1.0 / math.sqrt(DH)))[0]
+    ident = 0x3F803F80
+
+    out = ttnn.allocate_tensor_on_device(ttnn.Shape([B, NQH, Sq, DH]), ttnn.bfloat16, ttnn.TILE_LAYOUT, device, out_mc)
+
+    # CB ids in ttnn's allocation order (no mask / sink / paging CBs)
+    CB_Q, CB_K, CB_V, CB_ID, CB_COL, CB_RECIP, CB_QK, CB_OA, CB_OB, CB_MA, CB_MB, CB_SA, CB_SB, CB_EMD, CB_OUT = range(15)
+    q_tiles = Sq_chunk_t * DHt * (2 if max_chunks > 1 else 1)
+    sizes = {CB_Q: q_tiles, CB_K: Skt * DHt, CB_V: Skt * vDHt, CB_ID: 1, CB_COL: 1, CB_RECIP: 1,
+             CB_QK: Sq_chunk_t * Sk_chunk_t, CB_OA: Sq_chunk_t * vDHt, CB_OB: Sq_chunk_t * vDHt,
+             CB_MA: Sq_chunk_t, CB_MB: Sq_chunk_t, CB_SA: Sq_chunk_t, CB_SB: Sq_chunk_t, CB_EMD: Sq_chunk_t,
+             CB_OUT: out0_t}
+    cores = ttnn.CoreRangeSet([ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(gx - 1, gy - 1))])
+    cbs = [ttnn.CBDescriptor(total_size=nt * 2048, core_ranges=cores,
+                             format_descriptors=[ttnn.CBFormatDescriptor(buffer_index=i, data_format=ttnn.bfloat16,
+                                                                         page_size=2048)])
+           for i, nt in sizes.items()]
+
+    ta_out = list(ttnn.TensorAccessorArgs(out).get_compile_time_args())
+    num_cores = gx * gy
+    cp_ct = [B, NQH, NKH, Skt, DHt, vDHt, Sq_chunk_t, q_num_chunks, Sk_chunk_t, k_num_chunks,
+             DHt, qk_w, qk_h, Sq_chunk_t // qk_h, Sk_chunk_t // qk_w, 1,
+             Sk_chunk_t, o_w, o_h, Sq_chunk_t // o_h, vDHt // o_w, 1,
+             num_cores, 0, 0, 0, 0, scale_bits, 0, 0, 1, Skt, 0, 0,
+             CB_Q, CB_K, CB_V, _INACTIVE, _INACTIVE, CB_ID, CB_COL, _INACTIVE, CB_RECIP, CB_OUT, CB_QK,
+             CB_OA, CB_OB, CB_MA, CB_MB, CB_SA, CB_SB, CB_EMD] + ta_out
+    wr_ct = [B, NQH, NKH, Sqt, Sqt, Sk, DHt, vDHt, Sq_chunk_t, q_num_chunks, Sk_chunk_t, k_num_chunks,
+             ident, scale_bits, num_cores, 0, 0, 0, 0, 0, 0, 1, o_h, 0, 0, 0] + ta_out + [0, 0] + \
+            [_INACTIVE, CB_ID, CB_COL, _INACTIVE, CB_OUT, CB_Q]
+    rd_ct = [q_num_chunks, Sq_chunk_t, Sqt, Skt, DHt, CB_Q, CB_K, CB_V, NQH // NKH] + \
+        list(ttnn.TensorAccessorArgs(q).get_compile_time_args()) + \
+        list(ttnn.TensorAccessorArgs(k).get_compile_time_args()) + \
+        list(ttnn.TensorAccessorArgs(v).get_compile_time_args())
+
+    rd_rt, wr_rt, cp_rt = ttnn.RuntimeArgs(), ttnn.RuntimeArgs(), ttnn.RuntimeArgs()
+    for i in range(num_cores):
+        x, y = i % gx, i // gx
+        if i < n:
+            gs = i * base + min(i, extra)
+            gc = base + (1 if i < extra else 0)
+        else:
+            gs, gc = total, 0
+        rd_rt[x][y] = [q.buffer_address(), k.buffer_address(), v.buffer_address(), gs, gc]
+        wr_rt[x][y] = [out.buffer_address(), i, 1, 0, 0, 0, 0, 0, gs, gc, 0, 0]
+        cp_rt[x][y] = [i, 1, 0, 0, 0, gs, gc]
+    defines = [("STATS_GRANULARITY", str(_gran(Sq_chunk_t, dst))), ("SUB_EXP_GRANULARITY", str(_gran(Sk_chunk_t, dst))),
+               ("MUL_BCAST_GRANULARITY", str(_gran(Sq_chunk_t * Sk_chunk_t, dst))),
+               ("DHT_GRANULARITY", str(_gran(DHt, dst))), ("REDUCE_GRANULARITY", str(_gran(Sq_chunk_t, dst // 2))),
+               ("EXP_APPROX_MODE", "0")]
+    rd = ttnn.KernelDescriptor(
+        kernel_source=os.path.join(_KDIR, "fattn_reader.cpp"), source_type=ttnn.KernelDescriptor.SourceType.FILE_PATH,
+        core_ranges=cores, compile_time_args=rd_ct, runtime_args=rd_rt, config=ttnn.ReaderConfigDescriptor(),
+        defines=[(d, "1") for d in os.environ.get("MAST3R_FATTN_DIAG", "").split(",") if d])
+    wr = ttnn.KernelDescriptor(
+        kernel_source=_SDPA_K + "dataflow/writer_interleaved.cpp", source_type=ttnn.KernelDescriptor.SourceType.FILE_PATH,
+        core_ranges=cores, compile_time_args=wr_ct, runtime_args=wr_rt, defines=defines,
+        config=ttnn.WriterConfigDescriptor())
+    ccfg = ttnn.ComputeConfigDescriptor(math_fidelity=ttnn.MathFidelity.HiFi2, math_approx_mode=True,
+                                        fp32_dest_acc_en=False)
+    local = os.environ.get("MAST3R_FATTN_LOCAL", "0") == "1"
+    if local:
+        defines = defines + [("FSDPA_PROF", os.environ.get("MAST3R_FATTN_PROF", "0"))]
+    cp = ttnn.KernelDescriptor(
+        kernel_source=os.path.join(_KDIR, "fsdpa", "sdpa.cpp") if local else _SDPA_K + "compute/sdpa.cpp", source_type=ttnn.KernelDescriptor.SourceType.FILE_PATH,
+        core_ranges=cores, compile_time_args=cp_ct, runtime_args=cp_rt, defines=defines, config=ccfg)
+    prog = ttnn.ProgramDescriptor(kernels=[rd, wr, cp], semaphores=[], cbs=cbs)
+    ttnn.generic_op([q, k, v, out], prog)
+    return out

code/models/demos/mast3r/tt/fused.py

@@ -55,6 +55,171 @@ DEFAULT_TRACE_REGION_SIZE = 512 * 1024 * 1024  # bytes; measured: the 944-op tra
 
 _TRUE = ("1", "true", "yes", "on")
 
+#: Galaxy/p150 12x10 optimisation pass (2026-10-02, OPT_REPORT.md). ``MAST3R_OPT`` selects
+#: them: ``all`` (default) | ``none`` | comma list. Each is an independent A/B knob.
+#:   out   compact channel-first ROW_MAJOR head output (D2H 29 -> 3 ms)
+#:   mm    per-shape matmul program configs tuned on 12x10 + exact GELU fused into fc1
+#:   gelua approximate GELU in the fused fc1 epilogue (needs "mm"; accuracy-gated, default off)
+#: ``OPT_DEFAULT_OFF`` entries are only enabled when named explicitly (accuracy-gated A/B).
+#:   dpt   DPT layout cleanup: refinenet out_conv before the upsample, RM conv in/out (no tilize)
+#:   l1    transient per-block activations (LN / qkv / heads / RoPE / SDPA / fc1) in L1 interleaved
+#:   lofi  LoFi math fidelity for the transformer matmuls (accuracy-gated A/B, default off)
+#:   dptf  HiFi2 (fp32 acc) math fidelity for the DPT convs except the two head convs (was HiFi4)
+#:   gelut tanh-form GELU (UnaryOpType.GELU_TANH, fp32 SFPU tanh; |tanh-form - erf-form| <= 4.7e-4)
+#:         in the fused fc1 epilogue instead of the erf-form piecewise polynomial (needs "mm")
+#:   lnfold fold the per-block LayerNorm affine (gamma, beta) into the following linear
+#:         (W' = W * gamma, b' = b + W @ beta, fp32 on host, then bf16): affine-free LN on device
+#:   phase  DPT head tail (bilinear x2 -> 3x3 conv + relu -> 1x1 to 4) as 4 polyphase 3x3 convs on the
+#:         256^2 grid with the bilinear weights folded into the kernels (exact algebra, fp64 fold), the
+#:         2-pixel output ring recomputed exactly from 4 thin strips; no 512^2 x 128 intermediate.
+#:         Output [1, 1, 16, hw/4 + 2(H+W)] ROW_MAJOR; the host interleaves phases (needs "out", "dpt")
+#:   hostcol the per-call host cast of the two views writes the patch-embed im2col layout [2, N, 3*16*16]
+#:         bf16 directly (one strided copy, replaces cat + cast); the device skips its reshape/permute im2col
+#:   phase0 head.0 too: refinenet1's bilinear x2 + head.0 3x3 conv as 4 phase convs on the 128^2 grid,
+#:         phases interleaved on device (RM page concats) into the 256^2 head.0 output; the final
+#:         6-px ring is recomputed exactly by a thin-strip run of the original ops (needs "phase")
+#:   decb  decoder: the two branch blocks of a step share one B=2 split-heads / RoPE / SDPA / concat-heads
+#:         per attention (per-branch matmuls stay separate; outputs concatenated / sliced on batch)
+#:   resl1 encoder residual stream (patch-embed output and every residual matmul output) in L1 interleaved
+#:         (the fused residual op's ones vector follows it into L1); enc_norm output back in DRAM
+#:   ropes  RoPE trans_mat HEIGHT_SHARDED on every core (one 32x32 tile each, per-pass L1 copy) -> the
+#:         prefill-sharded rotary_embedding_llama factory reads it locally (bit-identical; 26.7 -> 21.3 us enc)
+#:   sym    return_pose (B-mode): one symmetric graph -- encoder once, decoder for both view orders, DPT
+#:         head 1 + head 2 of (1, 2) and head 1 of (2, 1); head 2 of (2, 1) (unused by the PairViewer pose)
+#:         skipped. Outputs bit-identical to the two-pass path. Server-side switch only.
+#:   tsplit with MAST3R_CQS=2: head 2's main rows and its ring strips are separate trace segments, so the
+#:         2 MB main-rows readback (CQ1) overlaps the ~0.5 ms strip computation (CQ0); host assembles both
+#:   sdpa2  per-shape SDPA chunks balanced for 120 cores (flat B*H*q-chunk scheduling): encoder (B*H=32) q160/k512
+#:         (7 q chunks/head -> 224 chunks <= 2 per core), decoder (B*H=24) q224/k512 (5 chunks/head -> 120 chunks,
+#:         1 per core); the kernel pads the partial last chunk. Probe (sdpa_probe3.py): 133 -> 94 us, 93 -> 67 us
+#:   hrope  model-local generic_op kernel (tt/heads_rope.py + tt/kernels/heads_rope_*.cpp): split heads + RoPE of q and k
+#:         in one program on all 120 cores, reading the per-branch qkv / q / kv matmul outputs directly (no decoder
+#:         batch concat); same compute sequence as rotary_embedding_llama -> bit-identical (heads_rope_probe.py:
+#:         enc 72.4 -> 26.5 us, dec self 74.3 -> 24.2 us, dec cross 75.3 -> 24.0 us). Needs rope=llama
+#:   hcat   model-local generic_op concat-heads kernel (data movement only, bit-exact) on all 120 cores; in the decoder
+#:         it writes the two branches' [1, N, D] outputs directly (no batch slices). Probe (heads_concat_probe.py):
+#:         enc 13.5 -> 9.6 us, dec 19.4 -> 8.3 us
+#:   mln    decoder: the two branches' affine-free LayerNorms of a step run as ONE program (stock LN kernels and
+#:         exact stock compile-time config via generic_op, each core's runtime args point at its branch's tensor):
+#:         bit-identical, 64 cores instead of 2 x 32 (multi_ln_probe.py: 30.9 -> 21.3 us per pair)
+#:   mc2d   plain linears (enc qkv, dec qkv / cq / ckv / embed) as full-grid 2-D mcast ttnn.linear with a per-shape
+#:         in0_block_w (8/16 for the K=768/1024/3072 decoder shapes) instead of minimal_matmul (mc2d_half_probe.py:
+#:         dec cq 22 -> 16 us, ckv 27 -> 22, qkv 33 -> 31, embed 27 -> 19); not bit-identical (different K blocking)
+#:   mc2dr  the residual linears (proj / cproj / fc2) as 2-D mcast linear + separate add instead of the fused dit op
+#:   mm32   mc2d / mc2dr linears with fp32 dest accumulation (out subblocks <= 4): with bf16 dest the larger
+#:         in0_block_w accumulates more K tiles in 16-bit dest (real-pair depth error 1.13 -> 1.43 %); fp32 dest is
+#:         more accurate than the original dit / minimal_matmul (resid_acc_probe.py) at the same speed
+#:   addln  every residual add that feeds a LayerNorm runs inside the LN program (stock FUSE_PRE_ADD reader/compute
+#:         path, model-local copy of layernorm.cpp that also packs the bf16 sum, model-local writer; cb_x kept bf16
+#:         so the LN sees the same bf16 sum ttnn.add produces): bit-identical, 2 programs -> 1 per residual
+#:         (multi_ln_probe.py add: dec pair 41.9 -> 24.6 us, enc 54.1 -> 31.2 us). Needs mc2dr, mln
+#:   cdbl   DPT 3x3 convs at <= 64x64: activation + weight double buffering (64^2 height-sharded, 32^2 / 16^2
+#:         block-sharded); bit-identical (conv_cfg_probe.py: 64^2 96 -> 76 us, 32^2 51 -> 40, 16^2 45 -> 39)
+#:   pil    phase0 interleave of the 4 head.0 phase convs by a model-local data-movement kernel on ROW_MAJOR pixel rows
+#:         (convs emit ROW_MAJOR; tt/kernels/phase_il_rm_*.cpp) instead of concat + transpose + 0/1 matmul + transpose
+#:         (exact; phase_il_probe.py: 30 us vs ~200 us per head)
+#:   til    head-tail phase interleave (rows (a,b,c) x (i,j) -> NCHW) by a model-local data-movement kernel
+#:         (tt/kernels/tail_il_*.cpp) instead of RM reshape + permute + reshape + tilize + 0/1 matmul + untilize (exact;
+#:         tail_il_probe.py: 35 us vs ~130 us per head)
+#:   dmm    decoder: the two branches' same-shape linears (qkv, proj, cq, ckv, cproj, fc1+GELU, fc2) as ONE program
+#:         each: ttnn's own 2-D mcast matmul descriptors placed on the top / bottom 12x5 halves (allowed_worker_cores)
+#:         and merged (ttnn.merge_program_descriptors), run through generic_op; bit-identical to full-grid linears
+#:         (dual_mm_probe.py rows: per pair 60 -> 51 us qkv, 30 -> 25 proj/cq/cproj, 44 -> 37 ckv, 154 -> 132 fc1,
+#:         72 -> 63 fc2). Needs mc2d, mc2dr, addln
+#:   dadd   DPT refinenet with skip: x + resConfUnit1(skip) and resConfUnit2's leading relu as ONE model-local kernel
+#:         (two bf16-rounded FPU adds in ttnn.add's order + SFPU relu; bit-identical; add3_probe.py: 128^2 165 -> 104 us,
+#:         64^2 56 -> 36, 32^2 21 -> 11)
+#:   addln2 addln at the stack ends too: the encoder's last add + enc_norm and the decoder's last add + dec_norm as the
+#:         fused add + LN program with gamma/beta (stock FUSE_GAMMA/FUSE_BETA path); the block-5/8 taps are copied from
+#:         the next step's fused add + LN sums instead of separate adds (bit-identical)
+#:   rchain DPT resConfUnit conv1 -> conv2 on ttnn's L1 conv path: conv1's sharded output feeds conv2 directly (no
+#:         S2I to DRAM + I2S back; no DRAM-path NHWC reshape copy at 16^2); >= 64^2 the relu input is height-sharded
+#:         into the conv's own spec and freed after halo (L1 CB budget at 128^2), < 64^2 it lives in L1 interleaved
+#:         (bit-identical; rchain_probe.py per resconv: 128^2 376 -> 322 us, 64^2 148 -> 128, 32^2 75 -> 68, 16^2 76 -> 47)
+#:   dlin   DPT refinenet out_conv (1x1) with its input (resConfUnit2's sum) in L1 and an explicit full-grid program
+#:         config with the auto config's in0_block_w / compute config (bit-identical); output in L1 below 128^2
+#:         (dlin_probe.py: 16384x256x256 96 -> 30 us, 4096: 33 -> 16, 1024: 15 -> 10). Needs rchain
+#:   dfront DPT front end: the 4 tap projections (dlin configs, auto in0_block_w) write L1, so the ConvTranspose /
+#:         ap3_down / layer_rn convs run on ttnn's L1 conv path (no DRAM-path NHWC reshape copies at 16^2, sharded
+#:         ap3_down output feeds layer4_rn); layer_rn outputs and the ap1 ConvTranspose output go to DRAM (l2_rn on the
+#:         sharded convT output would pick another config). Bit-identical. Needs dlin
+#:   dshard DPT >= 64^2: the refinenet adds read resConfUnit conv2's height shards in place and write relu(s) straight
+#:         into the height-shard spec of the next conv1 (model-local add3s / add2 kernels, same FPU adds -> bit-
+#:         identical); the leading relu of resConfUnit1 writes that spec directly. add3s_probe.py at 128^2: S2I + add3
+#:         + I2S 156 -> 72 us, S2I + add 77 -> 29 us. Needs rchain
+#:   ups1   refinenet upsample output staged in L1 (S2I to L1 + tilize L1 -> DRAM instead of S2I to DRAM + tilize
+#:         DRAM -> DRAM; exact)
+#:   tailf  head tail after the 4 per-phase 1x1s: phase sum + head.4 bias + transpose + slice + untilize + NCHW
+#:         interleave as ONE model-local program (tt/kernels/tailf_*.cpp; FPU adds, exact transpose and 16-bit
+#:         interleave; bit-identical; tailf_probe.py at 256^2: 113 -> 55 us). Needs til
+#:   pemm   patch-embed linear with an explicit full-grid 2-D config (auto config's in0_block_w / compute config:
+#:         bit-identical) writing the L1 residual directly (no 4 MB copy)
+#:   pcat   attention output projections (encoder proj, decoder self / cross proj) read the SDPA output [B, H, N, dh]
+#:         directly as the concatenated-heads matrix: ttnn's own 2-D mcast matmul descriptor with the in0 sender reader
+#:         swapped for a model-local copy that remaps the tile ids (tt/kernels/mm_in0_heads_reader.cpp); no
+#:         concat-heads program (pcat_probe.py: enc 39.9 -> 33.3 us, dec pair 31.1 -> 26.7 us; bit-identical).
+#:         Needs hcat, mc2d (and dmm for the decoder)
+#:   p1hs   refinenet1's out_conv writes the height-shard spec of head.0's phase convs directly; its DRAM copy for the
+#:         ring strips is one S2I (instead of a DRAM write + an I2S). Bit-identical. Needs dlin
+#:   hf3    HiFi3 (fp32 acc) instead of HiFi4 for the two output-resolution head convs (head.0 / head.2 phase convs and
+#:         their ring-strip convs). Not bit-identical (drops only the lo x lo mantissa partial product): real pair
+#:         gates unchanged to 4-5 digits (conf h1 0.99149 -> 0.99150 float); eth12 trace 28.72 -> 28.32 ms
+#:   gpoly  fc1 GELU epilogue (encoder fc1 and the decoder fc1 pairs) as a model-local SFPU erf-GELU minimax polynomial
+#:         (kernels/mast3r_gelu_poly.h: 0.5x + |x| s q(s^2), s = min(|x|, 4.25), 9 coefficients, tools_prof/gelu_fit.py)
+#:         in the GELU_TANH slot of a model-local copy of ttnn's matmul compute kernel (kernels/mm_gelu_compute.cpp, same
+#:         program otherwise). Max |err| vs the erf GELU of the reference 6e-5 (gelut: 4.7e-4); ~30 instead of ~40 SFPU
+#:         instructions per row (rejected_gelu/gelu_fast_probe.py EXP=5: enc fc1 198 -> 155 us). Needs gelut
+#:   qkx    decoder: n_b = LN(x_b) feeds branch b's qkv AND the other branch's cross k/v (norm_y(x_b) == n_b with the
+#:         affine folded), so each branch runs ONE linear with concatenated weights [qkv_b | ckv_other] (N = 3840, same
+#:         in0_block_w -> bit-identical); the self-attention RoPE reads q/k/v at row stride 5*Ht, the cross RoPE reads
+#:         k/v from the other branch's output (cols 3Ht..5Ht). One dual program per step instead of two. Needs dmm
+#:   sdbl   DPT ring-strip convs (head.0 / head.2 on the thin strips) with act + weight double buffering (height-sharded,
+#:         bit-identical; strip_conv_probe.py: 68 -> 60 us and 74 -> 69 us per strip conv)
+#:   pilhs  the phase0 interleave (pil) writes the head.2 phase convs' HEIGHT_SHARDED input spec directly (one contiguous
+#:         NoC write per 32-row unit into its shard; tt/kernels/phase_il_rm_writer_hs.cpp) instead of L1 interleaved + an
+#:         interleaved-to-sharded copy (~21 us per head). Exact. Needs pil
+#:   tups   DPT refinenet upsample (32^2 -> 64^2, 64^2 -> 128^2): bilinear x2 (half-pixel, clamped == ttnn.upsample) straight
+#:         from the TILE out_conv output to the TILE DRAM tensor by a model-local kernel (tt/kernels/ups2_*.cpp; ups2h_* for the
+#:         16-px-wide 16^2 -> 32^2 stage, two image rows per tile row): each output
+#:         tile = sum of <= 4 tile matmuls with constant interpolation tiles (entries exact in bf16, HiFi4, fp32 dest: exact
+#:         products and sums, one bf16 rounding). Replaces untilize + I2S + halo + upsample + S2I + tilize (tups_probe.py:
+#:         64^2 x 256 136 -> 43 us, 32^2 108 -> 20 us). Not bit-identical: the stock upsample rounds in a bf16 dest
+#:         (23 % of its outputs differ from the correctly rounded value, 3.5 % for this kernel)
+#:   sgat   DPT ring strips: their inputs (top/bottom 3 rows, left/right 3 columns of refinenet1's 128^2 output, the latter
+#:         already transposed) gathered straight from the height shards by a model-local data-movement kernel
+#:         (tt/kernels/strip_gather.cpp; 32-byte face-row reads) and carried to the strip segment, instead of an 8 MB S2I to
+#:         DRAM + untilize + 4 slices + 2 concats + permute. Exact. Needs p1hs
+#:   tapm   DPT front: the 4 tap projections (1x1 linears, DRAM-latency bound, 19-22 us each on 32-96 cores) as ONE merged
+#:         program on disjoint column blocks (6 / 3 / 2 / 1 columns for ap3 / ap2 / ap1 / ap0), dlin's in0_block_w ->
+#:         bit-identical. Needs dfront
+#:   obs    encoder qkv / proj / fc2 write BLOCK_SHARDED outputs (= the 2-D mcast config's per-core block: the matmul
+#:         packs into its own shard, no output NoC writes; r6_mm_sweep.py OUTM=bs: qkv 75 -> 60 us, proj 31 -> 26, fc2
+#:         92 -> 87). hrope reads the qkv shards through a TensorAccessor, the fused add + LN reads proj / fc2 through its
+#:         stock TensorAccessor. Bit-identical. Needs hrope, addln, addln2
+#:   obsd   decoder: the dual (top / bottom half) qkx, proj, cproj and fc2 linears write per-half BLOCK_SHARDED outputs
+#:         (qkx 74 -> 58 us); hrope takes up to two sharded layouts, the fused add + LN one reader kernel per layout.
+#:         Bit-identical. Needs dmm, qkx, pcat
+#:   resmm  encoder: the residual adds folded into the proj / fc2 matmul FUSE_BIAS epilogues (model-local compute kernel
+#:         mm_gelu_compute.cpp with MAST3R_RESID_CB: dst += resid tile via dest-reuse ELWADD; the residual stream stays
+#:         BLOCK_SHARDED in the obs output spec and is read through a CB on its shard). The LayerNorms read the sum only
+#:         (no add, no second write). Not bit-identical (one bf16 rounding of the sum instead of two). Needs obs, pcat,
+#:         mm32
+#:   resmmd decoder: as resmm for the dual proj / cproj / fc2 linears (per-half block-sharded residual streams, obsd
+#:         specs); the LayerNorms of both branches read the sums only. Needs obsd, pcat, mm32, qkx, dmm
+#:   esplit the final enc_norm LayerNorm writes the two views' encoder outputs as two tensors (rows of view 1 / view 2,
+#:         each core's writer points at its view's tensor) instead of one [2, N, 1024] tensor + two DRAM slices.
+#:         Bit-identical. Needs resmm
+#:   demb   the two views' decoder embeds as one dual (top / bottom 12x5) program writing decoder step 0's per-half
+#:         block-sharded residual specs directly (replaces 2 full-grid linears + 2 I2S). Bit-identical. Needs dmm, obsd,
+#:         resmmd, resl1d, decb
+#:   hrqk   hrope compute: q and k of a unit in one pass per stage (4 tiles per fp32 dest half instead of 2 x 2).
+#:         Same per-tile ops -> bit-identical (r9_hrope_probe.py: enc 28.9 -> 28.2 us)
+#:   resl1d decoder residual streams in L1 (decoder-embed output in L1; the block-5/8 taps copied to DRAM, dec_norm
+#:         outputs in DRAM, so nothing in L1 survives into the DPT heads)
+OPT_NAMES = ("out", "mm", "gelua", "dpt", "l1", "lofi", "dptf", "gelut", "lnfold", "phase", "hostcol", "phase0",
+             "decb", "resl1", "ropes", "sym", "tsplit", "sdpa2", "hrope", "hcat", "mln", "mc2d", "mc2dr", "resl1d", "addln", "cdbl", "pil", "til", "mm32", "dmm", "dadd", "addln2", "rchain", "dlin", "dfront", "dshard", "ups1", "tailf", "pemm", "pcat", "p1hs", "hf3", "gpoly", "qkx", "sdbl", "pilhs", "tups", "sgat", "tapm", "obs", "obsd", "resmm", "resmmd", "esplit", "demb", "hrqk")
+OPT_DEFAULT_OFF: Tuple[str, ...] = ("gelua", "lofi")
+
 
 def _truthy(v: Optional[str], default: bool) -> bool:
     if v is None or v.strip() == "":
@@ -94,6 +259,8 @@ class FusedConfig:
     ``MAST3R_TRACE``        ``1`` (default) | ``0``. Capture the whole device graph into one
                             metal trace (persistent input buffer, replay per call).
     ``MAST3R_TRACE_REGION`` trace region bytes for ``ttnn.open_device`` (default 512 MiB).
+    ``MAST3R_CQS``          ``1`` (default) | ``2``. With 2 command queues the trace is split after
+                            DPT head 1 and head 1's readback (CQ1) overlaps head 2 (CQ0).
     ``MAST3R_DPT_FUSE``     ``1`` (default) | ``0``. Exact DPT cleanups: TILE reshape of the
                             four taps, relu fused into the resconv conv1 / head2 convs.
 
@@ -110,6 +277,8 @@ class FusedConfig:
     trace_region_size: int = DEFAULT_TRACE_REGION_SIZE
     dpt_fuse: bool = False
     conv_weight_cache: bool = False
+    opts: Tuple[str, ...] = ()
+    num_cqs: int = 1
 
     @classmethod
     def from_env(cls, env: Optional[Mapping[str, str]] = None) -> "FusedConfig":
@@ -131,11 +300,20 @@ class FusedConfig:
         if region <= 0:
             raise ValueError(f"MAST3R_TRACE_REGION must be > 0, got {region}")
         dpt_fuse = _truthy(env.get("MAST3R_DPT_FUSE"), True)
+        opts = parse_opts(env.get("MAST3R_OPT"))
+        num_cqs = int(env.get("MAST3R_CQS") or 1)
+        if num_cqs not in (1, 2):
+            raise ValueError(f"MAST3R_CQS must be 1 or 2, got {num_cqs}")
         return cls(
             enabled=True, rope=rope, rope_lut=rope_lut, mm=mm, sdpa_chunks=sdpa_chunks,
             trace=trace, trace_region_size=region, dpt_fuse=dpt_fuse, conv_weight_cache=True,
+            opts=opts, num_cqs=num_cqs,
         )
 
+    def has(self, name: str) -> bool:
+        """Is optimisation ``name`` (one of :data:`OPT_NAMES`) enabled?"""
+        return name in self.opts
+
     @property
     def fused_rope(self) -> bool:
         """True when q/k weights, biases and LUTs must be channel-permuted."""
@@ -167,8 +345,36 @@ class FusedConfig:
         kw = dict(base)
         if self.enabled and self.trace:
             kw["trace_region_size"] = self.trace_region_size
+            if self.num_cqs == 2:
+                kw["num_command_queues"] = 2
         return kw
 
+    @property
+    def split_io(self) -> bool:
+        """``MAST3R_CQS=2`` (traced fused path): the forward is captured as two traces (A: encoder,
+        decoder, DPT head 1; B: DPT head 2); head 1 is read back on CQ1 (event-synced after A) while
+        CQ0 runs B. Needs a device opened with 2 command queues (Tensix dispatch on Galaxy)."""
+        return self.enabled and self.trace and self.num_cqs == 2
+
+
+def parse_opts(value: Optional[str]) -> Tuple[str, ...]:
+    """``MAST3R_OPT``: unset / ``all`` -> every :data:`OPT_NAMES`; ``none`` / ``0`` -> ();
+    otherwise a comma list (``-name`` removes one from ``all``: ``all,-gelu``)."""
+    v = (value or "all").strip().lower()
+    if v in ("none", "0", "off", ""):
+        return ()
+    sel = set()
+    for tok in (t.strip() for t in v.split(",") if t.strip()):
+        if tok == "all":
+            sel |= set(n for n in OPT_NAMES if n not in OPT_DEFAULT_OFF)
+        elif tok.startswith("-"):
+            sel.discard(tok[1:])
+        elif tok in OPT_NAMES:
+            sel.add(tok)
+        else:
+            raise ValueError(f"MAST3R_OPT: unknown optimisation {tok!r}; known {OPT_NAMES}")
+    return tuple(n for n in OPT_NAMES if n in sel)
+
 
 def parse_sdpa_chunks(value: Optional[str],
                       default: Optional[Tuple[int, int]] = None) -> Optional[Tuple[int, int]]:
@@ -219,6 +425,15 @@ def expand_head_perm(perm: torch.Tensor, heads: int) -> torch.Tensor:
     return torch.cat([perm + h * dh for h in range(heads)])
 
 
+def fold_ln_affine(weight: torch.Tensor, bias: Optional[torch.Tensor], gamma: torch.Tensor, beta: torch.Tensor):
+    """``Linear(LN(x) * gamma + beta)`` == ``Linear'(LN(x))`` with ``W' = W * gamma`` (scale the input
+    columns of the ``(out, in)`` weight) and ``b' = b + W @ beta``. Computed in fp32 (exact up to the
+    final bf16 rounding of ``W'`` / ``b'`` at upload)."""
+    w = weight.float()
+    b = torch.zeros(w.shape[0]) if bias is None else bias.float()
+    return (w * gamma.float()[None, :]).contiguous(), (b + w @ beta.float()).contiguous()
+
+
 def permute_out_rows(weight: torch.Tensor, bias: Optional[torch.Tensor], idx: torch.Tensor):
     """Row gather of an ``nn.Linear`` weight ``(out, in)`` and bias ``(out,)``: exact."""
     w = weight[idx].contiguous()

code/models/demos/mast3r/tt/heads_rope.py

@@ -0,0 +1,1138 @@
+"""Model-local fused split-heads + RoPE kernel (``ttnn.generic_op``; tt-metal untouched).
+
+``heads_rope(srcs, cos, sin, trans_mat, H, S, dh, device, ...) -> (q, k, v)``, each ``[B, H, S, dh]`` TILE bf16.
+
+``srcs`` lists, per output batch index b, ``(q_tensor, q_base, q_w, k_tensor, k_base, k_w, v_tensor, v_base, v_w)``
+where ``*_base`` is the tile offset of head 0 / seq-tile 0 of that batch and ``*_w`` the row width in tiles of the
+source tensor (all sources TILE bf16 interleaved, same buffer type). One program replaces
+``nlp_create_qkv_heads`` (+ the batch concat of per-branch sources in the decoder) + 2 x ``rotary_embedding_llama``;
+the compute kernel runs the stock rotary op's exact op sequence (bit-identical at the same compute config).
+"""
+from __future__ import annotations
+
+import math
+import os
+
+import ttnn
+
+_KDIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "kernels")
+_TILE = 2048
+
+
+def _cb(idx, ntiles, cores):
+    return ttnn.CBDescriptor(
+        total_size=ntiles * _TILE, core_ranges=cores,
+        format_descriptors=[ttnn.CBFormatDescriptor(buffer_index=idx, data_format=ttnn.bfloat16, page_size=_TILE)])
+
+
+def _is_dram(t):
+    return t.memory_config().buffer_type == ttnn.BufferType.DRAM
+
+
+def heads_rope(srcs, cos, sin, trans_mat, H, S, dh, device, out_mc=None, ncores=None, qk2=False):
+    out_mc = out_mc or ttnn.L1_MEMORY_CONFIG
+    B = len(srcs)
+    St, Wt = S // 32, dh // 32
+    grid = device.compute_with_storage_grid_size()
+    gx, gy = grid.x, grid.y
+    n = ncores or gx * gy
+    U = B * St * H
+    base, extra = divmod(U, n)
+    cores = ttnn.CoreRangeSet([ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(gx - 1, gy - 1))])
+    shape = ttnn.Shape([B, H, S, dh])
+    q = ttnn.allocate_tensor_on_device(shape, ttnn.bfloat16, ttnn.TILE_LAYOUT, device, out_mc)
+    k = ttnn.allocate_tensor_on_device(shape, ttnn.bfloat16, ttnn.TILE_LAYOUT, device, out_mc)
+    v = ttnn.allocate_tensor_on_device(shape, ttnn.bfloat16, ttnn.TILE_LAYOUT, device, out_mc)
+    src_dram = any(_is_dram(t) for s_ in srcs for t in (s_[0], s_[3], s_[6]))
+    # MAST3R_OPT "obs": sharded sources (block-sharded matmul outputs) are read through a TensorAccessor; up to two
+    # distinct sharded layouts (kind 1 / 2), interleaved sources keep the pow2 address generator (kind 0)
+    layouts = []
+
+    def kind(t):
+        if not t.memory_config().is_sharded():
+            assert _is_dram(t) == src_dram
+            return 0
+        ta = list(ttnn.TensorAccessorArgs(t).get_compile_time_args())
+        if ta not in layouts:
+            layouts.append(ta)
+        return 1 + layouts.index(ta)
+    rd_rt, wr_rt, cp_rt = ttnn.RuntimeArgs(), ttnn.RuntimeArgs(), ttnn.RuntimeArgs()
+    src_args = []
+    for (qt, qb, qw, kt, kb, kw, vt, vb, vw) in srcs:
+        src_args += [qt.buffer_address(), qb, qw, kind(qt), kt.buffer_address(), kb, kw, kind(kt),
+                     vt.buffer_address(), vb, vw, kind(vt)]
+    assert len(layouts) <= 2
+    ta_src = [v_ for l_ in layouts for v_ in l_]
+    start = 0
+    for i in range(gx * gy):
+        x, y = i % gx, i // gx
+        cnt = (base + (1 if i < extra else 0)) if i < n else 0
+        rd_rt[x][y] = [start, cnt]
+        wr_rt[x][y] = [start, cnt]
+        cp_rt[x][y] = [start, cnt]
+        start += cnt
+    assert start == U
+    ct = [H, St, Wt]
+    diag = [(d_, "1") for d_ in os.environ.get("MAST3R_HROPE_DIAG", "").split(",") if d_]   # probes only
+    rd = ttnn.KernelDescriptor(
+        kernel_source=os.path.join(_KDIR, "heads_rope_reader.cpp"), source_type=ttnn.KernelDescriptor.SourceType.FILE_PATH,
+        core_ranges=cores, compile_time_args=ct + [B] + ta_src, runtime_args=rd_rt,
+        common_runtime_args=[cos.buffer_address(), sin.buffer_address(), trans_mat.buffer_address()] + src_args,
+        defines=[("COS_DRAM", "true" if _is_dram(cos) else "false"), ("TRANS_DRAM", "true" if _is_dram(trans_mat) else "false"),
+                 ("SRC_DRAM", "true" if src_dram else "false")] + [("NTA", str(len(layouts)))] + diag,
+        config=ttnn.ReaderConfigDescriptor())
+    wr = ttnn.KernelDescriptor(
+        kernel_source=os.path.join(_KDIR, "heads_rope_writer.cpp"), source_type=ttnn.KernelDescriptor.SourceType.FILE_PATH,
+        core_ranges=cores, compile_time_args=ct, runtime_args=wr_rt,
+        common_runtime_args=[q.buffer_address(), k.buffer_address(), v.buffer_address()],
+        defines=[("OUT_DRAM", "true" if out_mc.buffer_type == ttnn.BufferType.DRAM else "false")] + diag,
+        config=ttnn.WriterConfigDescriptor())
+    cp = ttnn.KernelDescriptor(
+        kernel_source=os.path.join(_KDIR, "heads_rope_compute.cpp"), source_type=ttnn.KernelDescriptor.SourceType.FILE_PATH,
+        core_ranges=cores, compile_time_args=ct, runtime_args=cp_rt, defines=diag + ([("HROPE_QK2", "1")] if qk2 else []),
+        config=ttnn.ComputeConfigDescriptor(math_fidelity=ttnn.MathFidelity.HiFi4, math_approx_mode=False,
+                                            fp32_dest_acc_en=True))
+    m = 2 if qk2 else 1   # MAST3R_OPT "hrqk": q + k per stage -> the stage CBs hold 2 * Wt tiles
+    cbs = [_cb(0, 4 * Wt, cores), _cb(1, 2 * Wt, cores), _cb(2, 2 * Wt, cores), _cb(3, 1, cores), _cb(4, 2 * Wt, cores),
+           _cb(16, 4 * Wt, cores), _cb(24, m * Wt, cores), _cb(25, m * Wt, cores), _cb(26, m * Wt, cores)]
+    prog = ttnn.ProgramDescriptor(kernels=[rd, wr, cp], semaphores=[], cbs=cbs)
+    ins = []
+    for s in srcs:
+        for t in (s[0], s[3], s[6]):
+            if all(t is not u for u in ins):
+                ins.append(t)
+    ttnn.generic_op(ins + [cos, sin, trans_mat, q, k, v], prog)
+    return q, k, v
+
+
+def heads_concat(ctx, outs, H, S, dh, device, ub=2):
+    """Concatenate heads: ``ctx`` ``[B, H, S, dh]`` TILE bf16 interleaved -> per batch b the rows of
+    ``outs[b] = (tensor, base_tile, row_w_tiles)`` (pre-allocated, e.g. ``[1, S, H*dh]`` per branch, or one
+    ``[B, S, H*dh]`` with base b*St*row_w). Pure data movement (bit-exact), all cores."""
+    B = len(outs)
+    St, Wt = S // 32, dh // 32
+    grid = device.compute_with_storage_grid_size()
+    gx, gy = grid.x, grid.y
+    n = gx * gy
+    U = B * St * H
+    base, extra = divmod(U, n)
+    cores = ttnn.CoreRangeSet([ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(gx - 1, gy - 1))])
+    rd_rt, wr_rt = ttnn.RuntimeArgs(), ttnn.RuntimeArgs()
+    o_args = []
+    for (t, b0, w) in outs:
+        o_args += [t.buffer_address(), b0, w]
+    start = 0
+    for i in range(n):
+        x, y = i % gx, i // gx
+        cnt = base + (1 if i < extra else 0)
+        rd_rt[x][y] = [ctx.buffer_address(), start, cnt]
+        wr_rt[x][y] = [start, cnt] + o_args
+        start += cnt
+    ct = [H, St, Wt, ub]
+    rd = ttnn.KernelDescriptor(
+        kernel_source=os.path.join(_KDIR, "heads_concat_reader.cpp"), source_type=ttnn.KernelDescriptor.SourceType.FILE_PATH,
+        core_ranges=cores, compile_time_args=ct, runtime_args=rd_rt,
+        defines=[("IN_DRAM", "true" if _is_dram(ctx) else "false")], config=ttnn.ReaderConfigDescriptor())
+    wr = ttnn.KernelDescriptor(
+        kernel_source=os.path.join(_KDIR, "heads_concat_writer.cpp"), source_type=ttnn.KernelDescriptor.SourceType.FILE_PATH,
+        core_ranges=cores, compile_time_args=ct + [B], runtime_args=wr_rt,
+        defines=[("OUT_DRAM", "true" if _is_dram(outs[0][0]) else "false")], config=ttnn.WriterConfigDescriptor())
+    prog = ttnn.ProgramDescriptor(kernels=[rd, wr], semaphores=[], cbs=[_cb(0, 2 * ub * Wt, cores)])
+    ins = []
+    for o in outs:
+        if all(o[0] is not u for u in ins):
+            ins.append(o[0])
+    ttnn.generic_op([ctx] + ins, prog)
+    return [o[0] for o in outs]
+
+
+_LN_K = "ttnn/cpp/ttnn/operations/normalization/layernorm/device/kernels/"
+
+
+# probes only (tools_prof/r9_ln_probe.py): model-local reader for the plain (non-fused, affine-free) multi_layernorm
+_LN_READER_OVERRIDE = [os.environ.get("MAST3R_LN_READER") or None]
+_LN_READER_DEFS = [[(d_.split("=")[0], d_.split("=")[1]) for d_ in os.environ.get("MAST3R_LN_READER_DEFS", "").split(",") if d_]]
+
+
+def multi_layernorm(xs, eps, device, out_mc=None, residual_ys=None, res_mc=None, gamma=None, beta=None, fast=False, split=False,
+                    out_split=1):
+    """Affine-free LayerNorm of several same-shape TILE bf16 interleaved tensors in ONE program: the stock
+    ttnn.layer_norm kernels (reader_unary_interleaved_ln / layernorm.cpp / writer_unary_interleaved_start_id_blocked)
+    with the exact compile-time config ttnn's LayerNormMultiCoreProgramFactory builds for this case (default
+    compute config HiFi4 + fp32 dest, non-Welford, non-large-tensor), one tile row per core, but each core's runtime
+    args point at the tensor its row belongs to. Bit-identical to len(xs) separate ttnn.layer_norm calls.
+
+    ``residual_ys``: fused residual add -- normalises s_i = xs[i] + residual_ys[i] (stock FUSE_PRE_ADD path, the sum
+    kept in fp32 inside the kernel) and ALSO writes s_i (bf16) with the model-local ln_add_compute / ln_add_writer
+    kernels. Returns (normed list, sums list). Replaces ttnn.add + ttnn.layer_norm (not bit-identical: the LN sees
+    the unrounded sum).
+
+    ``out_split`` = k (one input, no residual): the normed output is written as k tensors of shape
+    ``[shape[0] // k, *shape[1:]]`` (consecutive row ranges; MAST3R_OPT "esplit": the two views' encoder outputs
+    without a batched tensor + slices). Same kernels and per-core math; only each writer's base address / start tile
+    differs -> bit-identical to slicing the single output."""
+    out_mc = out_mc or ttnn.L1_MEMORY_CONFIG
+    fuse = residual_ys is not None
+    x0 = xs[0]
+    shape = [int(d) for d in x0.shape]
+    W = int(shape[-1])
+    rows_per = 1
+    for d in shape[:-1]:
+        rows_per *= int(d)
+    rows_per //= 32
+    Wt = W // 32
+    grid = device.compute_with_storage_grid_size()
+    gx, gy = grid.x, grid.y
+    total = rows_per * len(xs)
+    assert total <= gx * gy, "one tile row per core"
+    if out_split > 1:
+        assert len(xs) == 1 and not fuse and shape[0] % out_split == 0 and rows_per % out_split == 0
+        oshape = [shape[0] // out_split] + shape[1:]
+        outs = [ttnn.allocate_tensor_on_device(ttnn.Shape(oshape), ttnn.bfloat16, ttnn.TILE_LAYOUT, device, out_mc)
+                for _ in range(out_split)]
+        rows_out = rows_per // out_split
+    else:
+        outs = [ttnn.allocate_tensor_on_device(ttnn.Shape(shape), ttnn.bfloat16, ttnn.TILE_LAYOUT, device, out_mc) for _ in xs]
+        rows_out = rows_per
+    sums = []
+    if fuse:
+        res_mc = res_mc or x0.memory_config()
+        sums = [ttnn.allocate_tensor_on_device(ttnn.Shape(shape), ttnn.bfloat16, ttnn.TILE_LAYOUT, device, res_mc)
+                for _ in xs]
+    block = 4
+    import struct
+    eps_bits = struct.unpack("<I", struct.pack("<f", eps))[0]
+    packed_one = 0x3F803F80
+    rd_rt, wr_rt, cp_rt = ttnn.RuntimeArgs(), ttnn.RuntimeArgs(), ttnn.RuntimeArgs()
+    rd_vals, wr_vals = {}, {}
+    for i in range(total):
+        x, y = i % gx, i // gx
+        t, r = divmod(i, rows_per)
+        rd_vals[(x, y)] = [xs[t].buffer_address(), 1, Wt, r, packed_one, eps_bits,
+                       gamma.buffer_address() if gamma is not None else 0,
+                       beta.buffer_address() if beta is not None else 0,
+                       residual_ys[t].buffer_address() if fuse else 0]
+        to, ro = (divmod(r, rows_out) if out_split > 1 else (t, r))
+        wr_rt[x][y] = wr_vals[(x, y)] = [outs[to].buffer_address(), Wt, 1, ro * Wt] + ([sums[t].buffer_address()] if fuse else [])
+        cp_rt[x][y] = [1]
+    ranges = []
+    full_rows, rem = divmod(total, gx)
+    if full_rows:
+        ranges.append(ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(gx - 1, full_rows - 1)))
+    if rem:
+        ranges.append(ttnn.CoreRange(ttnn.CoreCoord(0, full_rows), ttnn.CoreCoord(rem - 1, full_rows)))
+    cores = ttnn.CoreRangeSet(ranges)
+    ta_in = ttnn.TensorAccessorArgs(x0).get_compile_time_args()
+    ta_b = ttnn.TensorAccessorArgs(residual_ys[0]).get_compile_time_args() if fuse else [0, 0]
+    ta_out = ttnn.TensorAccessorArgs(outs[0]).get_compile_time_args()
+    named = {"cb_in": 0, "cb_inb": 1, "cb_scaler": 2, "cb_eps": 3, "cb_gamma": 5, "cb_beta": 6, "cb_out": 16, "cb_ex": 18,
+             "cb_ex2": 19, "cb_xmm2": 20, "cb_ex2pe": 21, "cb_fusion": 22, "cb_x": 23, "cb_xmm": 24, "cb_reciprocals": 25,
+             "cb_accumulate": 26, "cb_in_rm": 27, "cb_out_rm": 28, "cb_x_welford": 23 if fuse else 0,
+             "welford_fp32_alias": 0, "cb_ex_welford": 18, "cb_ex2_welford": 19, "welford_state_fp32_alias": 0}
+    named = list(named.items())
+    pre = [("FUSE_PRE_ADD", "1")] if fuse else []
+    ta_g = ttnn.TensorAccessorArgs(gamma).get_compile_time_args() if gamma is not None else [0, 0]
+    ta_be = ttnn.TensorAccessorArgs(beta).get_compile_time_args() if beta is not None else [0, 0]
+    rdef = list(pre) + ([("FUSE_GAMMA", "1")] if gamma is not None else []) + ([("FUSE_BETA", "1")] if beta is not None else [])
+    # MAST3R_OPT "obsd": the residual_ys may be block-sharded on different core grids (dual linears on the two grid
+    # halves) -> one reader kernel per distinct residual layout, each on the cores of its tensor's rows
+    groups = []      # [(ta_b, [core coords])]
+    for i in range(total):
+        t = i // rows_per
+        tb = list(ttnn.TensorAccessorArgs(residual_ys[t]).get_compile_time_args()) if fuse else list(ta_b)
+        tb = (tuple(ttnn.TensorAccessorArgs(xs[t]).get_compile_time_args()), tuple(tb))
+        if not groups or groups[-1][0] != tb:
+            if any(g_[0] == tb for g_ in groups):
+                gi = [g_[0] for g_ in groups].index(tb)
+                groups[gi][1].append((i % gx, i // gx))
+                continue
+            groups.append((tb, []))
+        groups[-1][1].append((i % gx, i // gx))
+    split = split and fuse
+    rds, wrs = [], []
+    for (ta_x, tb), crs in groups:
+        ta_x, tb = list(ta_x), list(tb)
+        rd_rt = ttnn.RuntimeArgs()
+        for (x_, y_) in crs:
+            rd_rt[x_][y_] = rd_vals[(x_, y_)]
+        rcores = cores if len(groups) == 1 else ttnn.CoreRangeSet(
+            [ttnn.CoreRange(ttnn.CoreCoord(x_, y_), ttnn.CoreCoord(x_, y_)) for (x_, y_) in crs])
+        rds.append(ttnn.KernelDescriptor(
+            kernel_source=os.path.join(_KDIR, "ln_reader_split.cpp") if split else
+            (_LN_READER_OVERRIDE[0] if (_LN_READER_OVERRIDE[0] and not fuse and gamma is None and beta is None
+                                        and Wt % block == 0) else _LN_K + "dataflow/reader_unary_interleaved_ln.cpp"),
+            source_type=ttnn.KernelDescriptor.SourceType.FILE_PATH,
+            core_ranges=rcores, compile_time_args=[block, 0, W] + ta_x + tb + ta_g + ta_be + [2048],
+            named_compile_time_args=named, defines=rdef + ([("LN_SPLIT_B", "1")] if split else []) + _LN_READER_DEFS[0],
+            runtime_args=rd_rt, config=ttnn.ReaderConfigDescriptor()))
+        if split:
+            # MAST3R_OPT "lnsplit": the writer reads b on the other RISC / NoC (kernels/ln_add_writer_rb.cpp)
+            w_rt = ttnn.RuntimeArgs()
+            for (x_, y_) in crs:
+                w_rt[x_][y_] = wr_vals[(x_, y_)] + [rd_vals[(x_, y_)][8]]
+            wrs.append(ttnn.KernelDescriptor(
+                kernel_source=os.path.join(_KDIR, "ln_add_writer_rb.cpp"),
+                source_type=ttnn.KernelDescriptor.SourceType.FILE_PATH,
+                core_ranges=rcores, compile_time_args=[block] + tb, runtime_args=w_rt,
+                defines=[("OUT_DRAM", "true" if _is_dram(outs[0]) else "false"),
+                         ("RES_DRAM", "true" if _is_dram(sums[0]) else "false")],
+                config=ttnn.WriterConfigDescriptor()))
+    if split:
+        wr = None
+    elif fuse:
+        wr = ttnn.KernelDescriptor(
+            kernel_source=os.path.join(_KDIR, "ln_add_writer.cpp"), source_type=ttnn.KernelDescriptor.SourceType.FILE_PATH,
+            core_ranges=cores, compile_time_args=[block], runtime_args=wr_rt,
+            defines=[("OUT_DRAM", "true" if _is_dram(outs[0]) else "false"),
+                     ("RES_DRAM", "true" if _is_dram(sums[0]) else "false")],
+            config=ttnn.WriterConfigDescriptor())
+    else:
+        wr = ttnn.KernelDescriptor(
+            kernel_source=_LN_K + "dataflow/writer_unary_interleaved_start_id_blocked.cpp",
+            source_type=ttnn.KernelDescriptor.SourceType.FILE_PATH, core_ranges=cores, compile_time_args=[block] + ta_out,
+            named_compile_time_args=named, runtime_args=wr_rt, config=ttnn.WriterConfigDescriptor())
+    ccfg = ttnn.ComputeConfigDescriptor(math_fidelity=ttnn.MathFidelity.HiFi4, math_approx_mode=False, fp32_dest_acc_en=True)
+    utd = [ttnn.UnpackToDestMode.Default] * 64
+    utd[26] = ttnn.UnpackToDestMode.UnpackToDestFp32
+    ccfg.unpack_to_dest_mode = utd
+    fast = fast and fuse and gamma is None and beta is None and Wt % block == 0
+    if fast:
+        # MAST3R_OPT "lnf": model-local fewer-pass add + LN compute kernel (kernels/ln_fast_compute.cpp)
+        ccfg = ttnn.ComputeConfigDescriptor(math_fidelity=ttnn.MathFidelity.HiFi4, math_approx_mode=False,
+                                            fp32_dest_acc_en=True)
+        inv_w_bits = struct.unpack("<I", struct.pack("<f", 1.0 / W))[0]
+        cp = ttnn.KernelDescriptor(
+            kernel_source=os.path.join(_KDIR, "ln_fast_compute.cpp"), source_type=ttnn.KernelDescriptor.SourceType.FILE_PATH,
+            core_ranges=cores, compile_time_args=[Wt, block, inv_w_bits, eps_bits], runtime_args=cp_rt, config=ccfg,
+            defines=[("LNF_PROF", "1")] if os.environ.get("MAST3R_LNF_PROF") else [])
+        f32, b16 = ttnn.float32, ttnn.bfloat16
+
+        def cbf(idx, n, fmt, page):
+            return ttnn.CBDescriptor(total_size=n * page, core_ranges=cores, format_descriptors=[
+                ttnn.CBFormatDescriptor(buffer_index=idx, data_format=fmt, page_size=page)])
+        cbs = [cbf(0, 2 * block, b16, 2048), cbf(1, 2 * block, b16, 2048), cbf(2, 2, b16, 2048), cbf(3, 2, b16, 2048),
+               cbf(16, 2 * block, b16, 2048), cbf(17, 2 * block, b16, 2048), cbf(23, Wt, b16, 2048),
+               cbf(24, 1, b16, 2048), cbf(18, 1, f32, 4096), cbf(20, 1, f32, 4096), cbf(21, 1, f32, 4096)]
+        prog = ttnn.ProgramDescriptor(kernels=rds + (wrs if split else [wr]) + [cp], semaphores=[], cbs=cbs)
+        ttnn.generic_op(list(xs) + list(residual_ys) + outs + sums, prog)
+        return outs, sums
+    cp = ttnn.KernelDescriptor(
+        kernel_source=os.path.join(_KDIR, "ln_add_compute.cpp") if fuse else _LN_K + "compute/layernorm.cpp",
+        source_type=ttnn.KernelDescriptor.SourceType.FILE_PATH,
+        core_ranges=cores, compile_time_args=[Wt, block, int(gamma is not None), int(beta is not None), 1, 1, 0, W, 32],
+        named_compile_time_args=named,
+        defines=pre, runtime_args=cp_rt, config=ccfg)
+
+    def cb(idx, n, fmt, page):
+        return ttnn.CBDescriptor(total_size=n * page, core_ranges=cores,
+                                 format_descriptors=[ttnn.CBFormatDescriptor(buffer_index=idx, data_format=fmt, page_size=page)])
+    f32, b16 = ttnn.float32, ttnn.bfloat16
+    Wnb = -(-Wt // block) * block
+    cbs = [cb(0, 2 * block if fuse else Wnb, b16, 2048), cb(16, 2 * block, b16, 2048), cb(18, 2, f32, 4096),
+           cb(2, 2, b16, 2048), cb(3, 2, b16, 2048), cb(19, 2, f32, 4096), cb(24, Wnb, f32, 4096),
+           cb(20, Wnb, f32, 4096), cb(21, 8, f32, 4096)]
+    if gamma is not None or beta is not None:
+        cbs += [cb(22, 2 * block, f32, 4096)]
+    if gamma is not None:
+        cbs += [cb(5, Wnb, b16, 2048)]
+    if beta is not None:
+        cbs += [cb(6, Wnb, b16, 2048)]
+    if fuse:
+        # cb_x in bf16 (stock: fp32): the LN then sees exactly the bf16 sum ttnn.add would produce -> bit-identical
+        cbs += [cb(23, Wnb, b16, 2048), cb(1, 2 * block, b16, 2048), cb(17, 2 * block, b16, 2048)]
+    prog = ttnn.ProgramDescriptor(kernels=rds + (wrs if split else [wr]) + [cp], semaphores=[], cbs=cbs)
+    aff = [t for t in (gamma, beta) if t is not None]
+    ttnn.generic_op(list(xs) + (list(residual_ys) if fuse else []) + aff + outs + sums, prog)
+    return (outs, sums) if fuse else outs
+
+
+def _in0_dummy(x):
+    """A sharded in0 (rejected obsf: tools_prof/obsf_probe.py) is read by the stock interleaved in0 sender through a TensorAccessor built for
+    the sharded tensor: the descriptor is created for a shape-equal L1-interleaved placeholder, then patched."""
+    if not x.memory_config().is_sharded():
+        return None
+    return ttnn.allocate_tensor_on_device(x.shape, x.dtype, ttnn.TILE_LAYOUT, x.device(), ttnn.L1_MEMORY_CONFIG)
+
+
+def _patch_in0_ta(desc, dummy, x):
+    ta_d = list(ttnn.TensorAccessorArgs(dummy).get_compile_time_args())
+    ta_x = list(ttnn.TensorAccessorArgs(x).get_compile_time_args())
+    assert not list(ttnn.TensorAccessorArgs(x).get_common_runtime_args())
+    n = 0
+    for k in desc.kernels:
+        if str(k.kernel_source).endswith(_IN0_SENDER):
+            cta = list(k.compile_time_args)
+            assert cta[28:28 + len(ta_d)] == ta_d, (cta[28:40], ta_d)
+            k.compile_time_args = cta[:28] + ta_x + cta[28 + len(ta_d):]
+            ra = k.runtime_args
+            for cr in k.core_ranges.ranges():
+                for xx in range(cr.start.x, cr.end.x + 1):
+                    for yy in range(cr.start.y, cr.end.y + 1):
+                        a = list(ra[xx][yy])
+                        if a:
+                            assert a[0] == dummy.buffer_address(), (a[0], dummy.buffer_address())
+                            a[0] = x.buffer_address()
+                            ra[xx][yy] = a
+            n += 1
+    assert n == 1, n
+
+
+_RESMM_SRC = os.path.join(_KDIR, "mm_gelu_compute.cpp")
+
+
+def colmajor_ta_args(x):
+    """MAST3R_OPT "obsc": compile-time TensorAccessor args of a "virtual transposed" view of the BLOCK_SHARDED (row-major
+    shard grid) TILE tensor ``x`` [.., M, K] whose shards hold their tiles in COLUMN-major order (what a matmul with
+    out_subblock_w == 1 packs into its own shard): pages [Kt, Mt], shard [sk, sm] (row-major inside the virtual shard
+    == column-major inside the real one), shard (i, j) on the core of real shard (j, i). Page id of tile (r, k) = k * Mt + r.
+    Format (tt_metal/impl/buffers/tensor_accessor_args.cpp): [cfg, page, rank, nbanks, tshape.., sshape.., coords..]."""
+    a = list(ttnn.TensorAccessorArgs(x).get_compile_time_args())
+    cfg, page, rank, nb = a[:4]
+    assert rank == 2, a[:8]
+    t = a[4:6]
+    sh = a[6:8]
+    co = a[8:]
+    nbanks = nb & 0xFFFF
+    assert len(co) == (nbanks + 1) // 2, (len(co), nbanks)
+    coords = []
+    for i in range(nbanks):
+        w = co[i // 2]
+        coords.append(((w >> 8) & 0xFF, w & 0xFF) if i % 2 == 0 else ((w >> 24) & 0xFF, (w >> 16) & 0xFF))
+    nr, nc = -(-t[0] // sh[0]), -(-t[1] // sh[1])        # real shard grid (rows = M blocks, cols = K blocks)
+    assert nr * nc == nbanks, (nr, nc, nbanks)
+    vco = [coords[j * nc + i] for i in range(nc) for j in range(nr)]   # virtual shard (i, j) -> real shard (j, i)
+    packed = []
+    for i in range(0, nbanks, 2):
+        (x1, y1) = vco[i]
+        if i + 1 < nbanks:
+            (x2, y2) = vco[i + 1]
+            packed.append(((x2 & 0xFF) << 24) | ((y2 & 0xFF) << 16) | ((x1 & 0xFF) << 8) | (y1 & 0xFF))
+        else:
+            packed.append(((x1 & 0xFF) << 8) | (y1 & 0xFF))
+    return [cfg, page, rank, nb, t[1], t[0], sh[1], sh[0]] + packed
+
+
+def _patch_in0_colmajor(desc, dummy, x, Mt, Kt):
+    """MAST3R_OPT "obsc": the stock interleaved in0 sender reads a column-major-shard BLOCK_SHARDED in0 (colmajor_ta_args):
+    tile (r, k) is virtual page k * Mt + r -> strides w = Mt, h = 1, next K block = in0_block_w * Mt, next h block =
+    in0_block_h; start tile r0 * Kt -> r0."""
+    ta_d = list(ttnn.TensorAccessorArgs(dummy).get_compile_time_args())
+    ta_x = colmajor_ta_args(x)
+    n = 0
+    for k in desc.kernels:
+        if str(k.kernel_source).endswith(_IN0_SENDER):
+            cta = list(k.compile_time_args)
+            assert cta[28:28 + len(ta_d)] == ta_d, (cta[28:40], ta_d)
+            assert cta[0] == 1 and cta[1] == Kt, cta[:4]
+            ibw, ibh = cta[4], cta[5]
+            assert cta[2] == ibw and cta[3] == ibh * Kt, cta[:6]
+            cta[0], cta[1], cta[2], cta[3] = Mt, 1, ibw * Mt, ibh
+            k.compile_time_args = cta[:28] + ta_x + cta[28 + len(ta_d):]
+            ra = k.runtime_args
+            for cr in k.core_ranges.ranges():
+                for xx in range(cr.start.x, cr.end.x + 1):
+                    for yy in range(cr.start.y, cr.end.y + 1):
+                        a = list(ra[xx][yy])
+                        if a:
+                            assert a[0] == dummy.buffer_address(), (a[0], dummy.buffer_address())
+                            assert a[1] % Kt == 0, (a[1], Kt)
+                            a[0] = x.buffer_address()
+                            a[1] = a[1] // Kt
+                            ra[xx][yy] = a
+            n += 1
+    assert n == 1, n
+
+
+def _add_resid(desc, resid, out, pc):
+    """MAST3R_OPT "resmm": fold ``out = ... + resid`` into a 2-D mcast linear's FUSE_BIAS epilogue. ``resid`` must be
+    BLOCK_SHARDED with exactly the output's shard spec (== the per-core output block); the compute kernel becomes the
+    model-local copy with MAST3R_RESID_CB and reads resid through a CB globally allocated on its shard."""
+    rmc, omc = resid.memory_config(), out.memory_config()
+    assert rmc.is_sharded() and omc.is_sharded() and rmc == omc, (rmc, omc)
+    assert pc.out_block_h == pc.per_core_M and pc.out_block_w == pc.per_core_N
+    assert pc.out_subblock_h == 1 or pc.out_subblock_w == pc.per_core_N, "row-major packing of the sharded output"
+    assert [int(d) for d in resid.shape] == [int(d) for d in out.shape] and resid.dtype == ttnn.bfloat16
+    used = set()
+    for c in desc.cbs:
+        for f in c.format_descriptors:
+            used.add(int(f.buffer_index))
+    idx = next(i for i in range(8, 32) if i not in used)
+    desc.cbs = list(desc.cbs) + [ttnn.cb_descriptor_from_sharded_tensor(idx, resid)]
+    n = 0
+    for k in desc.kernels:
+        src = str(k.kernel_source)
+        if src.endswith("bmm_large_block_zm_fused_bias_activation.cpp") or src == _RESMM_SRC:
+            assert any(d[0] == "FUSE_BIAS" for d in k.defines), "resmm needs a fused bias"
+            k.kernel_source = _RESMM_SRC
+            k.defines = list(k.defines) + [("MAST3R_RESID_CB", str(idx))]
+            n += 1
+    assert n == 1, n
+
+
+def linear_descriptor(x, w, b, pc, ck, out_mc, compute_source=None, extra_defines=(), resid=None, in0_colmajor=False):
+    """ttnn.linear(x, w, bias=b, program_config=pc (2-D mcast), compute_kernel_config=ck) built from ttnn's own
+    MatmulMultiCoreReuseMcast2DProgramFactory.create_descriptor and run through generic_op. With
+    ``compute_source`` the compute kernel file is swapped for a model-local copy (same CBs / args / defines)."""
+    m = ttnn._ttnn.operations.matmul
+    p = m.MatmulParams()
+    p.program_config = pc
+    p.bcast_batch = True
+    p.output_mem_config = out_mc
+    p.output_dtype = ttnn.bfloat16
+    p.compute_kernel_config = ck
+    p.untilize_out = False
+    p.transpose_a = False
+    p.transpose_b = False
+    i = m.MatmulInputs()
+    dummy = _in0_dummy(x)
+    i.input_tensors = [dummy if dummy is not None else x, w]
+    i.optional_input_tensors = [b]
+    shp = [int(d) for d in x.shape][:-1] + [int(w.shape[-1])]
+    outs = [ttnn.allocate_tensor_on_device(ttnn.Shape(shp), ttnn.bfloat16, ttnn.TILE_LAYOUT, x.device(), out_mc)]
+    i.optional_output_tensors = [outs[0]]
+    desc = m.MatmulMultiCoreReuseMcast2DProgramFactory.create_descriptor(p, i, outs)
+    if dummy is not None:
+        if in0_colmajor:
+            ps = [int(d) for d in x.padded_shape]
+            _patch_in0_colmajor(desc, dummy, x, math.prod(ps[:-1]) // 32, ps[-1] // 32)
+        else:
+            _patch_in0_ta(desc, dummy, x)
+        ttnn.deallocate(dummy)
+    else:
+        assert not in0_colmajor
+    if compute_source is not None:
+        n = 0
+        for k in desc.kernels:
+            if str(k.kernel_source).endswith("bmm_large_block_zm_fused_bias_activation.cpp"):
+                k.kernel_source = compute_source
+                if extra_defines:
+                    k.defines = list(k.defines) + list(extra_defines)
+                n += 1
+        assert n == 1, n
+    if resid is not None:
+        _add_resid(desc, resid, outs[0], pc)
+    ttnn.generic_op([x, w, b] + ([resid] if resid is not None else []) + list(outs), desc)
+    return outs[0]
+
+
+def phase_interleave_rm(ys, h0, w0, C, device, out_mc=None, pb=32):
+    """ROW_MAJOR variant of phase_interleave: ys are [1, 1, h0*w0, C] ROW_MAJOR bf16 interleaved (page = pixel row),
+    output [1, 1, 4*h0*w0, C] ROW_MAJOR (pixel-row pages). Exact."""
+    out_mc = out_mc or ttnn.L1_MEMORY_CONFIG
+    row = C * 2
+    lr = row.bit_length() - 1
+    assert 1 << lr == row and (2 * w0) % pb == 0
+    out = ttnn.allocate_tensor_on_device(ttnn.Shape([1, 1, 4 * h0 * w0, C]), ttnn.bfloat16, ttnn.ROW_MAJOR_LAYOUT, device,
+                                         out_mc)
+    hs = out_mc.is_sharded()
+    if hs:
+        # MAST3R_OPT "pilhs": write the HEIGHT_SHARDED spec directly (shard rows a multiple of the unit size pb)
+        spec = out_mc.shard_spec
+        shard_rows = int(spec.shape[0])
+        assert out_mc.memory_layout == ttnn.TensorMemoryLayout.HEIGHT_SHARDED and shard_rows % pb == 0 and int(spec.shape[1]) == C
+        assert spec.orientation == ttnn.ShardOrientation.ROW_MAJOR
+        coords = []
+        for rng in spec.grid.ranges():
+            for yy in range(rng.start.y, rng.end.y + 1):
+                for xx in range(rng.start.x, rng.end.x + 1):
+                    coords.append((xx, yy))
+        coords.sort(key=lambda c: (c[1], c[0]))      # row-major shard order
+        assert len(coords) * shard_rows >= 4 * h0 * w0
+        noc = []
+        for (xx, yy) in coords:
+            v = device.worker_core_from_logical_core(ttnn.CoreCoord(xx, yy))
+            noc += [v.x, v.y]
+    U = 4 * h0 * w0 // pb
+    grid = device.compute_with_storage_grid_size()
+    gx, gy = grid.x, grid.y
+    n = gx * gy
+    base, extra = divmod(U, n)
+    cores = ttnn.CoreRangeSet([ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(gx - 1, gy - 1))])
+    rd_rt, wr_rt = ttnn.RuntimeArgs(), ttnn.RuntimeArgs()
+    start = 0
+    for i in range(n):
+        x, y = i % gx, i // gx
+        cnt = base + (1 if i < extra else 0)
+        rd_rt[x][y] = [t.buffer_address() for t in ys] + [start, cnt]
+        wr_rt[x][y] = [start, cnt] if hs else [out.buffer_address(), start, cnt]
+        start += cnt
+    rd = ttnn.KernelDescriptor(
+        kernel_source=os.path.join(_KDIR, "phase_il_rm_reader.cpp"), source_type=ttnn.KernelDescriptor.SourceType.FILE_PATH,
+        core_ranges=cores, compile_time_args=[w0, pb, lr], runtime_args=rd_rt,
+        defines=[("SRC_DRAM", "true" if _is_dram(ys[0]) else "false")], config=ttnn.ReaderConfigDescriptor())
+    if hs:
+        wr = ttnn.KernelDescriptor(
+            kernel_source=os.path.join(_KDIR, "phase_il_rm_writer_hs.cpp"), source_type=ttnn.KernelDescriptor.SourceType.FILE_PATH,
+            core_ranges=cores, compile_time_args=[pb, lr, shard_rows // pb], runtime_args=wr_rt,
+            common_runtime_args=[out.buffer_address()] + noc, config=ttnn.WriterConfigDescriptor())
+    else:
+        wr = ttnn.KernelDescriptor(
+            kernel_source=os.path.join(_KDIR, "phase_il_rm_writer.cpp"), source_type=ttnn.KernelDescriptor.SourceType.FILE_PATH,
+            core_ranges=cores, compile_time_args=[pb, lr], runtime_args=wr_rt,
+            defines=[("OUT_DRAM", "true" if _is_dram(out) else "false")], config=ttnn.WriterConfigDescriptor())
+    cb = ttnn.CBDescriptor(total_size=2 * pb * row, core_ranges=cores,
+                           format_descriptors=[ttnn.CBFormatDescriptor(buffer_index=0, data_format=ttnn.bfloat16,
+                                                                       page_size=pb * row)])
+    prog = ttnn.ProgramDescriptor(kernels=[rd, wr], semaphores=[], cbs=[cb])
+    ttnn.generic_op(list(ys) + [out], prog)
+    return out
+
+
+def tail_interleave(t_rm, hl, wl, device, out_mc=None):
+    """Head-tail polyphase interleave: t_rm [1, 1, 16, hl*wl] ROW_MAJOR bf16 (row (2a+b)*4 + c, col i*wl + j) ->
+    [1, 1, 4*hl, 4*wl] ROW_MAJOR with page (c, i) = hi-res rows 2i, 2i+1 of channel c (NCHW [4, 2hl, 2wl]).
+    Model-local data-movement kernel; exact."""
+    out_mc = out_mc or ttnn.L1_MEMORY_CONFIG
+    trow = hl * wl * 2
+    lt = trow.bit_length() - 1
+    page = 4 * wl * 2
+    lp = page.bit_length() - 1
+    assert 1 << lt == trow and 1 << lp == page
+    out = ttnn.allocate_tensor_on_device(ttnn.Shape([1, 1, 4 * hl, 4 * wl]), ttnn.bfloat16, ttnn.ROW_MAJOR_LAYOUT, device,
+                                         out_mc)
+    U = 4 * hl
+    grid = device.compute_with_storage_grid_size()
+    gx, gy = grid.x, grid.y
+    n = gx * gy
+    base, extra = divmod(U, n)
+    cores = ttnn.CoreRangeSet([ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(gx - 1, gy - 1))])
+    rd_rt, wr_rt = ttnn.RuntimeArgs(), ttnn.RuntimeArgs()
+    start = 0
+    for i in range(n):
+        x, y = i % gx, i // gx
+        cnt = base + (1 if i < extra else 0)
+        rd_rt[x][y] = [t_rm.buffer_address(), start, cnt]
+        wr_rt[x][y] = [out.buffer_address(), start, cnt]
+        start += cnt
+    rd = ttnn.KernelDescriptor(
+        kernel_source=os.path.join(_KDIR, "tail_il_reader.cpp"), source_type=ttnn.KernelDescriptor.SourceType.FILE_PATH,
+        core_ranges=cores, compile_time_args=[hl, wl, lt], runtime_args=rd_rt,
+        defines=[("SRC_DRAM", "true" if _is_dram(t_rm) else "false")], config=ttnn.ReaderConfigDescriptor())
+    wr = ttnn.KernelDescriptor(
+        kernel_source=os.path.join(_KDIR, "tail_il_writer.cpp"), source_type=ttnn.KernelDescriptor.SourceType.FILE_PATH,
+        core_ranges=cores, compile_time_args=[wl, lp], runtime_args=wr_rt,
+        defines=[("OUT_DRAM", "true" if _is_dram(out) else "false")], config=ttnn.WriterConfigDescriptor())
+    cbs = [ttnn.CBDescriptor(total_size=2 * 4 * wl * 2, core_ranges=cores,
+                             format_descriptors=[ttnn.CBFormatDescriptor(buffer_index=0, data_format=ttnn.bfloat16,
+                                                                         page_size=4 * wl * 2)]),
+           ttnn.CBDescriptor(total_size=2 * page, core_ranges=cores,
+                             format_descriptors=[ttnn.CBFormatDescriptor(buffer_index=1, data_format=ttnn.bfloat16,
+                                                                         page_size=page)])]
+    prog = ttnn.ProgramDescriptor(kernels=[rd, wr], semaphores=[], cbs=cbs)
+    ttnn.generic_op([t_rm, out], prog)
+    return out
+
+
+def _mm_params(pc, ck, out_mc):
+    m = ttnn._ttnn.operations.matmul
+    p = m.MatmulParams()
+    p.program_config = pc
+    p.bcast_batch = True
+    p.output_mem_config = out_mc
+    p.output_dtype = ttnn.bfloat16
+    p.compute_kernel_config = ck
+    p.untilize_out = False
+    p.transpose_a = False
+    p.transpose_b = False
+    return p
+
+
+def _swap_compute(desc, compute_source):
+    n = 0
+    for k in desc.kernels:
+        if str(k.kernel_source).endswith("bmm_large_block_zm_fused_bias_activation.cpp"):
+            k.kernel_source = compute_source
+            n += 1
+    assert n == 1, n
+
+
+def dual_linear(jobs, ck, out_mc=None, compute_source=None, resids=None):
+    """Several independent ``x_i @ w_i + b_i`` (2-D mcast) in ONE program on disjoint core blocks:
+    ``jobs`` = [(x, w, b, pc)] where each pc is a MatmulMultiCoreReuseMultiCastProgramConfig with its own
+    ``allowed_worker_cores`` block (ttnn's factory places the mcast grid there). The per-job descriptors come from
+    ttnn's MatmulMultiCoreReuseMcast2DProgramFactory.create_descriptor and are combined with
+    ttnn.merge_program_descriptors; each job computes exactly what ttnn.linear with that pc computes."""
+    m = ttnn._ttnn.operations.matmul
+    out_mc = out_mc or ttnn.L1_MEMORY_CONFIG
+    descs, outs, ios = [], [], []
+    mcs = out_mc if isinstance(out_mc, (list, tuple)) else [out_mc] * len(jobs)      # per-job (MAST3R_OPT "obsd")
+    for (x, w, b, pc), out_mc in zip(jobs, mcs):
+        p = _mm_params(pc, ck, out_mc)
+        i = m.MatmulInputs()
+        dummy = _in0_dummy(x)
+        i.input_tensors = [dummy if dummy is not None else x, w]
+        i.optional_input_tensors = [b]
+        shp = [int(d) for d in x.shape][:-1] + [int(w.shape[-1])]
+        o = ttnn.allocate_tensor_on_device(ttnn.Shape(shp), ttnn.bfloat16, ttnn.TILE_LAYOUT, x.device(), out_mc)
+        i.optional_output_tensors = [o]
+        d = m.MatmulMultiCoreReuseMcast2DProgramFactory.create_descriptor(p, i, [o])
+        if dummy is not None:
+            _patch_in0_ta(d, dummy, x)
+            ttnn.deallocate(dummy)
+        if compute_source is not None:
+            _swap_compute(d, compute_source)
+        r = resids[len(outs)] if resids is not None else None
+        if r is not None:
+            _add_resid(d, r, o, pc)
+        descs.append(d)
+        outs.append(o)
+        ios += [x, w, b] + ([r] if r is not None else [])
+    prog = ttnn.merge_program_descriptors(descs)
+    ttnn.generic_op(ios + outs, prog)
+    return outs
+
+
+def add3_relu(x, skip, c2, device, out_mc=None, fp32_dest=False):
+    """DPT refinenet plumbing: s = x + (skip + c2) (two bf16-rounded FPU adds, the order of ttnn.add(x, ttnn.add(skip,
+    c2))) and relu(s) in one program (model-local kernel). All [1, 1, N, C] TILE bf16 interleaved. Returns (s, relu(s))."""
+    out_mc = out_mc or ttnn.DRAM_MEMORY_CONFIG
+    shp = [int(d) for d in x.shape]
+    s_ = ttnn.allocate_tensor_on_device(ttnn.Shape(shp), ttnn.bfloat16, ttnn.TILE_LAYOUT, device, out_mc)
+    r_ = ttnn.allocate_tensor_on_device(ttnn.Shape(shp), ttnn.bfloat16, ttnn.TILE_LAYOUT, device, out_mc)
+    U = 1
+    for d in shp:
+        U *= d
+    U //= 1024
+    grid = device.compute_with_storage_grid_size()
+    gx, gy = grid.x, grid.y
+    n = gx * gy
+    base, extra = divmod(U, n)
+    cores = ttnn.CoreRangeSet([ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(gx - 1, gy - 1))])
+    rd_rt, wr_rt, cp_rt = ttnn.RuntimeArgs(), ttnn.RuntimeArgs(), ttnn.RuntimeArgs()
+    start = 0
+    for i in range(n):
+        xx, yy = i % gx, i // gx
+        cnt = base + (1 if i < extra else 0)
+        rd_rt[xx][yy] = [start, cnt]
+        wr_rt[xx][yy] = [start, cnt]
+        cp_rt[xx][yy] = [cnt]
+        start += cnt
+    rd = ttnn.KernelDescriptor(
+        kernel_source=os.path.join(_KDIR, "add3_reader.cpp"), source_type=ttnn.KernelDescriptor.SourceType.FILE_PATH,
+        core_ranges=cores, compile_time_args=[], runtime_args=rd_rt,
+        common_runtime_args=[x.buffer_address(), skip.buffer_address(), c2.buffer_address()],
+        defines=[("IN_DRAM", "true" if _is_dram(x) else "false")], config=ttnn.ReaderConfigDescriptor())
+    wr = ttnn.KernelDescriptor(
+        kernel_source=os.path.join(_KDIR, "add3_writer.cpp"), source_type=ttnn.KernelDescriptor.SourceType.FILE_PATH,
+        core_ranges=cores, compile_time_args=[], runtime_args=wr_rt,
+        common_runtime_args=[s_.buffer_address(), r_.buffer_address()],
+        defines=[("OUT_DRAM", "true" if _is_dram(s_) else "false")], config=ttnn.WriterConfigDescriptor())
+    cp = ttnn.KernelDescriptor(
+        kernel_source=os.path.join(_KDIR, "add3_compute.cpp"), source_type=ttnn.KernelDescriptor.SourceType.FILE_PATH,
+        core_ranges=cores, compile_time_args=[], runtime_args=cp_rt,
+        config=ttnn.ComputeConfigDescriptor(math_fidelity=ttnn.MathFidelity.HiFi4, math_approx_mode=False,
+                                            fp32_dest_acc_en=fp32_dest))
+    cbs = [_cb(0, 2, cores), _cb(1, 2, cores), _cb(2, 2, cores), _cb(24, 2, cores), _cb(16, 2, cores), _cb(17, 2, cores)]
+    ttnn.generic_op([x, skip, c2, s_, r_], ttnn.ProgramDescriptor(kernels=[rd, wr, cp], semaphores=[], cbs=cbs))
+    return s_, r_
+
+
+def _hs_ranges(t, device):
+    """Per-core tile ranges of a HEIGHT_SHARDED TILE tensor whose shard grid is the first n cores of the compute grid in
+    row-major order (ttnn.num_cores_to_corerangeset(n, grid, row_wise=True), ShardOrientation.ROW_MAJOR). Returns
+    (CoreRangeSet, [(x, y, start, count)]) or None if the spec is not of that form."""
+    mc = t.memory_config()
+    if mc.memory_layout != ttnn.TensorMemoryLayout.HEIGHT_SHARDED or t.layout != ttnn.TILE_LAYOUT:
+        return None
+    ss = mc.shard_spec
+    if ss is None or ss.orientation != ttnn.ShardOrientation.ROW_MAJOR:
+        return None
+    g = device.compute_with_storage_grid_size()
+    n = ss.grid.num_cores()
+    if ss.grid != ttnn.num_cores_to_corerangeset(n, g, row_wise=True):
+        return None
+    sh, sw = ss.shape
+    C = int(t.shape[-1])
+    if sw != C or sh % 32 or C % 32:
+        return None
+    shp = [int(d) for d in t.shape]
+    rows = 1
+    for d in shp[:-1]:
+        rows *= d
+    U = rows * C // 1024
+    per = (sh // 32) * (C // 32)
+    out = []
+    for k in range(n):
+        s = k * per
+        out.append((k % g.x, k // g.x, s, max(0, min(U, s + per) - s)))
+    return ss.grid, out
+
+
+def add3_relu_s(x, skip, c2, device, out_mc, r_mc=None):
+    """add3_relu with c2 HEIGHT_SHARDED (read from each core's own shard: no S2I) and optionally relu(s) written
+    straight into a height-sharded tensor with c2's spec (r_mc; the next conv's input: no I2S). x / skip interleaved
+    (same buffer type), s interleaved (out_mc). Same compute kernel as add3_relu -> bit-identical."""
+    hr = _hs_ranges(c2, device)
+    assert hr is not None
+    cores, rng = hr
+    shp = [int(d) for d in x.shape]
+    s_ = ttnn.allocate_tensor_on_device(ttnn.Shape(shp), ttnn.bfloat16, ttnn.TILE_LAYOUT, device, out_mc)
+    r_shard = r_mc is not None and r_mc.is_sharded()
+    if r_shard:
+        assert r_mc == c2.memory_config()
+    r_ = ttnn.allocate_tensor_on_device(ttnn.Shape(shp), ttnn.bfloat16, ttnn.TILE_LAYOUT, device, r_mc or out_mc)
+    rd_rt, wr_rt, cp_rt = ttnn.RuntimeArgs(), ttnn.RuntimeArgs(), ttnn.RuntimeArgs()
+    for xx, yy, st, cnt in rng:
+        rd_rt[xx][yy] = [st, cnt]
+        wr_rt[xx][yy] = [st, cnt]
+        cp_rt[xx][yy] = [cnt]
+    rd = ttnn.KernelDescriptor(
+        kernel_source=os.path.join(_KDIR, "add3s_reader.cpp"), source_type=ttnn.KernelDescriptor.SourceType.FILE_PATH,
+        core_ranges=cores, compile_time_args=[], runtime_args=rd_rt,
+        common_runtime_args=[x.buffer_address(), skip.buffer_address(), c2.buffer_address()],
+        defines=[("IN_DRAM", "true" if _is_dram(x) else "false"), ("C_SHARD", "1"), ("C_DRAM", "false"), ("N_IN", "3")],
+        config=ttnn.ReaderConfigDescriptor())
+    wr = ttnn.KernelDescriptor(
+        kernel_source=os.path.join(_KDIR, "add3s_writer.cpp"), source_type=ttnn.KernelDescriptor.SourceType.FILE_PATH,
+        core_ranges=cores, compile_time_args=[], runtime_args=wr_rt,
+        common_runtime_args=[s_.buffer_address(), r_.buffer_address()],
+        defines=[("OUT_DRAM", "true" if _is_dram(s_) else "false"), ("N_OUT", "2"), ("R_SHARD", "1" if r_shard else "0"),
+                 ("R_DRAM", "true" if _is_dram(r_) else "false")], config=ttnn.WriterConfigDescriptor())
+    cp = ttnn.KernelDescriptor(
+        kernel_source=os.path.join(_KDIR, "add3_compute.cpp"), source_type=ttnn.KernelDescriptor.SourceType.FILE_PATH,
+        core_ranges=cores, compile_time_args=[], runtime_args=cp_rt,
+        config=ttnn.ComputeConfigDescriptor(math_fidelity=ttnn.MathFidelity.HiFi4, math_approx_mode=False,
+                                            fp32_dest_acc_en=False))
+    cbs = [_cb(0, 2, cores), _cb(1, 2, cores), _cb(2, 2, cores), _cb(24, 2, cores), _cb(16, 2, cores), _cb(17, 2, cores)]
+    ttnn.generic_op([x, skip, c2, s_, r_], ttnn.ProgramDescriptor(kernels=[rd, wr, cp], semaphores=[], cbs=cbs))
+    return s_, r_
+
+
+def add2_s(a, c2, device, out_mc):
+    """a + c2 (ttnn.add's single FPU add, bit-identical) with c2 HEIGHT_SHARDED read from each core's own shard (no
+    S2I); a interleaved, output interleaved (out_mc)."""
+    hr = _hs_ranges(c2, device)
+    assert hr is not None
+    cores, rng = hr
+    shp = [int(d) for d in a.shape]
+    o_ = ttnn.allocate_tensor_on_device(ttnn.Shape(shp), ttnn.bfloat16, ttnn.TILE_LAYOUT, device, out_mc)
+    rd_rt, wr_rt, cp_rt = ttnn.RuntimeArgs(), ttnn.RuntimeArgs(), ttnn.RuntimeArgs()
+    for xx, yy, st, cnt in rng:
+        rd_rt[xx][yy] = [st, cnt]
+        wr_rt[xx][yy] = [st, cnt]
+        cp_rt[xx][yy] = [cnt]
+    rd = ttnn.KernelDescriptor(
+        kernel_source=os.path.join(_KDIR, "add3s_reader.cpp"), source_type=ttnn.KernelDescriptor.SourceType.FILE_PATH,
+        core_ranges=cores, compile_time_args=[], runtime_args=rd_rt,
+        common_runtime_args=[a.buffer_address(), a.buffer_address(), c2.buffer_address()],
+        defines=[("IN_DRAM", "true" if _is_dram(a) else "false"), ("C_SHARD", "1"), ("C_DRAM", "false"), ("N_IN", "2")],
+        config=ttnn.ReaderConfigDescriptor())
+    wr = ttnn.KernelDescriptor(
+        kernel_source=os.path.join(_KDIR, "add3s_writer.cpp"), source_type=ttnn.KernelDescriptor.SourceType.FILE_PATH,
+        core_ranges=cores, compile_time_args=[], runtime_args=wr_rt,
+        common_runtime_args=[o_.buffer_address(), o_.buffer_address()],
+        defines=[("OUT_DRAM", "true" if _is_dram(o_) else "false"), ("N_OUT", "1"), ("R_SHARD", "0"), ("R_DRAM", "false")],
+        config=ttnn.WriterConfigDescriptor())
+    cp = ttnn.KernelDescriptor(
+        kernel_source=os.path.join(_KDIR, "add2_compute.cpp"), source_type=ttnn.KernelDescriptor.SourceType.FILE_PATH,
+        core_ranges=cores, compile_time_args=[], runtime_args=cp_rt,
+        config=ttnn.ComputeConfigDescriptor(math_fidelity=ttnn.MathFidelity.HiFi4, math_approx_mode=False,
+                                            fp32_dest_acc_en=False))
+    cbs = [_cb(0, 2, cores), _cb(2, 2, cores), _cb(16, 2, cores)]
+    ttnn.generic_op([a, c2, o_], ttnn.ProgramDescriptor(kernels=[rd, wr, cp], semaphores=[], cbs=cbs))
+    return o_
+
+
+def tail_fused(zs, bias_full, hl, wl, device, out_mc=None):
+    """Fused head tail (MAST3R_OPT "tailf"): zs = the 4 per-phase 1x1 outputs [1, 1, hl*wl, 32] TILE (phase p nonzero
+    only in columns 4p..4p+3), bias_full = [1, 1, 32, 32] TILE with every row = the 16 head.4 biases (+ zeros).
+    Returns [1, 1, 4*hl, 4*wl] ROW_MAJOR NCHW rows (= tail_interleave(untilize(slice(transpose(z0+z1+z2+z3) + bcol))),
+    bit-identical): one program instead of 3 adds + transpose + bias add + slice + untilize + interleave."""
+    out_mc = out_mc or ttnn.L1_MEMORY_CONFIG
+    assert wl % 64 == 0
+    tph = wl // 64            # tiles per unit (unit = half a low-res image row)
+    page = 4 * wl * 2
+    lp = page.bit_length() - 1
+    assert 1 << lp == page
+    out = ttnn.allocate_tensor_on_device(ttnn.Shape([1, 1, 4 * hl, 4 * wl]), ttnn.bfloat16, ttnn.ROW_MAJOR_LAYOUT, device,
+                                         out_mc)
+    grid = device.compute_with_storage_grid_size()
+    gx, gy = grid.x, grid.y
+    n = gx * gy
+    base, extra = divmod(2 * hl, n)
+    cores = ttnn.CoreRangeSet([ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(gx - 1, gy - 1))])
+    rd_rt, wr_rt, cp_rt = ttnn.RuntimeArgs(), ttnn.RuntimeArgs(), ttnn.RuntimeArgs()
+    start = 0
+    for i in range(n):
+        x, y = i % gx, i // gx
+        cnt = base + (1 if i < extra else 0)
+        rd_rt[x][y] = [start, cnt]
+        wr_rt[x][y] = [start, cnt, out.buffer_address()]
+        cp_rt[x][y] = [cnt * tph]
+        start += cnt
+    rd = ttnn.KernelDescriptor(
+        kernel_source=os.path.join(_KDIR, "tailf_reader.cpp"), source_type=ttnn.KernelDescriptor.SourceType.FILE_PATH,
+        core_ranges=cores, compile_time_args=[tph], runtime_args=rd_rt,
+        common_runtime_args=[z.buffer_address() for z in zs] + [bias_full.buffer_address()],
+        defines=[("Z_DRAM", "true" if _is_dram(zs[0]) else "false"), ("B_DRAM", "true" if _is_dram(bias_full) else "false")],
+        config=ttnn.ReaderConfigDescriptor())
+    wr = ttnn.KernelDescriptor(
+        kernel_source=os.path.join(_KDIR, "tailf_writer.cpp"), source_type=ttnn.KernelDescriptor.SourceType.FILE_PATH,
+        core_ranges=cores, compile_time_args=[tph, wl, hl, lp], runtime_args=wr_rt,
+        defines=[("OUT_DRAM", "true" if _is_dram(out) else "false")], config=ttnn.WriterConfigDescriptor())
+    cp = ttnn.KernelDescriptor(
+        kernel_source=os.path.join(_KDIR, "tailf_compute.cpp"), source_type=ttnn.KernelDescriptor.SourceType.FILE_PATH,
+        core_ranges=cores, compile_time_args=[tph], runtime_args=cp_rt,
+        config=ttnn.ComputeConfigDescriptor(math_fidelity=ttnn.MathFidelity.HiFi4, math_approx_mode=False,
+                                            fp32_dest_acc_en=False))
+    cbs = [_cb(0, 2, cores), _cb(1, 2, cores), _cb(2, 2, cores), _cb(3, 2, cores), _cb(4, 1, cores),
+           _cb(24, 1, cores), _cb(25, 1, cores), _cb(26, 1, cores), _cb(27, 1, cores), _cb(16, 2, cores),
+           ttnn.CBDescriptor(total_size=2 * 8 * tph * 128, core_ranges=cores,
+                             format_descriptors=[ttnn.CBFormatDescriptor(buffer_index=8, data_format=ttnn.bfloat16,
+                                                                         page_size=page)])]
+    ttnn.generic_op(list(zs) + [bias_full, out], ttnn.ProgramDescriptor(kernels=[rd, wr, cp], semaphores=[], cbs=cbs))
+    return out
+
+
+_IN0_SENDER = "reader_bmm_tile_layout_in0_sender_padding.cpp"
+
+
+def _patch_in0_heads(desc, dummy_addr, ctx, heads, nt, dt, b0):
+    """MAST3R_OPT "pcat": swap the 2-D mcast matmul's in0 sender reader for the model-local remapping copy and point its
+    in0 address (runtime arg 0) from the shape-only dummy at the SDPA output ``ctx``."""
+    n = 0
+    for k in desc.kernels:
+        if str(k.kernel_source).endswith(_IN0_SENDER):
+            k.kernel_source = os.path.join(_KDIR, "mm_in0_heads_reader.cpp")
+            k.defines = list(k.defines) + [("MAST3R_HEADS", str(heads)), ("MAST3R_NT", str(nt)), ("MAST3R_DT", str(dt)),
+                                           ("MAST3R_B0", str(b0))]
+            ra = k.runtime_args
+            for cr in k.core_ranges.ranges():
+                for x in range(cr.start.x, cr.end.x + 1):
+                    for y in range(cr.start.y, cr.end.y + 1):
+                        a = list(ra[x][y])
+                        if a:
+                            assert a[0] == dummy_addr, (a[0], dummy_addr)
+                            a[0] = ctx.buffer_address()
+                            ra[x][y] = a
+            n += 1
+    assert n == 1, n
+
+
+def linear_from_heads(jobs, ck, out_mc=None, resids=None):
+    """MAST3R_OPT "pcat": x_i @ w_i + b_i where x_i = concat_heads(ctx_i)[branch b0_i] is read straight from the SDPA
+    output ctx_i [B, H, N, dh] (TILE interleaved) by a remapping in0 reader -- no concat-heads program. jobs =
+    [(ctx, b0, nb, w, b, pc)]: rows of batches b0 .. b0 + nb - 1; several jobs (disjoint allowed_worker_cores) are
+    merged into one program as in dual_linear. Bit-identical to heads_concat + ttnn.linear with the same pc."""
+    m = ttnn._ttnn.operations.matmul
+    out_mc = out_mc or ttnn.L1_MEMORY_CONFIG
+    descs, outs, ios, dummies = [], [], [], []
+    mcs = out_mc if isinstance(out_mc, (list, tuple)) else [out_mc] * len(jobs)      # per-job (MAST3R_OPT "obsd")
+    for (ctx, b0, nb, w, b, pc), out_mc in zip(jobs, mcs):
+        Bc, H, N, dh = (int(d) for d in ctx.shape)
+        K = H * dh
+        dev = ctx.device()
+        dummy = ttnn.allocate_tensor_on_device(ttnn.Shape([nb, N, K]), ttnn.bfloat16, ttnn.TILE_LAYOUT, dev,
+                                               ctx.memory_config())
+        p = _mm_params(pc, ck, out_mc)
+        i = m.MatmulInputs()
+        i.input_tensors = [dummy, w]
+        i.optional_input_tensors = [b]
+        o = ttnn.allocate_tensor_on_device(ttnn.Shape([nb, N, int(w.shape[-1])]), ttnn.bfloat16, ttnn.TILE_LAYOUT, dev,
+                                           out_mc)
+        i.optional_output_tensors = [o]
+        d = m.MatmulMultiCoreReuseMcast2DProgramFactory.create_descriptor(p, i, [o])
+        _patch_in0_heads(d, dummy.buffer_address(), ctx, H, N // 32, dh // 32, b0)
+        r = resids[len(outs)] if resids is not None else None
+        if r is not None:
+            _add_resid(d, r, o, pc)
+        descs.append(d)
+        outs.append(o)
+        ios += [ctx, w, b] + ([r] if r is not None else [])
+        dummies.append(dummy)
+    for t in dummies:
+        ttnn.deallocate(t)
+    prog = descs[0] if len(descs) == 1 else ttnn.merge_program_descriptors(descs)
+    ttnn.generic_op(ios + outs, prog)
+    return outs
+
+
+def ups2_tables():
+    """MAST3R_OPT "tups": the 12 constant interpolation tiles [12 * 32, 32] (row block 6 * wi + kind; wi 0 -> vertical
+    weight 0.25, 1 -> 0.75; kind 0 E_EDGE, 1 E_L, 2 E_R, 3 O_L, 4 O_R, 5 O_EDGE). Half-pixel bilinear x2 with clamping
+    (== ttnn.upsample): out X = 2k -> 0.25 x[k-1] + 0.75 x[k], X = 2k+1 -> 0.75 x[k] + 0.25 x[k+1]; out tile t' = 2m covers
+    the in-tile m (E_R) and pixel 31 of in-tile m-1 (E_L); t' = 2m+1 in-tile m (O_L) and pixel 0 of in-tile m+1 (O_R); the
+    image edges clamp (E_EDGE / O_EDGE). Every entry is exact in bf16."""
+    import torch
+    M = torch.zeros(6, 32, 32, dtype=torch.float64)
+    for i in range(32):
+        k = i // 2                      # even t': in-tile local index of x[k]
+        if i % 2 == 0:
+            if k - 1 >= 0:
+                M[2, i, k - 1] += 0.25
+            else:
+                M[1, i, 31] += 0.25     # x[k-1] = pixel 31 of the previous in-tile
+            M[2, i, k] += 0.75
+        else:
+            M[2, i, k] += 0.75
+            M[2, i, k + 1] += 0.25
+        k = 16 + i // 2                 # odd t'
+        if i % 2 == 0:
+            M[3, i, k - 1] += 0.25
+            M[3, i, k] += 0.75
+        else:
+            M[3, i, k] += 0.75
+            if k + 1 <= 31:
+                M[3, i, k + 1] += 0.25
+            else:
+                M[4, i, 0] += 0.25      # x[k+1] = pixel 0 of the next in-tile
+    M[0] = M[2]
+    M[0, 0, 0] += 0.25                  # left edge: x[-1] clamps to x[0]
+    M[5] = M[3]
+    M[5, 31, 31] += 0.25                # right edge: x[W] clamps to x[W-1]
+    A = torch.cat([M * 0.25, M * 0.75], 0).reshape(12 * 32, 32)
+    return A
+
+
+def upsample2x_tile(x, H, W, C, device, a_tiles, out_mc=None):
+    """MAST3R_OPT "tups": bilinear x2 upsample (half-pixel, clamped == ttnn.upsample) of x [1, 1, H*W, C] TILE bf16
+    interleaved (pixel rows, row-major image), straight to [1, 1, 4*H*W, C] TILE (no untilize / reshard / halo / tilize).
+    W must be a multiple of 32. Exact products and fp32 sums (see ups2_compute.cpp)."""
+    out_mc = out_mc or ttnn.DRAM_MEMORY_CONFIG
+    assert W % 32 == 0 and C % 32 == 0
+    NX, CT = W // 32, C // 32
+    out = ttnn.allocate_tensor_on_device(ttnn.Shape([1, 1, 4 * H * W, C]), ttnn.bfloat16, ttnn.TILE_LAYOUT, device, out_mc)
+    U = 2 * H * CT
+    grid = device.compute_with_storage_grid_size()
+    gx, gy = grid.x, grid.y
+    n = gx * gy
+    base, extra = divmod(U, n)
+    cores = ttnn.CoreRangeSet([ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(gx - 1, gy - 1))])
+    rd_rt, wr_rt, cp_rt = ttnn.RuntimeArgs(), ttnn.RuntimeArgs(), ttnn.RuntimeArgs()
+    start = 0
+    for i in range(n):
+        xx, yy = i % gx, i // gx
+        cnt = base + (1 if i < extra else 0)
+        rd_rt[xx][yy] = [start, cnt]
+        wr_rt[xx][yy] = [start, cnt]
+        cp_rt[xx][yy] = [start, cnt]
+        start += cnt
+    assert start == U
+    rd = ttnn.KernelDescriptor(
+        kernel_source=os.path.join(_KDIR, "ups2_reader.cpp"), source_type=ttnn.KernelDescriptor.SourceType.FILE_PATH,
+        core_ranges=cores, compile_time_args=[NX, CT, H], runtime_args=rd_rt,
+        common_runtime_args=[x.buffer_address(), a_tiles.buffer_address()],
+        defines=[("IN_DRAM", "true" if _is_dram(x) else "false"), ("A_DRAM", "true" if _is_dram(a_tiles) else "false")],
+        config=ttnn.ReaderConfigDescriptor())
+    wr = ttnn.KernelDescriptor(
+        kernel_source=os.path.join(_KDIR, "ups2_writer.cpp"), source_type=ttnn.KernelDescriptor.SourceType.FILE_PATH,
+        core_ranges=cores, compile_time_args=[NX, CT], runtime_args=wr_rt, common_runtime_args=[out.buffer_address()],
+        defines=[("OUT_DRAM", "true" if _is_dram(out) else "false")], config=ttnn.WriterConfigDescriptor())
+    cp = ttnn.KernelDescriptor(
+        kernel_source=os.path.join(_KDIR, "ups2_compute.cpp"), source_type=ttnn.KernelDescriptor.SourceType.FILE_PATH,
+        core_ranges=cores, compile_time_args=[NX, CT], runtime_args=cp_rt,
+        config=ttnn.ComputeConfigDescriptor(math_fidelity=ttnn.MathFidelity.HiFi4, math_approx_mode=False,
+                                            fp32_dest_acc_en=True))
+    cbs = [_cb(0, 4 * NX, cores), _cb(1, 12, cores), _cb(16, 4, cores)]
+    ttnn.generic_op([x, a_tiles, out], ttnn.ProgramDescriptor(kernels=[rd, wr, cp], semaphores=[], cbs=cbs))
+    return out
+
+
+def strip_gather(xs, h0, w0, C, device):
+    """MAST3R_OPT "sgat": the ring-strip inputs (tb [2, 3, w0, C], lr [2, 3, h0, C], ROW_MAJOR, DRAM) gathered straight from
+    the HEIGHT_SHARDED TILE tensor xs [1, 1, h0*w0, C] (see kernels/strip_gather.cpp). Exact."""
+    mc = xs.memory_config()
+    spec = mc.shard_spec
+    assert mc.memory_layout == ttnn.TensorMemoryLayout.HEIGHT_SHARDED and spec.orientation == ttnn.ShardOrientation.ROW_MAJOR
+    assert int(spec.shape[0]) % 32 == 0 and int(spec.shape[1]) == C and C % 32 == 0
+    PER, CT = int(spec.shape[0]) // 32, C // 32
+    coords = []
+    for rng in spec.grid.ranges():
+        for yy in range(rng.start.y, rng.end.y + 1):
+            for xx in range(rng.start.x, rng.end.x + 1):
+                coords.append((xx, yy))
+    coords.sort(key=lambda c: (c[1], c[0]))
+    noc = []
+    for (xx, yy) in coords:
+        v = device.worker_core_from_logical_core(ttnn.CoreCoord(xx, yy))
+        noc += [v.x, v.y]
+    tb = ttnn.allocate_tensor_on_device(ttnn.Shape([2, 3, w0, C]), ttnn.bfloat16, ttnn.ROW_MAJOR_LAYOUT, device,
+                                        ttnn.DRAM_MEMORY_CONFIG)
+    lr = ttnn.allocate_tensor_on_device(ttnn.Shape([2, 3, h0, C]), ttnn.bfloat16, ttnn.ROW_MAJOR_LAYOUT, device,
+                                        ttnn.DRAM_MEMORY_CONFIG)
+    P = 6 * (w0 + h0)
+    grid = device.compute_with_storage_grid_size()
+    gx, gy = grid.x, grid.y
+    n = gx * gy
+    base, extra = divmod(P, n)
+    maxp = base + (1 if extra else 0)
+    cores = ttnn.CoreRangeSet([ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(gx - 1, gy - 1))])
+    rt = ttnn.RuntimeArgs()
+    start = 0
+    for i in range(n):
+        xx, yy = i % gx, i // gx
+        cnt = base + (1 if i < extra else 0)
+        rt[xx][yy] = [start, cnt]
+        start += cnt
+    assert start == P
+    k = ttnn.KernelDescriptor(
+        kernel_source=os.path.join(_KDIR, "strip_gather.cpp"), source_type=ttnn.KernelDescriptor.SourceType.FILE_PATH,
+        core_ranges=cores, compile_time_args=[CT, PER, w0, h0], runtime_args=rt,
+        common_runtime_args=[xs.buffer_address(), tb.buffer_address(), lr.buffer_address()] + noc,
+        config=ttnn.ReaderConfigDescriptor())
+    page = CT * 64
+    cb = ttnn.CBDescriptor(total_size=max(1, maxp) * page, core_ranges=cores,
+                           format_descriptors=[ttnn.CBFormatDescriptor(buffer_index=0, data_format=ttnn.bfloat16, page_size=page)])
+    ttnn.generic_op([xs, tb, lr], ttnn.ProgramDescriptor(kernels=[k], semaphores=[], cbs=[cb]))
+    return tb, lr
+
+
+def ups2h_tables():
+    """MAST3R_OPT "tups", W = 16: 4 constant tiles [4 * 32, 32], index 2 * wi + h (wi 0 -> vertical weight 0.25, 1 -> 0.75;
+    h = which half of the input tile holds the source image row): out px i of the 32-px output row, k = i // 2:
+    even i -> 0.25 x[max(k-1, 0)] + 0.75 x[k], odd i -> 0.75 x[k] + 0.25 x[min(k+1, 15)], x at in-tile rows 16 h + k."""
+    import torch
+    M = torch.zeros(2, 32, 32, dtype=torch.float64)
+    for h in range(2):
+        for i in range(32):
+            k = i // 2
+            if i % 2 == 0:
+                M[h, i, 16 * h + max(k - 1, 0)] += 0.25
+                M[h, i, 16 * h + k] += 0.75
+            else:
+                M[h, i, 16 * h + k] += 0.75
+                M[h, i, 16 * h + min(k + 1, 15)] += 0.25
+    return torch.cat([M * 0.25, M * 0.75], 0).reshape(4 * 32, 32)
+
+
+def upsample2x_tile_half(x, H, C, device, a_tiles, out_mc=None):
+    """MAST3R_OPT "tups" for a 16-px-wide image (W = 16, two image rows per input tile row): [1, 1, 16*H, C] TILE ->
+    [1, 1, 64*H, C] TILE (32-px output rows). Same arithmetic as upsample2x_tile."""
+    out_mc = out_mc or ttnn.DRAM_MEMORY_CONFIG
+    assert C % 32 == 0 and H % 2 == 0
+    CT = C // 32
+    out = ttnn.allocate_tensor_on_device(ttnn.Shape([1, 1, 64 * H, C]), ttnn.bfloat16, ttnn.TILE_LAYOUT, device, out_mc)
+    U = 2 * H * CT
+    grid = device.compute_with_storage_grid_size()
+    gx, gy = grid.x, grid.y
+    n = gx * gy
+    base, extra = divmod(U, n)
+    cores = ttnn.CoreRangeSet([ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(gx - 1, gy - 1))])
+    rd_rt, wr_rt, cp_rt = ttnn.RuntimeArgs(), ttnn.RuntimeArgs(), ttnn.RuntimeArgs()
+    start = 0
+    for i in range(n):
+        xx, yy = i % gx, i // gx
+        cnt = base + (1 if i < extra else 0)
+        rd_rt[xx][yy] = [start, cnt]
+        wr_rt[xx][yy] = [start, cnt]
+        cp_rt[xx][yy] = [start, cnt]
+        start += cnt
+    assert start == U
+    rd = ttnn.KernelDescriptor(
+        kernel_source=os.path.join(_KDIR, "ups2h_reader.cpp"), source_type=ttnn.KernelDescriptor.SourceType.FILE_PATH,
+        core_ranges=cores, compile_time_args=[CT, H], runtime_args=rd_rt,
+        common_runtime_args=[x.buffer_address(), a_tiles.buffer_address()],
+        defines=[("IN_DRAM", "true" if _is_dram(x) else "false"), ("A_DRAM", "true" if _is_dram(a_tiles) else "false")],
+        config=ttnn.ReaderConfigDescriptor())
+    wr = ttnn.KernelDescriptor(
+        kernel_source=os.path.join(_KDIR, "ups2h_writer.cpp"), source_type=ttnn.KernelDescriptor.SourceType.FILE_PATH,
+        core_ranges=cores, compile_time_args=[], runtime_args=wr_rt, common_runtime_args=[out.buffer_address()],
+        defines=[("OUT_DRAM", "true" if _is_dram(out) else "false")], config=ttnn.WriterConfigDescriptor())
+    cp = ttnn.KernelDescriptor(
+        kernel_source=os.path.join(_KDIR, "ups2h_compute.cpp"), source_type=ttnn.KernelDescriptor.SourceType.FILE_PATH,
+        core_ranges=cores, compile_time_args=[CT, H], runtime_args=cp_rt,
+        config=ttnn.ComputeConfigDescriptor(math_fidelity=ttnn.MathFidelity.HiFi4, math_approx_mode=False,
+                                            fp32_dest_acc_en=True))
+    cbs = [_cb(0, 4, cores), _cb(1, 4, cores), _cb(16, 4, cores)]
+    ttnn.generic_op([x, a_tiles, out], ttnn.ProgramDescriptor(kernels=[rd, wr, cp], semaphores=[], cbs=cbs))
+    return out

code/models/demos/mast3r/tt/kernels/add2_compute.cpp

@@ -0,0 +1,34 @@
+// SPDX-License-Identifier: Apache-2.0
+// mast3r-p150 model-local kernel (ttnn.generic_op): out = a + c (CB 0 + CB 2 -> CB 16), one FPU add with the CB formats /
+// rounding of ttnn.add (same as add3_compute's adds).
+#include <cstdint>
+#include "api/compute/common.h"
+#include "api/compute/eltwise_binary.h"
+#include "api/compute/compute_kernel_hw_startup.h"
+#include "api/compute/cb_api.h"
+#include "api/compute/pack.h"
+
+constexpr uint32_t CB_A = 0, CB_C = 2, CB_OUT = 16;
+
+void kernel_main() {
+    const uint32_t u_count = get_arg_val<uint32_t>(0);
+    if (u_count == 0) {
+        return;
+    }
+    compute_kernel_hw_startup(CB_A, CB_C, CB_OUT);
+    add_init(CB_A, CB_C);
+    for (uint32_t u = 0; u < u_count; ++u) {
+        cb_wait_front(CB_A, 1);
+        cb_wait_front(CB_C, 1);
+        cb_reserve_back(CB_OUT, 1);
+        tile_regs_acquire();
+        add_tiles(CB_A, CB_C, 0, 0, 0);
+        tile_regs_commit();
+        tile_regs_wait();
+        pack_tile(0, CB_OUT);
+        tile_regs_release();
+        cb_push_back(CB_OUT, 1);
+        cb_pop_front(CB_A, 1);
+        cb_pop_front(CB_C, 1);
+    }
+}

code/models/demos/mast3r/tt/kernels/add3_compute.cpp

@@ -0,0 +1,62 @@
+// SPDX-License-Identifier: Apache-2.0
+// mast3r-p150 model-local kernel (ttnn.generic_op): t = skip + c2; sum = x + t; out_relu = relu(sum). FPU adds with the
+// same CB formats / rounding as two ttnn.add calls, relu on the SFPU (exact); writes sum (CB 16) and relu(sum) (CB 17).
+#include <cstdint>
+#include "api/compute/common.h"
+#include "api/compute/eltwise_binary.h"
+#include "api/compute/eltwise_unary/eltwise_unary.h"
+#include "api/compute/eltwise_unary/relu.h"
+#include "api/compute/compute_kernel_hw_startup.h"
+#include "api/compute/cb_api.h"
+#include "api/compute/pack.h"
+
+constexpr uint32_t CB_X = 0, CB_S = 1, CB_C = 2, CB_T = 24, CB_SUM = 16, CB_RELU = 17;
+
+void kernel_main() {
+    const uint32_t u_count = get_arg_val<uint32_t>(0);
+    if (u_count == 0) {
+        return;
+    }
+    compute_kernel_hw_startup(CB_S, CB_C, CB_T);
+    for (uint32_t u = 0; u < u_count; ++u) {
+        cb_wait_front(CB_S, 1);
+        cb_wait_front(CB_C, 1);
+        cb_reserve_back(CB_T, 1);
+        add_init(CB_S, CB_C);
+        tile_regs_acquire();
+        add_tiles(CB_S, CB_C, 0, 0, 0);
+        tile_regs_commit();
+        tile_regs_wait();
+        pack_tile(0, CB_T);
+        tile_regs_release();
+        cb_push_back(CB_T, 1);
+        cb_pop_front(CB_S, 1);
+        cb_pop_front(CB_C, 1);
+
+        cb_wait_front(CB_X, 1);
+        cb_wait_front(CB_T, 1);
+        cb_reserve_back(CB_SUM, 1);
+        add_init(CB_X, CB_T);
+        tile_regs_acquire();
+        add_tiles(CB_X, CB_T, 0, 0, 0);
+        tile_regs_commit();
+        tile_regs_wait();
+        pack_tile(0, CB_SUM);
+        tile_regs_release();
+        cb_push_back(CB_SUM, 1);
+
+        cb_reserve_back(CB_RELU, 1);
+        add_init(CB_X, CB_T);
+        tile_regs_acquire();
+        add_tiles(CB_X, CB_T, 0, 0, 0);
+        relu_tile_init();
+        relu_tile(0);
+        tile_regs_commit();
+        tile_regs_wait();
+        pack_tile(0, CB_RELU);
+        tile_regs_release();
+        cb_push_back(CB_RELU, 1);
+        cb_pop_front(CB_X, 1);
+        cb_pop_front(CB_T, 1);
+    }
+}

code/models/demos/mast3r/tt/kernels/add3_reader.cpp

@@ -0,0 +1,27 @@
+// SPDX-License-Identifier: Apache-2.0
+// mast3r-p150 model-local kernel (ttnn.generic_op): DPT refinenet "x + (skip + c2)" + relu, reader.
+// Reads tile u of x (CB 0), skip (CB 1), c2 (CB 2) for u in this core's range (all TILE bf16 interleaved, same shape).
+#include <stdint.h>
+#include "api/dataflow/dataflow_api.h"
+
+void kernel_main() {
+    constexpr uint32_t LOG2_TILE = 11;
+    constexpr uint32_t TB = 1u << LOG2_TILE;
+    const uint32_t u_start = get_arg_val<uint32_t>(0);
+    const uint32_t u_count = get_arg_val<uint32_t>(1);
+    const InterleavedPow2AddrGen<IN_DRAM> gx = {.bank_base_address = get_common_arg_val<uint32_t>(0), .log_base_2_of_page_size = LOG2_TILE};
+    const InterleavedPow2AddrGen<IN_DRAM> gs = {.bank_base_address = get_common_arg_val<uint32_t>(1), .log_base_2_of_page_size = LOG2_TILE};
+    const InterleavedPow2AddrGen<IN_DRAM> gc = {.bank_base_address = get_common_arg_val<uint32_t>(2), .log_base_2_of_page_size = LOG2_TILE};
+    for (uint32_t u = u_start; u < u_start + u_count; ++u) {
+        cb_reserve_back(0, 1);
+        cb_reserve_back(1, 1);
+        cb_reserve_back(2, 1);
+        noc_async_read(gx.get_noc_addr(u), get_write_ptr(0), TB);
+        noc_async_read(gs.get_noc_addr(u), get_write_ptr(1), TB);
+        noc_async_read(gc.get_noc_addr(u), get_write_ptr(2), TB);
+        noc_async_read_barrier();
+        cb_push_back(0, 1);
+        cb_push_back(1, 1);
+        cb_push_back(2, 1);
+    }
+}

code/models/demos/mast3r/tt/kernels/add3_writer.cpp

@@ -0,0 +1,23 @@
+// SPDX-License-Identifier: Apache-2.0
+// mast3r-p150 model-local kernel (ttnn.generic_op): writes sum (CB 16) and relu(sum) (CB 17) tile u.
+#include <stdint.h>
+#include "api/dataflow/dataflow_api.h"
+
+void kernel_main() {
+    constexpr uint32_t LOG2_TILE = 11;
+    constexpr uint32_t TB = 1u << LOG2_TILE;
+    const uint32_t u_start = get_arg_val<uint32_t>(0);
+    const uint32_t u_count = get_arg_val<uint32_t>(1);
+    const InterleavedPow2AddrGen<OUT_DRAM> g0 = {.bank_base_address = get_common_arg_val<uint32_t>(0), .log_base_2_of_page_size = LOG2_TILE};
+    const InterleavedPow2AddrGen<OUT_DRAM> g1 = {.bank_base_address = get_common_arg_val<uint32_t>(1), .log_base_2_of_page_size = LOG2_TILE};
+    for (uint32_t u = u_start; u < u_start + u_count; ++u) {
+        cb_wait_front(16, 1);
+        noc_async_write(get_read_ptr(16), g0.get_noc_addr(u), TB);
+        cb_wait_front(17, 1);
+        noc_async_write(get_read_ptr(17), g1.get_noc_addr(u), TB);
+        noc_async_writes_flushed();
+        cb_pop_front(16, 1);
+        cb_pop_front(17, 1);
+    }
+    noc_async_write_barrier();
+}

code/models/demos/mast3r/tt/kernels/add3s_reader.cpp

@@ -0,0 +1,40 @@
+// SPDX-License-Identifier: Apache-2.0
+// mast3r-p150 model-local kernel (ttnn.generic_op): DPT refinenet "x + (skip + c2)" + relu, reader, shard-aware.
+// Reads tile u of x (CB 0), skip (CB 1) from interleaved buffers and c2 (CB 2) either interleaved or (C_SHARD) from this
+// core's own height shard (the core's range [u_start, u_start + u_count) is exactly its shard's tile range).
+// N_IN = 3: x, skip, c2 ; N_IN = 2: x, c2 only (CB 0 and CB 2).
+#include <stdint.h>
+#include "api/dataflow/dataflow_api.h"
+
+void kernel_main() {
+    constexpr uint32_t LOG2_TILE = 11;
+    constexpr uint32_t TB = 1u << LOG2_TILE;
+    const uint32_t u_start = get_arg_val<uint32_t>(0);
+    const uint32_t u_count = get_arg_val<uint32_t>(1);
+    const InterleavedPow2AddrGen<IN_DRAM> gx = {.bank_base_address = get_common_arg_val<uint32_t>(0), .log_base_2_of_page_size = LOG2_TILE};
+    const InterleavedPow2AddrGen<IN_DRAM> gs = {.bank_base_address = get_common_arg_val<uint32_t>(1), .log_base_2_of_page_size = LOG2_TILE};
+    const uint32_t c_base = get_common_arg_val<uint32_t>(2);
+#if !C_SHARD
+    const InterleavedPow2AddrGen<C_DRAM> gc = {.bank_base_address = c_base, .log_base_2_of_page_size = LOG2_TILE};
+#endif
+    for (uint32_t u = u_start; u < u_start + u_count; ++u) {
+        cb_reserve_back(0, 1);
+        noc_async_read(gx.get_noc_addr(u), get_write_ptr(0), TB);
+#if N_IN == 3
+        cb_reserve_back(1, 1);
+        noc_async_read(gs.get_noc_addr(u), get_write_ptr(1), TB);
+#endif
+        cb_reserve_back(2, 1);
+#if C_SHARD
+        noc_async_read(get_noc_addr(c_base + ((u - u_start) << LOG2_TILE)), get_write_ptr(2), TB);
+#else
+        noc_async_read(gc.get_noc_addr(u), get_write_ptr(2), TB);
+#endif
+        noc_async_read_barrier();
+        cb_push_back(0, 1);
+#if N_IN == 3
+        cb_push_back(1, 1);
+#endif
+        cb_push_back(2, 1);
+    }
+}

code/models/demos/mast3r/tt/kernels/add3s_writer.cpp

@@ -0,0 +1,37 @@
+// SPDX-License-Identifier: Apache-2.0
+// mast3r-p150 model-local kernel (ttnn.generic_op): writer for add3s / add2s. N_OUT = 2: sum (CB 16, interleaved) and
+// relu (CB 17; R_SHARD: into this core's own height shard of the output); N_OUT = 1: CB 16 only.
+#include <stdint.h>
+#include "api/dataflow/dataflow_api.h"
+
+void kernel_main() {
+    constexpr uint32_t LOG2_TILE = 11;
+    constexpr uint32_t TB = 1u << LOG2_TILE;
+    const uint32_t u_start = get_arg_val<uint32_t>(0);
+    const uint32_t u_count = get_arg_val<uint32_t>(1);
+    const InterleavedPow2AddrGen<OUT_DRAM> g0 = {.bank_base_address = get_common_arg_val<uint32_t>(0), .log_base_2_of_page_size = LOG2_TILE};
+#if N_OUT == 2
+    const uint32_t r_base = get_common_arg_val<uint32_t>(1);
+#if !R_SHARD
+    const InterleavedPow2AddrGen<R_DRAM> g1 = {.bank_base_address = r_base, .log_base_2_of_page_size = LOG2_TILE};
+#endif
+#endif
+    for (uint32_t u = u_start; u < u_start + u_count; ++u) {
+        cb_wait_front(16, 1);
+        noc_async_write(get_read_ptr(16), g0.get_noc_addr(u), TB);
+#if N_OUT == 2
+        cb_wait_front(17, 1);
+#if R_SHARD
+        noc_async_write(get_read_ptr(17), get_noc_addr(r_base + ((u - u_start) << LOG2_TILE)), TB);
+#else
+        noc_async_write(get_read_ptr(17), g1.get_noc_addr(u), TB);
+#endif
+#endif
+        noc_async_writes_flushed();
+        cb_pop_front(16, 1);
+#if N_OUT == 2
+        cb_pop_front(17, 1);
+#endif
+    }
+    noc_async_write_barrier();
+}

code/models/demos/mast3r/tt/kernels/fattn_reader.cpp

@@ -0,0 +1,108 @@
+// SPDX-License-Identifier: Apache-2.0
+//
+// mast3r-p150 model-local SDPA reader (MAST3R_OPT "fattn"), used with ttnn's stock streaming SDPA compute kernel
+// (sdpa.cpp, single K chunk = the whole key sequence) and its stock writer through ttnn.generic_op.
+//
+// Differences to the stock reader_interleaved.cpp:
+//  * the whole K^T / V of a (batch, head) is loaded ONCE per core and stays resident: cb_k / cb_v hold exactly one
+//    chunk (= the full sequence), so after the compute pops it, reserve_back returns the same L1 slot and the reader
+//    re-publishes it (push_back) without any NoC traffic while the head does not change;
+//  * no KV chain forwarding (no inter-core coupling), no masks / paging / sinks;
+//  * flat B*H*q_chunk scheduling with per-core [g_start, g_start + g_count) ranges chosen by the host (balanced to
+//    the tile row, not to the stock q-chunk count).
+// The CB tile layouts are the stock ones: Q row-major [Sq_chunk_t x DHt], K transposed [DHt x Skt], V row-major
+// [Skt x DHt].
+
+#include <cstdint>
+#include "api/dataflow/dataflow_api.h"
+
+void kernel_main() {
+    constexpr uint32_t q_num_chunks = get_compile_time_arg_val(0);
+    constexpr uint32_t Sq_chunk_t = get_compile_time_arg_val(1);
+    constexpr uint32_t Sqt = get_compile_time_arg_val(2);
+    constexpr uint32_t Skt = get_compile_time_arg_val(3);
+    constexpr uint32_t DHt = get_compile_time_arg_val(4);
+    constexpr uint32_t cb_q = get_compile_time_arg_val(5);
+    constexpr uint32_t cb_k = get_compile_time_arg_val(6);
+    constexpr uint32_t cb_v = get_compile_time_arg_val(7);
+    constexpr uint32_t kv_heads_div = get_compile_time_arg_val(8);  // q head index -> kv head index divisor (1)
+    constexpr auto q_args = TensorAccessorArgs<9>();
+    constexpr auto k_args = TensorAccessorArgs<q_args.next_compile_time_args_offset()>();
+    constexpr auto v_args = TensorAccessorArgs<k_args.next_compile_time_args_offset()>();
+
+    const uint32_t q_addr = get_arg_val<uint32_t>(0);
+    const uint32_t k_addr = get_arg_val<uint32_t>(1);
+    const uint32_t v_addr = get_arg_val<uint32_t>(2);
+    const uint32_t g_start = get_arg_val<uint32_t>(3);
+    const uint32_t g_count = get_arg_val<uint32_t>(4);
+
+    constexpr uint32_t tile_bytes = get_tile_size(cb_q);
+    const auto q_rd = TensorAccessor(q_args, q_addr, tile_bytes);
+    const auto k_rd = TensorAccessor(k_args, k_addr, tile_bytes);
+    const auto v_rd = TensorAccessor(v_args, v_addr, tile_bytes);
+
+    constexpr uint32_t q_tiles = Sq_chunk_t * DHt;
+    constexpr uint32_t kv_tiles = Skt * DHt;
+
+    // NoC read throttle: with ~120 cores reading L1-interleaved tiles at once, more than ~2 outstanding tile reads per
+    // core congests the NoC (tools_prof/r7_read_bw.py: 2 -> 955 GB/s aggregate, unthrottled -> 270 GB/s)
+#ifndef BAR
+#define BAR 2
+#endif
+    uint32_t nb = 0;
+    uint32_t prev_kv = 0xFFFFFFFF;
+    for (uint32_t i = 0; i < g_count; ++i) {
+        const uint32_t g = g_start + i;
+        const uint32_t bh = g / q_num_chunks;  // flat (batch, q head)
+        const uint32_t qc = g - bh * q_num_chunks;
+        const uint32_t kvh = bh / kv_heads_div;
+#ifdef FATTN_DIAG_NOKV
+        const bool new_kv = false;
+#elif defined(FATTN_DIAG_FIRSTKV)
+        const bool new_kv = i == 0;
+#else
+        const bool new_kv = kvh != prev_kv;
+#endif
+
+        // K^T first (the compute waits on K before Q)
+        cb_reserve_back(cb_k, kv_tiles);
+        if (new_kv) {
+            const uint32_t base = get_write_ptr(cb_k);
+            uint32_t tid = kvh * Skt * DHt;
+            for (uint32_t r = 0; r < Skt; ++r) {
+                for (uint32_t c = 0; c < DHt; ++c) {
+                    noc_async_read(k_rd.get_noc_addr(tid++), base + (c * Skt + r) * tile_bytes, tile_bytes);
+                    if ((++nb & (BAR - 1)) == 0) noc_async_read_barrier();
+                }
+            }
+        }
+        // Q chunk
+        cb_reserve_back(cb_q, q_tiles);
+        {
+            uint32_t wp = get_write_ptr(cb_q);
+            uint32_t tid = (bh * Sqt + qc * Sq_chunk_t) * DHt;
+            for (uint32_t t = 0; t < q_tiles; ++t) {
+                noc_async_read(q_rd.get_noc_addr(tid++), wp, tile_bytes);
+                if ((++nb & (BAR - 1)) == 0) noc_async_read_barrier();
+                wp += tile_bytes;
+            }
+        }
+        noc_async_read_barrier();
+        cb_push_back(cb_k, kv_tiles);
+        cb_push_back(cb_q, q_tiles);
+
+        cb_reserve_back(cb_v, kv_tiles);
+        if (new_kv) {
+            uint32_t wp = get_write_ptr(cb_v);
+            uint32_t tid = kvh * Skt * DHt;
+            for (uint32_t t = 0; t < kv_tiles; ++t) {
+                noc_async_read(v_rd.get_noc_addr(tid++), wp, tile_bytes);
+                if ((++nb & (BAR - 1)) == 0) noc_async_read_barrier();
+                wp += tile_bytes;
+            }
+            noc_async_read_barrier();
+        }
+        cb_push_back(cb_v, kv_tiles);
+        prev_kv = kvh;
+    }
+}

code/models/demos/mast3r/tt/kernels/fsdpa/compute_common.hpp

@@ -0,0 +1,2484 @@
+// mast3r-p150 model-local copy of ttnn sdpa compute (tt-metal 8b98410e730); FSDPA_PROF=1 enables the stock profiling zones.
+// SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#pragma once
+
+#include <cstdint>
+
+#define REDUCE_OP (PoolType::MAX)
+#define REDUCE_DIM (ReduceDim::REDUCE_ROW)
+
+#include "api/debug/assert.h"
+#include "api/compute/compute_kernel_api.h"
+#include "api/compute/binary_max_min.h"
+#include "api/compute/eltwise_binary.h"
+#include "api/compute/eltwise_unary/exp.h"
+#include "api/compute/eltwise_unary/recip.h"
+#include "api/compute/eltwise_unary/softplus.h"
+#include "api/compute/eltwise_unary/negative.h"
+#include "api/compute/eltwise_unary/binop_with_scalar.h"
+#include "api/compute/bcast.h"
+#include "api/compute/tile_move_copy.h"
+#include "api/compute/matmul.h"
+#include "api/compute/reduce.h"
+#include "api/compute/reduce_custom.h"
+#include "api/dataflow/circular_buffer.h"
+#include "cpp/ttnn/operations/transformer/sdpa/device/kernels/q_chunk_remapping.hpp"
+#include "cpp/ttnn/operations/transformer/sdpa/device/kernels/dataflow/chunked_prefill_utils.hpp"
+#include "cpp/ttnn/kernel_lib/dest_helpers.hpp"
+#if defined(TRISC_MATH) || defined(TRISC_PACK)
+#include "experimental/llk_sfpu/ckernel_sfpu_sdpa.h"
+#endif
+
+ALWI void sdpa_reduce_copy_tile_to_dst_init_short(uint32_t cbid, uint32_t transpose = 0) {
+    UNPACK((llk_unpack_A_init<BroadcastType::NONE, false, EltwiseBinaryReuseDestType::NONE, UnpackToDestEn>(
+        transpose, true /*transpose within 16x16 face*/, cbid)));
+
+    MATH((llk_math_eltwise_unary_datacopy_init<
+          DataCopyType::A2D,
+          DST_ACCUM_MODE,
+          BroadcastType::NONE,
+          false,  // is_int_fpu_en
+          PackMode::Default>(cbid)));
+}
+
+/**
+ * in0_cb = max(in0_cb, in1_cb)
+ */
+template <uint32_t num_tiles>
+void max_block_inplace(uint32_t in0, uint32_t in1) {
+    CircularBuffer cb_in0(in0);
+    CircularBuffer cb_in1(in1);
+    // inputs come in full, outputs go out full
+    copy_tile_to_dst_init_short(in0);
+    copy_tile_to_dst_init_short(in1);
+    binary_max_tile_init();
+    constexpr uint32_t dst_reg_0 = 0;
+    constexpr uint32_t dst_reg_1 = 1;
+    cb_in0.wait_front(num_tiles);
+    cb_in1.wait_front(num_tiles);
+    for (uint32_t i = 0; i < num_tiles; ++i) {
+        tile_regs_acquire();
+        copy_tile(in0, i, dst_reg_0);
+        copy_tile(in1, i, dst_reg_1);
+        binary_max_tile(dst_reg_0, dst_reg_1, dst_reg_0, VectorMode::C);
+        tile_regs_commit();
+        tile_regs_wait();
+        pack_tile(dst_reg_0, in0);
+        tile_regs_release();
+    }
+    cb_in0.pop_front(num_tiles);
+    cb_in0.reserve_back(num_tiles);
+    cb_in0.push_back(num_tiles);
+}
+
+/**
+ * out_cb = eltwise_max(in0, in1)
+ */
+template <VectorMode vector_mode = VectorMode::RC>
+void max_block(uint32_t in0, uint32_t in1, uint32_t out_cb, uint32_t num_tiles) {
+    CircularBuffer cb_in0(in0);
+    CircularBuffer cb_in1(in1);
+    CircularBuffer cb_out(out_cb);
+    // inputs come in full, outputs go out full
+    copy_tile_to_dst_init_short(in0);
+    binary_max_tile_init();
+
+    constexpr uint32_t dst_reg_0 = 0;
+    constexpr uint32_t dst_reg_1 = 1;
+    cb_in0.wait_front(num_tiles);
+    cb_in1.wait_front(num_tiles);
+    cb_out.reserve_back(num_tiles);
+    for (uint32_t i = 0; i < num_tiles; ++i) {
+        tile_regs_acquire();
+        copy_tile(in0, i, dst_reg_0);
+        copy_tile(in1, i, dst_reg_1);
+        binary_max_tile(dst_reg_0, dst_reg_1, dst_reg_0, vector_mode);
+        tile_regs_commit();
+        tile_regs_wait();
+        pack_tile(dst_reg_0, out_cb, i);
+        tile_regs_release();
+    }
+    cb_out.push_back(num_tiles);
+}
+
+/**
+ * out_cb = reduce[MAX,SUM](in0_cb * scale_cb)
+ */
+template <
+    PoolType pool_type,
+    ReduceDim reduce_dim,
+    uint32_t in0_cb,
+    uint32_t scale_cb,
+    uint32_t rows,
+    uint32_t cols,
+    VectorMode vector_mode = VectorMode::C>
+void reduce_c(uint32_t out_cb, uint32_t prev_cb, bool do_eltwise_max = false) {
+    CircularBuffer cb_in0(in0_cb);
+    CircularBuffer cb_scale(scale_cb);
+    CircularBuffer cb_out(out_cb);
+    CircularBuffer cb_prev(prev_cb);
+    // Precondition: in0_cb has rows*cols produced. in0_cb has tiles in row-major order
+    // Precondition: scale_cb has 1 produced
+    // Precondition: out_cb has rows free
+    // Postcondition: in0_cb has rows*cols produced
+    // Precondition: scale_cb has 1 produced
+    // Postcondition: out_cb has rows produced
+    // If do_eltwise_max == true, prev_cb has rows produced.
+
+    constexpr uint32_t num_tiles = rows * cols;
+
+#if defined REDUCE_GRANULARITY
+    constexpr uint32_t dst_tiles = (rows < REDUCE_GRANULARITY) ? rows : REDUCE_GRANULARITY;
+    constexpr uint32_t granularity = (rows >= REDUCE_GRANULARITY) ? (rows / REDUCE_GRANULARITY) : 1;
+#else
+    constexpr uint32_t dst_tiles = 1;
+    constexpr uint32_t granularity = rows;
+#endif
+
+    cb_scale.wait_front(1);
+    cb_out.reserve_back(rows);
+
+    const uint32_t num_tiles_to_wait = dst_tiles * cols;
+    uint32_t in0_wait_tiles = num_tiles_to_wait;
+
+    uint32_t row_start_idx = 0;
+    for (uint32_t g = 0; g < granularity; g++) {
+        cb_in0.wait_front(in0_wait_tiles);
+        tile_regs_acquire();
+
+        if (do_eltwise_max) {
+            cb_prev.wait_front(g * dst_tiles);
+            /**
+             * Copy previous max values into DST register.
+             * Note that this special invocation of copy_tile is necessary to produce
+             * tiles in DST with transposed faces, as `reduce_block_max_row` expects.
+             */
+            reconfig_data_format_srca(prev_cb);
+            sdpa_reduce_copy_tile_to_dst_init_short(prev_cb);
+            for (uint32_t i = 0; i < dst_tiles; i++) {
+                const uint32_t cur_max_dst_idx = i;
+                copy_tile(prev_cb, (row_start_idx + i), cur_max_dst_idx);
+            }
+            reconfig_data_format_srca(in0_cb);
+        }
+
+        /**
+         * For `dst_tiles` number of rows, compute the max into the even indices of the DST register.
+         */
+        reduce_block_max_row_init<cols>(out_cb);
+        for (uint32_t i = 0; i < dst_tiles; i++) {
+            const uint32_t reduce_dst_idx = i;
+            reduce_block_max_row<cols>(in0_cb, scale_cb, (row_start_idx + i) * cols, reduce_dst_idx);
+        }
+        reduce_block_max_row_uninit(in0_cb);
+
+        tile_regs_commit();
+        tile_regs_wait();
+        pack_reconfig_data_format(out_cb);
+        for (uint32_t i = 0; i < dst_tiles; i++) {
+            const uint32_t cur_max_dst_idx = i;
+            pack_tile<true>(cur_max_dst_idx, out_cb, (row_start_idx + i));
+        }
+        tile_regs_release();
+
+        row_start_idx += dst_tiles;
+        in0_wait_tiles += num_tiles_to_wait;
+    }
+
+    cb_out.push_back(rows);
+}
+
+/**
+ * out_cb = reduce[MAX,SUM](in0_cb * scale_cb)
+ *
+ * In this version cols does not have to be a compile-time constant.
+ */
+template <
+    PoolType pool_type,
+    ReduceDim reduce_dim,
+    uint32_t in0_cb,
+    uint32_t scale_cb,
+    uint32_t rows,
+    VectorMode vector_mode = VectorMode::C>
+void reduce_c(uint32_t out_cb, uint32_t prev_cb, uint32_t cols, bool do_eltwise_max = false) {
+    CircularBuffer cb_in0(in0_cb);
+    CircularBuffer cb_scale(scale_cb);
+    CircularBuffer cb_out(out_cb);
+    // Precondition: in0_cb has rows*cols produced. in0_cb has tiles in row-major order
+    // Precondition: scale_cb has 1 produced
+    // Precondition: out_cb has rows free
+    // Postcondition: in0_cb has rows*cols produced
+    // Precondition: scale_cb has 1 produced
+    // Postcondition: out_cb has rows produced
+
+    uint32_t num_tiles = rows * cols;
+    cb_scale.wait_front(1);
+    cb_in0.wait_front(num_tiles);
+    cb_out.reserve_back(rows);
+
+    pack_reconfig_data_format(out_cb);
+
+    binary_max_tile_init();
+    constexpr uint32_t reduce_dst_idx = 0;
+    constexpr uint32_t prev_max_dst_idx = 1;
+
+    for (uint32_t i = 0; i < rows; i++) {
+        reconfig_data_format_srca(in0_cb);
+        tile_regs_acquire();
+        reduce_init<pool_type, reduce_dim>(in0_cb, scale_cb, out_cb);
+        for (uint32_t j = 0; j < cols; j++) {
+            reduce_tile<pool_type, reduce_dim>(in0_cb, scale_cb, i * cols + j, 0, reduce_dst_idx);
+        }
+        reduce_uninit();
+        if (do_eltwise_max) {
+            reconfig_data_format_srca(prev_cb);
+            copy_tile_to_dst_init_short(prev_cb);
+            copy_tile(prev_cb, i, prev_max_dst_idx);
+            binary_max_tile(reduce_dst_idx, prev_max_dst_idx, reduce_dst_idx, vector_mode);
+        }
+
+        tile_regs_commit();
+        tile_regs_wait();
+        pack_tile(reduce_dst_idx, out_cb);
+        tile_regs_release();
+    }
+
+    cb_out.push_back(rows);
+}
+
+#ifdef TRISC_MATH
+template <bool legacy_compat = true>
+void recip_tile_first_column(uint32_t idst) {
+    SFPU_UNARY_CALL(DST_SYNC_MODE, DST_ACCUM_MODE, calculate_recip_first_column, (legacy_compat), idst, VectorMode::C);
+}
+#endif
+
+/**
+ * in_cb = 1 / in_cb
+ */
+void recip_block_inplace(uint32_t in_cb, uint32_t num_tiles) {
+    CircularBuffer cb_in(in_cb);
+    // Precondition: in_cb has num_tiles produced
+    // Postcondition: in_cb has num_tiles produced
+    reconfig_data_format_srca(in_cb);
+    copy_tile_to_dst_init_short(in_cb);
+    recip_tile_init();
+    pack_reconfig_data_format(in_cb);
+
+    cb_in.wait_front(num_tiles);
+    for (uint32_t i = 0; i < num_tiles; ++i) {
+        tile_regs_acquire();
+        copy_tile(in_cb, i, 0);
+        MATH((recip_tile_first_column(0)));
+        tile_regs_commit();
+        tile_regs_wait();
+        pack_tile(0, in_cb);
+        tile_regs_release();
+    }
+    cb_in.pop_front(num_tiles);
+    cb_in.reserve_back(num_tiles);
+    cb_in.push_back(num_tiles);
+}
+
+/**
+ * in0_cb = exp((in0_cb - in1_cb) * scale_fp32)
+ */
+template <
+    uint32_t in0_cb,
+    uint32_t rows,
+    uint32_t scale_fp32,
+    bool write_result_inplace = true,
+    bool do_reduce = true,
+    VectorMode vector_mode = VectorMode::RC>
+void sub_exp_block_bcast_cols_inplace(uint32_t in1_cb, uint32_t reduce_cb, uint32_t cols) {
+    CircularBuffer cb_in0(in0_cb);
+    CircularBuffer cb_in1(in1_cb);
+    CircularBuffer cb_reduce(reduce_cb);
+    // Precondition: in0_cb has rows*cols produced
+    // Precondition: in1_cb has rows produced
+    // Postcondition: in0_cb has rows*cols produced
+    // Postcondition: in1_cb has rows produced
+    // llk_unpack_AB_init (inside sub_bcast_cols_init) validates the live
+    // unpacker configuration.  Reconfigure first: qk_im can be FP32 while the
+    // row maximum is BF16, notably after applying a windowed BF16 mask.
+    reconfig_data_format(in0_cb, in1_cb);
+    sub_bcast_cols_init(in0_cb, in1_cb);
+
+    // The exponential function uses InputClamping::None for better performance. This version
+    // produces incorrect outputs for inputs <~ -88, but those outputs are guaranteed to be negative.
+    // Enable packer ReLU to zero any negative values produced by the exponential approximation.
+    exp_tile_init<true /* approx */, scale_fp32, InputClamping::None>();
+    PACK((llk_pack_relu_config(ReluConfig::zero())));
+
+    cb_in0.wait_front(rows * cols);
+    cb_in1.wait_front(rows);
+    if constexpr (do_reduce) {
+        cb_reduce.reserve_back(rows);
+    }
+
+#ifdef SUB_EXP_GRANULARITY
+    uint32_t dst_tiles = (cols < SUB_EXP_GRANULARITY) ? cols : SUB_EXP_GRANULARITY;
+    uint32_t granularity = (cols >= SUB_EXP_GRANULARITY) ? (cols / SUB_EXP_GRANULARITY) : 1;
+#else
+    uint32_t dst_tiles = cols;
+    uint32_t granularity = 1;
+#endif
+
+    for (uint32_t i = 0; i < rows; ++i) {
+        for (uint32_t u = 0; u < granularity; u++) {
+            tile_regs_acquire();
+            for (uint32_t j = 0; j < dst_tiles; ++j) {
+                sub_tiles_bcast_cols(in0_cb, in1_cb, j, i, j);
+                constexpr int iterations = (vector_mode == VectorMode::RC) ? 32 /*ITER*/ : 8 /*ITER*/;
+                constexpr VectorMode vector_mode_exp = (vector_mode == VectorMode::RC) ? VectorMode::None : vector_mode;
+                exp_tile<true /* approx */, false /* scale_en */, InputClamping::None, iterations>(j, vector_mode_exp);
+            }
+            tile_regs_commit();
+
+            if constexpr (write_result_inplace) {
+                cb_in0.pop_front(dst_tiles);
+                cb_in0.reserve_back(dst_tiles);
+            }
+
+            tile_regs_wait();
+
+            if constexpr (write_result_inplace) {
+                pack_reconfig_data_format(in0_cb);
+                for (uint32_t j = 0; j < dst_tiles; ++j) {
+                    pack_tile(j, in0_cb);
+                }
+                // Granular write output to enable following matmul unpack to start early.
+                cb_in0.push_back(dst_tiles);
+            }
+
+            if constexpr (do_reduce) {
+                pack_reconfig_data_format(reduce_cb);
+                // While we have results in DST, take advantage of L1 accumulation
+                // to reduce row x cols tiles to rows x 1 tiles.
+                if (u > 0) {
+                    // If on the same row, keep accumulating
+                    PACK((llk_pack_reconfig_l1_acc(1)));
+                }
+                for (uint32_t j = 0; j < dst_tiles; ++j) {
+                    pack_tile<true>(j, reduce_cb, i);
+                    if (u == 0 && j == 0) {
+                        // If this was the first tile of a row, start accumulating
+                        PACK((llk_pack_reconfig_l1_acc(1)));
+                    }
+                }
+            }
+            tile_regs_release();
+            if constexpr (do_reduce) {
+                PACK((llk_pack_reconfig_l1_acc(0)));
+            }
+        }
+    }
+    if constexpr (do_reduce) {
+        cb_reduce.push_back(rows);
+    }
+
+    PACK((llk_pack_relu_config(ReluConfig::none())));
+}
+
+/**
+ * out_cb = in0_cb * in1_cb
+ * @tparam rows - Number of rows of tiles
+ * @tparam cols - Number of columns of tiles
+ * @tparam immediate_pop - If true, uses tile-by-tile processing with immediate CB pop after each tile.
+ *                         If false, uses batched processing with deferred CB pop, processing multiple tiles in
+ * parallel.
+ * @tparam pack_accumulate - If true, enables L1 accumulation to accumulate results onto existing tiles
+ *                           in out_cb. Only supported when immediate_pop=false.
+ */
+template <uint32_t rows, uint32_t cols, bool immediate_pop, bool pack_accumulate>
+void mul_block_bcast_cols(uint32_t in0_cb, uint32_t in1_cb, uint32_t out_cb) {
+    CircularBuffer cb_in0(in0_cb);
+    CircularBuffer cb_in1(in1_cb);
+    CircularBuffer cb_out(out_cb);
+    // Precondition: in0_cb has rows*cols produced
+    // Precondition: in1_cb has rows produced
+    // Precondition: out_cb has rows*cols produced
+    // Postcondition: in0_cb empty
+    // Postcondition: in1_cb empty
+    // Postcondition: out_cb has rows*cols produced
+
+    constexpr uint32_t num_tiles = rows * cols;
+
+    reconfig_data_format(in0_cb, in1_cb);
+    pack_reconfig_data_format(out_cb);
+    mul_bcast_cols_init(in0_cb, in1_cb);
+    cb_in0.wait_front(num_tiles);
+    cb_in1.wait_front(rows);
+
+    if constexpr (immediate_pop) {
+        static_assert(!pack_accumulate, "Unsupported parameter configuration");
+        for (uint32_t i = 0; i < rows; ++i) {
+            for (uint32_t j = 0; j < cols; ++j) {
+                tile_regs_acquire();
+                mul_tiles_bcast_cols(in0_cb, in1_cb, 0, i, 0);
+                tile_regs_commit();
+                cb_in0.pop_front(1);
+                cb_out.reserve_back(1);
+                tile_regs_wait();
+                pack_tile(0, out_cb);
+                tile_regs_release();
+                cb_out.push_back(1);
+            }
+        }
+        cb_in1.pop_front(rows);
+    } else {
+#ifdef DHT_GRANULARITY
+        constexpr uint32_t dst_tiles = (cols < DHT_GRANULARITY) ? cols : DHT_GRANULARITY;
+        constexpr uint32_t granularity = (cols >= DHT_GRANULARITY) ? (cols / DHT_GRANULARITY) : 1;
+#else
+        constexpr uint32_t dst_tiles = 1;
+        constexpr uint32_t granularity = cols;
+#endif
+        PACK((llk_pack_reconfig_l1_acc(pack_accumulate)));
+        if (!pack_accumulate) {
+            cb_out.reserve_back(num_tiles);
+        }
+        uint32_t in0_index = 0;
+        for (uint32_t i = 0; i < rows; ++i) {
+            for (uint32_t u = 0; u < granularity; ++u) {
+                tile_regs_acquire();
+                for (uint32_t j = 0; j < dst_tiles; ++j) {
+                    mul_tiles_bcast_cols(in0_cb, in1_cb, in0_index, i, j);
+                    in0_index++;
+                }
+                tile_regs_commit();
+                tile_regs_wait();
+                for (uint32_t j = 0; j < dst_tiles; ++j) {
+                    pack_tile(j, out_cb);
+                }
+                tile_regs_release();
+            }
+        }
+        cb_in1.pop_front(rows);
+        cb_in0.pop_front(num_tiles);
+        if (pack_accumulate) {
+            PACK((llk_pack_reconfig_l1_acc(false)));
+            cb_out.pop_front(num_tiles);
+            cb_out.reserve_back(num_tiles);
+            cb_out.push_back(num_tiles);
+        } else {
+            cb_out.push_back(num_tiles);
+        }
+    }
+}
+
+/**
+ * in0_cb *= in1_cb
+ */
+template <uint32_t rows, uint32_t cols>
+void mul_block_bcast_cols_inplace(uint32_t in0_cb, uint32_t in1_cb) {
+    CircularBuffer cb_in0(in0_cb);
+    CircularBuffer cb_in1(in1_cb);
+    // Precondition: in0_cb has rows*cols produced
+    // Precondition: in1_cb has rows produced
+    // Postcondition: in0_cb has rows*cols produced
+    // Postcondition: in1_cb has rows consumed
+
+    constexpr uint32_t num_tiles = rows * cols;
+
+#ifdef DHT_GRANULARITY
+    constexpr uint32_t dst_tiles = (cols < DHT_GRANULARITY) ? cols : DHT_GRANULARITY;
+    constexpr uint32_t granularity = (cols >= DHT_GRANULARITY) ? (cols / DHT_GRANULARITY) : 1;
+#else
+    constexpr uint32_t dst_tiles = 1;
+    constexpr uint32_t granularity = cols;
+#endif
+
+    reconfig_data_format(in0_cb, in1_cb);
+    mul_bcast_cols_init(in0_cb, in1_cb);
+    pack_reconfig_data_format(in0_cb);
+    cb_in0.wait_front(num_tiles);
+    cb_in1.wait_front(rows);
+    for (uint32_t i = 0; i < rows; ++i) {
+        for (uint32_t u = 0; u < granularity; ++u) {
+            tile_regs_acquire();
+            for (uint32_t j = 0; j < dst_tiles; ++j) {
+                mul_tiles_bcast_cols(in0_cb, in1_cb, j, i, j);
+            }
+            tile_regs_commit();
+            cb_in0.pop_front(dst_tiles);
+            cb_in0.reserve_back(dst_tiles);
+            tile_regs_wait();
+            for (uint32_t j = 0; j < dst_tiles; ++j) {
+                pack_tile(j, in0_cb);
+            }
+            cb_in0.push_back(dst_tiles);
+            tile_regs_release();
+        }
+    }
+    cb_in1.pop_front(rows);
+    reconfig_data_format_srcb(in0_cb);
+}
+
+template <uint32_t in1_scalar_cb, uint32_t num_tiles>
+void mul_block_bcast_scalar_inplace(uint32_t in0_cb) {
+    CircularBuffer cb_in0(in0_cb);
+    CircularBuffer cb_in1_scalar(in1_scalar_cb);
+    // Precondition: in0_cb has num_tiles produced
+    // Precondition: in1_scalar_cb has 1 produced
+    // Postcondition: in0_cb has num_tiles produced
+    // Postcondition: in1_scalar_cb has 1 produced
+
+#ifdef STATS_GRANULARITY
+    constexpr uint32_t dst_tiles = STATS_GRANULARITY;
+    constexpr uint32_t granularity = num_tiles / STATS_GRANULARITY;
+#else
+    constexpr uint32_t dst_tiles = 1;
+    constexpr uint32_t granularity = num_tiles;
+#endif
+
+    reconfig_data_format(in0_cb, in1_scalar_cb);
+    mul_bcast_scalar_init(in0_cb, in1_scalar_cb);
+    cb_in0.wait_front(num_tiles);
+    cb_in1_scalar.wait_front(1);
+    uint32_t in0_index = 0;
+    for (uint32_t g = 0; g < granularity; ++g) {
+        tile_regs_acquire();
+        for (uint32_t i = 0; i < dst_tiles; ++i) {
+            mul_tiles_bcast_scalar(in0_cb, in1_scalar_cb, in0_index, 0, i);
+            in0_index++;
+        }
+        tile_regs_commit();
+        tile_regs_wait();
+        for (uint32_t i = 0; i < dst_tiles; ++i) {
+            pack_tile(i, in0_cb);
+        }
+        tile_regs_release();
+    }
+    cb_in0.pop_front(num_tiles);
+    cb_in0.reserve_back(num_tiles);
+    cb_in0.push_back(num_tiles);
+}
+
+/**
+ * in0_cb += in1_cb
+ */
+template <bool pop_in1 = true>
+void add_block_inplace(uint32_t in0_cb, uint32_t in1_cb, uint32_t num_tiles) {
+    CircularBuffer cb_in0(in0_cb);
+    CircularBuffer cb_in1(in1_cb);
+    // Precondition: in0_cb and in1_cb have num_tiles produced
+    // Postcondition: in0_cb has num_tiles produced
+    // Postcondition: in1_cb has num_tiles consumed
+
+    reconfig_data_format(in0_cb, in1_cb);
+    pack_reconfig_data_format(in0_cb);
+    add_init(in0_cb, in1_cb);
+    cb_in0.wait_front(num_tiles);
+    cb_in1.wait_front(num_tiles);
+    for (uint32_t i = 0; i < num_tiles; i++) {
+        tile_regs_acquire();
+        add_tiles(in0_cb, in1_cb, i, i, 0);
+        tile_regs_commit();
+        tile_regs_wait();
+        pack_tile(0, in0_cb);
+        tile_regs_release();
+    }
+
+    cb_in0.pop_front(num_tiles);
+    if (pop_in1) {
+        cb_in1.pop_front(num_tiles);
+    }
+    cb_in0.reserve_back(num_tiles);
+    cb_in0.push_back(num_tiles);
+}
+
+void mul_tiles_bcast_cols_inplace(uint32_t in0_cb, uint32_t in1_cb, uint32_t num_tiles) {
+    CircularBuffer cb_in0(in0_cb);
+    CircularBuffer cb_in1(in1_cb);
+    /**
+     * Given in0_cb and in1_cb, multiply each tile of in0_cb by the corresponding tile of in1_cb
+     * and bcast cols of in1_cb.
+     */
+    // Precondition: in0_cb and in1_cb have num_tiles produced
+    // Postcondition: in0_cb has num_tiles produced
+    // Postcondition: in1_cb has num_tiles produced
+
+    reconfig_data_format(in0_cb, in1_cb);
+    mul_bcast_cols_init(in0_cb, in1_cb);
+    pack_reconfig_data_format(in0_cb);
+    cb_in0.wait_front(num_tiles);
+    cb_in1.wait_front(num_tiles);
+    for (uint32_t i = 0; i < num_tiles; i++) {
+        tile_regs_acquire();
+        mul_tiles_bcast_cols(in0_cb, in1_cb, 0, i, 0);
+        tile_regs_commit();
+        cb_in0.pop_front(1);
+        cb_in0.reserve_back(1);
+        tile_regs_wait();
+        pack_tile(0, in0_cb);
+        tile_regs_release();
+        cb_in0.push_back(1);
+    }
+}
+
+/**
+ * in0_cb *= in1_cb
+ */
+void mul_block_inplace(uint32_t in0_cb, uint32_t in1_cb, uint32_t num_tiles) {
+    CircularBuffer cb_in0(in0_cb);
+    CircularBuffer cb_in1(in1_cb);
+    // Precondition: in0_cb and in1_cb have num_tiles produced
+    // Postcondition: in0_cb has num_tiles produced
+    // Postcondition: in1_cb has num_tiles produced
+
+    mul_init(in0_cb, in1_cb);
+    cb_in0.wait_front(num_tiles);
+    cb_in1.wait_front(num_tiles);
+    for (uint32_t i = 0; i < num_tiles; i++) {
+        invalidate_l1_cache();
+        tile_regs_acquire();
+        mul_tiles(in0_cb, in1_cb, 0, i, 0);
+        tile_regs_commit();
+        cb_in0.pop_front(1);
+        cb_in0.reserve_back(1);
+        tile_regs_wait();
+        pack_tile(0, in0_cb);
+        tile_regs_release();
+        cb_in0.push_back(1);
+    }
+}
+
+#if defined(TRISC_MATH) || defined(TRISC_PACK)
+
+template <bool SDPA_EXP_APPROX_MODE, uint16_t scale_bf16>
+void exp_tile_first_column(uint32_t idst) {
+    SFPU_UNARY_CALL(
+        DST_SYNC_MODE,
+        DST_ACCUM_MODE,
+        calculate_exponential_first_column,
+        (SDPA_EXP_APPROX_MODE, scale_bf16),
+        idst,
+        VectorMode::C);
+}
+#endif  // defined(TRISC_MATH) || defined(TRISC_PACK)
+
+/**
+ * out_cb = exp((in0_cb - in1_cb) * scale_fp32)
+ */
+template <uint32_t scale_fp32>
+void sub_exp_block(uint32_t in0_cb, uint32_t in1_cb, uint32_t out_cb, uint32_t num_tiles) {
+    CircularBuffer cb_in0(in0_cb);
+    CircularBuffer cb_in1(in1_cb);
+    CircularBuffer cb_out(out_cb);
+    // Precondition: in0_cb and in1_cb have num_tiles produced
+    // Postcondition: out_cb has num_tiles produced
+    // Postcondition: in0_cb and in1_cb has num_tiles produced
+
+    sub_init(in0_cb, in1_cb);
+    exp_tile_init<EXP_APPROX_MODE>();
+    cb_in0.wait_front(num_tiles);
+    cb_in1.wait_front(num_tiles);
+    cb_out.reserve_back(num_tiles);
+
+    // Convert scale_fp32 to bf16 scale
+    constexpr uint16_t scale_bf16 = scale_fp32 >> 16;
+
+    for (uint32_t i = 0; i < num_tiles; i++) {
+        invalidate_l1_cache();
+        tile_regs_acquire();
+        sub_tiles(in0_cb, in1_cb, i, i, 0);
+        MATH((exp_tile_first_column<EXP_APPROX_MODE, scale_bf16>(0)));
+        tile_regs_commit();
+        tile_regs_wait();
+        pack_tile(0, out_cb);
+        tile_regs_release();
+        cb_out.push_back(1);
+    }
+}
+
+#ifdef TRISC_MATH
+template <VectorMode vector_mode = VectorMode::C>
+void fused_max_sub_exp_add_tile(uint32_t idst, int scale_bf16) {
+    SFPU_UNARY_CALL_NO_TEMPLATE_ARGS(
+        DST_SYNC_MODE, DST_ACCUM_MODE, calculate_fused_max_sub_exp_add_tile, idst, vector_mode, scale_bf16);
+}
+#endif
+
+template <uint32_t scale_fp32, VectorMode vector_mode = VectorMode::C>
+void correction_block(
+    uint32_t cb_worker_max,
+    uint32_t cb_worker_sum,
+    uint32_t cb_cur_max,
+    uint32_t cb_prev_max,
+    uint32_t cb_cur_sum,
+    uint32_t cb_prev_sum,
+    uint32_t cb_exp_max_diff,
+    uint32_t cb_exp_max_diff_2,
+    uint32_t num_head_tiles) {
+    CircularBuffer cb_worker_max_obj(cb_worker_max);
+    CircularBuffer cb_worker_sum_obj(cb_worker_sum);
+    CircularBuffer cb_cur_max_obj(cb_cur_max);
+    CircularBuffer cb_prev_max_obj(cb_prev_max);
+    CircularBuffer cb_cur_sum_obj(cb_cur_sum);
+    CircularBuffer cb_prev_sum_obj(cb_prev_sum);
+    CircularBuffer cb_exp_max_diff_obj(cb_exp_max_diff);
+    CircularBuffer cb_exp_max_diff_2_obj(cb_exp_max_diff_2);
+    cb_worker_max_obj.wait_front(num_head_tiles);
+    cb_worker_sum_obj.wait_front(num_head_tiles);
+    cb_prev_max_obj.wait_front(num_head_tiles);
+    cb_prev_sum_obj.wait_front(num_head_tiles);
+
+    cb_cur_max_obj.reserve_back(num_head_tiles);
+    cb_cur_sum_obj.reserve_back(num_head_tiles);
+    cb_exp_max_diff_obj.reserve_back(num_head_tiles);
+    cb_exp_max_diff_2_obj.reserve_back(num_head_tiles);
+
+    constexpr uint32_t dst_reg_0 = 0;  // dst_reg_0 is used for prev_max
+    constexpr uint32_t dst_reg_1 = 1;  // dst_reg_1 is used for worker_max
+    constexpr uint32_t dst_reg_2 = 2;  // dst_reg_2 is used for cur_max
+    constexpr uint32_t dst_reg_3 = 3;  // dst_reg_3 is used for prev_sum, returns cur_sum
+    constexpr uint32_t dst_reg_4 = 4;  // dst_reg_4 is used for worker_sum
+
+    // convert scale from fp32 to bf16
+    constexpr uint16_t scale_bf16 = scale_fp32 >> 16;
+
+    for (uint32_t i = 0; i < num_head_tiles; i++) {
+        tile_regs_acquire();
+        copy_tile_to_dst_init_short(cb_worker_max);
+        exp_tile_init<EXP_APPROX_MODE>();
+        copy_tile(cb_prev_max, i, dst_reg_0);
+        copy_tile(cb_worker_max, i, dst_reg_1);
+        copy_tile(cb_prev_sum, i, dst_reg_3);
+        copy_tile(cb_worker_sum, i, dst_reg_4);
+        MATH((fused_max_sub_exp_add_tile<vector_mode>(0, scale_bf16)));
+        tile_regs_commit();
+        tile_regs_wait();
+        pack_tile(dst_reg_0, cb_exp_max_diff);
+        pack_tile(dst_reg_1, cb_exp_max_diff_2);
+        pack_tile(dst_reg_2, cb_cur_max);
+        pack_tile(dst_reg_3, cb_cur_sum);
+        tile_regs_release();
+        cb_cur_max_obj.push_back(1);
+        cb_cur_sum_obj.push_back(1);
+        cb_exp_max_diff_obj.push_back(1);
+        cb_exp_max_diff_2_obj.push_back(1);
+    }
+    cb_prev_sum_obj.pop_front(num_head_tiles);
+    cb_worker_sum_obj.pop_front(num_head_tiles);
+}
+
+/**
+ * in_cb -> out_cb
+ */
+template <bool pop_in_cb>
+void move_block(uint32_t in_cb, uint32_t out_cb, uint32_t num_tiles) {
+    CircularBuffer cb_in(in_cb);
+    CircularBuffer cb_out(out_cb);
+    // Precondition: in_cb has num_tiles produced
+    // Precondition: out_cb has num_tiles free
+    // Postcondition: in_cb has num_tiles consumed
+    // Postcondition: out_cb has num_tiles produced
+
+    copy_tile_to_dst_init_short(in_cb);
+
+    cb_in.wait_front(num_tiles);
+    cb_out.reserve_back(num_tiles);
+
+#pragma GCC unroll 0
+    for (uint32_t i = 0; i < num_tiles; i++) {
+        tile_regs_acquire();
+        copy_tile(in_cb, i, 0 /*dst*/);
+        tile_regs_commit();
+        tile_regs_wait();
+        pack_tile(0, out_cb);
+        tile_regs_release();
+        cb_out.push_back(1);
+    }
+    if (pop_in_cb) {
+        cb_in.pop_front(num_tiles);
+    }
+}
+
+void copy_block(uint32_t in_cb, uint32_t out_cb, uint32_t num_tiles) {
+    CircularBuffer cb_in(in_cb);
+    CircularBuffer cb_out(out_cb);
+    // Precondition: in_cb has num_tiles produced
+    // Precondition: out_cb has num_tiles free
+    // Postcondition: in_cb has num_tiles consumed
+    // Postcondition: out_cb has num_tiles produced
+    copy_tile_to_dst_init_short(in_cb);
+    cb_in.wait_front(num_tiles);
+    cb_out.reserve_back(num_tiles);
+#pragma GCC unroll 0
+    for (uint32_t i = 0; i < num_tiles; i++) {
+        tile_regs_acquire();
+        copy_tile(in_cb, i, 0 /*dst*/);
+        tile_regs_commit();
+        tile_regs_wait();
+        pack_tile(0, out_cb);
+        tile_regs_release();
+        cb_out.push_back(1);
+    }
+    cb_in.pop_front(num_tiles);
+}
+
+void log_block(uint32_t in_cb, uint32_t out_cb, uint32_t num_tiles) {
+    reconfig_data_format_srca(in_cb);
+    pack_reconfig_data_format(out_cb);
+    CircularBuffer cb_in(in_cb);
+    CircularBuffer cb_out(out_cb);
+    copy_tile_to_dst_init_short(in_cb);
+    log_tile_init();
+    cb_in.wait_front(num_tiles);
+    cb_out.reserve_back(num_tiles);
+
+    for (uint32_t i = 0; i < num_tiles; i++) {
+        tile_regs_acquire();
+        copy_tile(in_cb, i, 0 /*dst*/);
+        log_tile(0);
+        tile_regs_commit();
+        tile_regs_wait();
+        pack_tile(0, out_cb);
+        tile_regs_release();
+        cb_out.push_back(1);
+    }
+}
+
+void sigmoid_sub(uint32_t in0_cb, uint32_t in1_cb, uint32_t out_cb, uint32_t num_tiles) {
+    CircularBuffer cb_in0(in0_cb);
+    CircularBuffer cb_in1(in1_cb);
+    CircularBuffer cb_out(out_cb);
+    // out_cb = sigmoid(in0_cb - in1_cb)
+    /**
+     * sigmoid(x) is accurately implemented as 1 / (1 + exp(-x))
+     * This function manually implements the composite, accurate sigmoid.
+     *
+     * Each input tile has only the first column containing valid data, so VectorMode::C is a useful optimization.
+     */
+    cb_in0.wait_front(num_tiles);
+    cb_in1.wait_front(num_tiles);
+    cb_out.reserve_back(num_tiles);
+    sub_init(in0_cb, in1_cb);
+    exp_tile_init<false>();
+    // recip_tile_first_column<false>() calls the scalar sfpu_reciprocal_iter path, so initialize exactly
+    // that SFPU state here. Blackhole needs vConstFloatPrgm0 = 2.0 for Newton-Raphson; Wormhole
+    // needs vConstFloatPrgm0/1/2 loaded with reciprocal polynomial coefficients.
+    // This init programs persistent SFPU constants, not per-tile data. It intentionally comes after
+    // exp_tile_init<false>() because the exp call below is the custom exp_tile_first_column<false>(),
+    // whose polynomial implementation loads the constants it consumes into LREGs in the tile body; it
+    // does not depend on vConstFloatPrgm* state from exp_tile_init. Conversely, that exp body also
+    // does not clobber the reciprocal constants, so one reciprocal init before the tile loop is enough.
+    MATH((ckernel::sfpu::sfpu_reciprocal_init<false>()));
+
+    for (uint32_t i = 0; i < num_tiles; i++) {
+        tile_regs_acquire();
+        sub_tiles(in0_cb, in1_cb, i, i, 0);
+        // exp_tile<false, true /*SCALE_EN*/>(0, (int)VectorMode::C, (uint16_t)0xBF80 /*bf16(-1.0) scale*/);
+        MATH((exp_tile_first_column<false /*APPROX_MODE*/, (uint16_t)0xBF80 /*bf16(-1.0) scale*/>(0)));
+        // add_unary_tile(0 /*dst_index*/, 0x3F800000); // Call the macro directly to get access to VectorMode argument
+        MATH(SFPU_UNARY_CALL(
+            DST_SYNC_MODE,
+            DST_ACCUM_MODE,
+            calculate_binop_with_scalar,
+            (APPROX, ADD_UNARY, 8 /* ITERATIONS */),
+            0 /*dst_index*/,
+            VectorMode::C,
+            0x3F800000 /*scalar*/));
+        // recip_tile<false>(0, (int)VectorMode::C);
+        MATH((recip_tile_first_column<false>(0 /*dst_index*/)));
+        tile_regs_commit();
+        tile_regs_wait();
+        pack_tile(0, out_cb);
+        tile_regs_release();
+    }
+    cb_out.push_back(num_tiles);
+}
+
+#ifdef TRISC_MATH
+void softplus_tile_first_column(uint32_t idst, uint beta, uint beta_reciprocal, uint threshold) {
+    SFPU_UNARY_CALL_NO_TEMPLATE_ARGS(
+        DST_SYNC_MODE,
+        DST_ACCUM_MODE,
+        calculate_softplus_first_column,
+        idst,
+        VectorMode::C,
+        beta,
+        beta_reciprocal,
+        threshold);
+}
+#endif
+
+void logsigmoid_sub(uint32_t in0_cb, uint32_t in1_cb, uint32_t out_cb, uint32_t num_tiles) {
+    CircularBuffer cb_in0(in0_cb);
+    CircularBuffer cb_in1(in1_cb);
+    CircularBuffer cb_out(out_cb);
+    // out_cb = logsigmoid(in0_cb - in1_cb)
+    // Implemented as softplus for numerical stability. logsigmoid(x) = -softplus(-x)
+    cb_in0.wait_front(num_tiles);
+    cb_in1.wait_front(num_tiles);
+    cb_out.reserve_back(num_tiles);
+    sub_init(in0_cb, in1_cb);
+    softplus_tile_init();
+    constexpr uint32_t const_1_fp32 = 0x3F800000;
+    constexpr uint32_t const_20_fp32 = 0x41A00000;
+
+    for (uint32_t i = 0; i < num_tiles; i++) {
+        tile_regs_acquire();
+        // Negate input to softplus by swapping inputs to sub
+        sub_tiles(in1_cb, in0_cb, i, i, 0);
+        // softplus_tile(0, 0x3F800000, 0x3F800000, 0x41A00000);  // beta, beta_reciprocal, threshold
+        // MATH((llk_math_eltwise_unary_sfpu_softplus<APPROX>(
+        //     0,
+        //     const_1_fp32 /*beta*/,
+        //     const_1_fp32 /*beta_reciprocal*/,
+        //     const_20_fp32 /*threshold*/,
+        //     (int)VectorMode::C)));
+
+        MATH((softplus_tile_first_column(0, const_1_fp32, const_1_fp32, const_20_fp32)));
+        // Negate the output of softplus
+        negative_tile(0);
+        tile_regs_commit();
+        tile_regs_wait();
+        pack_tile(0, out_cb);
+        tile_regs_release();
+    }
+    cb_out.push_back(num_tiles);
+}
+
+/**
+ * out_cb = in0_cb - in1_cb
+ * Compile with size optimization to prevent binary size exceeding the limit.
+ */
+__attribute__((optimize("Os"))) void sub_block(uint32_t in0_cb, uint32_t in1_cb, uint32_t out_cb, uint32_t num_tiles) {
+    CircularBuffer cb_in0(in0_cb);
+    CircularBuffer cb_in1(in1_cb);
+    CircularBuffer cb_out(out_cb);
+    cb_in0.wait_front(num_tiles);
+    cb_in1.wait_front(num_tiles);
+    cb_out.reserve_back(num_tiles);
+    sub_init(in0_cb, in1_cb);
+
+    for (uint32_t i = 0; i < num_tiles; i++) {
+        tile_regs_acquire();
+        sub_tiles(in0_cb, in1_cb, i, i, 0);
+        tile_regs_commit();
+        tile_regs_wait();
+        pack_tile(0, out_cb);
+        tile_regs_release();
+    }
+    cb_out.push_back(num_tiles);
+}
+
+/**
+ * out_cb = in0_cb @ in1_cb
+ */
+ALWI void matmul_blocks(
+    const uint32_t& in0_cb,
+    const uint32_t& in1_cb,
+    const uint32_t& out_cb,
+    const uint32_t& M,
+    const uint32_t& N,
+    const uint32_t& K,
+    const uint32_t& num_blocks,
+    const uint32_t& in0_num_subblocks,
+    const uint32_t& in1_num_subblocks,
+    const uint32_t& in0_block_w,
+    const uint32_t& subblock_h,
+    const uint32_t& subblock_w,
+    const bool& transpose,
+    const bool& add_mask = false,
+    const uint32_t& mask_cb = 0,
+    const uint32_t& zero_cb = 0) {
+    // precondition: in0_cb has M*K produced
+    // precondition: in1_cb has K*N produced
+    // postcondition: in0_cb is full, in1_cb is empty
+    // postcondition: out_cb has M*N produced
+
+    CircularBuffer cb_in0(in0_cb);
+    CircularBuffer cb_in1(in1_cb);
+    CircularBuffer cb_out(out_cb);
+    CircularBuffer cb_mask(mask_cb);
+    CircularBuffer cb_zero(zero_cb);
+
+    matmul_block_init(
+        in0_cb, in1_cb, transpose /*transpose*/, subblock_w /*ct_dim*/, subblock_h /*rt_dim*/, in0_block_w /*kt_dim*/);
+
+    const uint32_t output_num_tiles = M * N;
+    const uint32_t out_subblock_num_tiles = subblock_h * subblock_w;
+    const uint32_t in0_subblock_all_cols_num_tiles = subblock_h * N;
+
+    uint32_t in0_index_offset = 0;
+
+    const uint32_t in0_subblock_num_tiles = subblock_h * in0_block_w;
+    uint32_t in0_wait_tiles = in0_subblock_num_tiles;
+
+    reconfig_data_format(in1_cb, in0_cb);
+    cb_in1.wait_front(K * N);
+    cb_out.reserve_back(output_num_tiles);
+
+    for (uint32_t in0_subblock = 0; in0_subblock < in0_num_subblocks; ++in0_subblock) {
+        cb_in0.wait_front(in0_wait_tiles);
+        uint32_t in1_index_offset = 0;
+        for (uint32_t in1_subblock = 0; in1_subblock < in1_num_subblocks; ++in1_subblock) {
+            tile_regs_acquire();
+
+            uint32_t dst_index = 0;
+            uint32_t in0_index = in0_index_offset;
+            uint32_t in1_index = in1_index_offset;
+
+            for (uint32_t inner_dim = 0; inner_dim < in0_block_w; inner_dim++) {
+                matmul_block(
+                    in0_cb, in1_cb, in0_index, in1_index, dst_index, transpose, subblock_w, subblock_h, in0_block_w);
+                in0_index++;
+                in1_index += N;
+            }
+            if (add_mask) {
+                cb_mask.wait_front(out_subblock_num_tiles);
+                cb_zero.wait_front(1);
+                reconfig_data_format(zero_cb, mask_cb);
+                add_init(zero_cb, mask_cb, true);
+                for (uint32_t i = 0; i < out_subblock_num_tiles; i++) {
+                    add_tiles(zero_cb, mask_cb, 0, i, i);
+                }
+                reconfig_data_format(in1_cb, in0_cb);
+                matmul_block_init(in0_cb, in1_cb, transpose, subblock_w, subblock_h, in0_block_w);
+            }
+            tile_regs_commit();
+            tile_regs_wait();
+            uint32_t dst_idx = 0;
+            uint32_t out_col_offset = in1_subblock * subblock_w;
+            for (uint32_t r = 0; r < subblock_h; r++) {
+                uint32_t out_row_offset = r * N;
+                for (uint32_t c = 0; c < subblock_w; c++) {
+                    pack_tile<true>(dst_idx, out_cb, out_row_offset + out_col_offset + c);
+                    dst_idx++;
+                }
+            }
+            tile_regs_release();
+            in1_index_offset += subblock_w;
+        }
+        in0_index_offset += subblock_h * in0_block_w;
+        in0_wait_tiles += in0_subblock_num_tiles;
+        // Somewhat granularize the push of in0 subblocks
+        cb_out.push_back(in0_subblock_all_cols_num_tiles);
+    }
+    cb_in1.pop_front(K * N);
+}
+
+template <uint32_t M>
+void matmul_reduce(uint32_t in1_cb, const uint32_t& out_cb) {
+    CircularBuffer cb_in1(in1_cb);
+    CircularBuffer cb_out(out_cb);
+    // precondition: in0_cb has M*K produced
+    // precondition: in1_cb has K*N produced
+    // postcondition: in0_cb is full, in1_cb is empty
+    // postcondition: out_cb has M*N produced
+
+    constexpr uint32_t N = 1;  // Result of reduce is 1 column
+    constexpr uint32_t in0_block_w = N;
+    constexpr uint32_t subblock_w = N;
+    // Reuse the Sq_chunk_t granularity chosen for sub_exp_block
+#ifdef STATS_GRANULARITY
+    constexpr uint32_t subblock_h = STATS_GRANULARITY;
+    constexpr uint32_t in0_num_subblocks = M / STATS_GRANULARITY;
+#else
+    constexpr uint32_t subblock_h = 1;
+    constexpr uint32_t in0_num_subblocks = M;
+#endif
+
+    /**
+     * Use matmul on Mx1 input to reduce rows within tile to produce Mx1 output.
+     */
+
+    // matmul_block_init validates the live reverse-order unpacker setup
+    // (in1_cb -> SrcA, out_cb -> SrcB), so establish it before init.
+    reconfig_data_format(in1_cb, out_cb);
+    matmul_block_init(
+        out_cb, in1_cb, 0 /*transpose*/, subblock_w /*ct_dim*/, subblock_h /*rt_dim*/, in0_block_w /*kt_dim*/);
+
+    constexpr uint32_t output_num_tiles = M * N;
+    constexpr uint32_t out_subblock_num_tiles = subblock_h * subblock_w;
+
+    pack_reconfig_data_format(out_cb);
+    cb_in1.wait_front(N);
+    cb_out.wait_front(M);
+
+    for (uint32_t in0_subblock = 0; in0_subblock < in0_num_subblocks; ++in0_subblock) {
+        tile_regs_acquire();
+
+        uint32_t dst_index = 0;
+        uint32_t in0_index = 0;
+        uint32_t in1_index = 0;
+
+        matmul_block(out_cb, in1_cb, in0_index, in1_index, dst_index, 0, subblock_w, subblock_h, in0_block_w);
+
+        tile_regs_commit();
+        cb_out.pop_front(subblock_h);
+
+        tile_regs_wait();
+        for (uint32_t i = 0; i < subblock_h; i++) {
+            pack_tile(i, out_cb);
+        }
+        tile_regs_release();
+        cb_out.push_back(subblock_h);
+    }
+}
+
+/**
+ * Batch-stamp a single tile onto a range of positions in out_cb using L1 accumulate.
+ * Caller must have already called copy_tile_to_dst_init_short and llk_pack_reconfig_l1_acc(1).
+ *
+ * @tparam dst_batch  Max tiles per DST cycle (DST register capacity, typically 8 for fp16b half-sync).
+ */
+template <uint32_t dst_batch>
+void stamp_tile_range_l1_acc(
+    uint32_t src_cb, uint32_t src_tile_idx, uint32_t out_cb, uint32_t out_offset, uint32_t count) {
+    for (uint32_t base = 0; base < count; base += dst_batch) {
+        uint32_t batch = (count - base < dst_batch) ? (count - base) : dst_batch;
+        tile_regs_acquire();
+        for (uint32_t i = 0; i < batch; i++) {
+            copy_tile(src_cb, src_tile_idx, i);
+        }
+        tile_regs_commit();
+        tile_regs_wait();
+        for (uint32_t i = 0; i < batch; i++) {
+            pack_tile<true>(i, out_cb, out_offset + base + i);
+        }
+        tile_regs_release();
+    }
+}
+
+template <uint32_t dst_batch>
+void apply_padded_mask_lightweight_runtime(
+    uint32_t neginf_cb,
+    uint32_t neginf_tile_idx,
+    uint32_t out_cb,
+    uint32_t num_padded,
+    uint32_t num_cols,
+    uint32_t num_rows,
+    uint32_t row_base = 0) {  // first out_cb tile-row of this query band; nonzero when heads span >1 DEST band
+    uint32_t start = num_cols - num_padded;
+
+    reconfig_data_format_srca(neginf_cb);
+    pack_reconfig_data_format(out_cb);
+    copy_tile_to_dst_init_short(neginf_cb);
+    PACK((llk_pack_reconfig_l1_acc(1)));
+
+    for (uint32_t row = 0; row < num_rows; row++) {
+        stamp_tile_range_l1_acc<dst_batch>(
+            neginf_cb, neginf_tile_idx, out_cb, (row_base + row) * num_cols + start, num_padded);
+    }
+
+    PACK((llk_pack_reconfig_l1_acc(0)));
+}
+
+/**
+ * Lightweight partial mask: L1-accumulate a partial mask tile (0 for valid, -inf for padded columns)
+ * onto the boundary tile position in out_cb. The partial tile is permanently fronted in the CB.
+ *
+ * @param mask_cb          CB holding mask tiles, permanently fronted
+ * @param partial_tile_idx Index of the partial tile within the CB
+ * @param out_cb           QK intermediate CB (already wait-fronted)
+ * @param boundary_col     Column index within the chunk where the boundary tile is
+ * @param num_cols         Total K tiles per row (Sk_chunk_t)
+ * @param num_rows         Q tiles per chunk (Sq_chunk_t)
+ */
+void apply_partial_mask_lightweight(
+    uint32_t mask_cb,
+    uint32_t partial_tile_idx,
+    uint32_t out_cb,
+    uint32_t boundary_col,
+    uint32_t num_cols,
+    uint32_t num_rows,
+    uint32_t row_base = 0) {  // first out_cb tile-row of this query band; nonzero when heads span >1 DEST band
+    reconfig_data_format_srca(mask_cb);
+    pack_reconfig_data_format(out_cb);
+    copy_tile_to_dst_init_short(mask_cb);
+    PACK((llk_pack_reconfig_l1_acc(1)));
+
+    for (uint32_t row = 0; row < num_rows; row++) {
+        tile_regs_acquire();
+        copy_tile(mask_cb, partial_tile_idx, 0);
+        tile_regs_commit();
+        tile_regs_wait();
+        pack_tile<true>(0, out_cb, (row_base + row) * num_cols + boundary_col);
+        tile_regs_release();
+    }
+
+    PACK((llk_pack_reconfig_l1_acc(0)));
+}
+
+/**
+ * Lightweight causal mask: stamps neginf and diagonal tiles onto QKT using L1 accumulate.
+ *
+ * For Q tile-row i (0..num_rows-1) processing K chunk starting at k_start_tile:
+ *   diag_col = q_start_tile + i - k_start_tile
+ *   - diag_col < 0:           entire row above diagonal -> stamp neginf on all num_cols tiles
+ *   - 0 <= diag_col < num_cols: diagonal tile at col diag_col, neginf at cols diag_col+1..num_cols-1
+ *   - diag_col >= num_cols:     entire row below diagonal -> no mask needed
+ */
+template <uint32_t dst_batch>
+void apply_causal_mask_lightweight(
+    uint32_t mask_cb,
+    uint32_t neginf_idx,
+    uint32_t diag_idx,
+    uint32_t out_cb,
+    uint32_t q_start_tile,
+    uint32_t k_start_tile,
+    uint32_t num_rows,
+    uint32_t num_cols,
+    uint32_t straddle_col = 0,
+    uint32_t straddle_jump = 0) {
+    reconfig_data_format_srca(mask_cb);
+    pack_reconfig_data_format(out_cb);
+    copy_tile_to_dst_init_short(mask_cb);
+    PACK((llk_pack_reconfig_l1_acc(1)));
+
+    for (uint32_t row = 0; row < num_rows; row++) {
+        uint32_t row_offset = row * num_cols;
+        const int32_t q_pos = static_cast<int32_t>(q_start_tile + row);
+
+        if (straddle_col == 0) {
+            // Fast path: K coords contiguous across cols.
+            int32_t diag_col = q_pos - static_cast<int32_t>(k_start_tile);
+            if (diag_col < 0) {
+                // Entire row above diagonal -> stamp all neginf
+                stamp_tile_range_l1_acc<dst_batch>(mask_cb, neginf_idx, out_cb, row_offset, num_cols);
+            } else if (static_cast<uint32_t>(diag_col) < num_cols) {
+                // Stamp the diagonal tile
+                tile_regs_acquire();
+                copy_tile(mask_cb, diag_idx, 0);
+                tile_regs_commit();
+                tile_regs_wait();
+                pack_tile<true>(0, out_cb, row_offset + static_cast<uint32_t>(diag_col));
+                tile_regs_release();
+
+                // Stamp neginf tiles to the right of diagonal
+                uint32_t neginf_start = static_cast<uint32_t>(diag_col) + 1;
+                if (neginf_start < num_cols) {
+                    stamp_tile_range_l1_acc<dst_batch>(
+                        mask_cb, neginf_idx, out_cb, row_offset + neginf_start, num_cols - neginf_start);
+                }
+            }
+            // else: diag_col >= num_cols -> entire row below diagonal, no mask needed
+        } else {
+            // Chunked-prefill straddle: K coord jumps by straddle_jump at col >= straddle_col
+            // (the K-chunk crosses a slab boundary). Evaluate per-col.
+            for (uint32_t col = 0; col < num_cols; col++) {
+                int32_t k_pos = static_cast<int32_t>(k_start_tile) + static_cast<int32_t>(col);
+                if (col >= straddle_col) {
+                    k_pos += static_cast<int32_t>(straddle_jump);
+                }
+                if (k_pos > q_pos) {
+                    stamp_tile_range_l1_acc<dst_batch>(mask_cb, neginf_idx, out_cb, row_offset + col, 1);
+                } else if (k_pos == q_pos) {
+                    tile_regs_acquire();
+                    copy_tile(mask_cb, diag_idx, 0);
+                    tile_regs_commit();
+                    tile_regs_wait();
+                    pack_tile<true>(0, out_cb, row_offset + col);
+                    tile_regs_release();
+                }
+            }
+        }
+    }
+
+    PACK((llk_pack_reconfig_l1_acc(0)));
+}
+
+/**
+ * Context for lightweight mask application.
+ * All mask tiles reside in a single CB. This struct stores the pre-resolved mask metadata used when
+ * lightweight masking is enabled; enablement itself is controlled by the `lightweight_mask_enabled`
+ * template parameter(s), not by default-constructing this context.
+ */
+struct LightweightMaskContext {
+    bool is_causal = false;                       // Causal masking active for this context instance
+    uint32_t neginf_tile_idx = 0;                 // Index of -inf tile in the mask CB
+    uint32_t causal_diag_tile_idx = 0;            // Index of causal diagonal tile in the mask CB
+    uint32_t primary_diag_tile_idx = 0;           // Causal diagonal, or sliding-window trailing-primary tile
+    uint32_t sliding_leading_prev_tile_idx = 0;   // Index of previous sliding-window leading tile
+    uint32_t sliding_leading_tile_idx = 0;        // Index of current sliding-window leading tile in the mask CB
+    uint32_t sliding_trailing_next_tile_idx = 0;  // Index of next sliding-window trailing tile
+    uint32_t global_n_padded_tiles = 0;           // Fully padded K tile columns for global_n chunk
+    uint32_t local_n_padded_tiles = 0;            // Fully padded K tile columns for local_n chunk
+    uint32_t joint_n_padded_tiles = 0;            // Fully padded K tile columns for joint_l chunk
+    uint32_t global_n_partial_col = 0;            // Column within tile where global_n padding starts (0 = no partial)
+    uint32_t joint_l_partial_col = 0;             // Column within tile where joint_l padding starts (0 = no partial)
+    uint32_t global_n_partial_tile_idx = 0;       // Index of global_n partial tile in the mask CB
+    uint32_t joint_l_partial_tile_idx = 0;        // Index of joint_l partial tile in the mask CB
+    uint32_t straddle_num_padded_tiles = 0;       // Trailing -inf tiles on straddle chunk (0 = inactive)
+    uint32_t straddle_mask_chunk_id = 0;          // K chunk index where straddle mask applies
+
+    /**
+     * Resolve which mask type applies for a given K chunk and return pre-resolved params.
+     * Called once per K chunk; the resolved values are then used per subblock.
+     */
+    void resolve_for_chunk(
+        uint32_t Sk_chunk_t,
+        uint32_t k_chunk,
+        uint32_t num_local_k_chunks,
+        bool ring_iter_needs_global_n_mask,
+        bool ring_iter_needs_joint_n_mask,
+        bool local_n_needs_masking,
+        uint32_t global_n_mask_chunk_id,
+        uint32_t local_n_mask_chunk_id,
+        uint32_t joint_n_mask_chunk_id,
+        uint32_t& out_num_padded,
+        uint32_t& out_boundary_col,
+        uint32_t& out_partial_tile_idx,
+        bool& out_has_partial) const {
+        out_num_padded = 0;
+        out_has_partial = false;
+        out_boundary_col = 0;
+        out_partial_tile_idx = 0;
+
+        if (ring_iter_needs_global_n_mask && k_chunk == global_n_mask_chunk_id) {
+            out_num_padded = global_n_padded_tiles;
+            out_has_partial = (global_n_partial_col > 0);
+            out_boundary_col = Sk_chunk_t - out_num_padded - (out_has_partial ? 1 : 0);
+            out_partial_tile_idx = global_n_partial_tile_idx;
+        } else if (local_n_needs_masking && k_chunk == local_n_mask_chunk_id) {
+            out_num_padded = local_n_padded_tiles;
+        } else if (ring_iter_needs_joint_n_mask && (k_chunk - num_local_k_chunks) == joint_n_mask_chunk_id) {
+            out_num_padded = joint_n_padded_tiles;
+            out_has_partial = (joint_l_partial_col > 0);
+            out_boundary_col = Sk_chunk_t - out_num_padded - (out_has_partial ? 1 : 0);
+            out_partial_tile_idx = joint_l_partial_tile_idx;
+        } else if (straddle_num_padded_tiles > 0 && k_chunk == straddle_mask_chunk_id) {
+            out_num_padded = straddle_num_padded_tiles;
+        }
+    }
+
+    /**
+     * Apply the lightweight mask for a given K chunk (non-streaming path).
+     * When is_causal: applies causal stamp, then also applies padding stamp if this K chunk has padding.
+     * When !is_causal: applies padding stamp only.
+     * L1-accumulate stamps are additive, so applying both is safe (-inf + -inf = -inf, -inf + 0 = -inf).
+     *
+     * @tparam dst_size  DST register capacity (8 for fp16b half-sync, 4 for fp32_dest_acc).
+     */
+    template <uint32_t dst_size>
+    void apply(
+        uint32_t cb_mask_in,
+        uint32_t cb_qk_im,
+        uint32_t Sk_chunk_t,
+        uint32_t Sq_chunk_t,
+        uint32_t k_chunk,
+        uint32_t num_local_k_chunks,
+        bool ring_iter_needs_global_n_mask,
+        bool ring_iter_needs_joint_n_mask,
+        bool local_n_needs_masking,
+        uint32_t global_n_mask_chunk_id,
+        uint32_t local_n_mask_chunk_id,
+        uint32_t joint_n_mask_chunk_id,
+        uint32_t q_start_tile,
+        uint32_t k_start_tile,
+        uint32_t straddle_col = 0,
+        uint32_t straddle_jump = 0) const {
+        if (is_causal) {
+            apply_causal_mask_lightweight<dst_size>(
+                cb_mask_in,
+                neginf_tile_idx,
+                causal_diag_tile_idx,
+                cb_qk_im,
+                q_start_tile,
+                k_start_tile,
+                Sq_chunk_t,
+                Sk_chunk_t,
+                straddle_col,
+                straddle_jump);
+        }
+
+        // Apply padding stamp (also when is_causal — the causal stamp doesn't handle K padding).
+        uint32_t num_padded, boundary_col, partial_tile_idx;
+        bool has_partial;
+        resolve_for_chunk(
+            Sk_chunk_t,
+            k_chunk,
+            num_local_k_chunks,
+            ring_iter_needs_global_n_mask,
+            ring_iter_needs_joint_n_mask,
+            local_n_needs_masking,
+            global_n_mask_chunk_id,
+            local_n_mask_chunk_id,
+            joint_n_mask_chunk_id,
+            num_padded,
+            boundary_col,
+            partial_tile_idx,
+            has_partial);
+
+        if (has_partial) {
+            apply_partial_mask_lightweight(
+                cb_mask_in, partial_tile_idx, cb_qk_im, boundary_col, Sk_chunk_t, Sq_chunk_t);
+        }
+        if (num_padded > 0) {
+            apply_padded_mask_lightweight_runtime<dst_size>(
+                cb_mask_in, neginf_tile_idx, cb_qk_im, num_padded, Sk_chunk_t, Sq_chunk_t);
+        }
+    }
+};
+
+enum SDPAType {
+    STANDARD = 0,
+    JOINT = 1,
+    RING = 2,
+};
+
+/******************************************************************************
+ *                             SDPA INNER LOOP                                *
+ ******************************************************************************/
+/**
+ * Use the specialized wrapper functions below instead of calling this directly.
+ *
+ * Template Parameters:
+ * @tparam sdpa_type - SDPA variant: STANDARD, JOINT, or RING
+ * @tparam cb_qk_im - QK intermediate buffer
+ * @tparam cb_identity_scale_in - Identity scale buffer
+ * @tparam cb_attention_sink - Attention sink buffer
+ * @tparam cb_scale_in - Scale buffer
+ * @tparam Sq_chunk_t - Query chunk size in tiles
+ * @tparam Sk_chunk_t - Key chunk size in tiles
+ * @tparam DHt - Head dimension in tiles
+ * @tparam vDHt - Value head dimension in tiles
+ * @tparam use_attention_sink - Whether to use attention sink
+ * @tparam is_causal - Whether to use causal masking
+ * @tparam use_provided_mask - Whether to use user-provided mask
+ * @tparam use_padded_mask - Whether to use padding mask
+ * @tparam use_joint_mask - Whether to use joint mask
+ * @tparam is_chunked - Whether query is chunked
+ * @tparam scale_fp32 - FP32 scale factor
+ * @tparam sliding_window_size - Sliding window attention size
+ * @tparam lightweight_mask_enabled - Enables the lightweight mask path (compile-time gated)
+ *
+ * Runtime Parameters:
+ * @param Skt - Sequence length in tiles
+ * @param qk_in0_block_w - QK matmul block width
+ * @param qk_subblock_w - QK matmul subblock width
+ * @param qk_subblock_h - QK matmul subblock height
+ * @param qk_in0_num_subblocks - QK input0 subblocks
+ * @param qk_in1_num_subblocks - QK input1 subblocks
+ * @param qk_num_blocks - QK number of blocks
+ * @param out_in0_block_w - Output matmul block width
+ * @param out_subblock_w - Output matmul subblock width
+ * @param out_subblock_h - Output matmul subblock height
+ * @param out_in0_num_subblocks - Output input0 subblocks
+ * @param out_in1_num_subblocks - Output input1 subblocks
+ * @param out_num_blocks - Output number of blocks
+ * @param iter_q_start - Query iteration start
+ * @param iter_q_end - Query iteration end
+ * @param q_num_chunks - Total query chunks
+ * @param local_q_start - Local query start
+ * @param chunked_q_chunk_offset - Chunked query offset
+ * @param iter_k_chunk_start - Key chunk iteration start
+ * @param iter_k_chunk_end - Key chunk iteration end
+ * @param q_chunk_tiles - Query chunk tiles
+ * @param k_chunk_tiles - Key chunk tiles
+ * @param qk_chunk_tiles - QK chunk tiles
+ * @param out_chunk_tiles - Output chunk tiles
+ * @param mask_chunk_0 - First mask chunk index
+ * @param mask_chunk_1 - Second mask chunk index
+ * @param ring_iter - Ring iteration index
+ * @param ring_id - Ring ID
+ * @param num_local_k_chunks - Number of K chunks stored locally on this device (used in Ring SDPA)
+ * @param local_padded_Nt - Padded sequence length in tiles for local K/V chunks on this device
+ * @param logical_nt - Logical (unpadded) sequence length in tiles for K/V
+ * @param ring_iter_needs_global_n_mask - Whether current ring iteration requires global N masking
+ * @param ring_iter_needs_joint_n_mask - Whether current ring iteration requires joint N masking
+ * @param local_n_needs_masking - Whether local N dimension requires masking
+ * @param global_n_mask_chunk_id - K chunk index where global N mask should be applied
+ * @param local_n_mask_chunk_id - K chunk index where local N mask should be applied
+ * @param joint_n_mask_chunk_id - K chunk index where joint N mask should be applied (relative to joint chunks)
+ * @param cb_q_in - Query input buffer
+ * @param cb_k_in - Key input buffer
+ * @param cb_v_in - Value input buffer
+ * @param cb_mask_in - Mask input buffer
+ * @param cb_col_identity - Column identity buffer
+ * @param cb_out_im_A - Output intermediate buffer A
+ * @param cb_out_im_B - Output intermediate buffer B
+ * @param cb_max_A - Max buffer A
+ * @param cb_max_B - Max buffer B
+ * @param cb_sum_A - Sum buffer A
+ * @param cb_sum_B - Sum buffer B
+ * @param cb_exp_max_diff - Exp max diff buffer
+ * @param cb_lse_in - LSE input buffer
+ * @param cb_lse_out - LSE output buffer
+ * @param cb_prev_out - Previous output buffer
+ * @param cb_out - Output buffer
+ */
+template <
+    SDPAType sdpa_type,
+    uint32_t cb_qk_im,
+    uint32_t cb_identity_scale_in,
+    uint32_t cb_attention_sink,
+    uint32_t cb_scale_in,
+    uint32_t Sq_chunk_t,
+    uint32_t Sk_chunk_t,
+    uint32_t NH,
+    uint32_t DHt,
+    uint32_t vDHt,
+    bool use_attention_sink,
+    bool use_provided_mask,
+    bool use_padded_mask,
+    bool use_joint_mask,
+    bool is_chunked,
+    uint32_t scale_fp32,
+    uint32_t sliding_window_size,
+    bool lightweight_mask_enabled = false,
+    bool chunked_enabled = false,
+    uint32_t chunked_q_local_padded_Nt = 0,
+    uint32_t chunked_chunk_size_t = 0>
+void sdpa_inner_loop(
+    const uint32_t Skt,
+    const uint32_t qk_in0_block_w,
+    const uint32_t qk_subblock_w,
+    const uint32_t qk_subblock_h,
+    const uint32_t qk_in0_num_subblocks,
+    const uint32_t qk_in1_num_subblocks,
+    const uint32_t qk_num_blocks,
+    const uint32_t out_in0_block_w,
+    const uint32_t out_subblock_w,
+    const uint32_t out_subblock_h,
+    const uint32_t out_in0_num_subblocks,
+    const uint32_t out_in1_num_subblocks,
+    const uint32_t out_num_blocks,
+    const uint32_t iter_q_start,
+    const uint32_t iter_q_end,
+    const uint32_t q_num_chunks,
+    const uint32_t local_q_start,
+    const uint32_t chunked_q_chunk_offset,
+    const uint32_t iter_k_chunk_start,
+    const uint32_t iter_k_chunk_end,
+    const uint32_t q_chunk_tiles,
+    const uint32_t k_chunk_tiles,
+    const uint32_t v_chunk_tiles,
+    const uint32_t qk_chunk_tiles,
+    const uint32_t out_chunk_tiles,
+    const uint32_t mask_chunk_0,
+    const uint32_t mask_chunk_1,
+    const uint32_t ring_iter,
+    const uint32_t ring_id,
+    const uint32_t num_local_k_chunks,
+    const uint32_t local_padded_Nt,
+    const uint32_t logical_nt,
+    const bool ring_iter_needs_global_n_mask,
+    const bool ring_iter_needs_joint_n_mask,
+    const bool local_n_needs_masking,
+    const uint32_t global_n_mask_chunk_id,
+    const uint32_t local_n_mask_chunk_id,
+    const uint32_t joint_n_mask_chunk_id,
+    const uint32_t cb_q_in,
+    const uint32_t cb_k_in,
+    const uint32_t cb_v_in,
+    const uint32_t cb_mask_in,
+    const uint32_t cb_col_identity,
+    const uint32_t cb_out_im_A,
+    const uint32_t cb_out_im_B,
+    const uint32_t cb_max_A,
+    const uint32_t cb_max_B,
+    const uint32_t cb_sum_A,
+    const uint32_t cb_sum_B,
+    const uint32_t cb_exp_max_diff,
+    const uint32_t cb_lse_in,
+    const uint32_t cb_lse_out,
+    const uint32_t cb_prev_out,
+    const uint32_t cb_out,
+    const LightweightMaskContext& lw_mask = {},
+    const bool is_causal = false,
+    const bool is_balanced = false,
+    const bool use_zigzag_balancing = false,
+    const bool is_last_ring_iter = true,
+    const ChunkedContext& chunked = {}) {
+    // Parameter-stable CB locals. Aliases (cb_sum_A/B, cb_max_A/B, cb_out_im_A/B) are
+    // std::swap-mutated below, so they use inline CircularBuffer(alias).method() at the call
+    // sites instead. cb_out is constructed conditionally near its consumer (cur.out target).
+    CircularBuffer cb_q_in_obj(cb_q_in);
+    CircularBuffer cb_k_in_obj(cb_k_in);
+    CircularBuffer cb_v_in_obj(cb_v_in);
+    CircularBuffer cb_qk_im_obj(cb_qk_im);
+    CircularBuffer cb_attention_sink_obj(cb_attention_sink);
+    CircularBuffer cb_lse_in_obj(cb_lse_in);
+    CircularBuffer cb_prev_out_obj(cb_prev_out);
+    constexpr uint32_t dst_size = compute_kernel_lib::DEST_AUTO_LIMIT;
+    uint32_t KV_chunks_processed_in_iter = 0;
+    const uint32_t q_per_core = iter_q_end - iter_q_start;
+
+    for (uint32_t q_iter = iter_q_start; q_iter < iter_q_end; ++q_iter) {
+        uint32_t q_start_tile = 0;    // First tile of Q chunk (tile units, both STANDARD and RING)
+        uint32_t q_high_tile = 0;     // STANDARD: upper tile bound for K iteration
+        uint32_t causal_k_limit = 0;  // RING: K-chunk index beyond which all K is above the diagonal
+        if constexpr (sdpa_type == STANDARD) {
+            const uint32_t linear_q_chunk = local_q_start + (q_iter - iter_q_start);
+            // Mod is a no-op when the input is per-head ([0, q_num_chunks)) and extracts the
+            // per-head q_chunk when it's a flat global index (global Q scheduling spans heads).
+            uint32_t q_chunk = remap_q_index(linear_q_chunk, q_num_chunks, use_zigzag_balancing) % q_num_chunks;
+            // Get Q chunk
+            if constexpr (is_chunked) {
+                q_chunk = chunked_q_chunk_offset + q_chunk;
+            }
+            q_start_tile = q_chunk * Sq_chunk_t;
+            if (is_causal) {
+                // Clamp to total K-tile extent. Mirrors reader_interleaved's clamp; without
+                // both, the reader and compute disagree on K-chunk count when Q-chunk extends
+                // past total K (Sq_chunk_t > Skt) → CB deadlock.
+                const uint32_t q_high_unclamped = q_start_tile + Sq_chunk_t;
+                q_high_tile = q_high_unclamped < Skt ? q_high_unclamped : Skt;
+            } else {
+                q_high_tile = Skt;
+            }
+        } else if (sdpa_type == RING) {
+            uint32_t q_chunk = remap_q_index(q_iter, q_num_chunks, use_zigzag_balancing) % q_num_chunks;
+
+            if constexpr (chunked_enabled) {
+                // Absolute Q tile row. Diag stamp masks K past Q's range; logical_n skip handles K past the cache.
+                q_start_tile =
+                    chunked.q_start_idx_t + chunked.ring_index * chunked_q_local_padded_Nt + q_chunk * Sq_chunk_t;
+            } else if (is_causal) {
+                q_start_tile = q_chunk * Sq_chunk_t;
+                causal_k_limit = (q_start_tile + Sq_chunk_t + Sk_chunk_t - 1) / Sk_chunk_t;
+            }
+            if (is_balanced && (q_chunk < q_num_chunks / 2)) {
+                continue;
+            }
+        }  // If ring attention
+
+        // Set up ping pong buffers
+        uint32_t alias_prev_sum = cb_sum_A;
+        uint32_t alias_cur_sum = cb_sum_B;
+        uint32_t alias_prev_max = cb_max_A;
+        uint32_t alias_cur_max = cb_max_B;
+        uint32_t alias_mm2_prev_out = cb_out_im_A;
+        uint32_t alias_mm2_cur_out = cb_out_im_B;
+
+        uint32_t k_chunk_end;
+        if constexpr (sdpa_type == STANDARD) {
+            // loop while k_low < q_high => (k_chunk * Sk_chunk_t) < q_high_tile.
+            k_chunk_end = (q_high_tile + Sk_chunk_t - 1) / Sk_chunk_t;
+        } else {  // RING or JOINT.
+            k_chunk_end = iter_k_chunk_end;
+        }
+
+        uint32_t processed_k_chunks = 0;
+
+        for (uint32_t k_chunk = iter_k_chunk_start; k_chunk < k_chunk_end; ++k_chunk) {
+            uint32_t kv_global_start_tile = 0;  // RING only: abs K-tile index of this k_chunk's start
+            if constexpr (sdpa_type == RING) {
+                const bool kv_chunk_is_joint = k_chunk >= num_local_k_chunks;
+                // Global index into the padded KV tensor. Chunked: non-monotonic local→abs map.
+                if constexpr (chunked_enabled) {
+                    kv_global_start_tile =
+                        kv_global_tile_for_local<true, 0, chunked_chunk_size_t, chunked_q_local_padded_Nt>(
+                            ring_id, k_chunk * Sk_chunk_t);
+                } else {
+                    kv_global_start_tile = local_padded_Nt * ring_id + k_chunk * Sk_chunk_t;
+                }
+                if (!kv_chunk_is_joint && (kv_global_start_tile >= logical_nt)) {
+                    // This is a KV chunk on spatial input beyond the logical N, and not joint KV. Skip it.
+                    continue;
+                }
+            }
+
+            KV_chunks_processed_in_iter++;
+
+            // Chunked-prefill: never take this skip (local-frame causal_k_limit doesn't apply —
+            // the diag stamp uses absolute coords every k_chunk instead).
+            if (sdpa_type == RING && !chunked_enabled && k_chunk >= causal_k_limit && is_causal) {
+                cb_k_in_obj.wait_front(k_chunk_tiles);
+                cb_v_in_obj.wait_front(v_chunk_tiles);
+                cb_k_in_obj.pop_front(k_chunk_tiles);
+                cb_v_in_obj.pop_front(v_chunk_tiles);
+
+                continue;
+            }
+
+            /**
+             * QK = Q_CHUNK @ K_CHUNK
+             *
+             * matmul_blocks internally waits on both inputs
+             */
+            reconfig_data_format(cb_k_in, cb_q_in);
+            pack_reconfig_data_format(cb_qk_im);
+            matmul_blocks(
+                cb_q_in,
+                cb_k_in,
+                cb_qk_im,
+                Sq_chunk_t,
+                Sk_chunk_t,
+                DHt,
+                qk_num_blocks,
+                qk_in0_num_subblocks,
+                qk_in1_num_subblocks,
+                qk_in0_block_w,
+                qk_subblock_h,
+                qk_subblock_w,
+                true /*transpose*/);
+
+            /**
+             * Note
+             * Typically, scores is multiplied by a scalar here. We employed an optimization
+             * where we fuse the scaling into exp both in exp(x - max) and exp(prev_max - cur_max).
+             * This gives us scaling for free on the performance-critical exp(x - max) computation.
+             */
+
+            bool apply_mask = false;
+            bool needs_padding_mask =
+                (sdpa_type == RING) &&
+                ((ring_iter_needs_global_n_mask && k_chunk == global_n_mask_chunk_id) ||
+                 (local_n_needs_masking && k_chunk == local_n_mask_chunk_id) ||
+                 (ring_iter_needs_joint_n_mask && (k_chunk - num_local_k_chunks) == joint_n_mask_chunk_id) ||
+                 (lw_mask.straddle_num_padded_tiles > 0 && k_chunk == lw_mask.straddle_mask_chunk_id));
+            if (sdpa_type == RING && !is_causal) {
+                apply_mask = needs_padding_mask;
+            } else if (is_causal || sliding_window_size > 0) {
+                // Chunked-prefill (RING): use abs K so the causal-overlap test matches the diag stamp's frame.
+                uint32_t k_low_idx;
+                uint32_t k_high_idx;
+                if constexpr (sdpa_type == RING && chunked_enabled) {
+                    k_low_idx = kv_global_start_tile;
+                    k_high_idx = k_low_idx + Sk_chunk_t;
+                    // Straddle: chunk crosses a per-chunk slab boundary in the local cache. The
+                    // tail tiles land in the next slab, so the actual high global K is shifted by
+                    // chunk_size_t - q_local_padded_Nt.
+                    if (chunked_q_local_padded_Nt > 0) {
+                        const uint32_t local_start = k_chunk * Sk_chunk_t;
+                        const uint32_t slab_end_local =
+                            (local_start / chunked_q_local_padded_Nt + 1) * chunked_q_local_padded_Nt;
+                        if (local_start + Sk_chunk_t > slab_end_local) {
+                            k_high_idx += (chunked_chunk_size_t - chunked_q_local_padded_Nt);
+                        }
+                    }
+                } else {
+                    k_low_idx = k_chunk * Sk_chunk_t;
+                    k_high_idx = k_low_idx + Sk_chunk_t;
+                }
+                // Apply mask if causal overlap, sliding window, or this K chunk has padding
+                apply_mask = (q_start_tile < k_high_idx) || (sliding_window_size > 0) || needs_padding_mask;
+            } else if constexpr (use_provided_mask) {
+                apply_mask = true;
+            } else if constexpr (use_padded_mask) {
+                // Apply mask only on the last K chunk
+                apply_mask = (k_chunk == iter_k_chunk_end - 1);
+            } else if constexpr (use_joint_mask) {
+                // Apply mask for specific chunk combinations
+                apply_mask = (k_chunk == mask_chunk_0) || (k_chunk == mask_chunk_1);
+            }
+
+            if (apply_mask) {
+                /* QK += MASK */
+                reconfig_data_format(cb_qk_im, cb_mask_in);
+                if constexpr (lightweight_mask_enabled) {
+                    // Re-enter reserved state on cb_qk_im so the lightweight mask can be stamped in-place.
+                    // matmul_blocks above already pushed the QK tiles, so tiles_received has been bumped;
+                    // without the pop+push cycle below, reduce_c's wait-front would return immediately
+                    // and unpack could start before the mask stamps land in L1.
+                    // Safe because pop-front only moves rd_ptr (L1 data is untouched), and on a
+                    // single-buffered CB of size qk_chunk_tiles the re-reserved wr_ptr wraps back to the
+                    // same physical region holding the QK scores.
+                    // Warning: this won't work if cb_qk_im is double-buffered -- the stamps would land
+                    // in the other buffer, leaving the QK scores unmasked.
+                    cb_qk_im_obj.wait_front(Sk_chunk_t * Sq_chunk_t);
+                    cb_qk_im_obj.pop_front(Sk_chunk_t * Sq_chunk_t);
+                    cb_qk_im_obj.reserve_back(Sk_chunk_t * Sq_chunk_t);
+                    // Chunked-prefill: feed abs K (matches abs q_start_tile so diag stamp lines up).
+                    uint32_t k_start_tile_for_mask;
+                    uint32_t lw_straddle_col = 0;
+                    uint32_t lw_straddle_jump = 0;
+                    if constexpr (sdpa_type == RING && chunked_enabled) {
+                        k_start_tile_for_mask = kv_global_start_tile;
+                        // Chunked-prefill straddle: a K-chunk can begin in one per-chunk K
+                        // region and end in the next when k_chunk_size does not divide
+                        // q_local_padded_Nt. Global K is non-contiguous across the K-chunk in
+                        // that case (jumps by chunk_size_t - q_local_padded_Nt between regions),
+                        // so we signal the column boundary (straddle_col) and the jump
+                        // (straddle_jump) to the diag stamp — see same comment block in
+                        // compute_streaming.hpp:sdpa_inner_loop_step for the full picture.
+                        if (chunked_q_local_padded_Nt > 0) {
+                            const uint32_t local_start = k_chunk * Sk_chunk_t;
+                            const uint32_t slab_end_local =
+                                (local_start / chunked_q_local_padded_Nt + 1) * chunked_q_local_padded_Nt;
+                            if (local_start + Sk_chunk_t > slab_end_local) {
+                                lw_straddle_col = slab_end_local - local_start;
+                                lw_straddle_jump = chunked_chunk_size_t - chunked_q_local_padded_Nt;
+                            }
+                        }
+                    } else {
+                        k_start_tile_for_mask = k_chunk * Sk_chunk_t;
+                    }
+                    lw_mask.template apply<dst_size>(
+                        cb_mask_in,
+                        cb_qk_im,
+                        Sk_chunk_t,
+                        Sq_chunk_t,
+                        k_chunk,
+                        num_local_k_chunks,
+                        ring_iter_needs_global_n_mask,
+                        ring_iter_needs_joint_n_mask,
+                        local_n_needs_masking,
+                        global_n_mask_chunk_id,
+                        local_n_mask_chunk_id,
+                        joint_n_mask_chunk_id,
+                        q_start_tile,
+                        k_start_tile_for_mask,
+                        lw_straddle_col,
+                        lw_straddle_jump);
+                    cb_qk_im_obj.push_back(Sk_chunk_t * Sq_chunk_t);
+                } else {
+                    add_block_inplace(cb_qk_im, cb_mask_in, qk_chunk_tiles);
+                }
+            }
+
+            /**
+             * reduce_c can perform both reduce_max and eltwise max with previous result.
+             * if do_eltwise_max:
+             *  cur_max = eltwise_max(prev_max, max(qk, dim=-1))
+             * else:
+             *  cur_max = max(qk, dim=-1)
+             *
+             * Use the reduce_c overload with cols as a runtime arg which uses standard
+             * reduce_tile + binary_max_tile. The overload with cols as a template arg
+             * is bf16-only but cb_qk_im could be fp32.
+             */
+            reconfig_data_format(cb_qk_im, cb_identity_scale_in);
+            reduce_c<PoolType::MAX, ReduceDim::REDUCE_ROW, cb_qk_im, cb_identity_scale_in, Sq_chunk_t>(
+                alias_cur_max, alias_prev_max, Sk_chunk_t, processed_k_chunks > 0);
+
+            /**
+             * sub_exp fuses a few operations.
+             * In-place it performs `QK = exp((QK - cur_max) * scale)`
+             *
+             * It also partially performs reduce_sum on the output using L1 accumulation.
+             * `cur_sum = sum_tiles(exp((QK - cur_max) * scale), dim=-1)`
+             *
+             * Partial reduce_sum is used to push the final row_reduction within a tile
+             * outside of the loop over K chunks.
+             */
+            sub_exp_block_bcast_cols_inplace<cb_qk_im, Sq_chunk_t, scale_fp32, true>(
+                alias_cur_max, alias_cur_sum, Sk_chunk_t);
+
+            // Reconfigure unpackers: srcA (context 0) = cb_v_in, srcB (context 1) = cb_qk_im (operands are swapped in
+            // matmul)
+            reconfig_data_format(cb_v_in, cb_qk_im);
+            pack_reconfig_data_format(alias_mm2_cur_out);
+
+            /* OUT_IM = QK @ V_CHUNK */
+            matmul_blocks(
+                cb_qk_im,
+                cb_v_in,
+                alias_mm2_cur_out,
+                Sq_chunk_t,
+                vDHt,
+                Sk_chunk_t,
+                out_num_blocks,
+                out_in0_num_subblocks,
+                out_in1_num_subblocks,
+                out_in0_block_w,
+                out_subblock_h,
+                out_subblock_w,
+                false /*transpose*/);
+
+            cb_qk_im_obj.pop_front(qk_chunk_tiles);
+            reconfig_data_format(alias_prev_max, alias_cur_max);
+
+            /* OUT_ACC += OUT_IM */
+            if (processed_k_chunks > 0) {
+                /**
+                 * cb_exp_max_diff = torch.exp((cb_prev_max - cb_cur_max) * scale)
+                 * Scale is fused into exp again since max is the max of unscaled scores.
+                 */
+                sub_exp_block<scale_fp32>(alias_prev_max, alias_cur_max, cb_exp_max_diff, Sq_chunk_t);
+                CircularBuffer(alias_prev_max).pop_front(Sq_chunk_t);
+
+                /**
+                 * cb_prev_sum *= cb_exp_max_diff
+                 * This is a bcast_cols since max_diff is a column vector and prev_sum is a partial
+                 * reduction, containing the sum of tiles in dim=-1 of QK.
+                 */
+                mul_tiles_bcast_cols_inplace(alias_prev_sum, cb_exp_max_diff, Sq_chunk_t);
+
+                /* cb_cur_sum += cb_prev_sum */
+                add_block_inplace(alias_cur_sum, alias_prev_sum, Sq_chunk_t);
+
+                /**
+                 * alias_mm2_cur_out += alias_mm2_prev_out * cb_exp_max_diff
+                 * This uses L1 accumulation to accumulate onto mm2_cur_out.
+                 */
+                mul_block_bcast_cols<Sq_chunk_t, vDHt, false, true>(
+                    alias_mm2_prev_out, cb_exp_max_diff, alias_mm2_cur_out);
+            }
+
+            // Swap CB handles to prepare for next iteration
+            std::swap(alias_prev_sum, alias_cur_sum);
+            std::swap(alias_mm2_prev_out, alias_mm2_cur_out);
+            std::swap(alias_prev_max, alias_cur_max);
+
+            processed_k_chunks++;
+        }
+
+        /**
+         * Performs final row-reduction on the partial sum.
+         */
+        matmul_reduce<Sq_chunk_t>(cb_col_identity, alias_prev_sum);
+
+        /**
+         * Process attention sink as a virtual K chunk.
+         * The attention sink provides additional logits that are included in the softmax
+         * denominator but don't contribute to the output (no S @ V computation).
+         * This effectively allows some attention probability to be "absorbed" by the sink,
+         * reducing attention weights on actual tokens.
+         *
+         * Shape of attention_sink: [Sq_chunk_t, 1] tiles
+         * Each head has one sink logit value that is broadcast to all query positions in the chunk.
+         * The reader kernel replicates the per-head value across all Sq_chunk_t positions.
+         */
+        if constexpr (use_attention_sink) {
+            // Treat attention_sink as scores (already scaled)
+            // Shape: [Sq_chunk_t, 1] tiles - same per-head sink value broadcast to all query positions
+
+            // 1. Update running max: cur_max = max(prev_max, attention_sink)
+            //    This compares the previous max with the sink logit
+            reconfig_data_format(cb_attention_sink, cb_identity_scale_in);
+
+            reduce_c<PoolType::MAX, ReduceDim::REDUCE_ROW, cb_attention_sink, cb_identity_scale_in, Sq_chunk_t, 1>(
+                alias_cur_max, alias_prev_max, true);
+
+            // 2. Compute exp((prev_max - cur_max) * scale) to rescale previous statistics
+            sub_exp_block<scale_fp32>(alias_prev_max, alias_cur_max, cb_exp_max_diff, Sq_chunk_t);
+            CircularBuffer(alias_prev_max).pop_front(Sq_chunk_t);
+
+            // 3. Rescale previous sum: prev_sum *= exp(prev_max - cur_max)
+            mul_tiles_bcast_cols_inplace(alias_prev_sum, cb_exp_max_diff, Sq_chunk_t);
+            // 4. Compute exp((attention_sink - cur_max) * scale) and accumulate in cur_sum
+            //    This adds the attention sink's contribution to the softmax denominator
+            sub_exp_block_bcast_cols_inplace<cb_attention_sink, Sq_chunk_t, scale_fp32, false>(
+                alias_cur_max, alias_cur_sum, 1);
+
+            // 5. Add rescaled previous sum to current sum: cur_sum += prev_sum
+            add_block_inplace(alias_cur_sum, alias_prev_sum, Sq_chunk_t);
+
+            // 6. Update running statistics for final normalization
+            std::swap(alias_prev_sum, alias_cur_sum);
+            std::swap(alias_prev_max, alias_cur_max);
+
+            // 7. Rescale accumulated output: mm2_prev_out *= exp(prev_max - cur_max)
+            //    Note: We do NOT compute attention_sink @ V, so output only has real token contributions
+            //    But we need to rescale it due to the updated max
+            mul_block_bcast_cols<Sq_chunk_t, vDHt, false, false>(
+                alias_mm2_prev_out, cb_exp_max_diff, alias_mm2_cur_out);
+            std::swap(alias_mm2_prev_out, alias_mm2_cur_out);
+        }
+
+        if constexpr (sdpa_type == RING) {
+            log_block(alias_prev_sum, alias_cur_max, Sq_chunk_t);
+
+            // Scale prev_max by scale_fp32
+            mul_block_bcast_scalar_inplace<cb_scale_in, Sq_chunk_t>(alias_prev_max);
+            add_block_inplace(alias_prev_max, alias_cur_max, Sq_chunk_t);
+
+            /* cb_cur_sum = 1.0 / cb_cur_sum */
+            recip_block_inplace(alias_prev_sum, Sq_chunk_t);
+            /* cb_out_accumulate_im *= cb_cur_sum */
+            mul_block_bcast_cols_inplace<Sq_chunk_t, vDHt>(alias_mm2_prev_out, alias_prev_sum);
+
+            if (ring_iter > 0) {
+                // Update output according to previous and current LSE
+                /**
+                 * sig = torch.sigmoid(cur_lse - prev_lse)
+                 * out = prev_out - sig * (prev_out - cur_out)
+                 * lse = prev_lse - torch.logsigmoid(prev_lse - cur_lse)
+                 */
+                cb_lse_in_obj.wait_front(Sq_chunk_t);
+                cb_prev_out_obj.wait_front(out_chunk_tiles);
+
+                uint32_t alias_cur_lse = alias_prev_max;      // full
+                uint32_t alias_sig = alias_cur_max;           // empty
+                uint32_t alias_cur_out = alias_mm2_prev_out;  // full
+                uint32_t alias_sub = alias_mm2_cur_out;       // empty
+
+                // alias_sig = sigmoid(alias_cur_lse - cb_lse_in)
+                sigmoid_sub(alias_cur_lse, cb_lse_in, alias_sig, Sq_chunk_t);
+
+                // alias_sub = cb_prev_out - alias_cur_out
+                reconfig_data_format(cb_prev_out, alias_cur_out);
+                sub_block(cb_prev_out, alias_cur_out, alias_sub, out_chunk_tiles);
+                // alias_sub *= alias_sig
+                reconfig_data_format(alias_sub, alias_sig);
+                mul_block_bcast_cols_inplace<Sq_chunk_t, vDHt>(alias_sub, alias_sig);
+                // cb_out = cb_prev_out - alias_sub
+                reconfig_data_format(cb_prev_out, alias_sub);
+                pack_reconfig_data_format(cb_out);
+                sub_block(cb_prev_out, alias_sub, cb_out, out_chunk_tiles);
+                cb_prev_out_obj.pop_front(out_chunk_tiles);
+                CircularBuffer(alias_cur_out).pop_front(out_chunk_tiles);
+                CircularBuffer(alias_sub).pop_front(out_chunk_tiles);
+
+                // alias_sig = sigmoid(cb_lse_in - alias_cur_lse)
+                // alias_cur_lse = log(alias_sig)
+                // cb_lse_out = cb_lse_in - alias_cur_lse
+                pack_reconfig_data_format(alias_sig);
+                reconfig_data_format(cb_lse_in, alias_cur_lse);
+                logsigmoid_sub(cb_lse_in, alias_cur_lse, alias_sig, Sq_chunk_t);
+                sub_block(cb_lse_in, alias_sig, cb_lse_out, Sq_chunk_t);
+                CircularBuffer(alias_sig).pop_front(Sq_chunk_t);
+                CircularBuffer(alias_cur_lse).pop_front(Sq_chunk_t);
+                cb_lse_in_obj.pop_front(Sq_chunk_t);
+            } else {
+                pack_reconfig_data_format(cb_out);
+                copy_block(alias_mm2_prev_out, cb_out, out_chunk_tiles);
+
+                pack_reconfig_data_format(cb_lse_out);
+                copy_block(alias_prev_max, cb_lse_out, Sq_chunk_t);
+            }
+        } else {
+            /* cb_cur_sum = 1.0 / cb_cur_sum */
+            recip_block_inplace(alias_prev_sum, Sq_chunk_t);
+
+            /* cb_out_accumulate_im *= cb_cur_sum */
+            pack_reconfig_data_format(cb_out);
+            mul_block_bcast_cols<Sq_chunk_t, vDHt, false, false>(alias_mm2_prev_out, alias_prev_sum, cb_out);
+
+            // free up cb_prev_max after K chunks
+            CircularBuffer(alias_prev_max).pop_front(Sq_chunk_t);
+        }
+
+        // When q_per_core == 1, Q is identical across ring iterations so we keep it
+        // fronted in the CB and only pop on the last iteration to avoid redundant DRAM re-reads.
+        if (q_per_core > 1 || is_last_ring_iter) {
+            cb_q_in_obj.pop_front(q_chunk_tiles);
+        }
+
+        // Under global Q scheduling the reader pushes one cb_attention_sink slot per Q iter,
+        // so we must drain matching slots inside this loop. Pre-PR the push/pop pair was
+        // 1:1 outside the loop; the new cadence is 1:1 per iter.
+        if constexpr (use_attention_sink) {
+            cb_attention_sink_obj.pop_front(Sq_chunk_t);
+        }
+    }
+
+    if constexpr (sdpa_type == RING) {
+        if (KV_chunks_processed_in_iter % 2 == 0) {
+            cb_k_in_obj.wait_front(k_chunk_tiles);
+            cb_v_in_obj.wait_front(v_chunk_tiles);
+            cb_k_in_obj.pop_front(k_chunk_tiles);
+            cb_v_in_obj.pop_front(v_chunk_tiles);
+        }
+    }
+}
+
+/******************************************************************************
+ *                          SDPA WRAPPER FUNCTIONS                            *
+ ******************************************************************************/
+
+/**
+ * Standard SDPA with optional causal masking, attention sink, and sliding window.
+ */
+template <
+    uint32_t cb_qk_im,
+    uint32_t cb_identity_scale_in,
+    uint32_t cb_attention_sink,
+    uint32_t Sq_chunk_t,
+    uint32_t Sk_chunk_t,
+    uint32_t DHt,
+    uint32_t vDHt,
+    bool use_attention_sink,
+    bool is_causal,
+    bool use_provided_mask,
+    bool use_padded_mask,
+    bool is_chunked,
+    uint32_t scale_fp32,
+    uint32_t sliding_window_size,
+    bool lightweight_mask_enabled = false>
+void sdpa_standard(
+    const uint32_t Skt,
+    const uint32_t qk_in0_block_w,
+    const uint32_t qk_subblock_w,
+    const uint32_t qk_subblock_h,
+    const uint32_t qk_in0_num_subblocks,
+    const uint32_t qk_in1_num_subblocks,
+    const uint32_t qk_num_blocks,
+    const uint32_t out_in0_block_w,
+    const uint32_t out_subblock_w,
+    const uint32_t out_subblock_h,
+    const uint32_t out_in0_num_subblocks,
+    const uint32_t out_in1_num_subblocks,
+    const uint32_t out_num_blocks,
+    const uint32_t iter_q_start,
+    const uint32_t iter_q_end,
+    const uint32_t q_num_chunks,
+    const uint32_t local_q_start,
+    const uint32_t chunked_q_chunk_offset,
+    const uint32_t k_num_chunks,
+    const uint32_t q_chunk_tiles,
+    const uint32_t k_chunk_tiles,
+    const uint32_t v_chunk_tiles,
+    const uint32_t qk_chunk_tiles,
+    const uint32_t out_chunk_tiles,
+    const uint32_t cb_q_in,
+    const uint32_t cb_k_in,
+    const uint32_t cb_v_in,
+    const uint32_t cb_mask_in,
+    const uint32_t cb_col_identity,
+    const uint32_t cb_out_im_A,
+    const uint32_t cb_out_im_B,
+    const uint32_t cb_max_A,
+    const uint32_t cb_max_B,
+    const uint32_t cb_sum_A,
+    const uint32_t cb_sum_B,
+    const uint32_t cb_exp_max_diff,
+    const uint32_t cb_out,
+    const LightweightMaskContext& lw_mask = {},
+    const bool use_zigzag_balancing = false) {
+    sdpa_inner_loop<
+        STANDARD,
+        cb_qk_im,
+        cb_identity_scale_in,
+        cb_attention_sink,
+        0,  // cb_scale_in (not used)
+        Sq_chunk_t,
+        Sk_chunk_t,
+        0,  // NH (not used)
+        DHt,
+        vDHt,
+        use_attention_sink,
+        use_provided_mask,
+        use_padded_mask,
+        false,  // use_joint_mask (not used)
+        is_chunked,
+        scale_fp32,
+        sliding_window_size,
+        lightweight_mask_enabled>(
+        Skt,
+        qk_in0_block_w,
+        qk_subblock_w,
+        qk_subblock_h,
+        qk_in0_num_subblocks,
+        qk_in1_num_subblocks,
+        qk_num_blocks,
+        out_in0_block_w,
+        out_subblock_w,
+        out_subblock_h,
+        out_in0_num_subblocks,
+        out_in1_num_subblocks,
+        out_num_blocks,
+        iter_q_start,
+        iter_q_end,
+        q_num_chunks,
+        local_q_start,
+        chunked_q_chunk_offset,
+        0,             // iter_k_chunk_start
+        k_num_chunks,  // iter_k_chunk_end
+        q_chunk_tiles,
+        k_chunk_tiles,
+        v_chunk_tiles,
+        qk_chunk_tiles,
+        out_chunk_tiles,
+        0,      // mask_chunk_0 (not used)
+        0,      // mask_chunk_1 (not used)
+        0,      // ring_iter (not used)
+        0,      // ring_id (not used)
+        0,      // num_local_k_chunks (not used)
+        0,      // local_padded_Nt (not used)
+        0,      // logical_nt (not used)
+        false,  // ring_iter_needs_global_n_mask (not used)
+        false,  // ring_iter_needs_joint_n_mask (not used)
+        false,  // local_n_needs_masking (not used)
+        0,      // global_n_mask_chunk_id (not used)
+        0,      // local_n_mask_chunk_id (not used)
+        0,      // joint_n_mask_chunk_id (not used)
+        cb_q_in,
+        cb_k_in,
+        cb_v_in,
+        cb_mask_in,
+        cb_col_identity,
+        cb_out_im_A,
+        cb_out_im_B,
+        cb_max_A,
+        cb_max_B,
+        cb_sum_A,
+        cb_sum_B,
+        cb_exp_max_diff,
+        0,  // cb_lse_in (not used)
+        0,  // cb_lse_out (not used)
+        0,  // cb_prev_out (not used)
+        cb_out,
+        lw_mask,
+        is_causal,
+        false,  // is_balanced (not used)
+        use_zigzag_balancing);
+}
+
+/**
+ * Joint SDPA for multi-modal attention.
+ */
+template <
+    uint32_t cb_qk_im,
+    uint32_t cb_identity_scale_in,
+    uint32_t Sq_chunk_t,
+    uint32_t Sk_chunk_t,
+    uint32_t DHt,
+    bool use_joint_mask,
+    uint32_t scale_fp32>
+void sdpa_joint(
+    const uint32_t Skt,
+    const uint32_t qk_in0_block_w,
+    const uint32_t qk_subblock_w,
+    const uint32_t qk_subblock_h,
+    const uint32_t qk_in0_num_subblocks,
+    const uint32_t qk_in1_num_subblocks,
+    const uint32_t qk_num_blocks,
+    const uint32_t out_in0_block_w,
+    const uint32_t out_subblock_w,
+    const uint32_t out_subblock_h,
+    const uint32_t out_in0_num_subblocks,
+    const uint32_t out_in1_num_subblocks,
+    const uint32_t out_num_blocks,
+    const uint32_t local_q_start,
+    const uint32_t local_q_end,
+    const uint32_t k_num_chunks,
+    const uint32_t q_chunk_tiles,
+    const uint32_t k_chunk_tiles,
+    const uint32_t qk_chunk_tiles,
+    const uint32_t out_chunk_tiles,
+    const uint32_t mask_chunk_0,
+    const uint32_t mask_chunk_1,
+    const uint32_t cb_q_in,
+    const uint32_t cb_k_in,
+    const uint32_t cb_v_in,
+    const uint32_t cb_mask_in,
+    const uint32_t cb_col_identity,
+    const uint32_t cb_out_im_A,
+    const uint32_t cb_out_im_B,
+    const uint32_t cb_max_A,
+    const uint32_t cb_max_B,
+    const uint32_t cb_sum_A,
+    const uint32_t cb_sum_B,
+    const uint32_t cb_exp_max_diff,
+    const uint32_t cb_out) {
+    sdpa_inner_loop<
+        JOINT,
+        cb_qk_im,
+        cb_identity_scale_in,
+        0,  // cb_attention_sink (not used)
+        0,  // cb_scale_in (not used)
+        Sq_chunk_t,
+        Sk_chunk_t,
+        0,  // NH (not used)
+        DHt,
+        DHt,    // vDHt = DHt
+        false,  // use_attention_sink (not used)
+        false,  // use_provided_mask (not used)
+        false,  // use_padded_mask (not used)
+        use_joint_mask,
+        false,  // is_chunked (not used)
+        scale_fp32,
+        0>(  // sliding_window_size (not used)
+        Skt,
+        qk_in0_block_w,
+        qk_subblock_w,
+        qk_subblock_h,
+        qk_in0_num_subblocks,
+        qk_in1_num_subblocks,
+        qk_num_blocks,
+        out_in0_block_w,
+        out_subblock_w,
+        out_subblock_h,
+        out_in0_num_subblocks,
+        out_in1_num_subblocks,
+        out_num_blocks,
+        local_q_start,  // iter_q_start
+        local_q_end,    // iter_q_end
+        0,              // q_num_chunks (not used)
+        local_q_start,
+        0,             // chunked_q_chunk_offset (not used)
+        0,             // iter_k_chunk_start
+        k_num_chunks,  // iter_k_chunk_end
+        q_chunk_tiles,
+        k_chunk_tiles,
+        k_chunk_tiles,
+        qk_chunk_tiles,
+        out_chunk_tiles,
+        mask_chunk_0,
+        mask_chunk_1,
+        0,      // ring_iter (not used)
+        0,      // ring_id (not used)
+        0,      // num_local_k_chunks (not used)
+        0,      // local_padded_Nt (not used)
+        0,      // logical_nt (not used)
+        false,  // ring_iter_needs_global_n_mask (not used)
+        false,  // ring_iter_needs_joint_n_mask (not used)
+        false,  // local_n_needs_masking (not used)
+        0,      // global_n_mask_chunk_id (not used)
+        0,      // local_n_mask_chunk_id (not used)
+        0,      // joint_n_mask_chunk_id (not used)
+        cb_q_in,
+        cb_k_in,
+        cb_v_in,
+        cb_mask_in,
+        cb_col_identity,
+        cb_out_im_A,
+        cb_out_im_B,
+        cb_max_A,
+        cb_max_B,
+        cb_sum_A,
+        cb_sum_B,
+        cb_exp_max_diff,
+        0,  // cb_lse_in (not used)
+        0,  // cb_lse_out (not used)
+        0,  // cb_prev_out (not used)
+        cb_out);
+}
+
+/**
+ * Ring SDPA for distributed multi-device attention.
+ */
+template <
+    uint32_t cb_qk_im,
+    uint32_t cb_identity_scale_in,
+    uint32_t cb_scale_in,
+    uint32_t Sq_chunk_t,
+    uint32_t Sk_chunk_t,
+    uint32_t NH,
+    uint32_t DHt,
+    uint32_t vDHt,
+    uint32_t scale_fp32,
+    bool lightweight_mask_enabled = false,
+    bool chunked_enabled = false,
+    uint32_t chunked_q_local_padded_Nt = 0,
+    uint32_t chunked_chunk_size_t = 0>
+void sdpa_ring(
+    const uint32_t qk_in0_block_w,
+    const uint32_t qk_subblock_w,
+    const uint32_t qk_subblock_h,
+    const uint32_t qk_in0_num_subblocks,
+    const uint32_t qk_in1_num_subblocks,
+    const uint32_t qk_num_blocks,
+    const uint32_t out_in0_block_w,
+    const uint32_t out_subblock_w,
+    const uint32_t out_subblock_h,
+    const uint32_t out_in0_num_subblocks,
+    const uint32_t out_in1_num_subblocks,
+    const uint32_t out_num_blocks,
+    const uint32_t global_q_start,
+    const uint32_t global_q_end,
+    const uint32_t q_num_chunks,
+    const uint32_t iter_k_chunk_start,
+    const uint32_t iter_k_chunk_end,
+    const uint32_t q_chunk_tiles,
+    const uint32_t k_chunk_tiles,
+    const uint32_t v_chunk_tiles,
+    const uint32_t qk_chunk_tiles,
+    const uint32_t out_chunk_tiles,
+    const uint32_t ring_iter,
+    const uint32_t ring_id,
+    const uint32_t num_local_k_chunks,
+    const uint32_t local_padded_Nt,
+    const uint32_t logical_nt,
+    const bool ring_iter_needs_global_n_mask,
+    const bool ring_iter_needs_joint_n_mask,
+    const bool local_n_needs_masking,
+    const uint32_t global_n_mask_chunk_id,
+    const uint32_t local_n_mask_chunk_id,
+    const uint32_t joint_n_mask_chunk_id,
+    const uint32_t cb_q_in,
+    const uint32_t cb_k_in,
+    const uint32_t cb_v_in,
+    const uint32_t cb_mask_in,
+    const uint32_t cb_col_identity,
+    const uint32_t cb_out_im_A,
+    const uint32_t cb_out_im_B,
+    const uint32_t cb_max_A,
+    const uint32_t cb_max_B,
+    const uint32_t cb_sum_A,
+    const uint32_t cb_sum_B,
+    const uint32_t cb_exp_max_diff,
+    const uint32_t cb_lse_in,
+    const uint32_t cb_lse_out,
+    const uint32_t cb_prev_out,
+    const uint32_t cb_out,
+    const LightweightMaskContext& lw_mask,
+    const bool is_causal_ring_iter,
+    const bool skip_first_half_q,
+    const bool is_last_ring_iter,
+    const bool use_zigzag_balancing = false,
+    const ChunkedContext& chunked = {}) {
+    sdpa_inner_loop<
+        RING,
+        cb_qk_im,
+        cb_identity_scale_in,
+        0,  // cb_attention_sink (not used)
+        cb_scale_in,
+        Sq_chunk_t,
+        Sk_chunk_t,
+        NH,
+        DHt,
+        vDHt,
+        false,  // use_attention_sink (not used)
+        false,  // use_provided_mask (not used)
+        false,  // use_padded_mask (not used)
+        false,  // use_joint_mask (not used)
+        false,  // is_chunked (not used)
+        scale_fp32,
+        0,  // sliding_window_size (not used)
+        lightweight_mask_enabled,
+        chunked_enabled,
+        chunked_q_local_padded_Nt,
+        chunked_chunk_size_t>(
+        0,  // Skt (not used)
+        qk_in0_block_w,
+        qk_subblock_w,
+        qk_subblock_h,
+        qk_in0_num_subblocks,
+        qk_in1_num_subblocks,
+        qk_num_blocks,
+        out_in0_block_w,
+        out_subblock_w,
+        out_subblock_h,
+        out_in0_num_subblocks,
+        out_in1_num_subblocks,
+        out_num_blocks,
+        global_q_start,  // iter_q_start
+        global_q_end,    // iter_q_end
+        q_num_chunks,    // q_num_chunks (total per-head chunks: local + joint)
+        0,               // local_q_start (not used)
+        0,               // chunked_q_chunk_offset (not used)
+        0,
+        iter_k_chunk_end,
+        q_chunk_tiles,
+        k_chunk_tiles,
+        v_chunk_tiles,
+        qk_chunk_tiles,
+        out_chunk_tiles,
+        0,  // mask_chunk_0 (not used)
+        0,  // mask_chunk_1 (not used)
+        ring_iter,
+        ring_id,
+        num_local_k_chunks,
+        local_padded_Nt,
+        logical_nt,
+        ring_iter_needs_global_n_mask,
+        ring_iter_needs_joint_n_mask,
+        local_n_needs_masking,
+        global_n_mask_chunk_id,
+        local_n_mask_chunk_id,
+        joint_n_mask_chunk_id,
+        cb_q_in,
+        cb_k_in,
+        cb_v_in,
+        cb_mask_in,
+        cb_col_identity,
+        cb_out_im_A,
+        cb_out_im_B,
+        cb_max_A,
+        cb_max_B,
+        cb_sum_A,
+        cb_sum_B,
+        cb_exp_max_diff,
+        cb_lse_in,
+        cb_lse_out,
+        cb_prev_out,
+        cb_out,
+        lw_mask,
+        is_causal_ring_iter,
+        skip_first_half_q,
+        use_zigzag_balancing,
+        is_last_ring_iter,
+        chunked);
+}

code/models/demos/mast3r/tt/kernels/fsdpa/sdpa.cpp

@@ -0,0 +1,277 @@
+// mast3r-p150 model-local copy of ttnn sdpa compute (tt-metal 8b98410e730); FSDPA_PROF=1 enables the stock profiling zones.
+// SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include <cstdint>
+
+#define REDUCE_OP (PoolType::MAX)
+#define REDUCE_DIM (ReduceDim::REDUCE_ROW)
+
+#include "api/compute/compute_kernel_api.h"
+#include "api/compute/compute_kernel_hw_startup.h"
+#include "compute_common.hpp"
+#include "compute_streaming.hpp"
+
+void kernel_main() {
+    constexpr uint32_t B = get_compile_time_arg_val(0);
+    constexpr uint32_t NQH = get_compile_time_arg_val(1);
+    constexpr uint32_t NKH = get_compile_time_arg_val(2);
+    constexpr uint32_t Skt = get_compile_time_arg_val(3);
+    constexpr uint32_t DHt = get_compile_time_arg_val(4);
+    constexpr uint32_t vDHt = get_compile_time_arg_val(5);
+    constexpr uint32_t Sq_chunk_t = get_compile_time_arg_val(6);
+    constexpr uint32_t q_num_chunks = get_compile_time_arg_val(7);
+    constexpr uint32_t Sk_chunk_t = get_compile_time_arg_val(8);
+    constexpr uint32_t k_num_chunks = get_compile_time_arg_val(9);
+
+    constexpr uint32_t qk_in0_block_w = get_compile_time_arg_val(10);
+    constexpr uint32_t qk_subblock_w = get_compile_time_arg_val(11);
+    constexpr uint32_t qk_subblock_h = get_compile_time_arg_val(12);
+    constexpr uint32_t qk_in0_num_subblocks = get_compile_time_arg_val(13);
+    constexpr uint32_t qk_in1_num_subblocks = get_compile_time_arg_val(14);
+    constexpr uint32_t qk_num_blocks = get_compile_time_arg_val(15);
+    constexpr uint32_t out_in0_block_w = get_compile_time_arg_val(16);
+    constexpr uint32_t out_subblock_w = get_compile_time_arg_val(17);
+    constexpr uint32_t out_subblock_h = get_compile_time_arg_val(18);
+    constexpr uint32_t out_in0_num_subblocks = get_compile_time_arg_val(19);
+    constexpr uint32_t out_in1_num_subblocks = get_compile_time_arg_val(20);
+    constexpr uint32_t out_num_blocks = get_compile_time_arg_val(21);
+
+    constexpr uint32_t num_cores = get_compile_time_arg_val(22);
+
+    constexpr bool is_causal = get_compile_time_arg_val(23) == 1;
+    constexpr bool use_provided_mask = get_compile_time_arg_val(24) == 1;
+    constexpr bool use_padded_mask = get_compile_time_arg_val(25) == 1;
+    constexpr bool is_chunked = get_compile_time_arg_val(26) == 1;
+    constexpr uint32_t scale_fp32 = get_compile_time_arg_val(27);
+    constexpr uint32_t sliding_window_size = get_compile_time_arg_val(28);
+    constexpr bool use_attention_sink = get_compile_time_arg_val(29) == 1;
+    constexpr bool use_streaming_compute = get_compile_time_arg_val(30) == 1;
+    constexpr uint32_t valid_Skt = get_compile_time_arg_val(31);
+    constexpr uint32_t k_partial_col = get_compile_time_arg_val(32);
+    // Zigzag remap flag drives the external remap_q_index call on the flat B*NQH*q_num_chunks range.
+    constexpr bool use_zigzag_balancing = get_compile_time_arg_val(33) == 1;
+
+    const uint32_t core_id = get_arg_val<uint32_t>(0);
+    const uint32_t num_phases = get_arg_val<uint32_t>(1);
+    const uint32_t use_chunk_start_idx_tensor = get_arg_val<uint32_t>(2);
+    uint32_t chunked_q_chunk_offset_phase_1 = get_arg_val<uint32_t>(3);
+    uint32_t chunked_q_chunk_offset_phase_2 = 0;
+    if (num_phases == 2) {
+        chunked_q_chunk_offset_phase_2 = get_arg_val<uint32_t>(4);
+    }
+
+    // Global Q scheduling args follow phase_2 slot.
+    const uint32_t global_q_start = get_arg_val<uint32_t>(5);
+    const uint32_t global_q_count = get_arg_val<uint32_t>(6);
+
+    constexpr uint32_t q_chunk_tiles = Sq_chunk_t * DHt;
+    constexpr uint32_t k_chunk_tiles = Sk_chunk_t * DHt;
+    constexpr uint32_t v_chunk_tiles = Sk_chunk_t * vDHt;
+    constexpr uint32_t qk_chunk_tiles = Sq_chunk_t * Sk_chunk_t;
+    constexpr uint32_t out_chunk_tiles = Sq_chunk_t * vDHt;
+
+    constexpr uint32_t cb_arg_offset = 34;
+    constexpr uint32_t cb_q_in = get_compile_time_arg_val(cb_arg_offset + 0);
+    constexpr uint32_t cb_k_in = get_compile_time_arg_val(cb_arg_offset + 1);
+    constexpr uint32_t cb_v_in = get_compile_time_arg_val(cb_arg_offset + 2);
+    constexpr uint32_t cb_mask_in = get_compile_time_arg_val(cb_arg_offset + 3);
+    constexpr uint32_t cb_attention_sink = get_compile_time_arg_val(cb_arg_offset + 4);
+    constexpr uint32_t cb_identity_scale_in = get_compile_time_arg_val(cb_arg_offset + 5);
+    constexpr uint32_t cb_col_identity = get_compile_time_arg_val(cb_arg_offset + 6);
+    constexpr uint32_t cb_chunk_start_idx = get_compile_time_arg_val(cb_arg_offset + 7);
+    constexpr uint32_t cb_recip_scratch = get_compile_time_arg_val(cb_arg_offset + 8);
+    constexpr uint32_t cb_out = get_compile_time_arg_val(cb_arg_offset + 9);
+    constexpr uint32_t cb_qk_im = get_compile_time_arg_val(cb_arg_offset + 10);
+    constexpr uint32_t cb_out_im_A = get_compile_time_arg_val(cb_arg_offset + 11);
+    constexpr uint32_t cb_out_im_B = get_compile_time_arg_val(cb_arg_offset + 12);
+    constexpr uint32_t cb_max_A = get_compile_time_arg_val(cb_arg_offset + 13);
+    constexpr uint32_t cb_max_B = get_compile_time_arg_val(cb_arg_offset + 14);
+    constexpr uint32_t cb_sum_A = get_compile_time_arg_val(cb_arg_offset + 15);
+    constexpr uint32_t cb_sum_B = get_compile_time_arg_val(cb_arg_offset + 16);
+    constexpr uint32_t cb_exp_max_diff = get_compile_time_arg_val(cb_arg_offset + 17);
+    uint32_t chunked_q_chunk_offset = 0;
+    CircularBuffer cb_chunk_start_idx_obj(cb_chunk_start_idx);
+    CircularBuffer cb_identity_scale_in_obj(cb_identity_scale_in);
+    CircularBuffer cb_mask_in_obj(cb_mask_in);
+    compute_kernel_hw_startup<SrcOrder::Reverse>(cb_q_in, cb_k_in, cb_out);
+    matmul_init(cb_q_in, cb_k_in);
+
+    if constexpr (is_chunked) {
+        if (use_chunk_start_idx_tensor != 0) {
+            cb_chunk_start_idx_obj.wait_front(1);
+            uint32_t chunk_start_idx = ckernel::read_tile_value(cb_chunk_start_idx, 0, 0);
+            cb_chunk_start_idx_obj.pop_front(1);
+            const uint32_t q_chunk_size = Sq_chunk_t * TILE_HEIGHT;
+            chunked_q_chunk_offset_phase_1 = chunk_start_idx / q_chunk_size;
+            if (num_phases == 2) {
+                chunked_q_chunk_offset_phase_2 = chunked_q_chunk_offset_phase_1;
+            }
+        }
+    }
+
+    if constexpr (use_streaming_compute) {
+        // Streaming SDPA v2: direct cb_qkt_im writes via cb_push_back_hold_wr_ptr.
+        // No row buffers needed; a dedicated 1-tile CB is used as recip scratch.
+
+        // Wait once for identity scale; v2 removes per-call waits inside reduce_c_row_group
+        cb_identity_scale_in_obj.wait_front(1);
+
+        // Lightweight-mask context: writer pre-generates either [neginf, causal_diag, partial?]
+        // or, for sliding, [neginf, trailing_primary, leading_prev, leading_current, trailing_next, partial?].
+        // primary_diag_tile_idx is the per-layout tile used for the row-local diagonal stamp.
+        LightweightMaskContext lw_mask;
+        uint32_t lw_mask_tile_count = 1;
+        lw_mask.neginf_tile_idx = 0;
+        lw_mask.is_causal = is_causal;
+        if constexpr (sliding_window_size > 0) {
+            lw_mask.primary_diag_tile_idx = 1;
+            lw_mask.sliding_leading_prev_tile_idx = 2;
+            lw_mask.sliding_leading_tile_idx = 3;
+            lw_mask.sliding_trailing_next_tile_idx = 4;
+            lw_mask_tile_count = 5;
+        } else if constexpr (is_causal) {
+            lw_mask.causal_diag_tile_idx = lw_mask_tile_count++;
+            lw_mask.primary_diag_tile_idx = lw_mask.causal_diag_tile_idx;
+        }
+        if constexpr (k_partial_col > 0) {
+            lw_mask.global_n_partial_col = k_partial_col;
+            lw_mask.global_n_partial_tile_idx = lw_mask_tile_count++;
+            // global_n_padded_tiles = Sk_chunk_t - valid_tiles_in_last_chunk
+            constexpr uint32_t last_chunk_first_tile =
+                (valid_Skt > Sk_chunk_t) ? ((valid_Skt - 1) / Sk_chunk_t) * Sk_chunk_t : 0u;
+            constexpr uint32_t valid_tiles_in_last_chunk = valid_Skt - last_chunk_first_tile;
+            lw_mask.global_n_padded_tiles = Sk_chunk_t - valid_tiles_in_last_chunk;
+        }
+        // A user-provided dense mask is streamed per-chunk by the reader and consumed inside the
+        // inner loop — it does not use the writer-generated lightweight palette, so skip this wait.
+        if constexpr ((is_causal || sliding_window_size > 0 || k_partial_col > 0) && !use_provided_mask) {
+            cb_mask_in_obj.wait_front(lw_mask_tile_count);
+        }
+
+        // Global Q scheduling: sdpa_standard_v2 walks the per-core flat range over
+        // B*NQH*q_num_chunks chunks; the modulo inside its inner loop extracts the per-head q_chunk
+        // from each flat index. num_phases==1 is pinned for streaming, so the chunked offset comes
+        // from phase 1.
+        sdpa_standard_v2<
+            Sq_chunk_t,
+            Sk_chunk_t,
+            valid_Skt,
+            DHt,
+            vDHt,
+            scale_fp32,
+            qk_subblock_h,
+            qk_subblock_w,
+            out_subblock_h,
+            out_subblock_w,
+            use_padded_mask,
+            cb_q_in,
+            cb_k_in,
+            cb_v_in,
+            cb_qk_im,
+            cb_identity_scale_in,
+            cb_exp_max_diff,
+            cb_col_identity,
+            cb_recip_scratch,
+            cb_out,  // normalized output goes directly to output CB
+            cb_mask_in,
+            sliding_window_size,
+            is_causal,
+            use_attention_sink,
+            cb_attention_sink,
+            use_provided_mask>(
+            global_q_count,
+            k_num_chunks,
+            cb_out_im_A,
+            cb_out_im_B,
+            cb_max_A,
+            cb_max_B,
+            cb_sum_A,
+            cb_sum_B,
+            global_q_start,
+            chunked_q_chunk_offset_phase_1,
+            lw_mask,
+            q_num_chunks,
+            use_zigzag_balancing);
+    } else {
+        // Standard SDPA path (causal, masked, chunked, etc.)
+        constexpr bool use_lightweight_causal_mask = is_causal && !use_provided_mask && (sliding_window_size == 0);
+
+        LightweightMaskContext lw_mask;
+        if constexpr (use_lightweight_causal_mask) {
+            lw_mask.is_causal = true;
+            lw_mask.neginf_tile_idx = 0;
+            lw_mask.causal_diag_tile_idx = 1;
+            lw_mask.primary_diag_tile_idx = lw_mask.causal_diag_tile_idx;
+            cb_mask_in_obj.wait_front(2);
+        }
+
+        for (uint32_t phase = 0; phase < num_phases; ++phase) {
+            if (phase == 0) {
+                chunked_q_chunk_offset = chunked_q_chunk_offset_phase_1;
+            } else {
+                chunked_q_chunk_offset = chunked_q_chunk_offset_phase_2;
+            }
+
+            // Global Q scheduling: sdpa_standard walks the per-core flat range over
+            // B*NQH*q_num_chunks chunks; the modulo inside its inner loop extracts the per-head
+            // q_chunk from each flat index.
+            sdpa_standard<
+                cb_qk_im,
+                cb_identity_scale_in,
+                cb_attention_sink,
+                Sq_chunk_t,
+                Sk_chunk_t,
+                DHt,
+                vDHt,
+                use_attention_sink,
+                is_causal,
+                use_provided_mask,
+                use_padded_mask,
+                is_chunked,
+                scale_fp32,
+                sliding_window_size,
+                use_lightweight_causal_mask>(
+                Skt,
+                qk_in0_block_w,
+                qk_subblock_w,
+                qk_subblock_h,
+                qk_in0_num_subblocks,
+                qk_in1_num_subblocks,
+                qk_num_blocks,
+                out_in0_block_w,
+                out_subblock_w,
+                out_subblock_h,
+                out_in0_num_subblocks,
+                out_in1_num_subblocks,
+                out_num_blocks,
+                /*iter_q_start=*/0,
+                /*iter_q_end=*/global_q_count,
+                q_num_chunks,
+                /*local_q_start=*/global_q_start,
+                chunked_q_chunk_offset,
+                k_num_chunks,
+                q_chunk_tiles,
+                k_chunk_tiles,
+                v_chunk_tiles,
+                qk_chunk_tiles,
+                out_chunk_tiles,
+                cb_q_in,
+                cb_k_in,
+                cb_v_in,
+                cb_mask_in,
+                cb_col_identity,
+                cb_out_im_A,
+                cb_out_im_B,
+                cb_max_A,
+                cb_max_B,
+                cb_sum_A,
+                cb_sum_B,
+                cb_exp_max_diff,
+                cb_out,
+                lw_mask,
+                use_zigzag_balancing);
+        }
+    }
+}

code/models/demos/mast3r/tt/kernels/heads_concat_reader.cpp

@@ -0,0 +1,38 @@
+// SPDX-License-Identifier: Apache-2.0
+// mast3r-p150 model-local kernel (ttnn.generic_op): concatenate heads, reader.
+// Unit u = (b, s, h) over [B, St, H]: reads the Wt tiles of head h, seq-tile s, batch b of the [B, H, S, Wt*32]
+// TILE input (tile ((b*H + h)*St + s)*Wt) into CB 0. Units are processed in blocks of UB to batch NoC barriers.
+#include <stdint.h>
+#include "api/dataflow/dataflow_api.h"
+
+void kernel_main() {
+    constexpr uint32_t H = get_compile_time_arg_val(0);
+    constexpr uint32_t St = get_compile_time_arg_val(1);
+    constexpr uint32_t Wt = get_compile_time_arg_val(2);
+    constexpr uint32_t UB = get_compile_time_arg_val(3);
+    constexpr uint32_t LOG2_TILE = 11;
+    constexpr uint32_t tile_bytes = 1u << LOG2_TILE;
+    const uint32_t in_addr = get_arg_val<uint32_t>(0);
+    const uint32_t u_start = get_arg_val<uint32_t>(1);
+    const uint32_t u_count = get_arg_val<uint32_t>(2);
+    const InterleavedPow2AddrGen<IN_DRAM> g = {.bank_base_address = in_addr, .log_base_2_of_page_size = LOG2_TILE};
+    const uint32_t u_end = u_start + u_count;
+    for (uint32_t u0 = u_start; u0 < u_end; u0 += UB) {
+        const uint32_t n = (u_end - u0) < UB ? (u_end - u0) : UB;
+        cb_reserve_back(0, n * Wt);
+        uint32_t w = get_write_ptr(0);
+        for (uint32_t u = u0; u < u0 + n; ++u) {
+            const uint32_t b = u / (St * H);
+            const uint32_t r = u - b * (St * H);
+            const uint32_t s = r / H;
+            const uint32_t h = r - s * H;
+            const uint32_t it = ((b * H + h) * St + s) * Wt;
+            for (uint32_t j = 0; j < Wt; ++j) {
+                noc_async_read(g.get_noc_addr(it + j), w, tile_bytes);
+                w += tile_bytes;
+            }
+        }
+        noc_async_read_barrier();
+        cb_push_back(0, n * Wt);
+    }
+}

code/models/demos/mast3r/tt/kernels/heads_concat_writer.cpp

@@ -0,0 +1,43 @@
+// SPDX-License-Identifier: Apache-2.0
+// mast3r-p150 model-local kernel (ttnn.generic_op): concatenate heads, writer.
+// Writes unit (b, s, h) to output tensor of batch b at tile out_base[b] + s*out_w[b] + h*Wt.
+#include <stdint.h>
+#include "api/dataflow/dataflow_api.h"
+
+void kernel_main() {
+    constexpr uint32_t H = get_compile_time_arg_val(0);
+    constexpr uint32_t St = get_compile_time_arg_val(1);
+    constexpr uint32_t Wt = get_compile_time_arg_val(2);
+    constexpr uint32_t UB = get_compile_time_arg_val(3);
+    constexpr uint32_t NB = get_compile_time_arg_val(4);
+    constexpr uint32_t LOG2_TILE = 11;
+    constexpr uint32_t tile_bytes = 1u << LOG2_TILE;
+    uint32_t a = 0;
+    const uint32_t u_start = get_arg_val<uint32_t>(a++);
+    const uint32_t u_count = get_arg_val<uint32_t>(a++);
+    uint32_t o_addr[NB], o_base[NB], o_w[NB];
+    for (uint32_t b = 0; b < NB; ++b) {
+        o_addr[b] = get_arg_val<uint32_t>(a++); o_base[b] = get_arg_val<uint32_t>(a++); o_w[b] = get_arg_val<uint32_t>(a++);
+    }
+    const uint32_t u_end = u_start + u_count;
+    for (uint32_t u0 = u_start; u0 < u_end; u0 += UB) {
+        const uint32_t n = (u_end - u0) < UB ? (u_end - u0) : UB;
+        cb_wait_front(0, n * Wt);
+        uint32_t rd = get_read_ptr(0);
+        for (uint32_t u = u0; u < u0 + n; ++u) {
+            const uint32_t b = u / (St * H);
+            const uint32_t r = u - b * (St * H);
+            const uint32_t s = r / H;
+            const uint32_t h = r - s * H;
+            const InterleavedPow2AddrGen<OUT_DRAM> g = {.bank_base_address = o_addr[b], .log_base_2_of_page_size = LOG2_TILE};
+            const uint32_t ot = o_base[b] + s * o_w[b] + h * Wt;
+            for (uint32_t j = 0; j < Wt; ++j) {
+                noc_async_write(rd, g.get_noc_addr(ot + j), tile_bytes);
+                rd += tile_bytes;
+            }
+        }
+        noc_async_writes_flushed();
+        cb_pop_front(0, n * Wt);
+    }
+    noc_async_write_barrier();
+}

code/models/demos/mast3r/tt/kernels/heads_rope_compute.cpp

@@ -0,0 +1,180 @@
+// SPDX-License-Identifier: Apache-2.0
+// mast3r-p150 model-local kernel (ttnn.generic_op): fused split-heads + RoPE, compute.
+// Per unit, for q then k (Wt tiles each): the same op sequence as ttnn's rotary_embedding_llama compute kernel
+// (rotated = x @ trans_mat; sin_i = rotated * sin; cos_i = x * cos; out = cos_i + sin_i, each packed to bf16),
+// so the result is bit-identical to split_heads + rotary_embedding_llama with the same compute config.
+#include <cstdint>
+#include "api/compute/common.h"
+#include "api/compute/eltwise_binary.h"
+#include "api/compute/bcast.h"
+#include "api/compute/matmul.h"
+#include "api/compute/compute_kernel_hw_startup.h"
+#include "api/compute/reconfig_data_format.h"
+#include "api/compute/pack.h"
+#include "api/compute/cb_api.h"
+
+constexpr uint32_t CB_IN = 0, CB_COS = 1, CB_SIN = 2, CB_TRANS = 3;
+constexpr uint32_t CB_OUT = 16, CB_ROT = 24, CB_COSI = 25, CB_SINI = 26;
+
+#ifdef D_NOCOMP
+#define MATH_OP(x) ((void)0)
+#else
+#define MATH_OP(x) x
+#endif
+ALWI void ACQ() { tile_regs_acquire(); tile_regs_wait(); }
+ALWI void REL() { tile_regs_commit(); tile_regs_release(); }
+
+void kernel_main() {
+    constexpr uint32_t H = get_compile_time_arg_val(0);
+    constexpr uint32_t St = get_compile_time_arg_val(1);
+    constexpr uint32_t Wt = get_compile_time_arg_val(2);
+    const uint32_t u_start = get_arg_val<uint32_t>(0);
+    const uint32_t u_count = get_arg_val<uint32_t>(1);
+    if (u_count == 0) {
+        return;
+    }
+    compute_kernel_hw_startup<SrcOrder::Reverse>(CB_IN, CB_TRANS, CB_OUT);
+    compute_kernel_hw_startup(CB_COSI, CB_SINI, CB_OUT);
+    cb_wait_front(CB_TRANS, 1);
+    uint32_t last_s = 0xFFFFFFFF;
+    for (uint32_t u = u_start; u < u_start + u_count; ++u) {
+        const uint32_t b = u / (St * H);
+        const uint32_t r = u - b * (St * H);
+        const uint32_t s = r / H;
+        if (s != last_s) {
+            if (last_s != 0xFFFFFFFF) {
+                cb_pop_front(CB_SIN, Wt);
+                cb_pop_front(CB_COS, Wt);
+            }
+            cb_wait_front(CB_SIN, Wt);
+            cb_wait_front(CB_COS, Wt);
+            last_s = s;
+        }
+#ifdef HROPE_QK2
+        {
+            // MAST3R_OPT "hrqk": q and k of the unit in ONE pass per stage (2*Wt tiles per dest acquire; fp32 dest
+            // half = 4 tiles = 2*Wt for dh = 64). Same per-tile op sequence -> bit-identical.
+            constexpr uint32_t NT = 2 * Wt;
+            static_assert(NT <= 4, "q + k of one unit must fit one fp32 dest half (4 tiles)");
+            cb_wait_front(CB_IN, NT);
+            cb_reserve_back(CB_ROT, NT);
+            cb_reserve_back(CB_SINI, NT);
+            cb_reserve_back(CB_COSI, NT);
+            cb_reserve_back(CB_OUT, NT);
+
+            reconfig_data_format(CB_COSI, CB_TRANS, CB_SINI, CB_IN);
+            pack_reconfig_data_format(CB_OUT, CB_ROT);
+            matmul_init(CB_IN, CB_TRANS);
+            ACQ();
+            for (uint32_t j = 0; j < NT; ++j) {
+                MATH_OP(matmul_tiles(CB_IN, CB_TRANS, j, 0, j));
+                pack_tile(j, CB_ROT, j);
+            }
+            REL();
+            cb_push_back(CB_ROT, NT);
+            cb_wait_front(CB_ROT, NT);
+
+            reconfig_data_format(CB_TRANS, CB_ROT, CB_IN, CB_SIN);
+            pack_reconfig_data_format(CB_ROT, CB_SINI);
+            mul_init(CB_ROT, CB_SIN);
+            ACQ();
+            for (uint32_t j = 0; j < NT; ++j) {
+                MATH_OP(mul_tiles(CB_ROT, CB_SIN, j, j % Wt, j));
+                pack_tile(j, CB_SINI, j);
+            }
+            REL();
+            cb_push_back(CB_SINI, NT);
+            cb_pop_front(CB_ROT, NT);
+
+            reconfig_data_format(CB_ROT, CB_IN, CB_SIN, CB_COS);
+            pack_reconfig_data_format(CB_SINI, CB_COSI);
+            mul_init(CB_IN, CB_COS);
+            ACQ();
+            for (uint32_t j = 0; j < NT; ++j) {
+                MATH_OP(mul_tiles(CB_IN, CB_COS, j, j % Wt, j));
+                pack_tile(j, CB_COSI, j);
+            }
+            REL();
+            cb_push_back(CB_COSI, NT);
+            cb_pop_front(CB_IN, NT);
+
+            cb_wait_front(CB_SINI, NT);
+            cb_wait_front(CB_COSI, NT);
+            reconfig_data_format(CB_IN, CB_COSI, CB_COS, CB_SINI);
+            pack_reconfig_data_format(CB_COSI, CB_OUT);
+            add_init(CB_COSI, CB_SINI);
+            ACQ();
+            for (uint32_t j = 0; j < NT; ++j) {
+                MATH_OP(add_tiles(CB_COSI, CB_SINI, j, j, j));
+                pack_tile(j, CB_OUT, j);
+            }
+            REL();
+            cb_push_back(CB_OUT, NT);
+            cb_pop_front(CB_SINI, NT);
+            cb_pop_front(CB_COSI, NT);
+        }
+#else
+        for (uint32_t t = 0; t < 2; ++t) {  // q, k
+            cb_wait_front(CB_IN, Wt);
+            cb_reserve_back(CB_ROT, Wt);
+            cb_reserve_back(CB_SINI, Wt);
+            cb_reserve_back(CB_COSI, Wt);
+            cb_reserve_back(CB_OUT, Wt);
+
+            reconfig_data_format(CB_COSI, CB_TRANS, CB_SINI, CB_IN);
+            pack_reconfig_data_format(CB_OUT, CB_ROT);
+            matmul_init(CB_IN, CB_TRANS);
+            ACQ();
+            for (uint32_t j = 0; j < Wt; ++j) {
+                MATH_OP(matmul_tiles(CB_IN, CB_TRANS, j, 0, j));
+                pack_tile(j, CB_ROT, j);
+            }
+            REL();
+            cb_push_back(CB_ROT, Wt);
+            cb_wait_front(CB_ROT, Wt);
+
+            reconfig_data_format(CB_TRANS, CB_ROT, CB_IN, CB_SIN);
+            pack_reconfig_data_format(CB_ROT, CB_SINI);
+            mul_init(CB_ROT, CB_SIN);
+            ACQ();
+            for (uint32_t j = 0; j < Wt; ++j) {
+                MATH_OP(mul_tiles(CB_ROT, CB_SIN, j, j, j));
+                pack_tile(j, CB_SINI, j);
+            }
+            REL();
+            cb_push_back(CB_SINI, Wt);
+            cb_pop_front(CB_ROT, Wt);
+
+            reconfig_data_format(CB_ROT, CB_IN, CB_SIN, CB_COS);
+            pack_reconfig_data_format(CB_SINI, CB_COSI);
+            mul_init(CB_IN, CB_COS);
+            ACQ();
+            for (uint32_t j = 0; j < Wt; ++j) {
+                MATH_OP(mul_tiles(CB_IN, CB_COS, j, j, j));
+                pack_tile(j, CB_COSI, j);
+            }
+            REL();
+            cb_push_back(CB_COSI, Wt);
+            cb_pop_front(CB_IN, Wt);
+
+            cb_wait_front(CB_SINI, Wt);
+            cb_wait_front(CB_COSI, Wt);
+            reconfig_data_format(CB_IN, CB_COSI, CB_COS, CB_SINI);
+            pack_reconfig_data_format(CB_COSI, CB_OUT);
+            add_init(CB_COSI, CB_SINI);
+            ACQ();
+            for (uint32_t j = 0; j < Wt; ++j) {
+                MATH_OP(add_tiles(CB_COSI, CB_SINI, j, j, j));
+                pack_tile(j, CB_OUT, j);
+            }
+            REL();
+            cb_push_back(CB_OUT, Wt);
+            cb_pop_front(CB_SINI, Wt);
+            cb_pop_front(CB_COSI, Wt);
+        }
+#endif
+    }
+    cb_pop_front(CB_SIN, Wt);
+    cb_pop_front(CB_COS, Wt);
+    cb_pop_front(CB_TRANS, 1);
+}

code/models/demos/mast3r/tt/kernels/heads_rope_reader.cpp

@@ -0,0 +1,110 @@
+// SPDX-License-Identifier: Apache-2.0
+// mast3r-p150 model-local kernel (ttnn.generic_op): fused split-heads + RoPE, reader.
+// Work unit u = (b, s, h) in row-major order over [B, St, H]. For each unit it reads the Wt q tiles and
+// Wt k tiles of head h at seq-tile s of batch b (from per-batch source tensors with arbitrary row width and
+// column offset) into CB_IN, the Wt v tiles into CB_V, and (when s changes) the cos/sin rows into CB_COS/CB_SIN.
+#include <stdint.h>
+#include "api/dataflow/dataflow_api.h"
+
+constexpr uint32_t CB_IN = 0, CB_COS = 1, CB_SIN = 2, CB_TRANS = 3, CB_V = 4;
+
+void kernel_main() {
+    constexpr uint32_t H = get_compile_time_arg_val(0);
+    constexpr uint32_t St = get_compile_time_arg_val(1);
+    constexpr uint32_t Wt = get_compile_time_arg_val(2);
+    constexpr uint32_t NB = get_compile_time_arg_val(3);  // batch entries described in the rt args
+    constexpr uint32_t LOG2_TILE = 11;                     // bf16 32x32 tile = 2048 B
+
+    // per-core runtime args: unit range; common runtime args (same for every core): addresses / source layout
+    const uint32_t u_start = get_arg_val<uint32_t>(0);
+    const uint32_t u_count = get_arg_val<uint32_t>(1);
+    uint32_t a = 0;
+    const uint32_t cos_addr = get_common_arg_val<uint32_t>(a++);
+    const uint32_t sin_addr = get_common_arg_val<uint32_t>(a++);
+    const uint32_t trans_addr = get_common_arg_val<uint32_t>(a++);
+    // per batch entry and role (q, k, v): address, tile base, row width (tiles), source kind
+    // (kind 0 = interleaved, 1 / 2 = the 1st / 2nd TensorAccessor layout in the compile-time args: MAST3R_OPT "obs",
+    // block-sharded matmul outputs)
+    uint32_t q_addr[NB], q_base[NB], q_w[NB], k_addr[NB], k_base[NB], k_w[NB], v_addr[NB], v_base[NB], v_w[NB];
+    uint32_t q_kind[NB], k_kind[NB], v_kind[NB];
+    for (uint32_t b = 0; b < NB; ++b) {
+        q_addr[b] = get_common_arg_val<uint32_t>(a++); q_base[b] = get_common_arg_val<uint32_t>(a++); q_w[b] = get_common_arg_val<uint32_t>(a++);
+        q_kind[b] = get_common_arg_val<uint32_t>(a++);
+        k_addr[b] = get_common_arg_val<uint32_t>(a++); k_base[b] = get_common_arg_val<uint32_t>(a++); k_w[b] = get_common_arg_val<uint32_t>(a++);
+        k_kind[b] = get_common_arg_val<uint32_t>(a++);
+        v_addr[b] = get_common_arg_val<uint32_t>(a++); v_base[b] = get_common_arg_val<uint32_t>(a++); v_w[b] = get_common_arg_val<uint32_t>(a++);
+        v_kind[b] = get_common_arg_val<uint32_t>(a++);
+    }
+    if (u_count == 0) {
+        return;
+    }
+#if NTA >= 1
+    constexpr auto ta1_args = TensorAccessorArgs<4>();
+#endif
+#if NTA >= 2
+    constexpr auto ta2_args = TensorAccessorArgs<ta1_args.next_compile_time_args_offset()>();
+#endif
+    auto src_noc = [&](uint32_t kind, uint32_t addr, uint32_t tile) -> uint64_t {
+#if NTA >= 1
+        if (kind == 1) {
+            const auto ta = TensorAccessor(ta1_args, addr, 1u << LOG2_TILE);
+            return ta.get_noc_addr(tile);
+        }
+#endif
+#if NTA >= 2
+        if (kind == 2) {
+            const auto ta = TensorAccessor(ta2_args, addr, 1u << LOG2_TILE);
+            return ta.get_noc_addr(tile);
+        }
+#endif
+        const InterleavedPow2AddrGen<SRC_DRAM> g = {.bank_base_address = addr, .log_base_2_of_page_size = LOG2_TILE};
+        return g.get_noc_addr(tile);
+    };
+    const InterleavedPow2AddrGen<COS_DRAM> cos_g = {.bank_base_address = cos_addr, .log_base_2_of_page_size = LOG2_TILE};
+    const InterleavedPow2AddrGen<COS_DRAM> sin_g = {.bank_base_address = sin_addr, .log_base_2_of_page_size = LOG2_TILE};
+    const InterleavedPow2AddrGen<TRANS_DRAM> tr_g = {.bank_base_address = trans_addr, .log_base_2_of_page_size = LOG2_TILE};
+    constexpr uint32_t tile_bytes = 1u << LOG2_TILE;
+
+    cb_reserve_back(CB_TRANS, 1);
+    noc_async_read(tr_g.get_noc_addr(0), get_write_ptr(CB_TRANS), tile_bytes);
+    noc_async_read_barrier();
+    cb_push_back(CB_TRANS, 1);
+
+    uint32_t last_s = 0xFFFFFFFF;
+    for (uint32_t u = u_start; u < u_start + u_count; ++u) {
+        const uint32_t b = u / (St * H);
+        const uint32_t r = u - b * (St * H);
+        const uint32_t s = r / H;
+        const uint32_t h = r - s * H;
+        if (s != last_s) {
+            cb_reserve_back(CB_COS, Wt);
+            cb_reserve_back(CB_SIN, Wt);
+            uint32_t cw = get_write_ptr(CB_COS), sw = get_write_ptr(CB_SIN);
+            for (uint32_t j = 0; j < Wt; ++j) {
+                noc_async_read(cos_g.get_noc_addr(s * Wt + j), cw + j * tile_bytes, tile_bytes);
+                noc_async_read(sin_g.get_noc_addr(s * Wt + j), sw + j * tile_bytes, tile_bytes);
+            }
+        }
+        cb_reserve_back(CB_IN, 2 * Wt);
+        cb_reserve_back(CB_V, Wt);
+        uint32_t iw = get_write_ptr(CB_IN), vw = get_write_ptr(CB_V);
+        const uint32_t qt = q_base[b] + s * q_w[b] + h * Wt;
+        const uint32_t kt = k_base[b] + s * k_w[b] + h * Wt;
+        const uint32_t vt = v_base[b] + s * v_w[b] + h * Wt;
+#ifndef D_NOREAD
+        for (uint32_t j = 0; j < Wt; ++j) {
+            noc_async_read(src_noc(q_kind[b], q_addr[b], qt + j), iw + j * tile_bytes, tile_bytes);
+            noc_async_read(src_noc(k_kind[b], k_addr[b], kt + j), iw + (Wt + j) * tile_bytes, tile_bytes);
+            noc_async_read(src_noc(v_kind[b], v_addr[b], vt + j), vw + j * tile_bytes, tile_bytes);
+        }
+#endif
+        noc_async_read_barrier();
+        if (s != last_s) {
+            cb_push_back(CB_COS, Wt);
+            cb_push_back(CB_SIN, Wt);
+            last_s = s;
+        }
+        cb_push_back(CB_IN, 2 * Wt);
+        cb_push_back(CB_V, Wt);
+    }
+}

code/models/demos/mast3r/tt/kernels/heads_rope_writer.cpp

@@ -0,0 +1,57 @@
+// SPDX-License-Identifier: Apache-2.0
+// mast3r-p150 model-local kernel (ttnn.generic_op): fused split-heads + RoPE, writer.
+// Writes rotated q (Wt tiles), rotated k (Wt tiles) from CB_OUT and v (Wt tiles) from CB_V of unit (b, s, h)
+// to the [B, H, St*32, Wt*32] TILE outputs at tile ((b*H + h)*St + s)*Wt.
+#include <stdint.h>
+#include "api/dataflow/dataflow_api.h"
+
+constexpr uint32_t CB_V = 4, CB_OUT = 16;
+#ifdef D_NOWRITE
+#define WR(a, b, c) ((void)(b))
+#else
+#define WR(a, b, c) noc_async_write(a, b, c)
+#endif
+
+void kernel_main() {
+    constexpr uint32_t H = get_compile_time_arg_val(0);
+    constexpr uint32_t St = get_compile_time_arg_val(1);
+    constexpr uint32_t Wt = get_compile_time_arg_val(2);
+    constexpr uint32_t LOG2_TILE = 11;
+    constexpr uint32_t tile_bytes = 1u << LOG2_TILE;
+    const uint32_t q_addr = get_common_arg_val<uint32_t>(0);
+    const uint32_t k_addr = get_common_arg_val<uint32_t>(1);
+    const uint32_t v_addr = get_common_arg_val<uint32_t>(2);
+    const uint32_t u_start = get_arg_val<uint32_t>(0);
+    const uint32_t u_count = get_arg_val<uint32_t>(1);
+    const InterleavedPow2AddrGen<OUT_DRAM> qg = {.bank_base_address = q_addr, .log_base_2_of_page_size = LOG2_TILE};
+    const InterleavedPow2AddrGen<OUT_DRAM> kg = {.bank_base_address = k_addr, .log_base_2_of_page_size = LOG2_TILE};
+    const InterleavedPow2AddrGen<OUT_DRAM> vg = {.bank_base_address = v_addr, .log_base_2_of_page_size = LOG2_TILE};
+    for (uint32_t u = u_start; u < u_start + u_count; ++u) {
+        const uint32_t b = u / (St * H);
+        const uint32_t r = u - b * (St * H);
+        const uint32_t s = r / H;
+        const uint32_t h = r - s * H;
+        const uint32_t ot = ((b * H + h) * St + s) * Wt;
+        cb_wait_front(CB_V, Wt);
+        uint32_t rd = get_read_ptr(CB_V);
+        for (uint32_t j = 0; j < Wt; ++j) {
+            WR(rd + j * tile_bytes, vg.get_noc_addr(ot + j), tile_bytes);
+        }
+        cb_wait_front(CB_OUT, Wt);
+        rd = get_read_ptr(CB_OUT);
+        for (uint32_t j = 0; j < Wt; ++j) {
+            WR(rd + j * tile_bytes, qg.get_noc_addr(ot + j), tile_bytes);
+        }
+        noc_async_writes_flushed();
+        cb_pop_front(CB_OUT, Wt);
+        cb_wait_front(CB_OUT, Wt);
+        rd = get_read_ptr(CB_OUT);
+        for (uint32_t j = 0; j < Wt; ++j) {
+            WR(rd + j * tile_bytes, kg.get_noc_addr(ot + j), tile_bytes);
+        }
+        noc_async_writes_flushed();
+        cb_pop_front(CB_OUT, Wt);
+        cb_pop_front(CB_V, Wt);
+    }
+    noc_async_write_barrier();
+}

code/models/demos/mast3r/tt/kernels/ln_add_compute.cpp

@@ -0,0 +1,436 @@
+// SPDX-FileCopyrightText: © 2023 Tenstorrent USA, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+// mast3r-p150 model-local copy of ttnn's layernorm.cpp compute kernel (tt-metal 8b98410e730), used through
+// ttnn.generic_op. Only change: with FUSE_PRE_ADD the pre-add sum x = a + b is ALSO packed (bf16) into CB
+// "cb_res_out", which the model-local writer stores as the new residual stream. One program replaces
+// ttnn.add(residual, y) + ttnn.layer_norm(sum).
+
+#include <cstdint>
+
+#define BCAST_LLKOP EltwiseBinaryType::ELWMUL
+#define BCAST_DIM BroadcastType::COL
+
+#include "api/compute/compute_kernel_api.h"
+#include "api/compute/bcast.h"
+#include "api/compute/eltwise_binary.h"
+#include "api/compute/layernorm.h"
+#ifdef TILIZE_IN
+#include "api/compute/tilize.h"
+#endif
+#ifdef UNTILIZE_OUT
+#include "api/compute/pack_untilize.h"
+#endif
+#include <tt-metalium/constants.hpp>
+#include "ttnn/operations/normalization/kernel_util/compute/numeric.h"
+#include "ttnn/operations/normalization/kernel_util/generic/blocked_range.h"
+#include "ttnn/operations/normalization/kernel_util/generic/bit.h"
+#include "ttnn/operations/normalization/layernorm/device/kernels/layernorm_scaler_tiles.h"
+#include "api/compute/eltwise_unary/sfpu_split_includes.h"
+#include "api/compute/tile_move_copy.h"
+#include "api/compute/eltwise_unary/eltwise_unary.h"
+#include "api/dataflow/dataflow_buffer.h"
+#include "ttnn/operations/normalization/layernorm/device/kernels/compute/layernorm_compute_utils.h"
+
+namespace generic = norm::kernel_util::generic;
+namespace kutil = norm::kernel_util;
+namespace numeric = kutil::compute::numeric;
+namespace policies = kutil::compute::policies;
+
+constexpr uint32_t CB_RES_OUT = 17;
+
+void kernel_main() {
+    uint32_t NCHt = get_arg_val<uint32_t>(0);
+    constexpr uint32_t Wt = get_compile_time_arg_val(0);
+    constexpr uint32_t block_size = get_compile_time_arg_val(1);
+    constexpr uint32_t do_gamma = get_compile_time_arg_val(2);
+    constexpr uint32_t do_beta = get_compile_time_arg_val(3);
+    constexpr bool FLOAT32_DTYPE = get_compile_time_arg_val(4) == 1;
+    constexpr bool FLOAT32_REDUCTION = get_compile_time_arg_val(5) == 1;
+    constexpr bool LEGACY_RSQRT = get_compile_time_arg_val(6) == 1;
+    constexpr uint32_t W = get_compile_time_arg_val(7);
+    constexpr uint32_t tile_width = get_compile_time_arg_val(8);
+
+    // CB indices - configurable via named compile-time args for kernel chaining support
+    constexpr auto dfb_scaler_id = get_named_compile_time_arg_val("cb_scaler");  // single tile generated by the reader
+    constexpr auto dfb_eps_id = get_named_compile_time_arg_val("cb_eps");        // single tile generated by the reader
+    constexpr auto dfb_in_id = get_named_compile_time_arg_val("cb_in");    // input x or a for fused pre-add (x=a+b)
+    constexpr auto dfb_inb_id = get_named_compile_time_arg_val("cb_inb");  // input b for fused pre-add
+    constexpr auto dfb_out_id = get_named_compile_time_arg_val("cb_out");  // output
+    constexpr auto dfb_gamma_id = get_named_compile_time_arg_val("cb_gamma");
+    constexpr auto dfb_beta_id = get_named_compile_time_arg_val("cb_beta");
+#if defined RMSNORM and not defined FUSE_PRE_ADD
+    constexpr uint32_t dfb_xmm_id = dfb_in_id;  // x minus mean
+#else
+    constexpr uint32_t dfb_xmm_id = get_named_compile_time_arg_val("cb_xmm");  // x minus mean
+#endif
+    DataflowBuffer dfb_xmm(dfb_xmm_id);
+    constexpr auto dfb_ex_id = get_named_compile_time_arg_val("cb_ex");          // E[x]
+    constexpr auto dfb_ex2_id = get_named_compile_time_arg_val("cb_ex2");        // E[(x-E[x])^2]
+    constexpr auto dfb_xmm2_id = get_named_compile_time_arg_val("cb_xmm2");      // xmm^2
+    constexpr auto dfb_ex2pe_id = get_named_compile_time_arg_val("cb_ex2pe");    // E[(x-E[x])^2]+eps
+    constexpr auto dfb_fusion_id = get_named_compile_time_arg_val("cb_fusion");  // stream gamma/beta
+    DataflowBuffer dfb_eps(dfb_eps_id);
+    DataflowBuffer dfb_in(dfb_in_id);
+    DataflowBuffer dfb_inb(dfb_inb_id);
+    DataflowBuffer dfb_out(dfb_out_id);
+    DataflowBuffer dfb_gamma(dfb_gamma_id);
+    DataflowBuffer dfb_beta(dfb_beta_id);
+    DataflowBuffer dfb_ex(dfb_ex_id);
+    DataflowBuffer dfb_ex2(dfb_ex2_id);
+    DataflowBuffer dfb_xmm2(dfb_xmm2_id);
+    DataflowBuffer dfb_ex2pe(dfb_ex2pe_id);
+    DataflowBuffer dfb_fusion(dfb_fusion_id);
+    DataflowBuffer dfb_scaler(dfb_scaler_id);
+
+    constexpr auto dfb_in_rm_id =
+        get_named_compile_time_arg_val("cb_in_rm");  // input row-major (if row-major input, otherwise unused)
+    DataflowBuffer dfb_in_rm(dfb_in_rm_id);
+
+    constexpr int onetile = 1;
+    constexpr int dst0 = 0;
+    constexpr int dst1 = 1;
+    constexpr auto scaler0 = 0;
+
+#ifdef FUSE_PRE_ADD
+#ifdef RMSNORM
+    constexpr uint32_t dfb_x_id = dfb_xmm_id;
+#else
+    constexpr uint32_t dfb_x_id = get_named_compile_time_arg_val("cb_x");
+#endif
+#else
+    constexpr uint32_t dfb_x_id = dfb_in_id;
+#endif
+    DataflowBuffer dfb_x(dfb_x_id);
+
+#ifdef TILIZE_IN
+    compute_kernel_hw_startup(dfb_in_rm_id, dfb_in_rm_id, dfb_in_id);
+#elif defined(FUSE_PRE_ADD)
+    compute_kernel_hw_startup(dfb_in_id, dfb_inb_id, dfb_x_id);
+#elif defined(RMSNORM)
+    compute_kernel_hw_startup(dfb_xmm_id, dfb_xmm_id, dfb_xmm2_id);
+#else
+    compute_kernel_hw_startup(dfb_x_id, dfb_scaler_id, dfb_ex_id);
+#endif
+
+    dfb_eps.wait_front(1);  // comes from the reader
+
+    constexpr int dfb_im_or_out_id = (do_gamma | do_beta) ? dfb_fusion_id : dfb_out_id;
+    DataflowBuffer dfb_im_or_out(dfb_im_or_out_id);
+
+    // Intermediate buffers need to be reserved/pushed/popped
+    // in full blocks
+    const auto total_buffer_size = generic::blocks(Wt, block_size).total_with_remainder();
+
+    for (uint32_t ncht = 0; ncht < NCHt; ncht++) {
+#ifdef TILIZE_IN
+        tilize_all_blocks_to_cb<block_size>(dfb_in_rm, dfb_in, Wt);
+        // Re-init binary ops after tilize hardware reconfiguration.
+#ifdef FUSE_PRE_ADD
+        // TODO(#52395): compute_kernel_hw_startup is a call-once API; this mid-kernel re-init (preserving the
+        // pre-cleanup full-init behaviour) should become a targeted DST re-arm.
+        compute_kernel_hw_startup(dfb_in_id, dfb_inb_id, dfb_x_id);
+#elif defined(RMSNORM)
+        // TODO(#52395): compute_kernel_hw_startup is a call-once API; this mid-kernel re-init (preserving the
+        // pre-cleanup full-init behaviour) should become a targeted DST re-arm.
+        compute_kernel_hw_startup(dfb_xmm_id, dfb_xmm_id, dfb_xmm2_id);
+#else
+        // TODO(#52395): compute_kernel_hw_startup is a call-once API; this mid-kernel re-init (preserving the
+        // pre-cleanup full-init behaviour) should become a targeted DST re-arm.
+        compute_kernel_hw_startup(dfb_x_id, dfb_scaler_id, dfb_ex_id);
+#endif
+#endif
+/*
+ * X + Y
+ */
+#ifdef FUSE_PRE_ADD
+        reconfig_data_format(dfb_in_id, dfb_inb_id);
+        pack_reconfig_data_format(dfb_x_id);
+        add_init(dfb_in_id, dfb_inb_id);
+        for (auto block : generic::blocks(Wt, block_size)) {
+            // In/inb come from the reader and need to be
+            // synced on full block size. Keep cb_x_id aligned
+            // to full block size as well so pre-add/no-pre-add
+            // can be handled the same way.
+            dfb_in.wait_front(block.full_block_size());
+            dfb_inb.wait_front(block.full_block_size());
+
+            tile_regs_acquire();
+            for (auto i : block.local()) {
+                add_tiles(dfb_in_id, dfb_inb_id, i, i, i);
+            }
+            tile_regs_commit();
+
+            dfb_in.pop_front(block.full_block_size());
+            dfb_inb.pop_front(block.full_block_size());
+
+            dfb_x.reserve_back(block.full_block_size());
+            cb_reserve_back(CB_RES_OUT, block.full_block_size());
+
+            tile_regs_wait();
+            for (auto i : block.local()) {
+                pack_tile(i, dfb_x_id);
+            }
+            pack_reconfig_data_format(dfb_x_id, CB_RES_OUT);
+            for (auto i : block.local()) {
+                pack_tile(i, CB_RES_OUT);
+            }
+            pack_reconfig_data_format(CB_RES_OUT, dfb_x_id);
+            tile_regs_release();
+
+            dfb_x.push_back(block.full_block_size());  // push the sum into the same buffer
+            cb_push_back(CB_RES_OUT, block.full_block_size());
+        }
+#ifndef RMSNORM
+        reconfig_data_format(dfb_in_id, dfb_x_id, dfb_inb_id, dfb_scaler_id);
+#else
+        reconfig_data_format(dfb_in_id, dfb_x_id, dfb_inb_id, dfb_x_id);
+#endif
+        // by the end of this loop we should end up with Wt tiles in cb_x_id
+#else
+#ifdef RMSNORM
+        reconfig_data_format(dfb_in_id, dfb_in_id);
+        pack_reconfig_data_format(dfb_xmm2_id);
+#endif
+#endif
+
+#ifndef RMSNORM
+        // E[x]
+        numeric::
+            row_wise_mean<PoolType::SUM, ReduceDim::REDUCE_ROW, FLOAT32_REDUCTION, policies::FullBlockWithoutPopPolicy>(
+                dfb_x, dfb_scaler, dfb_ex, W, Wt, block_size, tile_width);
+
+        // x - E[x]
+        reconfig_data_format(dfb_x_id, dfb_ex_id);
+        dfb_xmm.reserve_back(total_buffer_size);
+        sub_bcast_cols_init(dfb_x_id, dfb_ex_id);
+        for (auto block : generic::blocks(Wt, block_size)) {
+            tile_regs_acquire();
+            for (auto i : block.local()) {
+                sub_tiles_bcast_cols(dfb_x_id, dfb_ex_id, i, 0, i);
+            }
+            tile_regs_commit();
+
+            dfb_x.pop_front(block.full_block_size());
+
+            tile_regs_wait();
+            for (auto i : block.local()) {
+                pack_tile(i, dfb_xmm_id);
+            }
+            tile_regs_release();
+
+            dfb_xmm.push_back(block.full_block_size());
+        }
+        dfb_ex.pop_front(1);
+
+        // model-local: cb_x is bf16 here (the model passes a bf16 cb_x so the LN sees exactly the bf16 sum),
+        // so srcA must be reconfigured to the fp32 xmm format as in the non-fused path
+        reconfig_data_format_srca(dfb_x_id, dfb_xmm_id);
+#endif
+
+        /* (x - E[x])^2
+         * compute temp = xmm*xmm = (x-E[x])^2
+         */
+        mul_init(dfb_xmm_id, dfb_xmm_id);
+        for (auto block : generic::blocks(Wt, block_size)) {
+#ifndef RMSNORM
+            dfb_xmm.wait_front(block.start() + block.size());
+#else
+            dfb_xmm.wait_front(block.start() + block.full_block_size());
+#endif
+            tile_regs_acquire();
+            for (auto i : block.local()) {
+                const auto global_i = block.to_global(i);
+                mul_tiles(dfb_xmm_id, dfb_xmm_id, global_i, global_i, i);
+            }
+            tile_regs_commit();
+
+            dfb_xmm2.reserve_back(block.full_block_size());
+
+            tile_regs_wait();
+            for (auto i : block.local()) {
+                pack_tile(i, dfb_xmm2_id);
+            }
+            tile_regs_release();
+
+            dfb_xmm2.push_back(block.full_block_size());
+        }
+#if defined RMSNORM and not defined FUSED_PRE_ADD
+        reconfig_data_format(dfb_xmm_id, dfb_xmm2_id, dfb_xmm_id, dfb_scaler_id);
+#endif
+
+        // Var[x]
+        numeric::
+            row_wise_mean<PoolType::SUM, ReduceDim::REDUCE_ROW, FLOAT32_REDUCTION, policies::FullBlockWithPopPolicy>(
+                dfb_xmm2, dfb_scaler, dfb_ex2, W, Wt, block_size, tile_width);
+
+        // Var[x] + eps
+        dfb_ex2.wait_front(1);
+        reconfig_data_format(dfb_ex2_id, dfb_eps_id);
+
+        tile_regs_acquire();
+        add_init(dfb_ex2_id, dfb_eps_id);
+        add_tiles(dfb_ex2_id, dfb_eps_id, 0, 0, dst0);
+        rsqrt_tile_init<LEGACY_RSQRT>();
+        rsqrt_tile<LEGACY_RSQRT>(dst0);
+        tile_regs_commit();
+
+        dfb_ex2.pop_front(1);
+
+        dfb_ex2pe.reserve_back(1);
+        pack_reconfig_data_format(dfb_ex2pe_id);
+
+        tile_regs_wait();
+        pack_tile(dst0, dfb_ex2pe_id);
+        tile_regs_release();
+
+        dfb_ex2pe.push_back(1);
+
+        // (x-E[x]) / sqrt(Var[x] + eps) * gamma + beta
+        dfb_ex2pe.wait_front(1);
+        for (auto block : generic::blocks(Wt, block_size)) {
+            reconfig_data_format(dfb_xmm_id, dfb_ex2pe_id);
+            if constexpr (do_gamma == 0 && do_beta == 0) {
+                pack_reconfig_data_format(dfb_out_id);
+            } else {
+                pack_reconfig_data_format(dfb_fusion_id);
+            }
+            dfb_im_or_out.reserve_back(block.full_block_size());
+#if defined RMSNORM and not defined FUSE_PRE_ADD
+            reconfig_data_format_srca(dfb_fusion_id, dfb_xmm_id);
+#endif
+            tile_regs_acquire();
+            mul_bcast_cols_init(dfb_xmm_id, dfb_ex2pe_id);
+            for (auto i : block.local()) {
+                mul_tiles_bcast_cols(dfb_xmm_id, dfb_ex2pe_id, block.to_global(i), 0, i);  // tile *= 1/(sum(exp(x)))
+#ifdef SFPU_OP_INIT_ACTIVATION
+                // Activation must be applied last. If do_gamma != 0 or do_beta != 0 then
+                // activation will be applied after the gamma/beta multiplication/addition.
+                // Otherwise, we can apply the activation here.
+                if constexpr (!(do_gamma == 1 || do_beta == 1)) {
+                    SFPU_OP_INIT_ACTIVATION
+                    SFPU_OP_FUNC_ACTIVATION
+                }
+#endif
+            }
+            tile_regs_commit();
+
+            tile_regs_wait();
+            for (auto i : block.local()) {
+                pack_tile(i, dfb_im_or_out_id);  // pack either to intermediate (dfb_fusion or out0)
+            }
+            tile_regs_release();
+
+            dfb_im_or_out.push_back(
+                block.full_block_size());  // if no gamma/beta are provided, this will be passed on to the writer
+
+            if constexpr (!(do_gamma == 0 && do_beta == 0)) {
+#if defined RMSNORM and not defined FUSE_PRE_ADD
+                reconfig_data_format_srca(dfb_xmm_id, dfb_fusion_id);
+#endif
+            }
+
+            if constexpr (do_gamma) {
+                if constexpr (do_beta == 0) {
+                    pack_reconfig_data_format(dfb_out_id);
+                }
+                reconfig_data_format_srcb(dfb_ex2pe_id, dfb_gamma_id);
+                uint32_t dfb_outg_id = do_beta ? dfb_fusion_id : dfb_out_id;
+                DataflowBuffer dfb_outg(dfb_outg_id);
+                dfb_gamma.wait_front(
+                    block.start() + block.full_block_size());  // we don't pop, TODO: only wait on first ht
+                dfb_fusion.wait_front(block.full_block_size());
+
+                tile_regs_acquire();
+                mul_bcast_rows_init(dfb_fusion_id, dfb_gamma_id);
+                for (auto i : block.local()) {
+                    mul_tiles_bcast_rows(
+                        dfb_fusion_id, dfb_gamma_id, i, block.to_global(i), i);  // tile *= 1/(sum(exp(x)))
+#ifdef SFPU_OP_INIT_ACTIVATION
+                    // Activation must be applied last. If do_beta != 0 then
+                    // activation will be applied after the beta addition.
+                    // Otherwise, we can apply the activation here.
+                    if constexpr (!(do_beta == 1)) {
+                        SFPU_OP_INIT_ACTIVATION
+                        SFPU_OP_FUNC_ACTIVATION
+                    }
+#endif
+                }
+                tile_regs_commit();
+
+                dfb_fusion.pop_front(block.full_block_size());
+                // We don't pop gamma since it's 1,1,1,Wt and we reuse it for all NCHt
+
+                dfb_outg.reserve_back(block.full_block_size());
+
+                tile_regs_wait();
+                for (auto i : block.local()) {
+                    pack_tile(i, dfb_outg_id);  // pack either to intermediate (dfb_fusion or out0)
+                }
+                tile_regs_release();
+
+                dfb_outg.push_back(block.full_block_size());
+            }
+            if constexpr (do_beta) {
+                pack_reconfig_data_format(dfb_out_id);
+                if constexpr (do_gamma) {
+                    reconfig_data_format_srcb(dfb_gamma_id, dfb_beta_id);
+                } else {
+                    reconfig_data_format_srcb(dfb_ex2pe_id, dfb_beta_id);
+                }
+                dfb_beta.wait_front(
+                    block.start() + block.full_block_size());  // TODO: optimization - only wait on first ht
+                dfb_fusion.wait_front(block.full_block_size());
+
+                tile_regs_acquire();
+                add_bcast_rows_init(dfb_fusion_id, dfb_beta_id);
+                for (auto i : block.local()) {
+                    add_tiles_bcast_rows(
+                        dfb_fusion_id, dfb_beta_id, i, block.to_global(i), i);  // tile *= 1/(sum(exp(x)))
+#ifdef SFPU_OP_INIT_ACTIVATION
+                    SFPU_OP_INIT_ACTIVATION
+                    SFPU_OP_FUNC_ACTIVATION
+#endif
+                }
+                tile_regs_commit();
+
+                dfb_fusion.pop_front(block.full_block_size());
+                // We don't pop beta since it's 1,1,1,Wt and we reuse it for all NCHt
+
+                dfb_out.reserve_back(block.full_block_size());
+
+                tile_regs_wait();
+                for (auto i : block.local()) {
+                    pack_tile(i, dfb_out_id);
+                }
+                tile_regs_release();
+
+                dfb_out.push_back(block.full_block_size());
+            }
+        }
+        dfb_ex2pe.pop_front(1);
+        dfb_xmm.pop_front(total_buffer_size);
+
+#ifdef UNTILIZE_OUT
+        constexpr auto dfb_out_rm_id = get_named_compile_time_arg_val("cb_out_rm");
+        DataflowBuffer dfb_out_rm(dfb_out_rm_id);
+        untilize_all_blocks_from_cb<block_size>(dfb_out, dfb_out_rm, Wt);
+#endif
+    }  // NCHt loop
+    // The reduce scaler is generated once by the reader and reused (waited inside row_wise_mean)
+    // across every NCHt iteration but never popped. Pop the producer's tile count once here to
+    // balance the CB. The reader pushes a second scaler tile only when the last column tile is
+    // partial (W not a multiple of tile_width), matching row_wise_mean's wait count.
+    //
+    // The reader generates the scalers using tt::constants::TILE_WIDTH; this kernel must use the
+    // same width for the count to match, so derive both from the shared helper. (tile_width is the
+    // tensor's tile width, which equals TILE_WIDTH for every supported layernorm config — see the
+    // partial-column handling in row_wise_mean above.)
+    static_assert(
+        tile_width == tt::constants::TILE_WIDTH,
+        "layernorm reader generates reduce scalers using TILE_WIDTH; compute must use the same tile "
+        "width or cb_scaler push/pop counts diverge (issue #48487)");
+    constexpr uint32_t num_scaler_tiles = norm::layernorm::reduce_scaler_tile_count(W, tile_width);
+    DataflowBuffer(dfb_scaler).pop_front(num_scaler_tiles);
+}

code/models/demos/mast3r/tt/kernels/ln_add_writer.cpp

@@ -0,0 +1,43 @@
+// SPDX-License-Identifier: Apache-2.0
+// mast3r-p150 model-local writer for ln_add_compute.cpp (generic_op): per tile row, first the residual sum
+// (CB 17, Wt tiles in blocks) to res_out, then the normalized row (CB cb_out) to out. Same blocking as ttnn's
+// writer_unary_interleaved_start_id_blocked.cpp.
+#include <stdint.h>
+#include "api/dataflow/dataflow_api.h"
+
+void kernel_main() {
+    const uint32_t dst_addr = get_arg_val<uint32_t>(0);
+    const uint32_t Wt = get_arg_val<uint32_t>(1);
+    const uint32_t num_tile_rows = get_arg_val<uint32_t>(2);
+    const uint32_t tile_offset = get_arg_val<uint32_t>(3);
+    const uint32_t res_addr = get_arg_val<uint32_t>(4);
+    constexpr uint32_t blk = get_compile_time_arg_val(0);
+    constexpr uint32_t CB_OUT = 16, CB_RES = 17;
+    constexpr uint32_t LOG2_TILE = 11;
+    constexpr uint32_t tile_bytes = 1u << LOG2_TILE;
+    const InterleavedPow2AddrGen<OUT_DRAM> og = {.bank_base_address = dst_addr, .log_base_2_of_page_size = LOG2_TILE};
+    const InterleavedPow2AddrGen<RES_DRAM> rg = {.bank_base_address = res_addr, .log_base_2_of_page_size = LOG2_TILE};
+    for (uint32_t h = 0; h < num_tile_rows; h++) {
+        const uint32_t row0 = tile_offset + h * Wt;
+        for (uint32_t w0 = 0; w0 < Wt; w0 += blk) {
+            const uint32_t n = (Wt - w0) < blk ? (Wt - w0) : blk;
+            cb_wait_front(CB_RES, blk);
+            uint32_t rd = get_read_ptr(CB_RES);
+            for (uint32_t i = 0; i < n; ++i) {
+                noc_async_write(rd + i * tile_bytes, rg.get_noc_addr(row0 + w0 + i), tile_bytes);
+            }
+            noc_async_write_barrier();
+            cb_pop_front(CB_RES, blk);
+        }
+        for (uint32_t w0 = 0; w0 < Wt; w0 += blk) {
+            const uint32_t n = (Wt - w0) < blk ? (Wt - w0) : blk;
+            cb_wait_front(CB_OUT, blk);
+            uint32_t rd = get_read_ptr(CB_OUT);
+            for (uint32_t i = 0; i < n; ++i) {
+                noc_async_write(rd + i * tile_bytes, og.get_noc_addr(row0 + w0 + i), tile_bytes);
+            }
+            noc_async_write_barrier();
+            cb_pop_front(CB_OUT, blk);
+        }
+    }
+}

code/models/demos/mast3r/tt/kernels/ln_add_writer_rb.cpp

@@ -0,0 +1,64 @@
+// SPDX-License-Identifier: Apache-2.0
+// mast3r-p150 model-local writer for the fused add + LN (MAST3R_OPT "lnsplit"): ln_add_writer.cpp plus the read of the
+// pre-add operand b (CB 1, in blocks, through a TensorAccessor: b may be block-sharded), so a (reader, one NoC) and b
+// (this RISC, the other NoC) stream in parallel. Per tile row: read b block i, then write residual-sum block i-1
+// (CB 17) as soon as the compute has produced it; then the normalized row (CB 16).
+#include <stdint.h>
+#include "api/dataflow/dataflow_api.h"
+
+void kernel_main() {
+    const uint32_t dst_addr = get_arg_val<uint32_t>(0);
+    const uint32_t Wt = get_arg_val<uint32_t>(1);
+    const uint32_t num_tile_rows = get_arg_val<uint32_t>(2);
+    const uint32_t tile_offset = get_arg_val<uint32_t>(3);
+    const uint32_t res_addr = get_arg_val<uint32_t>(4);
+    const uint32_t b_addr = get_arg_val<uint32_t>(5);
+    constexpr uint32_t blk = get_compile_time_arg_val(0);
+    constexpr auto b_args = TensorAccessorArgs<1>();
+    constexpr uint32_t CB_INB = 1, CB_OUT = 16, CB_RES = 17;
+    constexpr uint32_t LOG2_TILE = 11;
+    constexpr uint32_t tile_bytes = 1u << LOG2_TILE;
+    const InterleavedPow2AddrGen<OUT_DRAM> og = {.bank_base_address = dst_addr, .log_base_2_of_page_size = LOG2_TILE};
+    const InterleavedPow2AddrGen<RES_DRAM> rg = {.bank_base_address = res_addr, .log_base_2_of_page_size = LOG2_TILE};
+    const auto bg = TensorAccessor(b_args, b_addr, tile_bytes);
+    for (uint32_t h = 0; h < num_tile_rows; h++) {
+        const uint32_t row0 = tile_offset + h * Wt;
+        auto write_res = [&](uint32_t w0) {
+            const uint32_t n = (Wt - w0) < blk ? (Wt - w0) : blk;
+            cb_wait_front(CB_RES, blk);
+            uint32_t rd = get_read_ptr(CB_RES);
+            for (uint32_t i = 0; i < n; ++i) {
+                noc_async_write(rd + i * tile_bytes, rg.get_noc_addr(row0 + w0 + i), tile_bytes);
+            }
+            noc_async_write_barrier();
+            cb_pop_front(CB_RES, blk);
+        };
+        for (uint32_t w0 = 0; w0 < Wt; w0 += blk) {
+            const uint32_t n = (Wt - w0) < blk ? (Wt - w0) : blk;
+            cb_reserve_back(CB_INB, blk);
+            uint32_t wp = get_write_ptr(CB_INB);
+            for (uint32_t i = 0; i < n; ++i) {
+                noc_async_read(bg.get_noc_addr(row0 + w0 + i), wp + i * tile_bytes, tile_bytes);
+            }
+            noc_async_read_barrier();
+            cb_push_back(CB_INB, blk);
+            // Interleave: the residual-sum block of the previous step must be drained here. Reading ALL of b before
+            // writing any residual block DEADLOCKS (compute blocks on CB 17 while this kernel blocks on cb_inb):
+            // that variant hung chip 9 in round 7.
+            if (w0 > 0) {
+                write_res(w0 - blk);
+            }
+        }
+        write_res(((Wt - 1) / blk) * blk);
+        for (uint32_t w0 = 0; w0 < Wt; w0 += blk) {
+            const uint32_t n = (Wt - w0) < blk ? (Wt - w0) : blk;
+            cb_wait_front(CB_OUT, blk);
+            uint32_t rd = get_read_ptr(CB_OUT);
+            for (uint32_t i = 0; i < n; ++i) {
+                noc_async_write(rd + i * tile_bytes, og.get_noc_addr(row0 + w0 + i), tile_bytes);
+            }
+            noc_async_write_barrier();
+            cb_pop_front(CB_OUT, blk);
+        }
+    }
+}

code/models/demos/mast3r/tt/kernels/ln_fast_compute.cpp

@@ -0,0 +1,202 @@
+// SPDX-License-Identifier: Apache-2.0
+//
+// mast3r-p150 model-local fused residual-add + affine-free LayerNorm compute kernel (MAST3R_OPT "lnf"),
+// a drop-in for ln_add_compute.cpp with the same reader (ttnn's reader_unary_interleaved_ln, FUSE_PRE_ADD) and
+// writer (ln_add_writer.cpp) and the same per-row CB protocol (cb_in / cb_inb in blocks, cb_out + CB_RES_OUT in
+// blocks). Fewer passes than ttnn's layernorm.cpp (6 passes, 3 of them over fp32 tiles at HiFi4):
+//   A: s = a + b (fp32 dest) -> cb_x (bf16, the value the LN normalises, = ttnn.add's bf16 sum) + CB_RES_OUT
+//   M: mean_full = (sum_t s_t @ ONES) / W      (matmul with an all-ones tile: the row sum lands in all 32 columns,
+//      exact bf16 inputs, fp32 dest accumulation over the Wt tiles)
+//   B: d = s - mean_full; d^2 on the SFPU (fp32); the d^2 tiles are summed elementwise by the PACKER (L1
+//      accumulation into one fp32 tile) -> no xmm / xmm^2 round trips through L1
+//   V: rstd_full = rsqrt((acc @ ONES) / W + eps)  (one matmul, SFPU scale / add / rsqrt)
+//   C: y = (s - mean_full) * rstd_full  (FPU sub, then FPU mul with the DST operand reused as SrcA) -> cb_out
+// Precision: same operand precisions as the stock kernel for the subtraction and the final multiply (fp32 CB
+// operands enter SrcA/SrcB as tf32), the variance uses an exact fp32 SFPU square instead of a tf32 FPU square.
+
+#include <cstdint>
+
+#include "api/compute/compute_kernel_api.h"
+#include "api/compute/compute_kernel_hw_startup.h"
+#include "api/compute/eltwise_binary.h"
+#include "api/compute/matmul.h"
+#include "api/compute/pack.h"
+#include "api/compute/reconfig_data_format.h"
+#include "api/compute/eltwise_unary/eltwise_unary.h"
+#include "api/compute/eltwise_unary/fill.h"
+#include "api/compute/eltwise_unary/binop_with_scalar.h"
+#include "api/compute/eltwise_unary/rsqrt.h"
+#ifdef LNF_PROF
+#include "tools/profiler/kernel_profiler.hpp"
+#define LZ(n) DeviceZoneScopedN(n)
+#else
+#define LZ(n)
+#endif
+
+void kernel_main() {
+    const uint32_t NCHt = get_arg_val<uint32_t>(0);
+    constexpr uint32_t Wt = get_compile_time_arg_val(0);
+    constexpr uint32_t blk = get_compile_time_arg_val(1);
+    constexpr uint32_t inv_w_bits = get_compile_time_arg_val(2);
+    constexpr uint32_t eps_bits = get_compile_time_arg_val(3);
+
+    constexpr uint32_t CB_IN = 0, CB_INB = 1, CB_OUT = 16, CB_RES = 17, CB_MEAN = 18, CB_ACC2 = 20, CB_RSTD = 21,
+                       CB_X = 23, CB_ONES = 24;
+
+    compute_kernel_hw_startup(CB_IN, CB_INB, CB_X);
+
+    // all-ones bf16 tile (row-sum broadcast through a matmul)
+    cb_reserve_back(CB_ONES, 1);
+    tile_regs_acquire();
+    fill_tile_init();
+    fill_tile(0, 1.0f);
+    tile_regs_commit();
+    tile_regs_wait();
+    pack_reconfig_data_format(CB_ONES);
+    pack_tile(0, CB_ONES);
+    tile_regs_release();
+    cb_push_back(CB_ONES, 1);
+    cb_wait_front(CB_ONES, 1);
+
+    for (uint32_t row = 0; row < NCHt; ++row) {
+        {
+        LZ("LNF_A");
+        // ---- A: s = a + b ----
+        reconfig_data_format(CB_IN, CB_INB);
+        pack_reconfig_data_format(CB_X);
+        add_init(CB_IN, CB_INB);
+        for (uint32_t w0 = 0; w0 < Wt; w0 += blk) {
+            cb_wait_front(CB_IN, blk);
+            cb_wait_front(CB_INB, blk);
+            tile_regs_acquire();
+            for (uint32_t i = 0; i < blk; ++i) {
+                add_tiles(CB_IN, CB_INB, i, i, i);
+            }
+            tile_regs_commit();
+            cb_pop_front(CB_IN, blk);
+            cb_pop_front(CB_INB, blk);
+            cb_reserve_back(CB_X, blk);
+            cb_reserve_back(CB_RES, blk);
+            tile_regs_wait();
+            for (uint32_t i = 0; i < blk; ++i) {
+                pack_tile(i, CB_X);
+            }
+            for (uint32_t i = 0; i < blk; ++i) {
+                pack_tile(i, CB_RES);
+            }
+            tile_regs_release();
+            cb_push_back(CB_X, blk);
+            cb_push_back(CB_RES, blk);
+        }
+
+        }
+        {
+        LZ("LNF_M");
+        // ---- M: mean_full ----
+        cb_wait_front(CB_X, Wt);
+        reconfig_data_format(CB_ONES, CB_X);  // matmul: in1 (ONES) -> SrcA, in0 (X) -> SrcB
+        matmul_init(CB_X, CB_ONES);
+        tile_regs_acquire();
+        for (uint32_t t = 0; t < Wt; ++t) {
+            matmul_tiles(CB_X, CB_ONES, t, 0, 0);
+        }
+        binop_with_scalar_tile_init();
+        mul_unary_tile(0, inv_w_bits);
+        tile_regs_commit();
+        cb_reserve_back(CB_MEAN, 1);
+        pack_reconfig_data_format(CB_MEAN);
+        tile_regs_wait();
+        pack_tile(0, CB_MEAN);
+        tile_regs_release();
+        cb_push_back(CB_MEAN, 1);
+        cb_wait_front(CB_MEAN, 1);
+
+        }
+        {
+        LZ("LNF_B");
+        // ---- B: sum of (s - mean)^2, packer L1 accumulation into one fp32 tile ----
+        reconfig_data_format(CB_X, CB_MEAN);
+        sub_init(CB_X, CB_MEAN);
+        square_tile_init();
+        pack_reconfig_data_format(CB_ACC2);
+        cb_reserve_back(CB_ACC2, 1);
+        for (uint32_t w0 = 0; w0 < Wt; w0 += blk) {
+            tile_regs_acquire();
+            for (uint32_t i = 0; i < blk; ++i) {
+                sub_tiles(CB_X, CB_MEAN, w0 + i, 0, i);
+                square_tile(i);
+            }
+            tile_regs_commit();
+            tile_regs_wait();
+            for (uint32_t i = 0; i < blk; ++i) {
+                if (w0 + i == 0) {
+                    pack_reconfig_l1_acc(0);
+                }
+                pack_tile<true>(i, CB_ACC2, 0);
+                if (w0 + i == 0) {
+                    pack_reconfig_l1_acc(1);
+                }
+            }
+            tile_regs_release();
+        }
+        pack_reconfig_l1_acc(0);
+        cb_push_back(CB_ACC2, 1);
+
+        }
+        {
+        LZ("LNF_V");
+        // ---- V: rstd_full ----
+        cb_wait_front(CB_ACC2, 1);
+        reconfig_data_format(CB_ONES, CB_ACC2);
+        matmul_init(CB_ACC2, CB_ONES);
+        tile_regs_acquire();
+        matmul_tiles(CB_ACC2, CB_ONES, 0, 0, 0);
+        binop_with_scalar_tile_init();
+        mul_unary_tile(0, inv_w_bits);
+        add_unary_tile(0, eps_bits);
+        rsqrt_tile_init<false>();
+        rsqrt_tile<false>(0);
+        tile_regs_commit();
+        cb_pop_front(CB_ACC2, 1);
+        cb_reserve_back(CB_RSTD, 1);
+        pack_reconfig_data_format(CB_RSTD);
+        tile_regs_wait();
+        pack_tile(0, CB_RSTD);
+        tile_regs_release();
+        cb_push_back(CB_RSTD, 1);
+        cb_wait_front(CB_RSTD, 1);
+
+        }
+        {
+        LZ("LNF_C");
+        // ---- C: y = (s - mean) * rstd ----
+        pack_reconfig_data_format(CB_OUT);
+        for (uint32_t w0 = 0; w0 < Wt; w0 += blk) {
+            tile_regs_acquire();
+            reconfig_data_format(CB_X, CB_MEAN);
+            sub_init(CB_X, CB_MEAN);
+            for (uint32_t i = 0; i < blk; ++i) {
+                sub_tiles(CB_X, CB_MEAN, w0 + i, 0, i);
+            }
+            // SrcA takes the fp32 DST values: configure it for an fp32 operand (tf32 in SrcA, as the stock kernel's
+            // fp32 xmm CB) instead of the bf16 cb_x format
+            reconfig_data_format(CB_RSTD, CB_RSTD);
+            mul_reuse_dest_init<EltwiseBinaryReuseDestType::DEST_TO_SRCA>(CB_RSTD);
+            for (uint32_t i = 0; i < blk; ++i) {
+                mul_reuse_dest_tiles<EltwiseBinaryReuseDestType::DEST_TO_SRCA>(CB_RSTD, 0, i);
+            }
+            tile_regs_commit();
+            cb_reserve_back(CB_OUT, blk);
+            tile_regs_wait();
+            for (uint32_t i = 0; i < blk; ++i) {
+                pack_tile(i, CB_OUT);
+            }
+            tile_regs_release();
+            cb_push_back(CB_OUT, blk);
+        }
+        }
+        cb_pop_front(CB_X, Wt);
+        cb_pop_front(CB_MEAN, 1);
+        cb_pop_front(CB_RSTD, 1);
+    }
+}

code/models/demos/mast3r/tt/kernels/ln_reader_split.cpp

@@ -0,0 +1,209 @@
+// mast3r-p150 model-local copy of ttnn's reader_unary_interleaved_ln.cpp (tt-metal 8b98410e730). Only change: with
+// LN_SPLIT_B the second (pre-add) input b is NOT read here; the model-local writer ln_add_writer_rb.cpp reads it on
+// the other RISC / NoC, so the two operand streams of the fused add + LN load in parallel (MAST3R_OPT "lnsplit").
+// SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+// Merged reader kernel for layernorm / dit_rms_norm_unary_fused (standard / non-large-tensor path).
+//
+// Handles both TILE-layout input (default) and ROW_MAJOR input (#ifdef TILIZE_IN).
+// The loop structure, scaler/eps generation, gamma/beta reads, and FUSE_PRE_ADD reads
+// are shared between both paths.  Only the input accessor setup and the per-ncht input
+// read call branch on TILIZE_IN.
+//
+// Replaces the two separate files:
+//   - reader_unary_interleaved_ln.cpp          (TILE path, 105 lines)
+//   - reader_unary_interleaved_ln_rm_input.cpp (ROW_MAJOR path, 143 lines)
+//
+// Compile-time args:
+//   CTA[0]    = block_size
+//   CTA[1]    = use_welford  (0 for TILIZE_IN / RMSNORM)
+//   CTA[2]    = W                (logical width in elements)
+//   CTA[3..]  = TensorAccessorArgs for input a
+//   ...       = TensorAccessorArgs for b / residual  (may be null)
+//   ...       = TensorAccessorArgs for gamma          (may be null)
+//   ...       = TensorAccessorArgs for beta           (may be null)
+//   CTA[N]    = elem_size_bytes  (TILIZE_IN only; unused for TILE path)
+//
+// Runtime args:
+//   arg[0] = src_addr
+//   arg[1] = NCHt              (number of tile-rows assigned to this core)
+//   arg[2] = Wt                (width in tiles)
+//   arg[3] = start_tile_row    (tile-row index of first row for this core)
+//                              TILE:  previously passed as tile_offset = start_tile_row * Wt
+//                              RM:    previously passed as start_row;  start_tile_row = start_row / TILE_H
+//   arg[4] = packed_one_value  (legacy; unused, scaler is generated in-kernel)
+//   arg[5] = eps               (epsilon as bit-cast uint32)
+//   arg[6] = gamma_dram_addr
+//   arg[7] = beta_dram_addr
+//   arg[8] = b_dram_addr       (residual, unused if no FUSE_PRE_ADD)
+//   arg[9] = H_logical         (TILIZE_IN only: total valid rows; unused for TILE path)
+
+#include <stdint.h>
+#include "api/dataflow/dataflow_api.h"
+#include "ttnn/cpp/ttnn/kernel_lib/reduce_helpers_dataflow.hpp"
+#include "ttnn/kernel/dataflow/generate_bcast_scalar.hpp"
+#include "ttnn/operations/normalization/kernel_util/generic/blocked_range.h"
+#include "ttnn/operations/normalization/layernorm/device/kernels/layernorm_scaler_tiles.h"
+#include "ttnn/operations/normalization/layernorm/device/kernels/dataflow/layernorm_dataflow_utils.h"
+
+namespace generic = norm::kernel_util::generic;
+namespace layernorm_dataflow_utils = norm::layernorm::device::kernels::dataflow;
+
+void kernel_main() {
+    const uint32_t src_addr = get_arg_val<uint32_t>(0);
+    const uint32_t NCHt = get_arg_val<uint32_t>(1);
+    const uint32_t Wt = get_arg_val<uint32_t>(2);
+    const uint32_t start_tile_row = get_arg_val<uint32_t>(3);
+    const uint32_t gamma_addr = get_arg_val<uint32_t>(6);
+    const uint32_t beta_addr = get_arg_val<uint32_t>(7);
+    const uint32_t b_addr = get_arg_val<uint32_t>(8);
+#ifdef TILIZE_IN
+    const uint32_t H_logical = get_arg_val<uint32_t>(9);
+#endif
+
+    constexpr uint32_t dfb_id_in0 = get_named_compile_time_arg_val("cb_in");
+    // Welford-fp32 alias of cb_in (non-fused) or cb_x (fused). Shares SRAM with the
+    // primary CB but has its own read/write pointers, so we must push_back on it whenever we
+    // push to the primary CB. When welford_fp32_alias is 0, dfb_x_welford == cb_in.
+    constexpr uint32_t dfb_id_x_welford = get_named_compile_time_arg_val("cb_x_welford");
+    constexpr bool welford_fp32_alias = get_named_compile_time_arg_val("welford_fp32_alias") != 0;
+    constexpr uint32_t dfb_id_in1 = get_named_compile_time_arg_val("cb_inb");
+    constexpr uint32_t dfb_id_gamma = get_named_compile_time_arg_val("cb_gamma");
+    constexpr uint32_t dfb_id_beta = get_named_compile_time_arg_val("cb_beta");
+
+    Noc noc;
+    DataflowBuffer dfb_in0(dfb_id_in0);
+    DataflowBuffer dfb_x_welford(dfb_id_x_welford);
+#ifdef FUSE_PRE_ADD
+    DataflowBuffer dfb_in1(dfb_id_in1);
+#endif
+#ifdef FUSE_GAMMA
+    DataflowBuffer dfb_gamma(dfb_id_gamma);
+#endif
+#ifdef FUSE_BETA
+    DataflowBuffer dfb_beta(dfb_id_beta);
+#endif
+
+    constexpr uint32_t block_size = get_compile_time_arg_val(0);
+    constexpr bool use_welford = get_compile_time_arg_val(1) == 1;
+    constexpr uint32_t W = get_compile_time_arg_val(2);
+    constexpr auto src0_args = TensorAccessorArgs<3>();
+    [[maybe_unused]] constexpr auto src1_args = TensorAccessorArgs<src0_args.next_compile_time_args_offset()>();
+    [[maybe_unused]] constexpr auto gamma_args = TensorAccessorArgs<src1_args.next_compile_time_args_offset()>();
+    [[maybe_unused]] constexpr auto beta_args = TensorAccessorArgs<gamma_args.next_compile_time_args_offset()>();
+
+    constexpr uint32_t TILE_H = tt::constants::TILE_HEIGHT;
+    constexpr uint32_t TILE_W = tt::constants::TILE_WIDTH;
+
+#ifdef TILIZE_IN
+    // ROW_MAJOR path: input a is a row-major tensor.
+    // The compute kernel tilizes dfb_in_rm (c_27) → cb_in (c_0) before processing.
+    constexpr uint32_t elem_size_bytes = get_compile_time_arg_val(beta_args.next_compile_time_args_offset());
+
+    constexpr uint32_t rm_row_stride_bytes = block_size * TILE_W * elem_size_bytes;
+    constexpr uint32_t dfb_id_in_rm = get_named_compile_time_arg_val("cb_in_rm");
+    DataflowBuffer dfb_in_rm(dfb_id_in_rm);
+
+    const uint32_t src0_page_bytes = W * elem_size_bytes;
+#else
+    // TILE path: input a is already in tile layout.
+    const uint32_t src0_page_bytes = dfb_in0.get_tile_size();
+#endif
+
+    const auto src_a = TensorAccessor(src0_args, src_addr);
+
+#ifdef FUSE_GAMMA
+    const uint32_t gamma_tile_bytes = dfb_gamma.get_tile_size();
+    const auto addrg = TensorAccessor(gamma_args, gamma_addr);
+#endif
+#ifdef FUSE_BETA
+    const uint32_t beta_tile_bytes = dfb_beta.get_tile_size();
+    const auto addrb = TensorAccessor(beta_args, beta_addr);
+#endif
+#ifdef FUSE_PRE_ADD
+    const uint32_t src1_tile_bytes = dfb_in1.get_tile_size();
+    const auto src_b = TensorAccessor(src1_args, b_addr);
+#endif
+
+    // Generate constant tiles for layernorm compute
+    constexpr uint32_t dfb_scaler = get_named_compile_time_arg_val("cb_scaler");
+    constexpr uint32_t dfb_eps = get_named_compile_time_arg_val("cb_eps");
+
+    if constexpr (!use_welford) {
+        constexpr uint32_t partial_last_tile_cols = W % tt::constants::TILE_WIDTH;
+        // Push count shared with the compute kernel's cb_scaler pop count (issue #48487).
+        constexpr uint32_t num_scaler_tiles = norm::layernorm::reduce_scaler_tile_count(W, tt::constants::TILE_WIDTH);
+
+        dataflow_kernel_lib::calculate_and_prepare_reduce_scaler<
+            dfb_scaler,
+            ckernel::PoolType::SUM,
+            ckernel::ReduceDim::REDUCE_ROW,
+            dataflow_kernel_lib::SUM_AND_MAX_REDUCE_FACTOR>();
+
+        if constexpr (num_scaler_tiles == 2) {
+            dataflow_kernel_lib::calculate_and_prepare_reduce_scaler<
+                dfb_scaler,
+                ckernel::PoolType::SUM,
+                ckernel::ReduceDim::REDUCE_ROW,
+                dataflow_kernel_lib::SUM_AND_MAX_REDUCE_FACTOR>(partial_last_tile_cols);
+        }
+    }
+
+    const uint32_t eps = get_arg_val<uint32_t>(5);
+    generate_bcast_col_scalar(CircularBuffer(dfb_eps), eps);
+
+    for (uint32_t ncht = 0; ncht < NCHt; ncht++) {
+        const uint32_t curr_tile_row = start_tile_row + ncht;
+
+        // --- Input read: branches on layout ---
+#ifdef TILIZE_IN
+        // ROW_MAJOR: push one tile-row of row-major data into dfb_in_rm (block-by-block).
+        // The compute kernel's TILIZE_IN block converts dfb_in_rm → cb_in before processing.
+        layernorm_dataflow_utils::push_row_major_blocks_to_cb<decltype(src_a), TILE_W, TILE_H>(
+            noc, dfb_in_rm, src_a, Wt, block_size, curr_tile_row, elem_size_bytes, rm_row_stride_bytes, H_logical);
+
+#ifdef FUSE_PRE_ADD
+        for (auto block : generic::blocks(Wt, block_size)) {
+            layernorm_dataflow_utils::read_block_to_cb(
+                noc, dfb_in1, src_b, src1_tile_bytes, curr_tile_row * Wt + block.start(), block);
+        }
+#endif
+#else
+        // TILE: read input a and b (if present) interleaved per block.
+        for (auto block : generic::blocks(Wt, block_size)) {
+            const uint32_t flat_offset = curr_tile_row * Wt + block.start();
+            layernorm_dataflow_utils::read_block_to_cb(noc, dfb_in0, src_a, src0_page_bytes, flat_offset, block);
+#if defined(FUSE_PRE_ADD) && !defined(LN_SPLIT_B)
+            layernorm_dataflow_utils::read_block_to_cb(noc, dfb_in1, src_b, src1_tile_bytes, flat_offset, block);
+#elif !defined(FUSE_PRE_ADD)
+            // Non-fused welford-fp32 alias: dfb_x_welford shares dfb_in0's memory but has its own
+            // read/write pointers. After the data lands in dfb_in0, push
+            // dfb_x_welford by the same amount so compute can wait_front on the alias separately
+            // for welford reads. Skipped when no alias is active (dfb_x_welford == dfb_in0; the
+            // duplicate push would double-count dfb_in0's semaphore).
+            if constexpr (welford_fp32_alias) {
+                dfb_x_welford.reserve_back(block.full_block_size());
+                dfb_x_welford.push_back(block.full_block_size());
+            }
+#endif
+        }
+#endif
+
+        // --- Gamma / beta (shared): read once at ncht == 0 ---
+#if defined FUSE_GAMMA || defined FUSE_BETA
+        if (ncht == 0) {
+            for (auto block : generic::blocks(Wt, block_size)) {
+#ifdef FUSE_GAMMA
+                layernorm_dataflow_utils::read_block_to_cb(
+                    noc, dfb_gamma, addrg, gamma_tile_bytes, block.start(), block);
+#endif
+#ifdef FUSE_BETA
+                layernorm_dataflow_utils::read_block_to_cb(noc, dfb_beta, addrb, beta_tile_bytes, block.start(), block);
+#endif
+            }  // wt loop
+        }
+#endif
+    }  // ncht loop
+}

code/models/demos/mast3r/tt/kernels/mast3r_gelu_poly.h

@@ -0,0 +1,116 @@
+// SPDX-License-Identifier: Apache-2.0
+// mast3r-p150 model-local SFPU GELU for the fc1 matmul epilogue (MAST3R_OPT "gpoly"), used in the GELU_TANH slot of
+// the model-local mm_gelu_activation.hpp.
+//
+// The reference model (DUSt3R / CroCo Mlp) uses the exact erf GELU, gelu(x) = x Phi(x). Here
+//     gelu(x) = 0.5 x + |x| s q(s^2),   s = min(|x|, X),  X = 4.25,
+// where s q(s^2) is a degree-17 odd minimax polynomial (9 coefficients in s^2, weighted LP fit, tools_prof/gelu_fit.py)
+// of Phi(s) - 0.5 on [0, X] with s q(s^2) = 0.5 exactly at s = X (so |x| > X gives relu(x) up to |x| Phi(-X) <= 1e-5 |x|).
+// Max |error| vs the erf GELU on [-8, 8]: 6e-5 before the bf16 rounding (the tanh form, gelut, is 4.7e-4 away from
+// the erf GELU); fp32 SFPU arithmetic, one bf16 round-to-nearest at the end (as ttnn's gelu_tanh does without fp32
+// dest). ~30 SFPU instructions per row vs ~40 for ttnn's accurate gelu_tanh.
+#pragma once
+#include "api/compute/common_globals.h"
+#if defined(TRISC_MATH) || defined(TRISC_PACK)
+#include "llk_math_eltwise_unary_sfpu_macros.h"
+#include "sfpi.h"
+
+namespace ckernel::sfpu {
+
+constexpr float MG_X = 4.25f;
+constexpr int MG_N = 9;
+constexpr float MG_C[MG_N] = {3.989576101e-01f, -6.634455919e-02f, 9.786477312e-03f, -1.095801592e-03f, 9.108911036e-05f,
+                              -5.402808256e-06f, 2.137231121e-07f, -5.002082304e-09f, 5.199383632e-11f};
+
+#ifndef MAST3R_GPOLY_X1
+// Two dst rows per step: every coefficient is loaded once for both Horner chains (halves the SFPLOADI pairs) and the
+// two independent SFPMAD chains hide each other's latency. Same arithmetic per element as the one-row form.
+template <bool is_fp32_dest_acc_en, int ITERATIONS = 8>
+inline void calculate_mast3r_gelu_poly() {
+    // vConstFloatPrgm0 = X, vConstFloatPrgm1 = 0.5, vConstFloatPrgm2 = MG_C[MG_N - 1] (set in init)
+    static_assert(ITERATIONS % 2 == 0, "two rows per step");
+#pragma GCC unroll 4
+    for (int d = 0; d < ITERATIONS / 2; d++) {
+        sfpi::vFloat x0 = sfpi::dst_reg[0];
+        sfpi::vFloat x1 = sfpi::dst_reg[1];
+        sfpi::vFloat a0 = sfpi::abs(x0);
+        sfpi::vFloat a1 = sfpi::abs(x1);
+        sfpi::vFloat xb = sfpi::vConstFloatPrgm0;
+        sfpi::vFloat s0 = sfpi::min(a0, xb);
+        sfpi::vFloat s1 = sfpi::min(a1, xb);
+        sfpi::vFloat u0 = s0 * s0;
+        sfpi::vFloat u1 = s1 * s1;
+        sfpi::vFloat as0 = a0 * s0;
+        sfpi::vFloat as1 = a1 * s1;
+        sfpi::vFloat c = MG_C[MG_N - 2];
+        sfpi::vFloat p0 = sfpi::vConstFloatPrgm2 * u0 + c;
+        sfpi::vFloat p1 = sfpi::vConstFloatPrgm2 * u1 + c;
+#pragma GCC unroll 8
+        for (int k = MG_N - 3; k >= 0; --k) {
+            c = MG_C[k];
+            p0 = p0 * u0 + c;
+            p1 = p1 * u1 + c;
+        }
+        p0 = p0 * as0;
+        p1 = p1 * as1;
+        x0 = sfpi::dst_reg[0];
+        x1 = sfpi::dst_reg[1];
+        sfpi::vFloat y0 = x0 * sfpi::vConstFloatPrgm1 + p0;
+        sfpi::vFloat y1 = x1 * sfpi::vConstFloatPrgm1 + p1;
+        if constexpr (!is_fp32_dest_acc_en) {
+            y0 = sfpi::convert<sfpi::vFloat16b>(y0, sfpi::RoundMode::Nearest);
+            y1 = sfpi::convert<sfpi::vFloat16b>(y1, sfpi::RoundMode::Nearest);
+        }
+        sfpi::dst_reg[0] = y0;
+        sfpi::dst_reg[1] = y1;
+        sfpi::dst_reg += 2;
+    }
+}
+#else
+template <bool is_fp32_dest_acc_en, int ITERATIONS = 8>
+inline void calculate_mast3r_gelu_poly() {
+    // vConstFloatPrgm0 = X, vConstFloatPrgm1 = 0.5, vConstFloatPrgm2 = MG_C[MG_N - 1] (set in init)
+    sfpi::vFloat h0 = MG_C[0], h1 = MG_C[1];
+#pragma GCC unroll 8
+    for (int d = 0; d < ITERATIONS; d++) {
+        sfpi::vFloat x = sfpi::dst_reg[0];
+        sfpi::vFloat a = sfpi::abs(x);
+        sfpi::vFloat xb = sfpi::vConstFloatPrgm0;
+        sfpi::vFloat s = sfpi::min(a, xb);
+        sfpi::vFloat u = s * s;
+        sfpi::vFloat p = sfpi::vConstFloatPrgm2 * u + MG_C[MG_N - 2];
+#pragma GCC unroll 8
+        for (int k = MG_N - 3; k >= 2; --k) {
+            p = p * u + MG_C[k];
+        }
+        p = p * u + h1;
+        p = p * u + h0;
+        sfpi::vFloat as = a * s;
+        p = p * as;
+        sfpi::vFloat y = x * sfpi::vConstFloatPrgm1 + p;
+        if constexpr (!is_fp32_dest_acc_en) {
+            y = sfpi::convert<sfpi::vFloat16b>(y, sfpi::RoundMode::Nearest);
+        }
+        sfpi::dst_reg[0] = y;
+        sfpi::dst_reg++;
+    }
+}
+#endif
+
+template <bool is_fp32_dest_acc_en>
+inline void mast3r_gelu_poly_init() {
+    math::reset_counters(p_setrwc::SET_ABD_F);
+    sfpi::vConstFloatPrgm0 = MG_X;
+    sfpi::vConstFloatPrgm1 = 0.5f;
+    sfpi::vConstFloatPrgm2 = MG_C[MG_N - 1];
+}
+
+}  // namespace ckernel::sfpu
+#endif
+
+ALWI void mast3r_gelu_tanh_tile_init_pack() {
+    PACK(SFPU_UNARY_INIT_FN(gelu_tanh, sfpu::mast3r_gelu_poly_init, (DST_ACCUM_MODE)));
+}
+ALWI void mast3r_gelu_tanh_tile_pack(uint32_t idst) {
+    PACK(SFPU_UNARY_CALL(DST_SYNC_MODE, DST_ACCUM_MODE, calculate_mast3r_gelu_poly, (DST_ACCUM_MODE), idst, VectorMode::RC));
+}

code/models/demos/mast3r/tt/kernels/mm_gelu_activation.hpp

@@ -0,0 +1,123 @@
+// SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+#pragma once
+#include "ttnn/operations/matmul/shared_with_host/activation_type.hpp"
+#include "api/compute/compute_kernel_api.h"
+#include "api/compute/eltwise_unary/gelu.h"
+#include "api/compute/eltwise_unary/relu.h"
+#include "api/compute/eltwise_unary/activations.h"
+#include "api/compute/eltwise_unary/hardtanh.h"
+#include "api/compute/eltwise_unary/selu.h"
+#include "api/compute/eltwise_unary/softplus.h"
+#include "internal/risc_attribs.h"
+#include "mast3r_gelu_poly.h"  // mast3r-p150 model-local: GELU_TANH slot -> erf-GELU minimax polynomial (MAST3R_OPT "gpoly")
+#include <cstring>  // for memcpy
+
+using ttnn::operations::matmul::KernelActivation;
+
+// Helper templates to select activation variants based on parameters
+template <KernelActivation ACT, uint32_t PARAM0 = 0, uint32_t PARAM1 = 0>
+struct ActivationInitHelper {
+    // Compile-time validation
+    static_assert(
+        ACT == KernelActivation::NONE || ACT == KernelActivation::SILU || ACT == KernelActivation::TANH ||
+            ACT == KernelActivation::GELU || ACT == KernelActivation::GELU_TANH || ACT == KernelActivation::RELU6 ||
+            ACT == KernelActivation::SIGMOID || ACT == KernelActivation::HARDSIGMOID ||
+            ACT == KernelActivation::HARDTANH || ACT == KernelActivation::SELU || ACT == KernelActivation::SOFTPLUS,
+        "Unsupported KernelActivation type for fused activation init");
+
+    FORCE_INLINE static void init() {
+        if constexpr (ACT == KernelActivation::SILU) {
+            silu_tile_init_pack();
+        } else if constexpr (ACT == KernelActivation::TANH) {
+            // PARAM0: 0 = accurate, non-zero = fast
+            tanh_tile_init_pack<PARAM0 != 0>();
+        } else if constexpr (ACT == KernelActivation::GELU) {
+            // PARAM0: 0 = accurate, non-zero = fast
+            gelu_tile_init_pack<PARAM0 != 0>();
+        } else if constexpr (ACT == KernelActivation::GELU_TANH) {
+            mast3r_gelu_tanh_tile_init_pack();
+        } else if constexpr (ACT == KernelActivation::RELU6) {
+            relu_max_tile_init_pack();
+        } else if constexpr (ACT == KernelActivation::SIGMOID) {
+            // Enhanced: PARAM1 is fast_approximate flag
+            sigmoid_tile_init_pack<PARAM1 != 0>();
+        } else if constexpr (ACT == KernelActivation::HARDSIGMOID) {
+            hardsigmoid_tile_init_pack();
+        } else if constexpr (ACT == KernelActivation::HARDTANH) {
+            hardtanh_tile_init_pack();
+        } else if constexpr (ACT == KernelActivation::SELU) {
+            selu_tile_init_pack();
+        } else if constexpr (ACT == KernelActivation::SOFTPLUS) {
+            softplus_tile_init_pack();
+        }
+    }
+};
+
+template <KernelActivation ACT, uint32_t PARAM0 = 0, uint32_t PARAM1 = 0, uint32_t PARAM2 = 0>
+struct ActivationApplyHelper {
+    // Compile-time validation
+    static_assert(
+        ACT == KernelActivation::NONE || ACT == KernelActivation::SILU || ACT == KernelActivation::TANH ||
+            ACT == KernelActivation::GELU || ACT == KernelActivation::GELU_TANH || ACT == KernelActivation::RELU6 ||
+            ACT == KernelActivation::SIGMOID || ACT == KernelActivation::HARDSIGMOID ||
+            ACT == KernelActivation::HARDTANH || ACT == KernelActivation::SELU || ACT == KernelActivation::SOFTPLUS,
+        "Unsupported KernelActivation type for fused activation apply");
+
+    // Parameter-specific validation
+    static_assert(
+        ACT != KernelActivation::SOFTPLUS || PARAM0 != 0,
+        "SOFTPLUS PARAM0 (beta) must be non-zero to avoid division by zero");
+
+    FORCE_INLINE static void apply(uint32_t tile_index) {
+        if constexpr (ACT == KernelActivation::SILU) {
+            silu_tile_pack(tile_index);
+        } else if constexpr (ACT == KernelActivation::TANH) {
+            // PARAM0: 0 = accurate, non-zero = fast
+            tanh_tile_pack<PARAM0 != 0>(tile_index);
+        } else if constexpr (ACT == KernelActivation::GELU) {
+            // PARAM0: 0 = accurate, non-zero = fast
+            gelu_tile_pack<PARAM0 != 0>(tile_index);
+        } else if constexpr (ACT == KernelActivation::GELU_TANH) {
+            mast3r_gelu_tanh_tile_pack(tile_index);
+        } else if constexpr (ACT == KernelActivation::RELU6) {
+            // PARAM0 is the max value (as uint32_t bit pattern)
+            // Default to 6.0 if PARAM0 is 0
+            constexpr uint32_t max = (PARAM0 != 0) ? PARAM0 : 0x40c00000u;
+            relu_max_tile_pack(tile_index, max);
+        } else if constexpr (ACT == KernelActivation::SIGMOID) {
+            // Enhanced: PARAM0 is vector mode, PARAM1 is fast_approximate
+            constexpr VectorMode vec_mode = (PARAM0 == 1)   ? VectorMode::R
+                                            : (PARAM0 == 2) ? VectorMode::C
+                                                            : VectorMode::RC;
+            sigmoid_tile_pack<vec_mode, PARAM1 != 0>(tile_index);
+        } else if constexpr (ACT == KernelActivation::HARDSIGMOID) {
+            hardsigmoid_tile_pack(tile_index);
+        } else if constexpr (ACT == KernelActivation::HARDTANH) {
+            hardtanh_tile_pack(tile_index, PARAM0, PARAM1);
+        } else if constexpr (ACT == KernelActivation::SELU) {
+            // PARAM0 is alpha, PARAM1 is lambda
+            selu_tile_pack(tile_index, PARAM0, PARAM1);
+        } else if constexpr (ACT == KernelActivation::SOFTPLUS) {
+            // PARAM0 is beta, PARAM2 beta reciprocal, PARAM1 is threshold
+            softplus_tile_pack(tile_index, PARAM0, PARAM2, PARAM1);
+        }
+    }
+};
+
+template <KernelActivation ACT, uint32_t PARAM0 = 0, uint32_t PARAM1 = 0, uint32_t PARAM2 = 0>
+FORCE_INLINE void apply_activation_from_pack(uint32_t out_subblock_num_tiles) {
+    PACK(TTI_SEMWAIT(
+        p_stall::STALL_TDMA | p_stall::STALL_CFG, semaphore::t6_sem(semaphore::MATH_PACK), p_stall::STALL_ON_ZERO));
+
+    // Flip destination register offset for PACKER access
+    PACK(TT_SETC16(DEST_TARGET_REG_CFG_MATH_Offset_ADDR32, ckernel::packer::get_packer_dest_offset()));
+
+    for (uint32_t i = 0; i < out_subblock_num_tiles; i++) {
+        ActivationApplyHelper<ACT, PARAM0, PARAM1, PARAM2>::apply(i);
+    }
+
+    // Wait for SFPU completion before packing
+    PACK(TTI_STALLWAIT(p_stall::STALL_PACK, p_stall::WAIT_SFPU));
+}

code/models/demos/mast3r/tt/kernels/mm_gelu_compute.cpp

@@ -0,0 +1,657 @@
+// SPDX-FileCopyrightText: © 2023 Tenstorrent USA, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+// mast3r-p150 model-local copy of ttnn bmm_large_block_zm_fused_bias_activation.cpp (tt-metal 8b98410e730), unchanged
+// except that its "mm_gelu_activation.hpp" resolves to the model-local copy (GELU_TANH -> mast3r_gelu_poly.h, MAST3R_OPT "gpoly")
+// and the optional fused residual add (MAST3R_RESID_CB, MAST3R_OPT "resmm").
+
+#include <cstdint>
+
+#include "api/compute/matmul.h"
+#include "api/compute/compute_kernel_hw_startup.h"
+#include "api/compute/pack_untilize.h"
+#include "api/compute/tile_move_copy.h"
+#include "api/compute/transpose.h"
+#include "api/dataflow/dataflow_buffer.h"
+#include "internal/mod_div_lib.h"
+
+#ifdef FUSE_BIAS
+#include "api/compute/bcast.h"
+#endif
+
+#include "api/compute/eltwise_binary.h"
+#ifdef SFPU_ACTIVATION
+#include "mm_gelu_activation.hpp"
+#endif
+
+// Please update
+// tests/tt_metal/tt_metal/perf_microbenchmark/1_compute_mm/kernels/bmm_large_block_zm_fused_bias_activation_copy.cpp
+// when making any changes to this file.
+// Have to keep a copy because cannot import ttnn into tests/tt_metal.
+// With FUSE_BIAS: row_broadcast_bias (row-broadcast vs elementwise add_tiles) is compile-time arg 18 here;
+// the perf copy uses index 14 (different compile-time arg layout).
+
+/**
+ * @brief Transposes a block of tiles from one circular buffer to another.
+ *
+ * This function reads a block of tiles from the input circular buffer (cb), performs a width-height
+ * (WH) transpose on each tile, and writes the transposed tiles to the output circular buffer.
+ * The operation is performed in blocks of `block_size` tiles for efficiency, with a separate loop
+ * at the end to handle any leftover tiles when the total tile count is not divisible by
+ * `block_size`. The default block size is 4, since there are guaranteed to be 4 tiles in the dst
+ *               regs irrespective of dst sync mode or data format.
+ *
+ * @tparam in0_block_num_tiles The number of tiles in the block to be transposed.
+ * @tparam block_size The number of tiles in each block to be transposed.
+ * @param in0_transpose_dfb_id Circular buffer ID to read the original tiles from.
+ * @param in0_dfb_id Circular buffer ID to which the transposed tiles are written.
+ */
+template <uint32_t in0_block_num_tiles, uint32_t block_size = 4>
+FORCE_INLINE void transpose_tile_block(uint32_t in0_transpose_dfb_id, uint32_t in0_dfb_id) {
+    DataflowBuffer in0_transpose_dfb(in0_transpose_dfb_id);
+    DataflowBuffer in0_dfb(in0_dfb_id);
+    constexpr uint32_t num_blocks = in0_block_num_tiles / block_size;
+    constexpr uint32_t last_block_size = in0_block_num_tiles % block_size;
+    // Lets do 2 passes: One loop until last and one last for the left overs
+    for (uint32_t block_idx = 0; block_idx < num_blocks; ++block_idx) {
+        in0_transpose_dfb.wait_front(block_size);
+        tile_regs_acquire();
+        for (uint32_t tile_idx = 0; tile_idx < block_size; tile_idx++) {
+            transpose_tile(in0_transpose_dfb_id, tile_idx, tile_idx);
+        }
+        tile_regs_commit();
+        in0_transpose_dfb.pop_front(block_size);
+
+        in0_dfb.reserve_back(block_size);
+        tile_regs_wait();
+        for (uint32_t tile_idx = 0; tile_idx < block_size; tile_idx++) {
+            pack_tile(tile_idx, in0_dfb_id);
+        }
+        tile_regs_release();
+        in0_dfb.push_back(block_size);
+    }
+
+    if constexpr (last_block_size > 0) {
+        in0_transpose_dfb.wait_front(last_block_size);
+        tile_regs_acquire();
+        for (uint32_t tile_idx = 0; tile_idx < last_block_size; tile_idx++) {
+            transpose_tile(in0_transpose_dfb_id, tile_idx, tile_idx);
+        }
+        tile_regs_commit();
+        in0_transpose_dfb.pop_front(last_block_size);
+
+        in0_dfb.reserve_back(last_block_size);
+        tile_regs_wait();
+        for (uint32_t tile_idx = 0; tile_idx < last_block_size; tile_idx++) {
+            pack_tile(tile_idx, in0_dfb_id);
+        }
+        tile_regs_release();
+        in0_dfb.push_back(last_block_size);
+    }
+}
+
+FORCE_INLINE void reload_from_cb_to_dst(
+    uint32_t in0_dfb_id,
+    uint32_t in1_dfb_id,
+    uint32_t mm_partials_dfb_id,
+    uint32_t mm_partials_reload_dfb_id,
+    bool in1_transpose_tile,
+    uint32_t out_subblock_num_tiles,
+    uint32_t out_subblock_w,
+    uint32_t out_subblock_h,
+    uint32_t in0_block_w) {
+    DataflowBuffer mm_partials_dfb(mm_partials_dfb_id);
+    // mm_partials_reload_dfb_id is the CB view the reload copies through. It equals mm_partials_dfb_id
+    // unless the partials CB is also read as an FPU operand elsewhere (the fused bias add reads it via
+    // SrcA), in which case UnpackToDestFp32 cannot be set on it directly; instead a second buffer index
+    // aliases the same SRAM with UnpackToDestFp32 set, and the reload copies through that alias while the
+    // FPU consumer keeps the original view. The alias has its own read pointer, so align it with the
+    // partials CB's current read position before copying.
+    // Reconfigure input
+    copy_tile_to_dst_init_short_with_dt(in1_dfb_id, mm_partials_reload_dfb_id);
+    mm_partials_dfb.wait_front(out_subblock_num_tiles);
+
+    if (mm_partials_reload_dfb_id != mm_partials_dfb_id) {
+        // Only the unpacker owns cb_interface / the read pointer; keep this off the MATH/PACK threads.
+        UNPACK(
+            (get_local_cb_interface(mm_partials_reload_dfb_id).fifo_rd_ptr =
+                 get_local_cb_interface(mm_partials_dfb_id).fifo_rd_ptr));
+    }
+
+    uint32_t start_dst_index = 0;
+    uint32_t start_tile_index = 0;
+    copy_block(mm_partials_reload_dfb_id, start_tile_index, start_dst_index, out_subblock_num_tiles);
+
+    mm_partials_dfb.pop_front(out_subblock_num_tiles);
+    // Reconfigure srcA back
+    reconfig_data_format_srca(mm_partials_reload_dfb_id, in1_dfb_id);
+    matmul_block_init(in0_dfb_id, in1_dfb_id, in1_transpose_tile, out_subblock_w, out_subblock_h, in0_block_w);
+}
+
+template <uint32_t out_subblock_w, uint32_t out_block_w>
+inline void reblock_and_untilize(
+    uint32_t num_out_subblocks_in_col,
+    uint32_t out_subblock_num_tiles,
+    uint32_t out_subblock_h,
+    uint32_t interm_dfb_id,
+    uint32_t out_dfb_id) {
+    DataflowBuffer interm_dfb(interm_dfb_id);
+    DataflowBuffer out_dfb(out_dfb_id);
+    uint32_t num_tiles_in_row_of_subblocks = mulsi3(out_subblock_num_tiles, num_out_subblocks_in_col);
+    interm_dfb.wait_front(num_tiles_in_row_of_subblocks);
+
+    uint32_t within_block_index = 0;
+    for (uint32_t h = 0; h < out_subblock_h; h++) {
+        uint32_t block_offset = 0;
+
+        out_dfb.reserve_back(out_block_w);
+        for (uint32_t n = 0; n < num_out_subblocks_in_col; n++) {
+            tile_regs_acquire();
+            for (uint32_t w = 0; w < out_subblock_w; w++) {
+                uint32_t tile_index = block_offset + within_block_index + w;
+                copy_tile(interm_dfb_id, tile_index, w);
+            }
+            tile_regs_commit();
+            tile_regs_wait();
+            pack_untilize_dest<out_subblock_w, out_block_w>(out_dfb_id, 1, n);
+            tile_regs_release();
+            block_offset += out_subblock_num_tiles;
+        }
+        out_dfb.push_back(out_block_w);
+
+        within_block_index += out_subblock_w;
+    }
+    interm_dfb.pop_front(num_tiles_in_row_of_subblocks);
+}
+
+void kernel_main() {
+// RUNTIME ARGS
+#ifdef MATMUL_DRAM_SHARDED
+    const bool is_worker_core = get_arg_val<uint32_t>(0) == 1;
+    // if not worker core, skip
+    if (not is_worker_core) {
+        return;
+    }
+#endif
+
+    constexpr uint32_t in0_block_w = get_compile_time_arg_val(0);        // inner block size in tiles
+    constexpr uint32_t in0_num_subblocks = get_compile_time_arg_val(1);  // outer row block size (in inner row blocks)
+    constexpr uint32_t in0_block_num_tiles =
+        get_compile_time_arg_val(2);  // out_subblock_h*in0_block_w*in0_num_subblocks;
+    constexpr uint32_t in0_subblock_num_tiles = get_compile_time_arg_val(3);  // out_subblock_h*in0_block_w
+    constexpr uint32_t in1_num_subblocks =
+        get_compile_time_arg_val(4);  // outer column block size (in inner column blocks)
+    constexpr uint32_t in1_block_num_tiles =
+        get_compile_time_arg_val(5);                               // out_subblock_w*in0_block_w* in1_num_subblocks;
+    constexpr uint32_t in1_block_w = get_compile_time_arg_val(6);  // out_subblock_w*in1_num_subblocks
+    constexpr uint32_t num_blocks_inner_dim = get_compile_time_arg_val(7);     // outer inner dim (in inner dim blocks)
+    constexpr uint32_t num_blocks_w_dim = get_compile_time_arg_val(8);         // outer inner dim (in inner dim blocks)
+    constexpr uint32_t num_blocks_h_dim = get_compile_time_arg_val(9);         // outer inner dim (in inner dim blocks)
+    constexpr uint32_t out_subblock_h = get_compile_time_arg_val(10);          // inner row block size in tiles
+    constexpr uint32_t out_subblock_w = get_compile_time_arg_val(11);          // inner column block size in tiles
+    constexpr uint32_t out_subblock_num_tiles = get_compile_time_arg_val(12);  // out_subblock_h * out_subblock_w;
+    constexpr uint32_t batch = get_compile_time_arg_val(13);                   // batch dim
+    constexpr uint32_t out_block_num_tiles = get_compile_time_arg_val(14);     // number of tiles in out_block
+    constexpr bool untilize_out = get_compile_time_arg_val(15);                // untilize output
+    // This boolean is set when the number of batches is only known at runtime, typically based on a sparsity tensor.
+    constexpr bool get_batch_from_reader = (bool)get_compile_time_arg_val(16);
+    constexpr bool in0_transpose_tile = (bool)get_compile_time_arg_val(17);
+
+    constexpr uint32_t out_block_w = out_subblock_w * in1_num_subblocks;
+
+    constexpr uint32_t in0_dfb_id = in0_transpose_tile ? get_named_compile_time_arg_val("cb_in0_transposed")
+                                                       : get_named_compile_time_arg_val("cb_in0");
+    constexpr uint32_t in1_dfb_id = get_named_compile_time_arg_val("cb_in1");
+    constexpr uint32_t out_dfb_id = get_named_compile_time_arg_val("cb_out");
+    constexpr uint32_t mm_partials_dfb_id = get_named_compile_time_arg_val("cb_intermed0");
+    // CB view the cross-block reload copies through: the UnpackToDestFp32-marked alias of the partials
+    // CB when it is also read as an FPU operand (fused bias), otherwise the partials CB itself.
+#ifdef MM_PARTIALS_RELOAD_ALIAS_CB
+    // The partials CB is also read as an FPU operand (fused bias) and so cannot carry UnpackToDestFp32;
+    // the reload instead copies through this alias view of the same SRAM, which does carry the flag.
+    constexpr uint32_t mm_partials_reload_dfb_id = MM_PARTIALS_RELOAD_ALIAS_CB;
+#else
+    constexpr uint32_t mm_partials_reload_dfb_id = mm_partials_dfb_id;
+#endif
+    constexpr uint32_t untilize_mode_out_dfb_id = untilize_out ? mm_partials_dfb_id : out_dfb_id;
+    // When in0 needs to be transposed, the original data is read from cb_in0 (in0_transpose_dfb_id),
+    // transposed, and the result is written to cb_in0_transposed (in0_dfb_id), which is then used
+    // as input for the matmul call.
+    constexpr uint32_t in0_transpose_dfb_id = get_named_compile_time_arg_val("cb_in0");
+
+    DataflowBuffer in0_dfb(in0_dfb_id);
+    DataflowBuffer in1_dfb(in1_dfb_id);
+    DataflowBuffer mm_partials_dfb(mm_partials_dfb_id);
+    DataflowBuffer untilize_mode_out_dfb(untilize_mode_out_dfb_id);
+
+#ifdef FUSE_BIAS
+    constexpr uint32_t bias_dfb_id = get_named_compile_time_arg_val("cb_bias");
+    constexpr uint32_t bias_ntiles = get_named_compile_time_arg_val("bias_ntiles");
+    constexpr uint32_t mm_out_dfb_id = mm_partials_dfb_id;
+    // true: row-0 broadcast ([N] / [...,1,N]); false: elementwise add_tiles (bias has multiple M rows).
+    constexpr bool row_broadcast_bias = (bool)get_compile_time_arg_val(18);
+    DataflowBuffer bias_dfb(bias_dfb_id);
+#ifdef MAST3R_RESID_CB
+    // MAST3R_OPT "resmm": out = in0 @ in1 + bias + resid. The residual CB is globally allocated on a BLOCK_SHARDED
+    // tensor whose shard == this core's output block (row-major per_core_M x per_core_N tiles, same spec as the
+    // output). Only valid for batch == 1, one output block per core (asserted on the host).
+    constexpr uint32_t resid_dfb_id = MAST3R_RESID_CB;
+    DataflowBuffer resid_dfb(resid_dfb_id);
+#endif
+#else
+    constexpr uint32_t mm_out_dfb_id = untilize_mode_out_dfb_id;
+#endif
+    DataflowBuffer mm_out_dfb(mm_out_dfb_id);
+
+    // Number of valid in1 columns in the last in1 subblock. For the DRAM-sharded variant the
+    // planner may pad per_core_N_compute beyond per_core_N_in1_sender so that out_subblock_w can be
+    // larger; the reader only pushes per_core_N_in1_sender tiles per block into cb_in1. To avoid
+    // reading those non-existent (padded) cb_in1 tiles, the compute kernel narrows the matmul_block
+    // call on the last in1 subblock to last_subblock_w_valid lanes. When no padding occurs this
+    // equals out_subblock_w and the original full-width path is preserved.
+#ifdef MATMUL_DRAM_SHARDED
+    constexpr uint32_t last_subblock_w_valid = get_named_compile_time_arg_val("last_subblock_w_valid");
+#else
+    constexpr uint32_t last_subblock_w_valid = out_subblock_w;
+#endif
+    constexpr bool last_subblock_padded = last_subblock_w_valid < out_subblock_w;
+
+#ifdef SFPU_ACTIVATION
+    constexpr KernelActivation activation_type =
+        static_cast<KernelActivation>(get_named_compile_time_arg_val("activation_type"));
+    constexpr uint32_t activation_param0 = get_named_compile_time_arg_val("activation_param0");
+    constexpr uint32_t activation_param1 = get_named_compile_time_arg_val("activation_param1");
+    constexpr uint32_t activation_param2 = get_named_compile_time_arg_val("activation_param2");
+
+    ActivationInitHelper<activation_type, activation_param0, activation_param1>::init();
+#endif
+
+#ifdef IN1_TRANSPOSE_TILE
+    constexpr uint32_t in1_transpose_tile = true;
+#else
+    constexpr uint32_t in1_transpose_tile = false;
+#endif
+
+    constexpr bool spill = num_blocks_inner_dim > 1;
+
+    compute_kernel_hw_startup<SrcOrder::Reverse>(in0_dfb_id, in1_dfb_id, mm_partials_dfb_id);
+    matmul_block_init(in0_dfb_id, in1_dfb_id, in1_transpose_tile, out_subblock_w, out_subblock_h, in0_block_w);
+    for (uint32_t b = 0; b < batch; b++) {
+        if constexpr (get_batch_from_reader) {
+            // Check whether this batch is valid
+            bool is_batch_valid = false;
+            UNPACK(is_batch_valid = (bool)mailbox_read(ckernel::ThreadId::BriscThreadId);)
+            MATH(is_batch_valid = (bool)mailbox_read(ckernel::ThreadId::BriscThreadId);)
+            PACK(is_batch_valid = (bool)mailbox_read(ckernel::ThreadId::BriscThreadId);)
+            if (!is_batch_valid) {
+                continue;
+            }
+        }
+
+        for (uint32_t bh = 0; bh < num_blocks_h_dim; ++bh) {
+            for (uint32_t bw = 0; bw < num_blocks_w_dim; ++bw) {
+                bool enable_reload = false;
+
+#ifdef PACK_RELU
+                // for each batch we start with relu disabled so that intermediate results are not relu'd
+                if constexpr (batch > 1 || num_blocks_h_dim > 1 || num_blocks_w_dim > 1) {
+                    PACK((llk_pack_relu_config(ReluConfig::none())));
+                }
+#endif
+
+                if constexpr (batch > 1 || num_blocks_h_dim > 1 || num_blocks_w_dim > 1) {
+                    PACK((pack_reconfig_data_format(mm_partials_dfb_id)));
+                }
+
+                for (uint32_t block = 0; block < num_blocks_inner_dim; block++) {
+                    bool last_out = block == (num_blocks_inner_dim - 1);
+// Configure packer once for pack out without Bias
+#if not defined FUSE_BIAS and defined PACK_RELU
+                    if (last_out) {
+                        // if last block we pack the final result with relu enabled
+                        PACK((llk_pack_relu_config(ReluConfig::zero())));
+                    }
+#endif
+
+                    if constexpr (in0_transpose_tile) {
+                        reconfig_data_format_srca(in1_dfb_id, in0_transpose_dfb_id);
+                        transpose_init(in0_transpose_dfb_id);
+                        PACK((pack_reconfig_data_format(in0_dfb_id)));
+#ifdef PACKER_L1_ACC
+                        PACK((llk_pack_reconfig_l1_acc(0)));
+#endif
+                        transpose_tile_block<in0_block_num_tiles>(in0_transpose_dfb_id, in0_dfb_id);
+                        reconfig_data_format_srca(in0_transpose_dfb_id, in1_dfb_id);
+                        matmul_block_init(
+                            in0_dfb_id, in1_dfb_id, in1_transpose_tile, out_subblock_w, out_subblock_h, in0_block_w);
+                        PACK((pack_reconfig_data_format(mm_partials_dfb_id)));
+                    }
+
+                    in0_dfb.wait_front(in0_block_num_tiles);
+                    in1_dfb.wait_front(in1_block_num_tiles);
+
+                    int in0_index_subblock_offset = 0;
+                    for (uint32_t in0_subblock = 0; in0_subblock < in0_num_subblocks; in0_subblock++) {
+                        int in1_index_subblock_offset = 0;
+                        for (uint32_t in1_subblock = 0; in1_subblock < in1_num_subblocks; in1_subblock++) {
+                            // When last_subblock_padded is true the last in1 subblock has
+                            // (out_subblock_w - last_subblock_w_valid) padded lanes whose cb_in1 tiles were
+                            // never pushed by the reader. Narrow matmul_block so the unpacker only touches
+                            // tiles that exist; the padded dst lanes are left at whatever the previous
+                            // operation wrote there and the output writer (BRISC) drops those columns.
+                            const bool is_last_in1_subblock_padded =
+                                last_subblock_padded && (in1_subblock == in1_num_subblocks - 1);
+                            const uint32_t effective_subblock_w =
+                                is_last_in1_subblock_padded ? last_subblock_w_valid : out_subblock_w;
+
+                            tile_regs_acquire();
+                            if (enable_reload) {
+                                reload_from_cb_to_dst(
+                                    in0_dfb_id,
+                                    in1_dfb_id,
+                                    mm_partials_dfb_id,
+                                    mm_partials_reload_dfb_id,
+                                    in1_transpose_tile,
+                                    out_subblock_num_tiles,
+                                    out_subblock_w,
+                                    out_subblock_h,
+                                    in0_block_w);
+                            }
+
+#ifndef SKIP_COMPUTE
+                            // Compute output sub-block
+                            uint32_t dst_index =
+                                0;  // start at 0, each call to matmul_block internally increments dst_index
+                            uint32_t in0_index = in0_index_subblock_offset;  // offset into in0 block
+                            uint32_t in1_index = in1_index_subblock_offset;  // offset into in1 block
+                            // inner dim that we accumulate is the inner dim of in0/in1, which is in0_block_w
+                            for (uint32_t inner_dim_idx = 0; inner_dim_idx < in0_block_w; ++inner_dim_idx) {
+                                // matmul outer product of (out_subblock_h x out_subblock_w) tiles that fill dst
+                                // accumulation is done by iterating matmul_block across inner dim
+                                // in0_block_w is passed as innder dim (kt) to matmul_block, internally used to stride
+                                // in0
+                                matmul_block(
+                                    in0_dfb_id,
+                                    in1_dfb_id,
+                                    in0_index,
+                                    in1_index,
+                                    dst_index,
+                                    in1_transpose_tile,
+                                    effective_subblock_w,
+                                    out_subblock_h,
+                                    in0_block_w);
+                                in0_index++;               // stride right by 1
+                                in1_index += in1_block_w;  // to stride down by 1 need to stride by in_per_core_w
+                                                           // (should be called in1_block_w)
+                            }
+
+#endif  // SKIP_COMPUTE
+
+                            if (last_out) {
+                                tile_regs_commit();
+                                mm_out_dfb.reserve_back(out_subblock_num_tiles);
+
+#if defined SFPU_ACTIVATION and not defined FUSE_BIAS
+                                apply_activation_from_pack<
+                                    activation_type,
+                                    activation_param0,
+                                    activation_param1,
+                                    activation_param2>(out_subblock_num_tiles);
+#else
+                                tile_regs_wait();
+#endif
+
+#if defined FP32_DEST_ACC_EN or defined PACKER_L1_ACC
+                                PACK((pack_reconfig_data_format(mm_out_dfb_id)));
+#endif
+
+#ifdef PACKER_L1_ACC
+#ifdef FUSE_BIAS
+                                if (block == 0) {  // no accumulation for first iteration
+                                    PACK((llk_pack_reconfig_l1_acc(0)));
+                                } else {
+                                    PACK((llk_pack_reconfig_l1_acc(1)));
+                                }
+#else
+                                PACK((llk_pack_reconfig_l1_acc(0)));
+#endif
+#endif
+                                uint32_t start_dst_index = 0;
+                                pack_block(start_dst_index, mm_out_dfb_id, out_subblock_num_tiles);
+
+                                tile_regs_release();
+                                mm_out_dfb.push_back(out_subblock_num_tiles);
+
+                            } else {
+                                tile_regs_commit();
+                                mm_partials_dfb.reserve_back(out_subblock_num_tiles);
+                                tile_regs_wait();
+
+#ifdef PACKER_L1_ACC
+                                if (block == 0) {  // no accumulation for first iteration
+                                    PACK((llk_pack_reconfig_l1_acc(0)));
+                                } else if (block == 1) {
+                                    PACK((llk_pack_reconfig_l1_acc(1)));
+                                } else if (in0_transpose_tile) {
+                                    // For each block, l1_acc would have been enabled during the
+                                    // transpose stage. So let us put it back here.
+                                    PACK((llk_pack_reconfig_l1_acc(1)));
+                                }
+#endif
+
+                                uint32_t start_dst_index = 0;
+                                pack_block(start_dst_index, mm_partials_dfb_id, out_subblock_num_tiles);
+
+                                tile_regs_release();
+                                mm_partials_dfb.push_back(out_subblock_num_tiles);
+                            }
+
+                            in1_index_subblock_offset += out_subblock_w;
+                        }
+                        in0_index_subblock_offset += in0_subblock_num_tiles;
+                    }
+
+#ifdef PACKER_L1_ACC
+#ifdef FUSE_BIAS
+                    if (block < num_blocks_inner_dim - 1) {
+                        // Wait/pop in subblock-sized steps so the step size
+                        // matches the bias section's wait_front(out_subblock_num_tiles),
+                        // satisfying the CB API requirement that all wait_front
+                        // increments on a given CB are identical.
+                        for (uint32_t s = 0; s < out_block_num_tiles; s += out_subblock_num_tiles) {
+                            mm_partials_dfb.wait_front(out_subblock_num_tiles);
+                            mm_partials_dfb.pop_front(out_subblock_num_tiles);
+                        }
+                    }
+                    // never reload when with bias, bias uses intermediate buffer
+                    enable_reload = false;
+#else
+                    // Last iteration does spill and reload to output buffer
+                    if (block < num_blocks_inner_dim - 2) {
+                        for (uint32_t s = 0; s < out_block_num_tiles; s += out_subblock_num_tiles) {
+                            mm_partials_dfb.wait_front(out_subblock_num_tiles);
+                            mm_partials_dfb.pop_front(out_subblock_num_tiles);
+                        }
+                    }
+                    if (block == num_blocks_inner_dim - 2) {
+                        enable_reload = true;
+                    }  // reload when last iteration
+#endif
+#else
+                    if constexpr (spill) {
+                        enable_reload = true;
+                    }
+#endif
+
+                    in0_dfb.pop_front(in0_block_num_tiles);
+                    in1_dfb.pop_front(in1_block_num_tiles);
+                }
+
+#ifdef FUSE_BIAS
+#ifdef PACK_RELU
+                // if last block we pack the final result with relu enabled
+                PACK((llk_pack_relu_config(ReluConfig::zero())));
+#endif
+#if defined FP32_DEST_ACC_EN or defined PACKER_L1_ACC
+                PACK((pack_reconfig_data_format(out_dfb_id)));
+#endif
+#ifdef PACKER_L1_ACC
+                PACK((llk_pack_reconfig_l1_acc(0)));
+#endif
+                reconfig_data_format(in1_dfb_id, mm_partials_dfb_id, in0_dfb_id, bias_dfb_id);
+                if constexpr (row_broadcast_bias) {
+                    add_bcast_rows_init(mm_partials_dfb_id, bias_dfb_id);
+                } else {
+                    add_init(mm_partials_dfb_id, bias_dfb_id);
+                }
+                // Reader only pushes bias once when num_blocks_w_dim == 1;
+                // the tiles stay in the CB for reuse across bh/batch iterations.
+                if ((b == 0 && bh == 0) || num_blocks_w_dim > 1) {
+                    bias_dfb.wait_front(bias_ntiles);
+                }
+#ifdef MAST3R_RESID_CB
+                // nothing produces the sharded residual CB: publish its (already resident) tiles once
+                resid_dfb.reserve_back(out_block_num_tiles);
+                resid_dfb.push_back(out_block_num_tiles);
+                resid_dfb.wait_front(out_block_num_tiles);
+#endif
+                for (uint32_t in0_subblock = 0; in0_subblock < in0_num_subblocks; in0_subblock++) {
+                    int in1_index_subblock_offset = 0;
+                    for (uint32_t in1_subblock = 0; in1_subblock < in1_num_subblocks; in1_subblock++) {
+                        // See matmul stage: the last in1 subblock has padded lanes whose bias tile was
+                        // never pushed by the reader. Redirect those out-of-range bias_tile_idx reads to
+                        // tile 0 of cb_bias to keep them in-bounds; the resulting padded output columns
+                        // are dropped by the writer.
+                        const bool is_last_in1_subblock_padded =
+                            last_subblock_padded && (in1_subblock == in1_num_subblocks - 1);
+                        // Redundant wait since we know data was just pushed
+                        mm_partials_dfb.wait_front(out_subblock_num_tiles);
+                        tile_regs_acquire();
+                        for (uint32_t i = 0, j = 0; j < out_subblock_h; j++) {
+#ifdef BIAS_FULL_BLOCK
+                            // The bias CB holds a full [M, N] tile block. m_tile is this output tile's
+                            // row within that block; bias_tile_idx is the position of the matching bias
+                            // tile in the CB (row m_tile, column in1_index_subblock_offset). Only
+                            // matmul_multicore_reuse_optimized loads the full block; other callers load a
+                            // single bias row and use the N-only index below.
+                            const uint32_t m_tile = in0_subblock * out_subblock_h + j;
+                            uint32_t bias_tile_idx = m_tile * in1_block_w + in1_index_subblock_offset;
+#else
+                            uint32_t bias_tile_idx = in1_index_subblock_offset;
+#endif
+                            for (uint32_t k = 0; k < out_subblock_w; k++, i++) {
+                                const uint32_t safe_bias_tile_idx =
+                                    (is_last_in1_subblock_padded && k >= last_subblock_w_valid)
+                                        ? 0u              // Padded output columns with tile 0 of cb_bias added are
+                                        : bias_tile_idx;  // dropped by the writer.
+
+                                if constexpr (row_broadcast_bias) {
+                                    add_tiles_bcast_rows(mm_partials_dfb_id, bias_dfb_id, i, safe_bias_tile_idx, i);
+                                } else {
+                                    add_tiles(mm_partials_dfb_id, bias_dfb_id, i, safe_bias_tile_idx, i);
+                                }
+                                bias_tile_idx++;
+                            }
+                        }
+#ifdef MAST3R_RESID_CB
+                        // dst (fp32: partials + bias) += resid tile; DST -> SrcA, resid -> SrcB (bf16, same as bias)
+                        add_reuse_dest_init<EltwiseBinaryReuseDestType::DEST_TO_SRCA>(resid_dfb_id);
+                        for (uint32_t i = 0, j = 0; j < out_subblock_h; j++) {
+                            const uint32_t r0 = (in0_subblock * out_subblock_h + j) * in1_block_w + in1_index_subblock_offset;
+                            for (uint32_t k = 0; k < out_subblock_w; k++, i++) {
+                                add_reuse_dest_tiles<EltwiseBinaryReuseDestType::DEST_TO_SRCA>(resid_dfb_id, r0 + k, i);
+                            }
+                        }
+                        if constexpr (row_broadcast_bias) {
+                            add_bcast_rows_init(mm_partials_dfb_id, bias_dfb_id);
+                        } else {
+                            add_init(mm_partials_dfb_id, bias_dfb_id);
+                        }
+#endif
+                        tile_regs_commit();
+
+                        mm_partials_dfb.pop_front(out_subblock_num_tiles);
+
+                        // Pack out to output buffer
+                        untilize_mode_out_dfb.reserve_back(out_subblock_num_tiles);
+
+#ifdef SFPU_ACTIVATION
+                        PACK(TTI_SEMWAIT(
+                            p_stall::STALL_TDMA | p_stall::STALL_CFG,
+                            semaphore::t6_sem(semaphore::MATH_PACK),
+                            p_stall::STALL_ON_ZERO));
+
+                        // Flip destination register offset for PACKER access
+                        PACK(TT_SETC16(
+                            DEST_TARGET_REG_CFG_MATH_Offset_ADDR32, ckernel::packer::get_packer_dest_offset()));
+
+                        for (uint32_t i = 0; i < out_subblock_num_tiles; i++) {
+                            ActivationApplyHelper<activation_type, activation_param0, activation_param1>::apply(i);
+                        }
+
+                        PACK(TTI_STALLWAIT(p_stall::STALL_PACK, p_stall::WAIT_SFPU));
+#else
+                        tile_regs_wait();
+#endif
+                        for (uint32_t i = 0; i < out_subblock_num_tiles; i++) {
+                            pack_tile(i, untilize_mode_out_dfb_id);
+                        }
+                        tile_regs_release();
+                        untilize_mode_out_dfb.push_back(out_subblock_num_tiles);
+
+                        in1_index_subblock_offset += out_subblock_w;
+                    }
+                }
+                if constexpr (num_blocks_w_dim > 1) {
+                    bias_dfb.pop_front(bias_ntiles);
+                }
+#endif  // FUSE_BIAS
+                if constexpr (untilize_out) {
+#ifdef PACK_RELU
+                    PACK((llk_pack_relu_config(ReluConfig::none())));
+#endif  // PACK_RELU
+#ifndef FUSE_BIAS
+                    reconfig_data_format_srca(in1_dfb_id, mm_partials_dfb_id);
+#if defined FP32_DEST_ACC_EN or defined PACKER_L1_ACC
+                    PACK((pack_reconfig_data_format(out_dfb_id)));
+#endif
+#ifdef PACKER_L1_ACC
+                    PACK((llk_pack_reconfig_l1_acc(0)));
+#endif
+#endif  // FUSE_BIAS
+                    pack_untilize_dest_init<out_subblock_w, out_block_w>(out_dfb_id);
+                    copy_tile_to_dst_init_short(mm_partials_dfb_id);
+                    for (uint32_t in0_subblock_i = 0; in0_subblock_i < in0_num_subblocks; ++in0_subblock_i) {
+                        reblock_and_untilize<out_subblock_w, out_block_w>(
+                            in1_num_subblocks, out_subblock_num_tiles, out_subblock_h, mm_partials_dfb_id, out_dfb_id);
+                    }
+                    pack_untilize_uninit(mm_partials_dfb_id);
+                }
+                if constexpr (batch > 1 || num_blocks_w_dim > 1 || num_blocks_h_dim > 1) {
+#ifdef FUSE_BIAS
+                    // reconfigure unpacker df for src A and src B
+                    reconfig_data_format(mm_partials_dfb_id, in1_dfb_id, bias_dfb_id, in0_dfb_id);
+#else
+                    // reconfigure unpacker df for src A
+                    reconfig_data_format_srca(mm_partials_dfb_id, in1_dfb_id);
+#endif
+                    // reconfigure init for matmul
+                    matmul_block_init(
+                        in0_dfb_id, in1_dfb_id, in1_transpose_tile, out_subblock_w, out_subblock_h, in0_block_w);
+                }
+            }
+        }
+    }
+#ifdef FUSE_BIAS
+    // For num_blocks_w_dim == 1 the reader pushes bias once and the kernel holds it resident,
+    // reusing it across all batch/bh/block iterations without popping. Pop it once here, after the
+    // last use, so the CB is balanced. (For num_blocks_w_dim > 1 the per-block pop above already
+    // balances each re-pushed bias block.)
+    if constexpr (num_blocks_w_dim == 1) {
+        bias_dfb.pop_front(bias_ntiles);
+    }
+#ifdef MAST3R_RESID_CB
+    resid_dfb.pop_front(out_block_num_tiles);
+#endif
+#endif
+}

code/models/demos/mast3r/tt/kernels/mm_in0_heads_reader.cpp

@@ -0,0 +1,472 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+// mast3r-p150 model-local copy of ttnn's reader_bmm_tile_layout_in0_sender_padding.cpp (tt-metal 8b98410e730) for
+// MAST3R_OPT "pcat": the only change is the in0 page id remap below, so the projection matmul reads the SDPA output
+// [B, H, N, dh] (TILE, interleaved) directly as the concatenated-heads matrix [B*N, H*dh] (no concat-heads program).
+// in0 tile (m, k) of the [M = NB*N, K = H*dh] view -> tile ((b*H + h)*NT + n)*DT + d of the head-major tensor, with
+// b = m / NT + B0, n = m % NT, h = k / DT, d = k % DT.
+
+#include <stdint.h>
+
+#include "api/dataflow/dataflow_api.h"
+#include "api/debug/assert.h"
+#include "hostdevcommon/common_values.hpp"
+#include "ttnn/operations/ccl/kernel_common/worker_sync_utils.hpp"
+#include "ttnn/operations/kernel_helper_functions/pad_tile.hpp"
+#include "ckernel.h"
+#include "ckernel_defs.h"
+#include "api/dataflow/noc.h"
+#include "api/dataflow/dataflow_buffer.h"
+#include "api/dataflow/noc_semaphore.h"
+#include "api/tensor/noc_traits.h"
+#include "api/dataflow/endpoints.h"
+#include "api/core_local_mem.h"
+
+#ifndef MAST3R_HEADS
+#error "MAST3R_HEADS / MAST3R_NT / MAST3R_DT / MAST3R_B0 defines required"
+#endif
+static inline __attribute__((always_inline)) uint32_t mast3r_in0_remap(uint32_t t) {
+    constexpr uint32_t KT = MAST3R_HEADS * MAST3R_DT;
+    const uint32_t m = t / KT;
+    const uint32_t k = t - m * KT;
+    const uint32_t mb = m / MAST3R_NT;
+    const uint32_t n = m - mb * MAST3R_NT;
+    const uint32_t h = k / MAST3R_DT;
+    const uint32_t d = k - h * MAST3R_DT;
+    return (((mb + MAST3R_B0) * MAST3R_HEADS + h) * MAST3R_NT + n) * MAST3R_DT + d;
+}
+
+void kernel_main() {
+    uint32_t rt_args_idx = 0;
+    // in0 tensor args
+    const uint32_t in0_tensor_addr = get_arg_val<uint32_t>(rt_args_idx++);
+    uint32_t in0_tensor_start_tile_id = get_arg_val<uint32_t>(rt_args_idx++);
+    // in0 mcast args
+    const uint32_t in0_mcast_dest_noc_start_x = get_arg_val<uint32_t>(rt_args_idx++);
+    const uint32_t in0_mcast_dest_noc_start_y = get_arg_val<uint32_t>(rt_args_idx++);
+    const uint32_t in0_mcast_dest_noc_end_x = get_arg_val<uint32_t>(rt_args_idx++);
+    const uint32_t in0_mcast_dest_noc_end_y = get_arg_val<uint32_t>(rt_args_idx++);
+
+    // padding args
+    const uint32_t last_block_h = get_arg_val<uint32_t>(rt_args_idx++);
+    // sparsity args
+    const uint32_t sparsity_addr = get_arg_val<uint32_t>(rt_args_idx++);
+
+    // COMPILE TIME ARGS
+    // in0 tensor args
+    constexpr uint32_t in0_tensor_stride_w = get_compile_time_arg_val(0);
+    constexpr uint32_t in0_tensor_stride_h = get_compile_time_arg_val(1);
+    constexpr uint32_t in0_tensor_next_inner_dim_block_stride = get_compile_time_arg_val(2);
+    constexpr uint32_t in0_tensor_next_h_dim_block_stride = get_compile_time_arg_val(3);
+    // in0 block args
+    constexpr uint32_t in0_block_w = get_compile_time_arg_val(4);
+    constexpr uint32_t in0_block_h = get_compile_time_arg_val(5);
+    constexpr uint32_t in0_block_num_tiles = get_compile_time_arg_val(6);
+    constexpr uint32_t in0_last_ktile_w = get_compile_time_arg_val(7);
+    constexpr uint32_t in0_last_ktile_h = get_compile_time_arg_val(8);
+
+    constexpr bool extract_shard_sub_blocks = (bool)get_compile_time_arg_val(9);
+    constexpr uint32_t shard_width_in_tiles = get_compile_time_arg_val(10);
+    constexpr uint32_t shard_height_in_tiles = get_compile_time_arg_val(11);
+    // in0/in1 common args
+    constexpr uint32_t num_blocks_inner_dim = get_compile_time_arg_val(12);
+    constexpr uint32_t num_blocks_w_dim = get_compile_time_arg_val(13);
+    constexpr uint32_t num_blocks_h_dim = get_compile_time_arg_val(14);
+    // in0 mcast args
+    constexpr uint32_t in0_mcast_num_dests = get_compile_time_arg_val(17);
+    constexpr uint32_t in0_mcast_num_cores = get_compile_time_arg_val(18);
+    // batch args
+    constexpr uint32_t MtKt = get_compile_time_arg_val(19);  // if 0
+    constexpr uint32_t in0_B = get_compile_time_arg_val(20);
+    constexpr uint32_t in1_B = get_compile_time_arg_val(21);
+    constexpr uint32_t in0_reuse_in_CB = get_compile_time_arg_val(22);
+
+    // sparsity args
+
+    constexpr uint32_t batchB = get_compile_time_arg_val(23);
+    constexpr uint32_t sparsity_pagesize = get_compile_time_arg_val(24);
+    // Boolean that is set when input A is sparse. If set, both input A and B are assumed to be sparse.
+    // Based on the sparsity tensor, the corresponding batch in input A and B are skipped.
+    constexpr bool bcast_A = (bool)get_compile_time_arg_val(25);
+    // This boolean is set when the number of batches is only known at runtime, typically based on a sparsity tensor.
+    constexpr bool get_batch_from_reader = (bool)get_compile_time_arg_val(26);
+
+    constexpr bool fuse_op = (bool)get_compile_time_arg_val(27);
+
+    constexpr auto in0_args = TensorAccessorArgs<28>();
+
+    constexpr auto sparsity_args = TensorAccessorArgs<in0_args.next_compile_time_args_offset()>();
+
+    // Number of valid (non-zero sparsity) batches the receiver and compute kernels are configured to
+    // process. When nnz is supplied (get_batch_from_reader == false), those kernels loop exactly
+    // num_batch_compute times, while this sender only multicasts once per non-zero sparsity entry, i.e.
+    // count_nonzero(sparsity) times. The op silently requires count_nonzero(sparsity) == num_batch_compute;
+    // if they disagree, the receivers wait on multicasts that never come (or the sender waits on receivers
+    // that already finished) and the device deadlocks. count_nonzero(sparsity) is data-dependent and only
+    // known here at runtime, so we validate it on-device by counting the multicasts we actually issue and
+    // asserting the contract holds -- surfacing a loud assert (under watcher) instead of a silent hang.
+    // See https://github.com/tenstorrent/tt-metal/issues/45943.
+    [[maybe_unused]] constexpr uint32_t num_batch_compute =
+        get_compile_time_arg_val(sparsity_args.next_compile_time_args_offset());
+
+    // 0 is used to specify "INVALID" state, i.e. when the multicasted data has not been received by the receiver.
+    // 0x1 is used to specify "VALID" state, i.e. when the batch is valid.
+    // 0x2 is used to specify "IGNORE_BATCH" state, i.e. when the batch is not valid.
+    constexpr uint32_t IGNORE_BATCH = 0x2;
+
+    // When sparsity is disabled, we just loop once
+    constexpr uint32_t batchB_lim = batchB == 0 ? 1u : batchB;
+
+    // Indexed/gather mode: iterate only the num_active selected sparse groups (the ids the caller
+    // passed in the `indices` operand). Every iterated batch is valid -- no sparsity scan, no skip,
+    // no validity multicast -- so this sender never has to look at the id list itself: A is either
+    // broadcast (bcast_A) or already compact and advanced sequentially per iteration (!bcast_A).
+    // Every factory that builds this kernel passes "num_active"; only the sparse matmul factory ever
+    // sets it non-zero. 0 means not indexed, i.e. the unchanged dense sparsity-scan path.
+    constexpr uint32_t num_active = get_named_compile_time_arg_val("num_active");
+    constexpr bool use_indices = num_active > 0;
+    constexpr uint32_t batch_loop_lim = use_indices ? num_active : batchB_lim;
+
+    MatmulOpReceiver fused_op_receiver;
+    if constexpr (fuse_op) {
+        fused_op_receiver = MatmulOpReceiver(
+            true, /* wait_for_op_signal */
+            rt_args_idx,
+            num_blocks_inner_dim,
+            in0_block_w /* tiles_per_block (in the same dimension as tensor slice) */
+        );
+    }
+
+    constexpr uint32_t dfb_id_in0 = get_named_compile_time_arg_val("cb_in0");
+    constexpr uint32_t in0_single_tile_size_bytes = get_tile_size(dfb_id_in0);
+    // Tiles whose size is not a multiple of the DRAM alignment are padded to it in DRAM, and the
+    // interleaved in0 CB pages are sized to match (see the program factory). The NOC reads the
+    // unpadded tile of data into each padded slot, and tiles are laid out / multicast at the padded
+    // stride. No-op when already aligned. The sharded path keeps the natural (unpadded) stride.
+    constexpr uint32_t in0_aligned_tile_size_bytes =
+        (in0_single_tile_size_bytes + (DRAM_ALIGNMENT - 1)) & ~(DRAM_ALIGNMENT - 1);
+#ifdef IN0_SHARDED
+    constexpr uint32_t in0_block_size_bytes = in0_block_num_tiles * in0_single_tile_size_bytes;
+#else
+    constexpr uint32_t in0_block_size_bytes = in0_block_num_tiles * in0_aligned_tile_size_bytes;
+#endif
+
+    Noc noc;
+    DataflowBuffer dfb_in0(dfb_id_in0);
+    Semaphore<> sender_sem(get_compile_time_arg_val(15));
+    Semaphore<> receiver_sem(get_compile_time_arg_val(16));
+
+#ifdef IN0_SHARDED
+    // In case we need to send multiple blocks per shard, in0 sharded cb is cb2 and we extract the sub-blocks to cb0
+    constexpr uint32_t shard_read_stride = shard_width_in_tiles * in0_single_tile_size_bytes;
+    constexpr uint32_t shard_read_width = in0_single_tile_size_bytes * in0_block_w;
+    constexpr uint32_t shard_num_tiles = shard_width_in_tiles * shard_height_in_tiles;
+    constexpr uint32_t in0_tensor_next_h_dim_block_stride_bytes =
+        in0_tensor_next_h_dim_block_stride * in0_single_tile_size_bytes;
+
+    uint32_t noc_shard_read_start_addr = 0;
+    if constexpr (extract_shard_sub_blocks) {
+        constexpr uint32_t dfb_id_in2 =
+            get_named_compile_time_arg_val("cb_in0_sharded");  // in0 sharded cb if extract_shard_sub_blocks
+        DataflowBuffer dfb_in2(dfb_id_in2);
+        noc_shard_read_start_addr = dfb_in2.get_read_ptr();
+    }
+
+#else
+    const auto s0 = TensorAccessor(in0_args, in0_tensor_addr);
+#endif  // IN0_SHARDED
+
+    // sparsity accessor
+    constexpr uint32_t dfb_id_sparsity = get_named_compile_time_arg_val("cb_sparsity");
+    DataflowBuffer dfb_sparsity(dfb_id_sparsity);
+    const auto s_sparsity = TensorAccessor(sparsity_args, sparsity_addr);
+
+#ifndef SKIP_MCAST
+    // Set ur local VALID value, to be mcasted to destinations flag address after the data has been mcasted
+    receiver_sem.set(VALID);
+    // local address that will be atomically incremented by mcast receivers, to know when all receivers are ready
+    // to receive the mcast
+
+#ifdef IN0_SHARDED
+    uint32_t in0_start_address = dfb_in0.get_write_ptr();
+#endif  // IN0_SHARDED
+#endif  // SKIP_MCAST
+
+    uint32_t l1_write_addr_sparsity = 0;
+    if constexpr (batchB > 0 && !use_indices) {
+        dfb_sparsity.reserve_back(1);
+        l1_write_addr_sparsity = dfb_sparsity.get_write_ptr();
+    }
+
+    // Counts the in0 multicasts actually issued (one per non-zero sparsity entry). Used to validate
+    // count_nonzero(sparsity) == num_batch_compute when nnz is supplied (see num_batch_compute above).
+    [[maybe_unused]] uint32_t num_valid_batches = 0;
+
+    for (uint32_t b = 0; b < in0_B; ++b) {
+        if constexpr (batchB > 0 && !use_indices) {
+            noc.async_read(s_sparsity, dfb_sparsity, sparsity_pagesize, {.page_id = b}, {.offset_bytes = 0});
+            noc.async_read_barrier();
+        }
+
+        for (uint32_t bB = 0; bB < batch_loop_lim; ++bB) {
+            if constexpr (batchB > 0 && !use_indices) {
+                volatile auto is_batch_valid =
+                    ((reinterpret_cast<volatile tt_l1_ptr uint16_t*>(l1_write_addr_sparsity))[bB]) != 0;
+
+                if constexpr (get_batch_from_reader) {
+#ifndef SKIP_MCAST
+                    // First broadcast this to other cores
+                    sender_sem.wait(in0_mcast_num_dests);
+                    sender_sem.set(0);
+                    receiver_sem.set(is_batch_valid ? VALID : IGNORE_BATCH);
+                    receiver_sem.set_multicast(
+                        noc,
+                        in0_mcast_dest_noc_start_x,
+                        in0_mcast_dest_noc_start_y,
+                        in0_mcast_dest_noc_end_x,
+                        in0_mcast_dest_noc_end_y,
+                        in0_mcast_num_cores);
+                    noc.async_writes_flushed();
+                    // Reset the semaphore value to VALID
+                    receiver_sem.set(VALID);
+#endif  // SKIP_MCAST
+
+                    // We need to pass the value to compute cores regardless of the value of is_batch_valid
+                    ckernel::mailbox_write(ckernel::ThreadId::UnpackThreadId, is_batch_valid);
+                    ckernel::mailbox_write(ckernel::ThreadId::MathThreadId, is_batch_valid);
+                    ckernel::mailbox_write(ckernel::ThreadId::PackThreadId, is_batch_valid);
+                }
+
+                if (!is_batch_valid) {
+                    if constexpr (!bcast_A) {
+                        in0_tensor_start_tile_id += MtKt;
+                    }
+                    continue;
+                }
+
+                // This is a valid (non-zero) batch that we are about to multicast. When nnz was supplied,
+                // catch count_nonzero(sparsity) > num_batch_compute here, before the sender blocks below
+                // waiting on receivers that have already finished their num_batch_compute iterations.
+                if constexpr (!get_batch_from_reader) {
+                    ++num_valid_batches;
+                    ASSERT(num_valid_batches <= num_batch_compute);
+                }
+            }
+
+#ifdef IN0_SHARDED
+            uint32_t in0_tensor_current_h_dim_block_start_addr = noc_shard_read_start_addr;
+#endif  // IN0_SHARDED
+            uint32_t in0_tensor_current_h_dim_block_tile_id = in0_tensor_start_tile_id;
+            for (uint32_t bh = 0; bh < num_blocks_h_dim; ++bh) {
+                for (uint32_t bw = 0; bw < num_blocks_w_dim; ++bw) {
+#ifdef IN0_SHARDED
+                    uint32_t in0_tensor_current_inner_dim_block_start_addr = in0_tensor_current_h_dim_block_start_addr;
+#endif  // IN0_SHARDED
+                    uint32_t in0_tensor_current_inner_dim_block_start_tile_id = in0_tensor_current_h_dim_block_tile_id;
+                    for (uint32_t block = 0; block < num_blocks_inner_dim; ++block) {
+                        if constexpr (fuse_op) {
+                            fused_op_receiver.update_current_block_start_tile_id(
+                                block, in0_tensor_current_inner_dim_block_start_tile_id, in0_tensor_start_tile_id);
+                        }
+
+                        // Operand 0
+                        // Common for sharded and interleaved paths
+                        dfb_in0.reserve_back(in0_block_num_tiles);
+#ifndef IN0_SHARDED
+
+                        uint32_t in0_write_offset = 0;
+
+#ifndef SKIP_MCAST
+                        uint32_t in0_start_address =
+                            dfb_in0.get_write_ptr();  // copy start address of block, to be used for mcasting
+#endif                                                // SKIP_MCAST
+
+                        // Copy in0 block into CB, as the default kernel
+                        uint32_t in0_tensor_row_start_tile_id = in0_tensor_current_inner_dim_block_start_tile_id;
+                        for (uint32_t h = 0; h < in0_block_h; ++h) {
+                            uint32_t in0_tensor_tile_id = in0_tensor_row_start_tile_id;
+                            for (uint32_t w = 0; w < in0_block_w; ++w) {
+                                if (bh < num_blocks_h_dim - 1 || h < last_block_h) {
+                                    noc.async_read(
+                                        s0,
+                                        dfb_in0,
+                                        in0_single_tile_size_bytes,
+                                        {.page_id = mast3r_in0_remap(in0_tensor_tile_id)},
+                                        {.offset_bytes = in0_write_offset});
+                                }
+
+                                // Zero out padded regions for the very last tile
+                                if constexpr (in0_last_ktile_w > 0) {
+                                    if ((block == num_blocks_inner_dim - 1) && (w == in0_block_w - 1)) {
+                                        noc.async_read_barrier();
+                                        constexpr DataFormat in0_data_format = get_dataformat(dfb_id_in0);
+                                        pad_last_ktile<in0_data_format, in0_last_ktile_w>(
+                                            dfb_in0.get_write_ptr() + in0_write_offset);
+                                    }
+                                }
+                                if constexpr (in0_last_ktile_h > 0) {
+                                    if ((block == num_blocks_inner_dim - 1) && (w == in0_block_w - 1)) {
+                                        noc.async_read_barrier();
+                                        constexpr DataFormat in0_data_format = get_dataformat(dfb_id_in0);
+                                        pad_last_transposed_ktile<in0_data_format, in0_last_ktile_h>(
+                                            dfb_in0.get_write_ptr() + in0_write_offset);
+                                    }
+                                }
+
+                                in0_write_offset += in0_aligned_tile_size_bytes;
+                                in0_tensor_tile_id += in0_tensor_stride_w;
+                            }
+                            in0_tensor_row_start_tile_id += in0_tensor_stride_h;
+                        }
+                        in0_tensor_current_inner_dim_block_start_tile_id += in0_tensor_next_inner_dim_block_stride;
+
+                        // Barrier! make sure the reads are done
+                        noc.async_read_barrier();
+#else
+                        if constexpr (extract_shard_sub_blocks) {
+                            uint32_t l1_write_addr_in0 = dfb_in0.get_write_ptr();
+
+#ifndef SKIP_MCAST
+                            in0_start_address =
+                                l1_write_addr_in0;  // copy start address of block, to be used for mcasting
+#endif  // SKIP_MCAST
+
+                            UnicastEndpoint self_ep;
+                            uint32_t noc_shard_read_l1_addr = in0_tensor_current_inner_dim_block_start_addr;
+
+                            for (uint32_t i = 0; i < in0_block_h; i++) {
+                                noc.async_read(
+                                    self_ep,
+                                    CoreLocalMem<uint32_t>(l1_write_addr_in0),
+                                    shard_read_width,
+                                    {.noc_x = my_x[0], .noc_y = my_y[0], .addr = noc_shard_read_l1_addr},
+                                    {});
+
+                                l1_write_addr_in0 += shard_read_width;
+                                noc_shard_read_l1_addr += shard_read_stride;
+                            }
+
+                            in0_tensor_current_inner_dim_block_start_addr += shard_read_width;
+                            noc.async_read_barrier();
+                        }
+
+                        {
+                            constexpr DataFormat in0_data_format = get_dataformat(dfb_id_in0);
+                            uint32_t in0_pad_base_addr = dfb_in0.get_write_ptr();
+                            if constexpr (in0_last_ktile_w > 0) {
+                                if ((block == num_blocks_inner_dim - 1)) {
+                                    for (uint32_t h = 0; h < in0_block_h; ++h) {
+                                        auto ptr = in0_pad_base_addr +
+                                                   (h * in0_block_w + in0_block_w - 1) * in0_single_tile_size_bytes;
+                                        pad_last_ktile<in0_data_format, in0_last_ktile_w>(ptr);
+                                    }
+                                }
+                            }
+                            if constexpr (in0_last_ktile_h > 0) {
+                                if ((block == num_blocks_inner_dim - 1)) {
+                                    for (uint32_t w = 0; w < in0_block_w; ++w) {
+                                        auto ptr = in0_pad_base_addr +
+                                                   ((in0_block_h - 1) * in0_block_w + w) * in0_single_tile_size_bytes;
+                                        pad_last_transposed_ktile<in0_data_format, in0_last_ktile_h>(ptr);
+                                    }
+                                }
+                            }
+                        }
+#endif  // IN0_SHARDED
+
+#ifndef SKIP_MCAST
+                        // wait until all in0 mcast destinations have atomically incremented the in0 semaphore_addr
+                        // (i.e. its value should be in0_mcast_num_dests), then reset the semaphore_addr value back to
+                        // zero for the next block
+                        sender_sem.wait(in0_mcast_num_dests);
+                        sender_sem.set(0);
+
+                        // Now we have the block in the CB address, we can mcast to dests!
+                        MulticastEndpoint mcast_dst;
+                        // num_dests must not include source, since we are NOT really doing a local copy!
+                        noc.async_write_multicast(
+                            CoreLocalMem<uint32_t>(in0_start_address),
+                            mcast_dst,
+                            in0_block_size_bytes,
+                            in0_mcast_num_cores,
+                            {},
+                            {.noc_x_start = in0_mcast_dest_noc_start_x,
+                             .noc_y_start = in0_mcast_dest_noc_start_y,
+                             .noc_x_end = in0_mcast_dest_noc_end_x,
+                             .noc_y_end = in0_mcast_dest_noc_end_y,
+                             .addr = in0_start_address},
+                            true);
+
+                        // Note: no need for write barrier, since these two multicasts are done on the same noc id, same
+                        // vc, same cmd_buf Also, this only works because we are setting VCs statically (using
+                        // NOC_CMD_STATIC_VC).
+#ifdef ARCH_BLACKHOLE
+                        // On Blackhole the flush is needed because NoC latency is higher than L1 <-> RISCV
+                        // latency which means data could be changed before write is issued.
+                        noc.async_writes_flushed();
+#endif  // ARCH_BLACKHOLE
+
+                        // We should also multicast the flag to destinations
+                        // num_dests must not include source, since we are NOT really doing a local copy!
+                        receiver_sem.set_multicast(
+                            noc,
+                            in0_mcast_dest_noc_start_x,
+                            in0_mcast_dest_noc_start_y,
+                            in0_mcast_dest_noc_end_x,
+                            in0_mcast_dest_noc_end_y,
+                            in0_mcast_num_cores);
+#endif  // SKIP_MCAST
+
+                        // Common for sharded and interleaved paths
+                        dfb_in0.push_back(in0_block_num_tiles);
+                    }
+                }
+#ifdef IN0_SHARDED
+                in0_tensor_current_h_dim_block_start_addr += in0_tensor_next_h_dim_block_stride_bytes;
+#endif  // IN0_SHARDED
+                in0_tensor_current_h_dim_block_tile_id += in0_tensor_next_h_dim_block_stride;
+            }
+
+            if constexpr (!bcast_A) {
+                in0_tensor_start_tile_id += MtKt;
+            }
+        }
+
+        if constexpr (bcast_A) {
+            in0_tensor_start_tile_id += MtKt;
+        }
+
+        // this is an optimization for the case when in0 is [1, 1, M, K] and in1 is [1, H, K, N], i.e. when in0_B ==
+        // 1 and in1_B > 1 in this case we originally had to replicate the in0 block for each batch, but with this
+        // optimization we just read the tensor slice once into each core's L1 and keep it there for all weight
+        // batches since the needed in0 data is already in L1 after batch 0, we can just move read pointer for this
+        // CB so compute kernel thinks it has new data
+        if (in0_reuse_in_CB) {
+            for (uint32_t fake_batch = 0; fake_batch < in1_B - in0_B; ++fake_batch) {
+                for (uint32_t blk = 0; blk < num_blocks_inner_dim; ++blk) {
+                    dfb_in0.reserve_back(in0_block_num_tiles);
+                    dfb_in0.push_back(in0_block_num_tiles);
+                }
+            }
+        }
+    }
+    noc.async_write_barrier();
+
+    // When nnz was supplied, the receiver and compute kernels loop exactly num_batch_compute times.
+    // If we issued fewer multicasts than that (count_nonzero(sparsity) < num_batch_compute), those
+    // kernels are now waiting on multicasts that will never arrive and the device would deadlock.
+    // Fail loudly instead. See https://github.com/tenstorrent/tt-metal/issues/45943.
+    // (Indexed/gather mode never counts: it multicasts exactly num_batch_compute == num_active times
+    // by construction, since every iterated group is active.)
+    if constexpr (!get_batch_from_reader && batchB > 0 && !use_indices) {
+        ASSERT(num_valid_batches == num_batch_compute);
+    }
+
+    // For completeness, we empty the sparsity CB if it was reserved earlier
+    if constexpr (batchB > 0 && !use_indices) {
+        dfb_sparsity.push_back(1);
+        dfb_sparsity.wait_front(1);
+        dfb_sparsity.pop_front(1);
+    }
+}

code/models/demos/mast3r/tt/kernels/phase_il_rm_reader.cpp

@@ -0,0 +1,45 @@
+// SPDX-License-Identifier: Apache-2.0
+// mast3r-p150 model-local kernel (ttnn.generic_op): polyphase interleave on ROW_MAJOR pixel rows, reader.
+// X[(2i+a)*wl + 2j+b, :] = Y_ab[i*w0 + j, :]; one page = one pixel row (C bf16 = ROW_BYTES). Unit = a block of
+// PB consecutive output pixels of one output row. Pure data movement: exact.
+#include <stdint.h>
+#include "api/dataflow/dataflow_api.h"
+
+void kernel_main() {
+    constexpr uint32_t W0 = get_compile_time_arg_val(0);       // source row width (pixels)
+    constexpr uint32_t PB = get_compile_time_arg_val(1);       // pixels per unit (divides 2*W0)
+    constexpr uint32_t LOG2_ROW = get_compile_time_arg_val(2); // log2(row bytes)
+    constexpr uint32_t ROW = 1u << LOG2_ROW;
+    constexpr uint32_t BPR = (2 * W0) / PB;                    // units per output pixel row
+    uint32_t y_addr[4];
+    for (uint32_t p = 0; p < 4; ++p) {
+        y_addr[p] = get_arg_val<uint32_t>(p);
+    }
+    const uint32_t u_start = get_arg_val<uint32_t>(4);
+    const uint32_t u_count = get_arg_val<uint32_t>(5);
+    const InterleavedPow2AddrGen<SRC_DRAM> g0 = {.bank_base_address = y_addr[0], .log_base_2_of_page_size = LOG2_ROW};
+    const InterleavedPow2AddrGen<SRC_DRAM> g1 = {.bank_base_address = y_addr[1], .log_base_2_of_page_size = LOG2_ROW};
+    const InterleavedPow2AddrGen<SRC_DRAM> g2 = {.bank_base_address = y_addr[2], .log_base_2_of_page_size = LOG2_ROW};
+    const InterleavedPow2AddrGen<SRC_DRAM> g3 = {.bank_base_address = y_addr[3], .log_base_2_of_page_size = LOG2_ROW};
+    for (uint32_t u = u_start; u < u_start + u_count; ++u) {
+        const uint32_t y = u / BPR;
+        const uint32_t x0 = (u - y * BPR) * PB;
+        const uint32_t i = y >> 1;
+        const bool a = (y & 1) != 0;
+        cb_reserve_back(0, 1);
+        uint32_t dst = get_write_ptr(0);
+        for (uint32_t x = x0; x < x0 + PB; ++x) {
+            const uint32_t page = i * W0 + (x >> 1);
+            uint64_t src;
+            if (!a) {
+                src = (x & 1) ? g1.get_noc_addr(page) : g0.get_noc_addr(page);
+            } else {
+                src = (x & 1) ? g3.get_noc_addr(page) : g2.get_noc_addr(page);
+            }
+            noc_async_read(src, dst, ROW);
+            dst += ROW;
+        }
+        noc_async_read_barrier();
+        cb_push_back(0, 1);
+    }
+}

code/models/demos/mast3r/tt/kernels/phase_il_rm_writer.cpp

@@ -0,0 +1,25 @@
+// SPDX-License-Identifier: Apache-2.0
+// mast3r-p150 model-local kernel (ttnn.generic_op): polyphase interleave on ROW_MAJOR pixel rows, writer.
+#include <stdint.h>
+#include "api/dataflow/dataflow_api.h"
+
+void kernel_main() {
+    constexpr uint32_t PB = get_compile_time_arg_val(0);
+    constexpr uint32_t LOG2_ROW = get_compile_time_arg_val(1);
+    constexpr uint32_t ROW = 1u << LOG2_ROW;
+    const uint32_t out_addr = get_arg_val<uint32_t>(0);
+    const uint32_t u_start = get_arg_val<uint32_t>(1);
+    const uint32_t u_count = get_arg_val<uint32_t>(2);
+    const InterleavedPow2AddrGen<OUT_DRAM> g = {.bank_base_address = out_addr, .log_base_2_of_page_size = LOG2_ROW};
+    for (uint32_t u = u_start; u < u_start + u_count; ++u) {
+        cb_wait_front(0, 1);
+        uint32_t src = get_read_ptr(0);
+        for (uint32_t p = u * PB; p < (u + 1) * PB; ++p) {
+            noc_async_write(src, g.get_noc_addr(p), ROW);
+            src += ROW;
+        }
+        noc_async_writes_flushed();
+        cb_pop_front(0, 1);
+    }
+    noc_async_write_barrier();
+}

code/models/demos/mast3r/tt/kernels/phase_il_rm_writer_hs.cpp

@@ -0,0 +1,29 @@
+// SPDX-License-Identifier: Apache-2.0
+// mast3r-p150 model-local kernel (ttnn.generic_op): polyphase interleave on ROW_MAJOR pixel rows, writer straight into a
+// HEIGHT_SHARDED L1 output (MAST3R_OPT "pilhs": the head.2 phase convs' own input shard spec, no interleaved->sharded
+// copy). A unit = PB consecutive pixel rows (one CB page); shards hold UPS units, so a unit never straddles two shards and
+// is written with one contiguous NoC write. Common runtime args: [out_addr, noc_x0, noc_y0, noc_x1, noc_y1, ...] per shard.
+#include <stdint.h>
+#include "api/dataflow/dataflow_api.h"
+
+void kernel_main() {
+    constexpr uint32_t PB = get_compile_time_arg_val(0);
+    constexpr uint32_t LOG2_ROW = get_compile_time_arg_val(1);
+    constexpr uint32_t UPS = get_compile_time_arg_val(2);  // units per shard
+    constexpr uint32_t ROW = 1u << LOG2_ROW;
+    constexpr uint32_t UNIT_BYTES = PB * ROW;
+    const uint32_t u_start = get_arg_val<uint32_t>(0);
+    const uint32_t u_count = get_arg_val<uint32_t>(1);
+    const uint32_t out_addr = get_common_arg_val<uint32_t>(0);
+    for (uint32_t u = u_start; u < u_start + u_count; ++u) {
+        const uint32_t shard = u / UPS;
+        const uint32_t off = (u - shard * UPS) * UNIT_BYTES;
+        const uint32_t nx = get_common_arg_val<uint32_t>(1 + 2 * shard);
+        const uint32_t ny = get_common_arg_val<uint32_t>(2 + 2 * shard);
+        cb_wait_front(0, 1);
+        noc_async_write(get_read_ptr(0), get_noc_addr(nx, ny, out_addr + off), UNIT_BYTES);
+        noc_async_writes_flushed();
+        cb_pop_front(0, 1);
+    }
+    noc_async_write_barrier();
+}

code/models/demos/mast3r/tt/kernels/strip_gather.cpp

@@ -0,0 +1,56 @@
+// SPDX-License-Identifier: Apache-2.0
+// mast3r-p150 model-local kernel (ttnn.generic_op), MAST3R_OPT "sgat": gathers the DPT ring-strip inputs straight from
+// refinenet1's HEIGHT_SHARDED TILE output xs [H0*W0, 32*CT] (pixel rows): tb [2, 3, W0, C] (image rows 0..2 and H0-3..H0-1)
+// and lr [2, 3, H0, C] (columns 0..2 and W0-3..W0-1, already in the transposed strip order), ROW_MAJOR pixel pages in
+// DRAM. Per output pixel: 2 face-row reads (16 channels = 32 B each) per channel tile from the owning shard. Pure data
+// movement (exact). Replaces S2I to DRAM + untilize + 4 slices + 2 concats + permute.
+#include <stdint.h>
+#include "api/dataflow/dataflow_api.h"
+
+void kernel_main() {
+    constexpr uint32_t CT = get_compile_time_arg_val(0);    // channel tiles
+    constexpr uint32_t PER = get_compile_time_arg_val(1);   // tile rows per shard
+    constexpr uint32_t W0 = get_compile_time_arg_val(2);
+    constexpr uint32_t H0 = get_compile_time_arg_val(3);
+    constexpr uint32_t PAGE = CT * 64;                      // one pixel, all channels (bf16)
+    const uint32_t p_start = get_arg_val<uint32_t>(0);
+    const uint32_t p_count = get_arg_val<uint32_t>(1);
+    if (p_count == 0) {
+        return;
+    }
+    const uint32_t xs_addr = get_common_arg_val<uint32_t>(0);
+    const InterleavedAddrGen<true> gtb = {.bank_base_address = get_common_arg_val<uint32_t>(1), .page_size = PAGE};
+    const InterleavedAddrGen<true> glr = {.bank_base_address = get_common_arg_val<uint32_t>(2), .page_size = PAGE};
+    const uint32_t buf = get_write_ptr(0);
+    uint32_t l1 = buf;
+    for (uint32_t q = p_start; q < p_start + p_count; ++q) {
+        uint32_t pix;
+        if (q < 6 * W0) {
+            const uint32_t sj = q / W0, x = q - sj * W0, s = sj / 3, j = sj - s * 3;
+            pix = (s == 0 ? j : H0 - 3 + j) * W0 + x;
+        } else {
+            const uint32_t q2 = q - 6 * W0;
+            const uint32_t sj = q2 / H0, i = q2 - sj * H0, s = sj / 3, j = sj - s * 3;
+            pix = i * W0 + (s == 0 ? j : W0 - 3 + j);
+        }
+        const uint32_t R = pix >> 5, r = pix & 31;
+        const uint32_t shard = R / PER, lt = R - shard * PER;
+        const uint32_t nx = get_common_arg_val<uint32_t>(3 + 2 * shard);
+        const uint32_t ny = get_common_arg_val<uint32_t>(4 + 2 * shard);
+        const uint32_t foff = (r >= 16 ? 1024u : 0u) + (r & 15) * 32;
+        for (uint32_t c = 0; c < CT; ++c) {
+            const uint64_t src = get_noc_addr(nx, ny, xs_addr + (lt * CT + c) * 2048 + foff);
+            noc_async_read(src, l1 + c * 64, 32);
+            noc_async_read(src + 512, l1 + c * 64 + 32, 32);
+        }
+        l1 += PAGE;
+    }
+    noc_async_read_barrier();
+    l1 = buf;
+    for (uint32_t q = p_start; q < p_start + p_count; ++q) {
+        const uint64_t dst = q < 6 * W0 ? gtb.get_noc_addr(q) : glr.get_noc_addr(q - 6 * W0);
+        noc_async_write(l1, dst, PAGE);
+        l1 += PAGE;
+    }
+    noc_async_write_barrier();
+}

code/models/demos/mast3r/tt/kernels/tail_il_reader.cpp

@@ -0,0 +1,30 @@
+// SPDX-License-Identifier: Apache-2.0
+// mast3r-p150 model-local kernel (ttnn.generic_op): head-tail polyphase interleave, reader.
+// T [16, hl*wl] ROW_MAJOR (row (2a+b)*4 + c = phase (a,b) channel c, column i*wl + j) -> page (c, i) of the output
+// [4*hl, 2*W] ROW_MAJOR = hi-res rows 2i and 2i+1 of channel c. For page (c, i) this reads the 4 source segments
+// T[(2a+b)*4 + c][i*wl : (i+1)*wl] (a, b in {0,1}; wl*2 bytes each) into one CB slot, order (a, b).
+#include <stdint.h>
+#include "api/dataflow/dataflow_api.h"
+
+void kernel_main() {
+    constexpr uint32_t HL = get_compile_time_arg_val(0);
+    constexpr uint32_t WL = get_compile_time_arg_val(1);
+    constexpr uint32_t LOG2_TROW = get_compile_time_arg_val(2);   // log2(T row bytes) = log2(hl*wl*2)
+    constexpr uint32_t SEG = WL * 2;
+    const uint32_t t_addr = get_arg_val<uint32_t>(0);
+    const uint32_t u_start = get_arg_val<uint32_t>(1);
+    const uint32_t u_count = get_arg_val<uint32_t>(2);
+    const InterleavedPow2AddrGen<SRC_DRAM> g = {.bank_base_address = t_addr, .log_base_2_of_page_size = LOG2_TROW};
+    for (uint32_t u = u_start; u < u_start + u_count; ++u) {
+        const uint32_t c = u / HL;
+        const uint32_t i = u - c * HL;
+        cb_reserve_back(0, 1);
+        uint32_t dst = get_write_ptr(0);
+        for (uint32_t p = 0; p < 4; ++p) {
+            noc_async_read(g.get_noc_addr(p * 4 + c) + i * SEG, dst, SEG);
+            dst += SEG;
+        }
+        noc_async_read_barrier();
+        cb_push_back(0, 1);
+    }
+}

code/models/demos/mast3r/tt/kernels/tail_il_writer.cpp

@@ -0,0 +1,42 @@
+// SPDX-License-Identifier: Apache-2.0
+// mast3r-p150 model-local kernel (ttnn.generic_op): head-tail polyphase interleave, writer.
+// out[c][2i+a][2j+b] = T[(2a+b)*4 + c][i*wl + j]: 16-bit element interleave of the (a,0) / (a,1) segments, then one
+// page write (2 hi-res rows). Pure data movement: exact.
+#include <stdint.h>
+#include "api/dataflow/dataflow_api.h"
+
+void kernel_main() {
+    constexpr uint32_t WL = get_compile_time_arg_val(0);
+    constexpr uint32_t LOG2_PAGE = get_compile_time_arg_val(1);   // log2(out page bytes) = log2(4*wl*2)
+    constexpr uint32_t SEG = WL * 2;
+    const uint32_t out_addr = get_arg_val<uint32_t>(0);
+    const uint32_t u_start = get_arg_val<uint32_t>(1);
+    const uint32_t u_count = get_arg_val<uint32_t>(2);
+    const InterleavedPow2AddrGen<OUT_DRAM> g = {.bank_base_address = out_addr, .log_base_2_of_page_size = LOG2_PAGE};
+    uint32_t slot = 0;
+    for (uint32_t u = u_start; u < u_start + u_count; ++u) {
+        cb_wait_front(0, 1);
+        const uint32_t src = get_read_ptr(0);
+        // two output staging slots (CB 1 holds 2 pages): wait for the write issued two pages ago before reuse
+        if (u - u_start >= 2) {
+            noc_async_writes_flushed();
+        }
+        const uint32_t dst = get_write_ptr(1) + slot * (1u << LOG2_PAGE);
+        for (uint32_t a = 0; a < 2; ++a) {
+            const uint32_t* s0 = reinterpret_cast<const uint32_t*>(src + (2 * a) * SEG);
+            const uint32_t* s1 = reinterpret_cast<const uint32_t*>(src + (2 * a + 1) * SEG);
+            uint32_t* d = reinterpret_cast<uint32_t*>(dst + a * 2 * SEG);
+#pragma GCC unroll 4
+            for (uint32_t k = 0; k < WL / 2; ++k) {
+                const uint32_t w0 = s0[k];
+                const uint32_t w1 = s1[k];
+                d[2 * k] = (w0 & 0xFFFFu) | (w1 << 16);
+                d[2 * k + 1] = (w0 >> 16) | (w1 & 0xFFFF0000u);
+            }
+        }
+        cb_pop_front(0, 1);
+        noc_async_write(dst, g.get_noc_addr(u), 1u << LOG2_PAGE);
+        slot ^= 1;
+    }
+    noc_async_write_barrier();
+}

code/models/demos/mast3r/tt/kernels/tailf_compute.cpp

@@ -0,0 +1,57 @@
+// SPDX-License-Identifier: Apache-2.0
+// mast3r-p150 model-local kernel (ttnn.generic_op): fused head tail, compute. Per tile: ((((z0 + B) + z1) + z2) + z3),
+// FPU adds with ttnn.add's CB formats (bf16, 16-bit dest), then a WH transpose (exact) so each phase-channel is one tile
+// row. The z_p are nonzero only in their own 4 columns, so every used column gets exactly one rounding, z_p + bias --
+// the value of ttnn's add chain + transpose + bias add.
+#include <cstdint>
+#include "api/compute/common.h"
+#include "api/compute/eltwise_binary.h"
+#include "api/compute/transpose.h"
+#include "api/compute/compute_kernel_hw_startup.h"
+#include "api/compute/cb_api.h"
+#include "api/compute/pack.h"
+
+ALWI void add_pass(uint32_t ca, uint32_t cb, uint32_t co, bool pop_b) {
+    cb_wait_front(ca, 1);
+    cb_wait_front(cb, 1);
+    cb_reserve_back(co, 1);
+    add_init(ca, cb);
+    tile_regs_acquire();
+    add_tiles(ca, cb, 0, 0, 0);
+    tile_regs_commit();
+    tile_regs_wait();
+    pack_tile(0, co);
+    tile_regs_release();
+    cb_push_back(co, 1);
+    cb_pop_front(ca, 1);
+    if (pop_b) {
+        cb_pop_front(cb, 1);
+    }
+}
+
+void kernel_main() {
+    const uint32_t n_tiles = get_arg_val<uint32_t>(0);
+    if (n_tiles == 0) {
+        return;
+    }
+    compute_kernel_hw_startup(0, 4, 24);
+    cb_wait_front(4, 1);
+    for (uint32_t u = 0; u < n_tiles; ++u) {
+        add_pass(0, 4, 24, false);     // z0 + B (bias CB 4 stays resident)
+        add_pass(24, 1, 25, true);
+        add_pass(25, 2, 26, true);
+        add_pass(26, 3, 27, true);
+        cb_wait_front(27, 1);
+        cb_reserve_back(16, 1);
+        transpose_init(27);
+        tile_regs_acquire();
+        transpose_tile(27, 0, 0);
+        tile_regs_commit();
+        tile_regs_wait();
+        pack_tile(0, 16);
+        tile_regs_release();
+        cb_push_back(16, 1);
+        cb_pop_front(27, 1);
+    }
+    cb_pop_front(4, 1);
+}

code/models/demos/mast3r/tt/kernels/tailf_reader.cpp

@@ -0,0 +1,41 @@
+// SPDX-License-Identifier: Apache-2.0
+// mast3r-p150 model-local kernel (ttnn.generic_op): fused head tail, reader. Unit = half a low-res image row (TPH tiles
+// of 32 pixels). Reads tile u of the 4 per-phase 1x1 outputs z_p ([hl*wl, 32] TILE, interleaved) into CB 0..3, and the
+// row-replicated bias tile once into CB 4.
+#include <stdint.h>
+#include "api/dataflow/dataflow_api.h"
+
+void kernel_main() {
+    constexpr uint32_t TPH = get_compile_time_arg_val(0);
+    constexpr uint32_t LOG2_TILE = 11;
+    constexpr uint32_t TB = 1u << LOG2_TILE;
+    const uint32_t u_start = get_arg_val<uint32_t>(0);
+    const uint32_t u_count = get_arg_val<uint32_t>(1);
+    if (u_count == 0) {
+        return;
+    }
+    const InterleavedPow2AddrGen<Z_DRAM> g0 = {.bank_base_address = get_common_arg_val<uint32_t>(0), .log_base_2_of_page_size = LOG2_TILE};
+    const InterleavedPow2AddrGen<Z_DRAM> g1 = {.bank_base_address = get_common_arg_val<uint32_t>(1), .log_base_2_of_page_size = LOG2_TILE};
+    const InterleavedPow2AddrGen<Z_DRAM> g2 = {.bank_base_address = get_common_arg_val<uint32_t>(2), .log_base_2_of_page_size = LOG2_TILE};
+    const InterleavedPow2AddrGen<Z_DRAM> g3 = {.bank_base_address = get_common_arg_val<uint32_t>(3), .log_base_2_of_page_size = LOG2_TILE};
+    const InterleavedPow2AddrGen<B_DRAM> gb = {.bank_base_address = get_common_arg_val<uint32_t>(4), .log_base_2_of_page_size = LOG2_TILE};
+    cb_reserve_back(4, 1);
+    noc_async_read(gb.get_noc_addr(0), get_write_ptr(4), TB);
+    noc_async_read_barrier();
+    cb_push_back(4, 1);
+    for (uint32_t u = u_start * TPH; u < (u_start + u_count) * TPH; ++u) {
+        cb_reserve_back(0, 1);
+        cb_reserve_back(1, 1);
+        cb_reserve_back(2, 1);
+        cb_reserve_back(3, 1);
+        noc_async_read(g0.get_noc_addr(u), get_write_ptr(0), TB);
+        noc_async_read(g1.get_noc_addr(u), get_write_ptr(1), TB);
+        noc_async_read(g2.get_noc_addr(u), get_write_ptr(2), TB);
+        noc_async_read(g3.get_noc_addr(u), get_write_ptr(3), TB);
+        noc_async_read_barrier();
+        cb_push_back(0, 1);
+        cb_push_back(1, 1);
+        cb_push_back(2, 1);
+        cb_push_back(3, 1);
+    }
+}

code/models/demos/mast3r/tt/kernels/tailf_writer.cpp

@@ -0,0 +1,67 @@
+// SPDX-License-Identifier: Apache-2.0
+// mast3r-p150 model-local kernel (ttnn.generic_op): fused head tail, writer. Unit = half of a low-res image row i
+// (TPH tiles of 32 pixels). The compute kernel hands over the WH-transposed tile: row ch = (2a+b)*4 + c (phase (a,b),
+// channel c), column = pixel jj. out[c][2i+a][2(j0+jj)+b] = T[(2a+b)*4+c][jj]: 16-bit interleave of rows (a,0,c) and
+// (a,1,c) into the staged segments, then for each (c, a) one write of TPH*64 elements into output page (c, i) = hi-res
+// rows 2i, 2i+1 of channel c. Pure data movement: exact.
+#include <stdint.h>
+#include "api/dataflow/dataflow_api.h"
+
+void kernel_main() {
+    constexpr uint32_t TPH = get_compile_time_arg_val(0);
+    constexpr uint32_t WL = get_compile_time_arg_val(1);
+    constexpr uint32_t HL = get_compile_time_arg_val(2);
+    constexpr uint32_t LOG2_PAGE = get_compile_time_arg_val(3);   // log2(4*wl*2)
+    constexpr uint32_t SEG = TPH * 64 * 2;                         // bytes of one (c, a) segment
+    constexpr uint32_t SLOT = 8 * SEG;
+    const uint32_t u_start = get_arg_val<uint32_t>(0);
+    const uint32_t u_count = get_arg_val<uint32_t>(1);
+    const uint32_t out_addr = get_arg_val<uint32_t>(2);
+    if (u_count == 0) {
+        return;
+    }
+    const InterleavedPow2AddrGen<OUT_DRAM> g = {.bank_base_address = out_addr, .log_base_2_of_page_size = LOG2_PAGE};
+    uint32_t slot = 0;
+    for (uint32_t hu = u_start; hu < u_start + u_count; ++hu) {
+        const uint32_t i = hu >> 1;
+        const uint32_t h = hu & 1;
+        if (hu - u_start >= 1) {
+            noc_async_writes_flushed();    // 2 staging slots: the writes of unit hu-2 are done
+        }
+        const uint32_t stage = get_write_ptr(8) + slot * SLOT;   // CB 8: staging only (never pushed)
+        for (uint32_t t = 0; t < TPH; ++t) {
+            cb_wait_front(16, 1);
+            // transposed tile: row ch < 16 = face 0 row ch (pixels 0..15) + face 1 row ch (pixels 16..31)
+            const uint32_t* tsrc = reinterpret_cast<const uint32_t*>(get_read_ptr(16));
+            for (uint32_t a = 0; a < 2; ++a) {
+                for (uint32_t c = 0; c < 4; ++c) {
+                    const uint32_t ch0 = (2 * a) * 4 + c;
+                    const uint32_t ch1 = (2 * a + 1) * 4 + c;
+                    uint32_t* d = reinterpret_cast<uint32_t*>(stage + (c * 2 + a) * SEG + t * 64 * 2);
+#pragma GCC unroll 2
+                    for (uint32_t f = 0; f < 2; ++f) {
+                        const uint32_t* s0 = tsrc + f * 128 + ch0 * 8;   // 16-element face row = 8 words
+                        const uint32_t* s1 = tsrc + f * 128 + ch1 * 8;
+                        uint32_t* dd = d + f * 16;
+#pragma GCC unroll 8
+                        for (uint32_t k = 0; k < 8; ++k) {
+                            const uint32_t w0 = s0[k];
+                            const uint32_t w1 = s1[k];
+                            dd[2 * k] = (w0 & 0xFFFFu) | (w1 << 16);
+                            dd[2 * k + 1] = (w0 >> 16) | (w1 & 0xFFFF0000u);
+                        }
+                    }
+                }
+            }
+            cb_pop_front(16, 1);
+        }
+        for (uint32_t c = 0; c < 4; ++c) {
+            const uint64_t pa = g.get_noc_addr(c * HL + i);
+            for (uint32_t a = 0; a < 2; ++a) {
+                noc_async_write(stage + (c * 2 + a) * SEG, pa + (a * 2 * WL + h * TPH * 64) * 2, SEG);
+            }
+        }
+        slot ^= 1;
+    }
+    noc_async_write_barrier();
+}

code/models/demos/mast3r/tt/kernels/ups2_compute.cpp

@@ -0,0 +1,66 @@
+// SPDX-License-Identifier: Apache-2.0
+// mast3r-p150 model-local kernel (ttnn.generic_op), MAST3R_OPT "tups": bilinear x2 upsample in TILE layout as exact
+// tile matmuls. Output tile (Yr, t', c) = sum over the 2 source rows (vertical weight w in {0.25, 0.75}, folded into the
+// constant tile) and the <= 2 source x-tiles of A_kind (32 out px x 32 in px, entries in {0.25, 0.75, 1} * w) x X_tile.
+// in0 = A (SrcB, <= 4 significant bits), in1 = X (SrcA); HiFi4 + fp32 dest: every product and the <= 4-term sum are
+// exact in fp32, one bf16 round at the pack. Constant tiles in CB 1: index 6 * wi + kind, wi 0 -> 0.25, 1 -> 0.75;
+// kind 0 E_EDGE, 1 E_L, 2 E_R, 3 O_L, 4 O_R, 5 O_EDGE.
+#include <cstdint>
+#include "api/compute/common.h"
+#include "api/compute/matmul.h"
+#include "api/compute/compute_kernel_hw_startup.h"
+#include "api/compute/cb_api.h"
+#include "api/compute/pack.h"
+
+constexpr uint32_t CB_X = 0, CB_A = 1, CB_O = 16;
+constexpr uint32_t E_EDGE = 0, E_L = 1, E_R = 2, O_L = 3, O_R = 4, O_EDGE = 5;
+
+void kernel_main() {
+    constexpr uint32_t NX = get_compile_time_arg_val(0);
+    constexpr uint32_t CT = get_compile_time_arg_val(1);
+    const uint32_t u_start = get_arg_val<uint32_t>(0);
+    const uint32_t u_count = get_arg_val<uint32_t>(1);
+    if (u_count == 0) {
+        return;
+    }
+    compute_kernel_hw_startup<SrcOrder::Reverse>(CB_A, CB_X, CB_O);
+    matmul_init(CB_A, CB_X);
+    cb_wait_front(CB_A, 12);
+    for (uint32_t u = u_start; u < u_start + u_count; ++u) {
+        const uint32_t Yr = u / CT;
+        const uint32_t w0 = (Yr & 1) == 0 ? 0 : 6;   // row 0 weight: 0.25 (even Yr) / 0.75 (odd Yr)
+        const uint32_t w1 = (Yr & 1) == 0 ? 6 : 0;
+        cb_wait_front(CB_X, 2 * NX);
+        for (uint32_t tp = 0; tp < 2 * NX; ++tp) {
+            const uint32_t m = tp >> 1;
+            cb_reserve_back(CB_O, 1);
+            tile_regs_acquire();
+            for (uint32_t r = 0; r < 2; ++r) {
+                const uint32_t wb = r == 0 ? w0 : w1;
+                const uint32_t ro = r * NX;
+                if ((tp & 1) == 0) {
+                    if (m == 0) {
+                        matmul_tiles(CB_A, CB_X, wb + E_EDGE, ro, 0);
+                    } else {
+                        matmul_tiles(CB_A, CB_X, wb + E_L, ro + m - 1, 0);
+                        matmul_tiles(CB_A, CB_X, wb + E_R, ro + m, 0);
+                    }
+                } else {
+                    if (m == NX - 1) {
+                        matmul_tiles(CB_A, CB_X, wb + O_EDGE, ro + m, 0);
+                    } else {
+                        matmul_tiles(CB_A, CB_X, wb + O_L, ro + m, 0);
+                        matmul_tiles(CB_A, CB_X, wb + O_R, ro + m + 1, 0);
+                    }
+                }
+            }
+            tile_regs_commit();
+            tile_regs_wait();
+            pack_tile(0, CB_O);
+            tile_regs_release();
+            cb_push_back(CB_O, 1);
+        }
+        cb_pop_front(CB_X, 2 * NX);
+    }
+    cb_pop_front(CB_A, 12);
+}

code/models/demos/mast3r/tt/kernels/ups2_reader.cpp

@@ -0,0 +1,56 @@
+// SPDX-License-Identifier: Apache-2.0
+// mast3r-p150 model-local kernel (ttnn.generic_op), MAST3R_OPT "tups": bilinear x2 upsample (half-pixel, clamped, ==
+// ttnn.upsample) of a TILE [H*W, C] pixel-row tensor, reader. Work unit u = (output image row Yr, channel tile c),
+// u = Yr * CT + c. Per unit: the 2 * NX input tiles of the two source image rows (clamped) at channel tile c, in the
+// order [row0 x-tiles..., row1 x-tiles...] into CB 0. Once per core: the 12 constant interpolation tiles into CB 1.
+#include <stdint.h>
+#include "api/dataflow/dataflow_api.h"
+
+void kernel_main() {
+    constexpr uint32_t NX = get_compile_time_arg_val(0);   // input x-tiles per image row (W / 32)
+    constexpr uint32_t CT = get_compile_time_arg_val(1);   // channel tiles
+    constexpr uint32_t H = get_compile_time_arg_val(2);    // input image rows
+    constexpr uint32_t NA = 12;
+    constexpr uint32_t LOG2_TILE = 11;
+    constexpr uint32_t TB = 1u << LOG2_TILE;
+    const uint32_t u_start = get_arg_val<uint32_t>(0);
+    const uint32_t u_count = get_arg_val<uint32_t>(1);
+    if (u_count == 0) {
+        return;
+    }
+    const InterleavedPow2AddrGen<IN_DRAM> gx = {.bank_base_address = get_common_arg_val<uint32_t>(0), .log_base_2_of_page_size = LOG2_TILE};
+    const InterleavedPow2AddrGen<A_DRAM> ga = {.bank_base_address = get_common_arg_val<uint32_t>(1), .log_base_2_of_page_size = LOG2_TILE};
+    cb_reserve_back(1, NA);
+    uint32_t pa = get_write_ptr(1);
+    for (uint32_t i = 0; i < NA; ++i) {
+        noc_async_read(ga.get_noc_addr(i), pa, TB);
+        pa += TB;
+    }
+    noc_async_read_barrier();
+    cb_push_back(1, NA);
+    for (uint32_t u = u_start; u < u_start + u_count; ++u) {
+        const uint32_t Yr = u / CT;
+        const uint32_t c = u - Yr * CT;
+        const uint32_t y = Yr >> 1;
+        uint32_t y0, y1;
+        if ((Yr & 1) == 0) {
+            y0 = y > 0 ? y - 1 : 0;
+            y1 = y;
+        } else {
+            y0 = y;
+            y1 = y + 1 < H ? y + 1 : H - 1;
+        }
+        cb_reserve_back(0, 2 * NX);
+        uint32_t p = get_write_ptr(0);
+        for (uint32_t xt = 0; xt < NX; ++xt) {
+            noc_async_read(gx.get_noc_addr((y0 * NX + xt) * CT + c), p, TB);
+            p += TB;
+        }
+        for (uint32_t xt = 0; xt < NX; ++xt) {
+            noc_async_read(gx.get_noc_addr((y1 * NX + xt) * CT + c), p, TB);
+            p += TB;
+        }
+        noc_async_read_barrier();
+        cb_push_back(0, 2 * NX);
+    }
+}

code/models/demos/mast3r/tt/kernels/ups2_writer.cpp

@@ -0,0 +1,26 @@
+// SPDX-License-Identifier: Apache-2.0
+// mast3r-p150 model-local kernel (ttnn.generic_op), MAST3R_OPT "tups": writer. Per unit u = (Yr, c): the 2 * NX output
+// tiles of output image row Yr at channel tile c (output pixel-row tile R = Yr * 2 * NX + t').
+#include <stdint.h>
+#include "api/dataflow/dataflow_api.h"
+
+void kernel_main() {
+    constexpr uint32_t NX = get_compile_time_arg_val(0);
+    constexpr uint32_t CT = get_compile_time_arg_val(1);
+    constexpr uint32_t LOG2_TILE = 11;
+    constexpr uint32_t TB = 1u << LOG2_TILE;
+    const uint32_t u_start = get_arg_val<uint32_t>(0);
+    const uint32_t u_count = get_arg_val<uint32_t>(1);
+    const InterleavedPow2AddrGen<OUT_DRAM> g = {.bank_base_address = get_common_arg_val<uint32_t>(0), .log_base_2_of_page_size = LOG2_TILE};
+    for (uint32_t u = u_start; u < u_start + u_count; ++u) {
+        const uint32_t Yr = u / CT;
+        const uint32_t c = u - Yr * CT;
+        for (uint32_t tp = 0; tp < 2 * NX; ++tp) {
+            cb_wait_front(16, 1);
+            noc_async_write(get_read_ptr(16), g.get_noc_addr((Yr * 2 * NX + tp) * CT + c), TB);
+            noc_async_writes_flushed();
+            cb_pop_front(16, 1);
+        }
+    }
+    noc_async_write_barrier();
+}

code/models/demos/mast3r/tt/kernels/ups2h_compute.cpp

@@ -0,0 +1,53 @@
+// SPDX-License-Identifier: Apache-2.0
+// mast3r-p150 model-local kernel (ttnn.generic_op), MAST3R_OPT "tups" for W = 16: output tile (Yr, c) = A[w0, h0] x
+// X(y0) + A[w1, h1] x X(y1), where h = y & 1 selects the half of the input tile that holds image row y and w the vertical
+// weight (0.25 / 0.75, folded into the constant tile). Constant tile index 2 * wi + h. HiFi4 + fp32 dest (exact).
+#include <cstdint>
+#include "api/compute/common.h"
+#include "api/compute/matmul.h"
+#include "api/compute/compute_kernel_hw_startup.h"
+#include "api/compute/cb_api.h"
+#include "api/compute/pack.h"
+
+constexpr uint32_t CB_X = 0, CB_A = 1, CB_O = 16;
+
+void kernel_main() {
+    constexpr uint32_t CT = get_compile_time_arg_val(0);
+    constexpr uint32_t H = get_compile_time_arg_val(1);
+    const uint32_t u_start = get_arg_val<uint32_t>(0);
+    const uint32_t u_count = get_arg_val<uint32_t>(1);
+    if (u_count == 0) {
+        return;
+    }
+    compute_kernel_hw_startup<SrcOrder::Reverse>(CB_A, CB_X, CB_O);
+    matmul_init(CB_A, CB_X);
+    cb_wait_front(CB_A, 4);
+    for (uint32_t u = u_start; u < u_start + u_count; ++u) {
+        const uint32_t Yr = u / CT;
+        const uint32_t y = Yr >> 1;
+        uint32_t y0, y1, wa, wb;
+        if ((Yr & 1) == 0) {
+            y0 = y > 0 ? y - 1 : 0;
+            y1 = y;
+            wa = 0;   // 0.25
+            wb = 2;   // 0.75
+        } else {
+            y0 = y;
+            y1 = y + 1 < H ? y + 1 : H - 1;
+            wa = 2;
+            wb = 0;
+        }
+        cb_wait_front(CB_X, 2);
+        cb_reserve_back(CB_O, 1);
+        tile_regs_acquire();
+        matmul_tiles(CB_A, CB_X, wa + (y0 & 1), 0, 0);
+        matmul_tiles(CB_A, CB_X, wb + (y1 & 1), 1, 0);
+        tile_regs_commit();
+        tile_regs_wait();
+        pack_tile(0, CB_O);
+        tile_regs_release();
+        cb_push_back(CB_O, 1);
+        cb_pop_front(CB_X, 2);
+    }
+    cb_pop_front(CB_A, 4);
+}

code/models/demos/mast3r/tt/kernels/ups2h_reader.cpp

@@ -0,0 +1,48 @@
+// SPDX-License-Identifier: Apache-2.0
+// mast3r-p150 model-local kernel (ttnn.generic_op), MAST3R_OPT "tups" for W = 16 (an input tile row holds 2 image rows):
+// reader. Unit u = (output image row Yr, channel tile c) = Yr * CT + c: the input tiles holding the two (clamped) source
+// image rows y0, y1 at channel tile c into CB 0 (2 tiles); once per core the 4 constant tiles into CB 1.
+#include <stdint.h>
+#include "api/dataflow/dataflow_api.h"
+
+void kernel_main() {
+    constexpr uint32_t CT = get_compile_time_arg_val(0);
+    constexpr uint32_t H = get_compile_time_arg_val(1);
+    constexpr uint32_t NA = 4;
+    constexpr uint32_t LOG2_TILE = 11;
+    constexpr uint32_t TB = 1u << LOG2_TILE;
+    const uint32_t u_start = get_arg_val<uint32_t>(0);
+    const uint32_t u_count = get_arg_val<uint32_t>(1);
+    if (u_count == 0) {
+        return;
+    }
+    const InterleavedPow2AddrGen<IN_DRAM> gx = {.bank_base_address = get_common_arg_val<uint32_t>(0), .log_base_2_of_page_size = LOG2_TILE};
+    const InterleavedPow2AddrGen<A_DRAM> ga = {.bank_base_address = get_common_arg_val<uint32_t>(1), .log_base_2_of_page_size = LOG2_TILE};
+    cb_reserve_back(1, NA);
+    uint32_t pa = get_write_ptr(1);
+    for (uint32_t i = 0; i < NA; ++i) {
+        noc_async_read(ga.get_noc_addr(i), pa, TB);
+        pa += TB;
+    }
+    noc_async_read_barrier();
+    cb_push_back(1, NA);
+    for (uint32_t u = u_start; u < u_start + u_count; ++u) {
+        const uint32_t Yr = u / CT;
+        const uint32_t c = u - Yr * CT;
+        const uint32_t y = Yr >> 1;
+        uint32_t y0, y1;
+        if ((Yr & 1) == 0) {
+            y0 = y > 0 ? y - 1 : 0;
+            y1 = y;
+        } else {
+            y0 = y;
+            y1 = y + 1 < H ? y + 1 : H - 1;
+        }
+        cb_reserve_back(0, 2);
+        const uint32_t p = get_write_ptr(0);
+        noc_async_read(gx.get_noc_addr((y0 >> 1) * CT + c), p, TB);
+        noc_async_read(gx.get_noc_addr((y1 >> 1) * CT + c), p + TB, TB);
+        noc_async_read_barrier();
+        cb_push_back(0, 2);
+    }
+}

code/models/demos/mast3r/tt/kernels/ups2h_writer.cpp

@@ -0,0 +1,20 @@
+// SPDX-License-Identifier: Apache-2.0
+// mast3r-p150 model-local kernel (ttnn.generic_op), MAST3R_OPT "tups" for W = 16: writer (one output tile per unit, the
+// 32-pixel output image row Yr at channel tile c = tile Yr * CT + c).
+#include <stdint.h>
+#include "api/dataflow/dataflow_api.h"
+
+void kernel_main() {
+    constexpr uint32_t LOG2_TILE = 11;
+    constexpr uint32_t TB = 1u << LOG2_TILE;
+    const uint32_t u_start = get_arg_val<uint32_t>(0);
+    const uint32_t u_count = get_arg_val<uint32_t>(1);
+    const InterleavedPow2AddrGen<OUT_DRAM> g = {.bank_base_address = get_common_arg_val<uint32_t>(0), .log_base_2_of_page_size = LOG2_TILE};
+    for (uint32_t u = u_start; u < u_start + u_count; ++u) {
+        cb_wait_front(16, 1);
+        noc_async_write(get_read_ptr(16), g.get_noc_addr(u), TB);
+        noc_async_writes_flushed();
+        cb_pop_front(16, 1);
+    }
+    noc_async_write_barrier();
+}