tt-model push diffusion-planner-p150 (container)

.gitattributes

@@ -57,3 +57,16 @@ media/dp_nuscenes_scene-0103_kf12_cam_front_tt_NC.jpg filter=lfs diff=lfs merge=
 media/dp_nuscenes_scene-0757_kf11_bev_tt_NC.png filter=lfs diff=lfs merge=lfs -text
 media/dp_nuscenes_scene-0916_kf11_bev_tt_NC.png filter=lfs diff=lfs merge=lfs -text
 media/dp_straight_road_tt_vs_cpu.png filter=lfs diff=lfs merge=lfs -text
+image/blobs/sha256/0505b13139e86b5156c958e759de0b39da1b9d9ae252a64b2f6e7a05a9a3e350 filter=lfs diff=lfs merge=lfs -text
+image/blobs/sha256/1315e7e97a55bd5693b68844d50689456daceaab85d9d103c19e578551169924 filter=lfs diff=lfs merge=lfs -text
+image/blobs/sha256/20475a01a7476e219fb670109972618cc19d3b4b136ef19e1d2ac110271457a2 filter=lfs diff=lfs merge=lfs -text
+image/blobs/sha256/312f7c835935adbd4db8a7fc44867f1fb75883fa686a3c95177fe7f6dc43e34b filter=lfs diff=lfs merge=lfs -text
+image/blobs/sha256/354c8aad2b5a35d4683b0174c5443906ec2a5cf58b6cb3f5540d3b9cacaef0ab filter=lfs diff=lfs merge=lfs -text
+image/blobs/sha256/3678157ea229048016d1006e8432030308540aaa7ddfb0520bdeeba9c833dd93 filter=lfs diff=lfs merge=lfs -text
+image/blobs/sha256/575a85c6e189c4e5adb52123d0a57f0f63352c99d7d3ec20a8a821f27f930709 filter=lfs diff=lfs merge=lfs -text
+image/blobs/sha256/82bcf016508f0384e4aca1660cda09ea6f63d7dec3223a9d02f08cf6be718663 filter=lfs diff=lfs merge=lfs -text
+image/blobs/sha256/9b8844999af8c55ff4923e14d78b61b81f93b1387154462deb0e7f2290abbcfd filter=lfs diff=lfs merge=lfs -text
+image/blobs/sha256/be7e87294626da08f944fc039cac2c322108d9c2dd2335ebf5f63f27c2c6335a filter=lfs diff=lfs merge=lfs -text
+image/blobs/sha256/c2ce3023f2ddc0bc464cfbabb86f6fa277861c23662a8b447909ff5057e1caa2 filter=lfs diff=lfs merge=lfs -text
+image/blobs/sha256/cab9ef9b426e995ed54516bbbff3ed9429b565b12b361fe88f327d1a6aa0a6c9 filter=lfs diff=lfs merge=lfs -text
+image/blobs/sha256/ff074d9d763b545567421043a31b67b5b0a9c8eb4a968e40ecc74476d54d8419 filter=lfs diff=lfs merge=lfs -text

OPT_BASELINE.md

@@ -20,6 +20,32 @@ power had a median of 76 W and a max of 86 W, at 63-68 °C. The profiler CSV hea
 shared (AMD EPYC-Rome, 8 vCPUs, up to 8 other agents, load average 8-12 during these runs), so host-side stages are
 quoted as p50 / p99 (min).
 
+## After optimization (the current release, 2026-10-11)
+
+Everything below this section is the baseline of the first release and stays as measured. The optimized release
+(rounds 1-5 of `OPT_REPORT.md`, verified in `VERIFICATION_OPT_2026-10-11.md`, re-measured for the republish on
+`bc966b4` with ttaw 0.23.2; same configuration: ETH dispatch, 1 CQ, 12×10, AICLK 1350 MHz, the stage bench above):
+
+| | baseline (`5541833`, 2026-10-08) | current (2026-10-11) |
+|---|---|---|
+| device latency per plan (one trace replay + sync), shipped sample | 102.13 ms p50 | **20.05 ms** p50 (p99 20.08) |
+| back-to-back replays, shipped sample (88 neighbours) | 102.04 ms = 9.80 plans/s | **20.00 ms = 50.0 plans/s** (5.1×) |
+| back-to-back replays by agent bucket (r32 / r64 / r96 / r128 / r192 / full) | 102.04 ms for every scene | 17.44 / 19.25 / 20.00 / 21.11 / 22.30 / 26.65 ms (the last three from `OPT_REPORT.md` round 5) |
+| `model(inputs=arrays)` end to end, shipped sample | 117.90 ms p50 (p99 134.56) | **26.73 ms** p50 (p99 27.48) |
+| `model(inputs=<.npz path>)` | 124.76 ms | 32.27 ms |
+| host pre · pack · host tensors · H2D · D2H · host post | 4.75 · 0.45 · 2.41 · 1.14 · 0.58 · 4.32 ms | 2.45 · 0.52 · 1.57 · 0.76 · 0.22 · 1.39 ms |
+| programs per plan (traced replay profile) | 6,282 | 1,418 at r96 (kernel 19.64 ms, op-to-op gaps 0.89 ms, span 20.54 ms) |
+| traces / trace buffers | 1 / 74.6 MB | 6 (full capacity + 5 agent buckets) / 60.3 MB |
+| `from_pretrained` load, empty / warm JIT cache | 315 s / 8.6 s | 352 s / 12.9 s |
+| served `/predict` `timing_ms.total` (uvicorn on the host) | 123.1 ms median | 37.2 ms median |
+| accuracy: 44 device tests; 99 scenes worst ego max / mean (gates 1.0 / 0.3 m), turn, neighbours (gate 1.5 m) | 44 / 44; 0.313 / 0.143 m, 99 / 99, ≤ 0.086 m | 44 / 44; 0.346 / 0.154 m, 99 / 99, ≤ 0.094 m (two gated precision changes: `OPT_REPORT.md` r2.3, r3.2) |
+
+Where the 82 ms went (same-session back-to-back A/B deltas of `OPT_REPORT.md`, largest steps): the encoder's batched channel matmuls on 4-8 cores (round 1,
+−23.4 ms), the precision cost recovered with custom `generic_op` kernels (fused fp32 LayerNorm −15.5 ms, K-concatenated
+split matmul −7.0 ms, fused fp32 attention −5.4 ms, softmax and glue fusions; rounds 2-3), L1 residency (round 4,
+−3.5 ms) and exact agent-bucket compaction (round 5, −6.6 ms on this scene). Logs: `logs/diffusion-planner/republish/`
+(`bench_recheck.json`, `device_suite.log`, `load_*.json`, `served_bench.json`).
+
 ## Summary
 
 | | baseline |

OPT_REPORT.md

@@ -1,14 +1,71 @@
 # diffusion-planner-p150 optimization report (p150, ETH dispatch, 12×10 grid)
 
-**Status: baseline port; optimization pending.** This is the first public release: the functional port, measured
-once (`OPT_BASELINE.md`, baseline commit `5541833`, 2026-10-08) and not optimized yet. No optimization round has run,
-so there is no step, no rejected attempt and no hang to report, and every number below is the baseline. The whole
-plan (encoder, the 11 DiT evaluations with the DPM-Solver++(2M) updates, the turn-indicator head) runs as one metal
-trace of 6,282 programs: 102.04 ms per back-to-back replay (9.80 plans/s of device throughput), 117.9 ms per
-synchronous `model()` call on the shipped sample. The trace is kernel-bound (op-to-op gaps 4.4 ms of a 103.7 ms span).
-Two sinks of similar size come first: the **numerics defaults** that the end-to-end gates need (split hi / lo matmuls,
-fp32 LayerNorm, fp32 matmul attention) cost **33.4 ms** over the first device round, and the encoder's channel-MLP and
-pre-projection matmuls run on only 4-8 cores (**25.8 ms**, unrelated to precision).
+**Status: optimization rounds 1-5 done and independently verified (`VERIFICATION_OPT_2026-10-11.md`, PASS); the
+optimized release is packaged for re-publishing (2026-10-11, section "Republish (2026-10-11)" below). Round 6
+(map-entity compaction, the cross attention, the MK-E / MK-D megakernels) is in progress on a branch and is NOT in
+this release.**
+Device time per plan, back to back, on the bench scene (`kashiwanoha_dense`, 88 neighbours: the r96 bucket):
+**102.03 (published baseline) -> 70.12 (round 1) -> 40.89 (round 2) -> 30.16 (round 3) -> 26.62 (round 4) -> 20.00 ms
+(round 5)**, 50.0 plans/s; one blocking plan 20.06 ms; e2e `model()` p50 **26.84 ms** (113.1 at the baseline, 81.2
+after round 1, 52.0 after round 2, 40.8 after round 3, 37.4 after round 4). Since round 5 the device time depends on
+the scene (the agent bucket): 17.44 ms (r32, ≤ 31 neighbours: 53 of the 99 gate scenes; 44 in r64, 2 in r96) to 26.65 ms (full capacity,
+> 191 neighbours); over the 99 gate scenes the served device time (`timing_ms.device`) is p50 20.83 ms (29.61 after
+round 4). Accuracy: 44 / 44 device tests with the gates read-only in every step; every round-5 step is bit-identical
+(the 99 per-scene numbers are those of round 3).
+
+Round 5 added four kept commits, all bit-identical on the rows that are read back, each with its own knob and a
+same-session alternating A/B (3 pairs, 3 scenes):
+
+- the decoder on the first R rows of an agent bucket (`COMPACT=1`, buckets 32 / 64 / 96 / 128 / 192, one trace
+  each, picked per request): 26.62 -> 23.25 ms on the bench scene (r96), 21.85 ms at r32;
+- the encoder's neighbour trunk and head on the same R neighbours, the other neighbour tokens zero (`COMPACT=2`):
+  23.25 -> 19.99 ms (r96), 17.42 ms at r32;
+- the host pre- and post-processing vectorised, bit-exact (`HOST_FAST`): e2e 30.15 -> 27.31 ms (host pre 3.77 ->
+  2.49 ms, post 3.02 -> 1.44 ms);
+- `cs` uploaded as one tile column and the all-zero inputs not re-uploaded (`INPUT_TRIM`): e2e 27.22 -> 26.84 ms,
+  device +0.014 ms (one concat program).
+
+Round 4 added four kept commits, **all bit-identical** to their predecessor (same raw readback, same 99 per-scene
+numbers), each with its own `DIFFUSION_PLANNER_*` knob and a same-session alternating A/B:
+
+- the K = 1024 split operands (mlp fc2, final p4) written L1 block-sharded by the operand build and read in place
+  by a sharded-in0 2-D multicast matmul (`KCAT_L1`): −1.63 ms;
+- the fused attention's inputs L1-resident: decoder qkv / cross q outputs, the hoisted cross K / V heads, the
+  masks (`ATTN_L1`): −0.41 ms, then −0.39 ms more by keeping the fusion q / kv linears in DRAM (`ATTN_L1=2`: the
+  stock linear with an L1 output ran on 6 cores, found in the round-4 profile);
+- the decoder blocks' intermediates in L1 (split-LN outputs and stream, attention outputs, linear outputs:
+  `DEC_L1=1`): −1.01 ms; plus the pre-projection and final layer (`DEC_L1=2`): −0.09 ms.
+
+Measured and not kept: the operand build's GELU once per tile (`KCAT_ACT_ONCE`, bit-identical, +0.06 ms), the mixer
+intermediates in L1 (`ENC_L1`, +1.67 ms and not bit-identical) and the fusion intermediates in L1 (`FUS_L1`, +0.39
+ms and not bit-identical): stock ops pick other programs for L1 outputs.
+
+Round 3 added five kept commits (each with its own `DIFFUSION_PLANNER_*` knob, default on, and a same-session
+alternating A/B); four are **bit-identical** to their predecessor (same raw readback, same 99 per-scene numbers):
+
+- the attention score scale + mask + softmax as one `generic_op` (`ATTN_SMSM`): −1.56 ms;
+- the encoder's split linears (pre-projections, ego / neighbour island) in the K-concatenated form (`ENC_KCAT`):
+  −1.70 ms, **the round's one precision change**: 99-scene worst ego max 0.327 -> 0.346 m, mean 0.149 -> 0.154 m,
+  neighbours 0.104 -> 0.094 m (gates 1.0 / 0.3 / 1.5 m; 47 scenes worse, 45 better);
+- the whole fp32 matmul attention (head split, Q·Kᵀ, scale / mask / softmax, P·V, head merge) as one `generic_op`
+  with the scores in L1 (`ATTN_FUSED`): −5.43 ms;
+- the split operand of the next K-concatenated linear written by its producer (split LN, fused attention:
+  `KCAT_EMIT`): −0.95 ms;
+- the two transposes around each mixer token-mixing MLP inside the fused LN programs (`LN_TR`): −1.08 ms.
+
+Measured and not kept: the mixer fc1 GELU in the matmul epilogue (`LIN_ACT`, −0.13 ms but 99-scene ego mean 0.149 ->
+0.174 m) and the softmax scale by an immediate (`ATTN_SMSM=2`, bit-identical, −0.03 ms).
+
+Round 2 (L2 `generic_op` kernels, six of seven steps bit-identical): fused fp32 LayerNorm (`LN_KERNEL=2`, −15.54
+ms), residual adds in the LN (`LN_RESID`, −0.21), K-concatenated decoder split matmuls (`SPLIT_KCAT=2`, −6.98, the
+precision change of round 2: ego max 0.313 -> 0.327 m), score scale + mask (`ATTN_SMASK`, −2.83), GELU in the
+operand build (`KCAT_ACT`, −0.60), split-row decoder LN (`LN_SPLIT`, −0.83), SFPU statistics broadcast
+(`LN_SFPU_BCAST`, −2.24).
+
+Round 1 (stock-op program configs, all bit-identical): the encoder's batched linears as 2-D matmuls (`ENC_CH2D`,
+−23.37 ms), P·V on one output tile per core (`ATTN_FAST=1`, −3.35 ms), explicit configs for the decoder's 352-row
+matmuls (`DEC_MMCFG`, −5.18 ms). Two faster attention variants of round 1 stay rejected (not bit-identical, raise
+the 99-scene error).
 
 All numbers: `code/scripts/bench.py` (median of 100 warm iterations, batch 1,
 `code/tt_diffusion_planner/samples/kashiwanoha_dense.npz`), ETH dispatch, 1 CQ, 12×10 grid, the pinned numerics
@@ -18,19 +75,19 @@ device profiler (one traced replay between signposts). Accuracy gates: `OPT_BASE
 
 ## Summary
 
-| | baseline `5541833` (2026-10-08) | **final (= baseline: no round yet)** |
-|---|---|---|
-| device trace, one blocking plan | 102.13 ms (p99 104.81) | same |
-| back-to-back traces | 102.04 ms (9.80 plans/s) | same |
-| e2e `model(inputs=arrays)` p50 / p99 | 117.90 / 134.56 ms | same |
-| e2e `model(inputs=<.npz path>)` p50 | 124.76 ms | same |
-| host pre-processing / pack / host tensors / H2D / D2H / post-processing | 4.76 / 0.45 / 2.41 / 1.14 / 0.58 / 4.32 ms | same |
-| device programs per plan (unique programs) | 6,282 (309) | same |
-| kernel sum / op-to-op gaps / span (profile) | 99.27 / 4.39 / 103.66 ms | same |
-| encoder / 11 decoder evaluations (+ solver updates) | 51.6 / 51.4 ms (4.67 ms per evaluation) | same |
-| accuracy gates (PCC / agreement vs the fp32 CPU reference) | 44 / 44 device tests: module PCC ≥ 0.999952, encoding 0.999978, decoder evaluation ≥ 0.9999995; 99 scenes: ego max 0.313 m / mean 0.143 m (gates 1.0 / 0.3 m), turn command 99 / 99, neighbours ≤ 0.086 m (gate 1.5 m) | same |
-| served `/predict` `timing_ms.total`, median of 50 (uvicorn on the host, the shipped sample) | 123.1 ms, measured on `c0d84f9` (same device code; the release verification) | same |
-| `from_pretrained` load: empty JIT cache / warm cache | 315 s / 8.6 s (`c0d84f9`) | same |
+| | baseline `5541833` (2026-10-08) | after round 1 (`dd6f249`, 2026-10-10) | after round 2 (`c42c540`, 2026-10-10) | after round 3 (`0b750e3`, 2026-10-10) | after round 4 (`93428bb`, 2026-10-10) | **after round 5 (`e933166`, 2026-10-10; bench scene = bucket r96)** |
+|---|---|---|---|---|---|---|
+| device trace, one blocking plan | 102.13 ms (p99 104.81) | **70.18 ms** (same-session baseline 102.09) | **40.96 ms** (p99 41.01) | **30.22 ms** (p99 30.42) | **26.68 ms** (p99 26.71) | **20.06 ms** (r32: 17.50, r64: 19.31; full capacity 26.7) |
+| back-to-back traces | 102.04 ms (9.80 plans/s) | **70.12 ms (14.26 plans/s)** (same-session baseline 102.03) | **40.90 ms (24.45 plans/s)** | **30.16 ms (33.16 plans/s)** | **26.62 ms (37.56 plans/s)** | **20.00 ms (50.0 plans/s)**; buckets r32 / r64 / r96 / r128 / r192 / full: 17.44 / 19.25 / 20.00 / 21.11 / 22.30 / 26.65 ms |
+| e2e `model(inputs=arrays)` p50 / p99 | 117.90 / 134.56 ms | **81.2 ms** p50 (same-session baseline 113.1; host load lower than on 2026-10-08) | **51.98 / 56.39 ms** | **40.84 / 41.52 ms** | **37.36 / 38.03 ms** | **26.84 ms** p50 (straight_road r32: 23.36, scene-0103_kf14 r64: 25.72) |
+| e2e `model(inputs=<.npz path>)` p50 | 124.76 ms | 86.7 ms (same-session baseline 118.7) | 57.9 ms | 46.6 ms | 43.3 ms | 32.5 ms |
+| host pre-processing / pack / host tensors / H2D / D2H / post-processing | 4.76 / 0.45 / 2.41 / 1.14 / 0.58 / 4.32 ms | 3.84 / 0.38 / 1.98 / 0.92 / 0.30 / 3.50 ms (unchanged code; quieter host) | 3.91 / 0.39 / 1.99 / 0.94 / 0.35 / 3.54 ms (unchanged code) | 3.85 / 0.39 / 1.96 / 0.92 / 0.33 / 3.49 ms (unchanged code) | 3.88 / 0.39 / 1.94 / 0.92 / 0.33 / 3.47 ms (unchanged code) | 2.49 / 0.52 / 1.58 / 0.76 / 0.25 / 1.44 ms (`HOST_FAST`, `INPUT_TRIM`) |
+| device programs per plan (unique programs) | 6,282 (309) | same | 2,155 (293 program-cache entries); trace buffer 16.8 MB (44.6 before) | 1,405 (195 program-cache entries); trace buffer 12.3 MB | 1,411 (199 program-cache entries); trace buffer 12.3 MB | 1,418 at r96 (480 program-cache entries for the 6 traces); trace buffers 60.3 MB (6 traces) |
+| kernel sum / op-to-op gaps / span (profile) | 99.27 / 4.39 / 103.66 ms | same | 39.99 / 1.47 / 41.46 ms | 29.71 / 0.88 / 30.59 ms | 26.00 / 0.89 / 26.89 ms | 19.64 / 0.89 / 20.54 ms (r96) |
+| encoder / 11 decoder evaluations (+ solver updates) | 51.6 / 51.4 ms (4.67 ms per evaluation) | same | 15.7 / 23.9 ms (2.17 ms per evaluation; kernel sums of the profile) | 12.6 / 17.1 ms (decoder 1.55 ms per evaluation; kernel sums of the profile) | 12.4 / 13.4 ms (decoder 1.22 ms per evaluation; kernel sums of the profile) | 9.4 / 10.0 ms at r96 (decoder blocks 9.49 ms: 0.86 ms per evaluation; kernel sums of the profile) |
+| accuracy gates (PCC / agreement vs the fp32 CPU reference) | 44 / 44 device tests: module PCC ≥ 0.999952, encoding 0.999978, decoder evaluation ≥ 0.9999995; 99 scenes: ego max 0.313 m / mean 0.143 m (gates 1.0 / 0.3 m), turn command 99 / 99, neighbours ≤ 0.086 m (gate 1.5 m) | identical (bit-identical readback; 44 / 44, all 99 per-scene numbers equal) | 44 / 44; 99 scenes: ego max 0.327 m / mean 0.149 m, turn 99 / 99, neighbours ≤ 0.104 m (the r2.3 precision change; every other step bit-identical) | 44 / 44; 99 scenes: ego max 0.346 m / mean 0.154 m, turn 99 / 99, neighbours ≤ 0.094 m (the r3.2 precision change; every other step bit-identical) | 44 / 44; 99 scenes identical to round 3 (every round-4 step bit-identical) | 44 / 44; 99 scenes identical to round 3 (every round-5 step bit-identical on the read rows) |
+| served `/predict` `timing_ms.total`, median of 50 (uvicorn on the host, the shipped sample) | 123.1 ms, measured on `c0d84f9` (same device code; the release verification) | same | not re-measured (no release in this round) | not re-measured | not re-measured | **37.2 ms** (min 36.7; decode 8.6, `model()` 26.8; republish re-measure on `bc966b4`, 2026-10-11) |
+| `from_pretrained` load: empty JIT cache / warm cache | 315 s / 8.6 s (`c0d84f9`) | same | – / 7.9 s (warm, bench) | – / 7.1 s (warm, bench) | – / 7.8 s (warm, bench) | 352 s / 12.9 s (6 traces; republish re-measure, 2026-10-11; 11.4 s in the bench process) |
 
 The host is shared with other agents' jobs, so the host-side rows move with its load (several ms); the device rows
 repeat to ±0.05 ms. AICLK 1350 MHz (min 1343) during the stage bench. The release re-check on `c0d84f9`
@@ -41,6 +98,35 @@ repeat to ±0.05 ms. AICLK 1350 MHz (min 1343) during the stage bench. The relea
 | step | commit | what | trace / b2b ms | e2e ms | accuracy | revert switch |
 |---|---|---|---|---|---|---|
 | baseline | `5541833` | first correct port: the whole plan in one trace, ETH 1CQ 12×10, the numerics defaults of PORT_LOG decision 11, host pre- and post-processing | 102.13 / 102.04 | 117.90 | all gates (above) | – |
+| r1.0 | `0f557b9` | re-vendor ttaw 0.20.0 -> 0.23.0 (no planner module changed) | – | – | 44 / 44, 99 scenes identical | – |
+| r1.1 | `9a885b3` | encoder linears on batched `[1, E, T, C]` activations as one 2-D matmul: free `[1, 1, E·64, 128]` view for the MixerBlock channel MLPs (auto 2-D multicast), a 1-D in1-multicast config with `fuse_batch=True` for the pre-projections (T = 6 / 20 / 40) | 102.09 -> 78.72 / 102.03 -> **78.66** (−23.37) | 113.1 -> 89.9 | bit-identical readback; 44 / 44; 99 scenes identical | `DIFFUSION_PLANNER_ENC_CH2D=0` |
+| r1.2 | `ef6a763` | fp32 matmul attention (`tt/attention.py`): P·V on `MatmulMultiCoreReuseProgramConfig` with one output tile per core and the whole K in one block (88 cores instead of 11; sweep: 73 -> 21 µs self, 71 -> 32 µs cross, 81 -> 42 µs fusion) | 78.72 -> 75.37 / 78.66 -> **75.31** (−3.35) | 89.8 -> 86.5 | bit-identical; 44 / 44; 99 scenes identical | `DIFFUSION_PLANNER_ATTN_FAST=0` |
+| r1.3 | `dd6f249` | decoder 352-row matmuls (every split pass): explicit `MatmulMultiCoreReuseMultiCastProgramConfig` per (K, N, pass), the fastest bit-identical candidate of a device sweep (`tt/layers.py` `DEC_MM_CONFIGS`; e.g. mlp fc2 30 -> 12 µs) | 75.37 -> 70.18 / 75.30 -> **70.12** (−5.18) | 86.5 -> 81.2 | bit-identical; 44 / 44; 99 scenes identical | `DIFFUSION_PLANNER_DEC_MMCFG=0` |
+| r2.1 | `2fc9af8` | fused fp32 LayerNorm (`LN_KERNEL=2`): the 9-program `layer_norm_fp32` decomposition as one `generic_op` (`tt/ln_kernel.py`, `kernels/ln32_*.cpp`), the same SFPU LLK calls in the same order (accurate `ttnn.mean` fold + `sfpu_reduce`, binary_ng fp32 SFPU ops, `rsqrt` Default) | 70.18 -> 54.64 / 70.12 -> **54.58** (−15.54) | 81.6 -> 65.6 | bit-identical readback; 44 / 44; 99 scenes identical | `DIFFUSION_PLANNER_LN_KERNEL=0` |
+| r2.2 | `6ef00e6` | residual adds fused into the LN program (`LN_RESID`): `h = x + r (* gate)` then `LN(h)`, writing both (mixers: x + token-mix, x + channel-MLP; decoder: the four stream adds of a block and the final one) | 54.64 -> 54.43 / 54.58 -> **54.37** (−0.21) | 65.5 -> 65.2 | bit-identical; 44 / 44; 99 scenes identical | `DIFFUSION_PLANNER_LN_RESID=0` |
+| r2.3 | `d7f84e1` | decoder split matmuls as ONE K-concatenated fp32 matmul (`SPLIT_KCAT=2`): `[x_hi \| x_hi \| x_lo \| 1 \| 0..] @ [w_hi; w_lo; w_hi; b_hi, b_lo; 0..]`, operand built by one `generic_op` (`kernels/kcat_*.cpp`), K' padded for K = 1024, per-shape 2-D multicast configs from a sweep | 54.44 -> 47.45 / 54.37 -> **47.39** (−6.98) | 65.4 -> 58.3 | **not bit-identical** (one fp32 accumulation, exact bias): final_x0 rel 1.1e-3; 44 / 44 green; 99 scenes ego max 0.313 -> 0.327 m, mean 0.143 -> 0.149 m, neighbours 0.086 -> 0.104 m (57 scenes worse, 41 better) | `DIFFUSION_PLANNER_SPLIT_KCAT=0` |
+| r2.4 | `2da6100` | attention score scale + mask in one pass (`ATTN_SMASK`, `kernels/smask_*.cpp`); reproduces the TF32 truncation of the stock fp32 + bf16 add (finding below) | 47.45 -> 44.63 / 47.39 -> **44.56** (−2.83) | 58.3 -> 55.6 | bit-identical to r2.3 (readback, 99 / 99 scenes); 44 / 44 | `DIFFUSION_PLANNER_ATTN_SMASK=0` |
+| r2.5 | `4699447` | the GELU between two K-concatenated linears applied inside the next operand build (`KCAT_ACT`; 77 unary programs per plan) | 44.62 -> 44.02 / 44.56 -> **43.96** (−0.60) | 55.5 -> 54.9 | bit-identical to r2.3; 44 / 44 | `DIFFUSION_PLANNER_KCAT_ACT=0` |
+| r2.6 | `98ca6c7` | decoder LayerNorms with each tile row over 8 cores (`LN_SPLIT`, `kernels/ln32s_*.cpp`: gather to a root core, ordered fold, broadcast; semaphores): 88 cores instead of 11 | 44.02 -> 43.19 / 43.96 -> **43.13** (−0.83) | 54.8 -> 54.3 | bit-identical to r2.3; 44 / 44 | `DIFFUSION_PLANNER_LN_SPLIT=0` |
+| r2.7 | `c42c540` | LN statistics broadcast by the SFPU row reduce itself (`LN_SFPU_BCAST`, `kernels/ln32_sfpu.h`: the stock `sfpu_reduce` row-sum arithmetic with the result stored to every column) instead of a RISC-V 1024-store column fill per statistic | 43.19 -> 40.95 / 43.13 -> **40.89** (−2.24) | 53.9 -> 51.9 | bit-identical to r2.3; 44 / 44 | `DIFFUSION_PLANNER_LN_SFPU_BCAST=0` |
+| r3.1 | `c03adc3` | attention score scale + mask + softmax as one `generic_op` (`ATTN_SMSM=1`, `kernels/smsm_*.cpp`): the smask SFPU sequence into an L1 row, then the stock `ttnn.softmax(numeric_stable=True)` compute with the same `kernel_lib` calls | 40.96 -> 39.39 / 40.89 -> **39.32** (−1.56) | 51.9 -> 50.3 | bit-identical to r2.7 (readback, 99 / 99 scenes); 44 / 44 | `DIFFUSION_PLANNER_ATTN_SMSM=0` |
+| r3.2 | `c6edfaf` | encoder split linears (`enc.pre.*`, `enc.island.*`, any K: weight blocks padded to whole tiles) as one K-concatenated fp32 matmul (`ENC_KCAT`), 1-D `fuse_batch` config for the `[1, E, T, K']` operands | 39.38 -> 37.68 / 39.32 -> **37.62** (−1.70) | 50.3 -> 48.4 | **not bit-identical** (one fp32 accumulation, exact bias; final_x0 rel 2.9e-3); 44 / 44 green; 99 scenes ego max 0.327 -> 0.346 m (scene-0103_kf14), mean 0.149 -> 0.154 m, neighbours 0.104 -> 0.094 m (47 scenes worse, 45 better, mean change −0.2 mm) | `DIFFUSION_PLANNER_ENC_KCAT=0` |
+| r3.3 | `8d402d1` | (not kept, knob default 0) mixer fc1 GELU in the matmul epilogue (`LIN_ACT`): an explicit copy of the stock auto config (bit-identical without the activation) + `fused_activation` | 39.38 -> 39.25 / 39.32 -> 39.19 (−0.13) | 50.3 -> 50.1 | not bit-identical; 44 / 44 green; ego max 0.327 -> 0.392 m, mean 0.149 -> 0.174 m (54 worse, 38 better) | default 0 |
+| r3.4 | `4acde70` | the whole fp32 matmul attention as one `generic_op` (`ATTN_FUSED`, `kernels/fattn_*.cpp`): per (head, query tile row) Q·Kᵀ into DEST (`matmul_block`, in1 transposed), the smsm phases on the row in L1, P·V accumulated in DEST in key order; Q / K / V read in place (no head split), output written as the merged heads (no head merge) | 37.68 -> 32.25 / 37.62 -> **32.19** (−5.43) | 48.9 -> 43.2 | bit-identical to r3.2 (readback, 99 / 99 scenes); 44 / 44 | `DIFFUSION_PLANNER_ATTN_FUSED=0` |
+| r3.5 | `ad03767` | the decoder's split-row LNs and the fused attention write the split operand `[x_hi \| x_hi \| x_lo \| 1]` of the next K-concatenated linear (`KCAT_EMIT`; `layers.KcatOperand`): 198 `kcat_operand` programs per plan less | 32.25 -> 31.30 / 32.19 -> **31.24** (−0.95) | 43.1 -> 42.4 | bit-identical to r3.4; 44 / 44 | `DIFFUSION_PLANNER_KCAT_EMIT=0` |
+| r3.6 | `0b750e3` | the two transposes around each mixer token-mixing MLP inside the fused LN programs (`LN_TR`): n1 writes its output per-entity transposed, n2 reads the token-mixing output transposed (`transpose_tile`, the stock transpose LLK with fp32 unpack-to-dest: exact); 48 transpose programs less | 31.30 -> 30.22 / 31.24 -> **30.16** (−1.08) | 42.1 -> 41.0 | bit-identical to r3.5; 44 / 44 | `DIFFUSION_PLANNER_LN_TR=0` |
+| r4.1 | `7e662d3` | the K = 1024 split operands X' `[352, 3328]` (mlp fc2) / `[352, 3168]` (final p4, K' padded to 11 slices) written **L1 block-sharded** by the operand build (its writer's TensorAccessor resolves the shards) and read in place by a 2-D multicast matmul with in0 sharded (`KCAT_L1`, `layers.KCAT_L1_CONFIGS`): no 4.7 MB DRAM round trip (kernel check: operand + matmul 62.6 -> 41.3 µs, 62.3 -> 41.8 µs) | 30.22 -> 28.59 / 30.16 -> **28.53** (−1.63) | 40.8 -> 39.3 | bit-identical to r3.6 (readback, 99 / 99 scenes); 44 / 44 | `DIFFUSION_PLANNER_KCAT_L1=0` |
+| r4.2 | `310bd38` | (not kept, knob default 0) the deferred GELU of the operand build on one DEST copy + `copy_dest_values` (`KCAT_ACT_ONCE`) | 28.62 -> 28.66 / 28.53 -> 28.59 (+0.06) | (host-noisy) | bit-identical | default 0 |
+| r4.3 | `f6cc333` | the fused attention's inputs L1-resident (`ATTN_L1=1`): decoder qkv / cross q outputs and fusion q / kv outputs in L1, the hoisted cross K / V heads moved to L1 once per plan, the self / cross / fusion masks in L1 (kernel check: self 51.0 -> 35.1 µs, cross 61.0 -> 54.2, fusion 106.2 -> 88.4) | 28.59 -> 28.18 / 28.53 -> **28.12** (−0.41) | 39.3 -> 38.9 | bit-identical to r4.1; 44 / 44 | `DIFFUSION_PLANNER_ATTN_L1=0` |
+| r4.4 | `e5fb151` | the decoder blocks' intermediates in L1 interleaved (`DEC_L1=1`): split-LN outputs and stream `h`, fused attention outputs, out / mlp / cross-out linear outputs (kernel check: split LN 24.1 -> 20.6 µs, K = 256 matmuls −2 to −4 µs, attention −1.7 µs) | 28.18 -> 27.16 / 28.12 -> **27.11** (−1.01) | 38.8 -> 37.9 | bit-identical to r4.3; 44 / 44 | `DIFFUSION_PLANNER_DEC_L1=0` |
+| r4.5 | `60eaef1` | (not kept, default 0) the mixer blocks' intermediates in L1 (`ENC_L1`) | 27.16 -> 28.84 / 27.11 -> 28.77 (+1.67) | 37.7 -> 39.5 | **not** bit-identical (final_x0 rel 8e-3; ego max 0.379 m): the stock matmul / unary programs change with an L1 output; 44 / 44 | default 0 |
+| r4.6 | `dab96fb` | `DEC_L1=2`: the decoder pre-projection and final layer intermediates in L1 too | 27.17 -> 27.07 / 27.11 -> **27.02** (−0.09) | 37.9 -> 37.9 | bit-identical to r4.4; 44 / 44 | `DIFFUSION_PLANNER_DEC_L1=1` |
+| r4.7 | `c3b6a5a` | (not kept, default 0) the fusion blocks' intermediates in L1 (`FUS_L1`) | 27.07 -> 27.47 / 27.02 -> 27.40 (+0.39) | 37.8 -> 38.2 | **not** bit-identical (ego max 0.364 m; `ttnn.layer_norm` with an L1 output); 44 / 44 | default 0 |
+| r4.8 | `93428bb` | `ATTN_L1=2`: the fusion q / kv linears back to DRAM outputs (the stock linear with an L1 output ran on 6 cores, 54 µs, with a separate bias add: the fusion stage 1.41 -> 1.74 ms in the first round-4 profile); masks and decoder inputs stay in L1 | 27.08 -> 26.69 / 27.02 -> **26.62** (−0.39) | 37.8 -> 37.7 | bit-identical to r4.6; 44 / 44 | `DIFFUSION_PLANNER_ATTN_L1=1` |
+| r5.1 | `9fe29a8` | exact decoder compaction (`COMPACT=1`): one trace per agent bucket (`AGENT_BUCKETS` 32, 64, 96, 128, 192) besides the full 352-row plan; the decoder runs on the first R rows (the ego, every valid self-attention key, every emitted neighbour), `y0` / `cs` / the key row sliced in the trace; per-R agent rows and cross mask; the matmul configs keep their K blocking (`per_core_M` fitted to the grid); the bucket picked per request (`TtDiffusionPlanner.variant_for`) | 26.68 -> 23.31 / 26.62 -> **23.25** (−3.37; r32 21.85, r64 22.91) | 38.3 -> 33.8 | bit-identical on the needed rows vs the full plan (15 scenes, `compact_check.py`); 44 / 44; 99 scenes identical | `DIFFUSION_PLANNER_COMPACT=0` |
+| r5.2 | `268ff32` | `COMPACT=2`: the encoder's neighbour trunk and head on the same first R neighbours, the other neighbour tokens a zero block (invalid entities: `token_valid` zeroes them anyway); the bucket also covers every valid neighbour token | 23.30 -> 20.05 / 23.25 -> **19.99** (−3.26; r32 17.42, r64 19.23) | 33.2 -> 30.1 | bit-identical on the needed rows (15 scenes); 44 / 44; 99 scenes identical | `DIFFUSION_PLANNER_COMPACT=1` |
+| r5.3 | `3ce8ea9` | vectorised host pre- / post-processing (`HOST_FAST`): normalisation of the non-small rows only, neighbour features on the 6 kept rows, the trajectory loops as array operations in the same float64 / float32 order, denormalisation and predicted paths on the emitted rows only | device unchanged (19.99) | 30.15 -> **27.31** | bit-exact on 134 inputs (`host_equiv.py`, both directions of the knob); 44 / 44; 99 scenes identical | `DIFFUSION_PLANNER_HOST_FAST=0` |
+| r5.4 | `e933166` | trimmed uploads (`INPUT_TRIM`): `cs` as one tile column `[352, 32]` widened in the trace by a tile-aligned concat; an all +0.0 input (`y0`, `static_x`) not re-uploaded while its buffer holds this model's last zeros | 20.05 -> 20.06 / 19.99 -> 20.00 (+0.014) | 27.22 -> **26.84** | readback bit-identical (bench dumps, 3 scenes x 6 runs); 44 / 44; 99 scenes identical | `DIFFUSION_PLANNER_INPUT_TRIM=0` |
 
 Later commits up to the first release changed no model numerics and no device code: the baseline documentation
 (`a6e5bbf`), the ttaw 0.19.0 re-vendor with the `serve.env` pins of the numerics knobs, the quickstart picture and two
@@ -49,11 +135,33 @@ identical gate values and per-scene numbers (`VERIFICATION_2026-10-08.md`).
 
 ## Round 1
 
-Not started. The optimization phase follows the publication of every model of the collection (`research/PLAN.md`
-§4.1, §5): each step is one commit with one `DIFFUSION_PLANNER_*` A/B knob, keeps the frozen gates and is re-checked on
-all 99 end-to-end scenes (the plan's sensitivity to device numerics is chaotic per scene: PORT_LOG known issues). The
-precision policy (PLAN.md §9.3) may be relaxed per module only with gate evidence; the precision cost is to be
-recovered with fused kernels that keep fp32-level accuracy, not by dropping precision.
+Date 2026-10-10, plan `OPT_PLAN.md` §7 round 1 (items 1, 2a, 4a). The device was shared with a METEOR debugging job
+that hung the chip three times during the round (05:22, 05:45 and 07:35 UTC, FAULT markers written by its own devrun
+jobs). None of this round's jobs ran on a faulted chip: every job script checks for the marker after taking the lock,
+and the queued jobs aborted with exit 75 and were re-queued after the orchestrator cleared the marker. No job of this
+round hung.
+
+Method, per step (`logs/diffusion-planner/opt_r1/`; scripts in `scripts/`):
+
+- `ab.sh <tag> <KNOB> <off> <on>`: the device suite with the new value (`TTAW_GATES_READONLY=1`, both device test
+  files, `DIFFUSION_PLANNER_E2E_REPORT` per-scene JSON), then 3 alternating pairs of `bench.py` (`--iters 30`, b2b 3 ×
+  50) in one devrun window. Every bench run writes its raw packed readback (`--dump`), and `cmp.py` compares on / off
+  bit for bit and the per-scene report with the baseline's `e2e_device.json`. The other knobs are pinned so each A/B
+  isolates one change; jobs run from a snapshot of the code (`snap.sh`), so later edits cannot leak into a queued job.
+- `mm_sweep.py` / `attn_sweep.py`: device micro-benches (20-call traces) of the decoder matmul shapes (three split
+  passes each) and the attention pieces under the auto config and candidate program configs, outputs compared bit for
+  bit (`mm_sweep.json`, `attn_sweep.log`).
+
+| A/B (same session, 3 pairs unless noted, kashiwanoha_dense) | off: b2b / trace / e2e ms | on: b2b / trace / e2e ms | Δ b2b | outputs | kept |
+|---|---|---|---:|---|---|
+| `ENC_CH2D` 0 / 1 (ATTN_FAST=0, DEC_MMCFG=0) | 102.03 / 102.09 / 113.1 | 78.66 / 78.72 / 89.9 | −23.37 | bit-identical | yes |
+| `ATTN_FAST` 0 / 1 (ENC_CH2D=1, DEC_MMCFG=0) | 78.66 / 78.72 / 89.8 | 75.31 / 75.37 / 86.5 | −3.35 | bit-identical | yes |
+| `DEC_MMCFG` 0 / 1 (ENC_CH2D=1, ATTN_FAST=1) | 75.30 / 75.37 / 86.5 | 70.12 / 70.18 / 81.2 | −5.18 | bit-identical | yes |
+| `ATTN_FAST` 0 / 3: + scale as an SFPU pre-activation of the mask add (ENC_CH2D=1, DEC_MMCFG=0; the first form of the knob, then a bool) | 78.66 / 78.72 / 89.6 | 72.81 / 72.88 / 83.8 | −5.85 (−2.5 over mode 1) | final_x0 rel 4.1e-3; 99 scenes: ego max 0.313 -> 0.385 m, mean 0.143 -> 0.173 m; 44 / 44 green | no |
+| `ATTN_FAST` 1 / 2: + scale on Q before Q·Kᵀ (2 pairs) | 70.12 / 70.18 / 81.2 | 68.15 / 68.21 / 79.5 | −1.97 | final_x0 rel 3.0e-3; ego max 0.313 -> 0.379 m, mean 0.143 -> 0.172 m, neighbours 0.086 -> 0.098 m; 44 / 44 green | no |
+
+AICLK 1350 MHz (min 1343) in every run. The device rows repeat to ±0.01 ms within a session; e2e moves with the host
+load by ±0.5 ms.
 
 ### Findings (measured)
 
@@ -74,19 +182,332 @@ The baseline profile (`OPT_BASELINE.md` "Device profile") is the starting point:
   is equal (the 12th column hardly matters to this graph yet), so ETH stays the default (D14); 2 CQs cost +9.8 ms on ETH,
   so 1 CQ stays pinned.
 
+Round 1 findings:
+
+- **The 4-D batched linear is the stock trap.** `ttnn.linear` on `[1, E, T, C] @ [C, N]` auto-configures with M per
+  batch element (`fuse_batch=False`, `matmul_program_config.cpp` `create_simple_matmul_program_config`), so the 6 ×
+  320 entity blocks ran on 8 cores. A free reshape (T a multiple of 32) or a program config with `fuse_batch=True`
+  gives the whole grid the work: same products, same fp32 DEST order, bit-identical output. 23.4 ms, the largest
+  single gain of the plan so far.
+- **K-blocking decides bit-identity, the core grid does not.** In both sweeps every candidate with the same
+  `in0_block_w` as the auto config (the whole K in one block for P·V; 8 for the decoder linears) was bit-identical;
+  smaller `in0_block_w` changed the P·V output by up to 7e-4 (the partial sums spill between K blocks).
+- **The 99-scene plan error reacts to any rounding change.** Both attention variants that only move one fp32 rounding
+  (where the 1/√32 is applied) raised the worst ego error by ~20 % (0.31 -> 0.38 m), with every gate still green. This
+  is the chaotic sensitivity of PORT_LOG known issues; bit-identical rewrites were preferred for that reason.
+- **`input_tensor_a_activations` is not a free fusion for fp32.** The scale pre-activation on the mask add is not
+  bit-identical to `multiply` + `add` (rel 4e-3 on `final_x0`), although both look like an fp32 SFPU multiply then
+  the same add.
+
 ### Megakernel / fusion work
 
-None yet. The candidates (D19) are in the backlog (item 6): the persistent DiT-evaluation kernel first.
+No megakernel yet; the candidates (D19) are in the backlog (item 6): the persistent DiT-evaluation kernel first. Round
+2's kernels are its building blocks (the fp32 LN with the residual phase, the split-row LN with semaphore gather /
+broadcast, the K-concatenated split operand, the score scale + mask). Round 3 added the first block-level fusion: the
+whole fp32 attention (Q·Kᵀ, scale / mask / softmax, P·V, head split / merge) as one bit-identical program
+(`ATTN_FUSED`), the attention phase of MK-D, plus producer-side operand emission (`KCAT_EMIT`) and the mixer
+transposes folded into the LN (`LN_TR`), which are the MK-E / MK-D data-layout pieces. The megakernel attempt keeps its
+own phase.
+
+## Round 2
+
+Date 2026-10-10, plan `OPT_PLAN.md` §7 round 2 (L2 `generic_op` kernels: items 3, 2, 4b, 5). The device was shared
+with METEOR hang-debugging jobs (two FAULT markers from their jobs, 09:19 and 10:05 UTC) and other agents' optimization
+jobs; none of this round's jobs hung, timed out or ran on a faulted chip (every job checks the marker after taking the
+lock). Every custom kernel went through the hang protocol before it touched the plan: smallest shape under
+`TT_METAL_WATCHER=2`, eager twice, then a traced replay, then the real shapes (`logs/diffusion-planner/opt_r2/`).
+
+Method, per step (`logs/diffusion-planner/opt_r2/`; scripts in `scripts/`): a device micro-check of the kernel against
+the stock programs it replaces (`ln_check.py`, `kcat_check.py`, `smask_check.py`: bit comparison, eager repeat,
+trace replay, µs per call over a 20-call trace), then `ab.sh <tag> <KNOB> <off> <on>` as in round 1 (device suite
+with the gates read-only, 3 alternating `bench.py` pairs with raw readback dumps, `cmp.py` bit comparison and the
+99-scene report). Every job runs from a snapshot of the code (`snap.sh`; scripts copied into it).
+
+| A/B (same session, 3 pairs, kashiwanoha_dense) | off: b2b / trace / e2e ms | on: b2b / trace / e2e ms | Δ b2b | outputs | kept |
+|---|---|---|---:|---|---|
+| `LN_KERNEL` 0 / 2 (LN_RESID=0, SPLIT_KCAT=0) | 70.12 / 70.18 / 81.6 | 54.58 / 54.64 / 65.6 | −15.54 | bit-identical | yes |
+| `LN_RESID` 0 / 1 (LN_KERNEL=2, SPLIT_KCAT=0) | 54.58 / 54.64 / 65.5 | 54.37 / 54.43 / 65.2 | −0.21 | bit-identical | yes |
+| `SPLIT_KCAT` 0 / 2 (LN_KERNEL=2, LN_RESID=1) | 54.37 / 54.44 / 65.4 | 47.39 / 47.45 / 58.3 | −6.98 | final_x0 rel 1.1e-3; 99 scenes ego max 0.313 -> 0.327 m, mean 0.143 -> 0.149 m, nb 0.086 -> 0.104 m; 44 / 44 green | yes |
+| `ATTN_SMASK` 0 / 1 (KCAT_ACT=0, LN_SPLIT=0) | 47.39 / 47.45 / 58.3 | 44.56 / 44.63 / 55.6 | −2.83 | bit-identical | yes |
+| `KCAT_ACT` 0 / 1 (ATTN_SMASK=1, LN_SPLIT=0) | 44.56 / 44.62 / 55.5 | 43.96 / 44.02 / 54.9 | −0.60 | bit-identical | yes |
+| `LN_SPLIT` 0 / 1 (ATTN_SMASK=1, KCAT_ACT=1) | 43.96 / 44.02 / 54.8 | 43.13 / 43.19 / 54.3 | −0.83 | bit-identical | yes |
+| `LN_SFPU_BCAST` 0 / 1 (all of the above on) | 43.13 / 43.19 / 53.9 | 40.89 / 40.95 / 51.9 | −2.24 | bit-identical | yes |
+
+AICLK 1350 MHz (min 1343) in every run; the device rows repeat to ±0.01 ms within a session.
+
+Kernel micro-benchmarks (µs per call, 20-call trace; stock programs -> the kernel; all bit-identical except the
+K-concatenated matmul, which is compared with its stock-op operand build):
+
+| kernel | shape | stock | kernel |
+|---|---|---:|---:|
+| fp32 LN (`LN_KERNEL`, one core per tile row) | decoder [352, 256] + affine / neighbour mixer [320·64, 128] / lane [140·64, 128] | 47.8 / 510.8 / 238.3 | 27.6 / 126.5 / 73.1 |
+| + residual and gate (`LN_RESID`) | decoder / neighbour mixer | 61.7 / 680.5 | 38.5 / 191.9 |
+| split-row LN (`LN_SPLIT`, 8 cores per row) | decoder, plain / + residual and gate | 47.9 / 61.7 | 26.8 / 29.4 |
+| split linear: 3 passes (`DEC_MMCFG`) vs K-concatenated (operand kernel + one matmul) | qkv 256->768 / out 256->256 / fc1 256->1024 / fc2 1024->256 | 62.1 / 50.0 / 70.4 / 70.7 | 27.2 / 22.4 / 31.2 / 62.5 |
+| score scale + mask (`ATTN_SMASK`) | 8 x 352 x 352 / 8 x 352 x 576 / 8 x 576 x 576 | 74.9 / 73.8 / 118.9 | 27.0 / 46.1 / 64.7 |
+
+### Findings (measured)
+
+- **Bit-identical fused kernels are possible for the precision decompositions** when the kernel issues the same LLK
+  calls in the same order as the stock programs and keeps every intermediate fp32 in L1 / DST: the fp32 LayerNorm
+  (accurate `ttnn.mean` = `add_binary_tile` fold + `sfpu_reduce` + `mul_unary_tile(1/W)`; binary_ng fp32 SFPU ops;
+  `rsqrt_tile<Default>`), the residual / gate adds, the GELUs and the attention scale reproduce the published numerics
+  bit for bit. That made 5 of the 6 kept steps exact rewrites (same readback, same 99 per-scene numbers).
+- **The stock fp32 + bf16 add is not an fp32 add.** `ttnn.add(scores_fp32, mask_bf16)` (the key mask of every fp32
+  matmul attention of the plan) runs on binary_ng's FPU path; its SrcA truncates the fp32 scores to TF32 (stock ==
+  exact & 0xFFFFE000 on every finite element, `smask_diag.log`; with an fp32 mask the add is exact). The published
+  graph therefore feeds TF32-truncated scores to the softmax. `ATTN_SMASK` reproduces that (an SFPU AND on the raw
+  bits); an exact variant (mask in fp32, or no truncation) would be a precision change to gate separately.
+- **One K-concatenated fp32 matmul replaces the three split passes** (and their typecast / subtract / add programs):
+  the products are the same (bf16-exact hi parts, TF32-truncated lo parts), only the accumulation order changes and
+  the bias becomes exact (two K rows instead of the packer epilogue, which rounds the fp32 output: TransFusion T4).
+  Per linear the rms error vs the fp64 product goes 3.2e-4 -> 3.1e-4; over the 99 scenes the worst scene moves
+  +4 % (0.313 -> 0.327 m), the others both ways (57 worse, 41 better, mean change +0.3 mm on the ego max). Kept: the
+  gates have 3x (max) / 2x (mean) margin, and it is the largest non-exact gain of the round (−7.0 ms).
+- **The K block must divide K'.** 3·Kt + 1 is prime for K = 1024 (97 tiles): the auto config then uses a 1-tile K
+  block (90 µs); padding K' to 104 = 8 x 13 with zero tiles and a 13-tile block gives 41 µs.
+- **The mixer LNs are DRAM-bound**, not compute-bound: the neighbour LN with its residual moves 42 MB (x, residual, h,
+  y; fp32) in ~190 µs. Removing that traffic needs the entity block resident in L1 (MK-E, item 7).
+- **The decoder LN is latency-bound** on 11 cores (one tile row each): 8 tiles per core of SFPU work in sequence. The
+  split-row form (8 cores per row, gather / fold / broadcast with semaphores) is bit-identical and 24 % faster with
+  the residual (38.5 -> 29.4 µs); its own floor is the root's two ordered 8-tile folds plus two broadcast rounds.
+- **A RISC-V fill loop is not free.** Broadcasting one fp32 statistic tile over its columns with 1,024 volatile L1
+  stores on a data-movement core sat on the critical path of every LN row twice. Storing the SFPU reduce's result
+  (already replicated across the 8 lanes by its butterfly) to all four column groups instead costs three more SFPU
+  stores: −2.24 ms over the plan (decoder split LN 29.4 -> 21.6 µs, neighbour mixer LN 191.9 -> 158.3 µs).
+
+### Profile (after r2.3, `d7f84e1`, one traced replay; `profile_r2/`)
+
+2,304 programs (from 6,282), kernel sum 46.52 ms, span 48.06 ms (b2b 47.39 ms). By op code: Matmul 16.28 ms (759
+programs), GenericOp (the new kernels) 13.34 ms (540), BinaryNg 7.80 ms (353), Softmax 3.48 ms (72), Unary 2.03 ms
+(191), Transpose 1.36 ms (84). Largest single items: the decoder LN programs on 11 cores (5.6 ms incl. the residual
+form; r2.6 then spread them over 88 cores), the mixer LNs (3.5 ms on 120 cores, DRAM-bound), the score scale + mask
+programs (4.8 ms decoder + 0.7 ms fusion; r2.4 removed half of that), the softmax (3.5 ms), the K = 1024 K-concatenated
+matmul (2.5 ms) and its operand build (1.6 ms, writes a 4.7 MB operand), the attention matmuls (~3.4 ms), the mixer
+transposes (1.3 ms) and the encoder pre-projection split glue (1.2 ms).
+
+## Round 3
+
+Date 2026-10-10, plan `OPT_PLAN.md` §0.2 (the ranked list after round 2: softmax fusion, encoder split linears, then
+the attention and operand traffic the round-3 profile pointed at). The device was shared with METEOR debugging jobs
+and other agents' optimization jobs; none of this round's jobs hung, timed out or ran on a faulted chip (every job
+checks the FAULT marker after taking the lock). Every new kernel or kernel mode went through the hang protocol before
+it touched the plan: smallest shapes under `TT_METAL_WATCHER=2`, eager twice, then the real shapes with a traced
+replay (`logs/diffusion-planner/opt_r3/*_tiny_watcher.log`, `*_real.log`).
+
+Method, per step (`logs/diffusion-planner/opt_r3/`; scripts in `scripts/`): a device micro-check of the kernel against
+the stock programs it replaces (`smsm_check.py`, `fattn_check.py`, `kemit_check.py`, `lntr_check.py`,
+`linact_check.py`: bit comparison, eager repeat, trace replay, µs per call over a 20-call trace), then `ab.sh <tag>
+<KNOB> <off> <on>` (device suite with the gates read-only and the per-scene report, 3 alternating `bench.py` pairs with
+raw readback dumps, `cmp.py` bit comparison and 99-scene comparison). Each A/B ran on the previous step's defaults
+(the knob of a not-yet-committed step set by the chain script); every job ran from a snapshot of the code.
+
+| A/B (same session, 3 pairs, kashiwanoha_dense) | off: b2b / trace / e2e ms | on: b2b / trace / e2e ms | Δ b2b | outputs | kept |
+|---|---|---|---:|---|---|
+| `ATTN_SMSM` 0 / 1 (round-2 defaults) | 40.89 / 40.96 / 51.9 | 39.32 / 39.39 / 50.3 | −1.56 | bit-identical | yes |
+| `ENC_KCAT` 0 / 1 | 39.32 / 39.38 / 50.3 | 37.62 / 37.68 / 48.4 | −1.70 | final_x0 rel 2.9e-3; 99 scenes ego max 0.327 -> 0.346 m, mean 0.149 -> 0.154 m, nb 0.104 -> 0.094 m; 44 / 44 green | yes |
+| `LIN_ACT` 0 / 1 (ENC_KCAT=0) | 39.32 / 39.38 / 50.3 | 39.19 / 39.25 / 50.1 | −0.13 | final_x0 rel 3.1e-3; ego max 0.327 -> 0.392 m, mean 0.149 -> 0.174 m; 44 / 44 green | no |
+| `ATTN_SMSM` 1 / 2 (ENC_KCAT=1) | 37.62 / 37.68 / 48.6 | 37.59 / 37.65 / 48.9 | −0.03 | bit-identical | no (noise level) |
+| `ATTN_FUSED` 0 / 1 | 37.62 / 37.68 / 48.9 | 32.19 / 32.25 / 43.2 | −5.43 | bit-identical | yes |
+| `KCAT_EMIT` 0 / 1 | 32.19 / 32.25 / 43.1 | 31.24 / 31.30 / 42.4 | −0.95 | bit-identical | yes |
+| `LN_TR` 0 / 1 (KCAT_EMIT=1) | 31.24 / 31.30 / 42.1 | 30.16 / 30.22 / 41.0 | −1.08 | bit-identical | yes |
+
+AICLK 1350 MHz (min 1343) in every run; the device rows repeat to ±0.01 ms within a session. Final check of the
+shipped defaults (`0b750e3`, no knob set; `logs/diffusion-planner/opt_r3/final/`): 44 / 44, readback and the 99
+per-scene numbers identical to the `LN_TR` on side; bench b2b 30.16 ms, trace p50 30.22 ms, e2e p50 40.84 ms.
+
+Kernel micro-benchmarks (µs per call, 20-call trace; every row bit-identical to the stock chain it replaces):
+
+| kernel | shape | stock chain | kernel |
+|---|---|---:|---:|
+| scale + mask + softmax (`ATTN_SMSM`) vs smask + `ttnn.softmax` | 8 x 352 x 352 / 8 x 352 x 576 / 8 x 576 x 576 | 62.0 / 101.6 / 149.9 | 50.6 / 73.2 / 106.1 |
+| fused attention (`ATTN_FUSED`) vs head split + Q·Kᵀ + smsm + P·V + head merge | decoder self (qkv in place) / cross (hoisted K, V heads) / fusion (q, kv in place) | 110.6 / 144.6 / 214.6 | 51.0 / 63.0 / 106.1 |
+| fused attention emitting the out-linear operand (`KCAT_EMIT`) vs fused attention + `kcat_operand` | decoder self | 57.9 | 51.7 |
+| split-row LN emitting the operand vs split-row LN + `kcat_operand` | [352, 256] | 26.2 | 22.1 |
+| LN writing its output per-entity transposed (`LN_TR`) vs LN + `ttnn.transpose` | neighbour [320, 64, 128] / lane [140, 64, 128] | 149.5 / 80.6 | 102.6 / 61.8 |
+| LN (+ residual) reading the residual transposed vs `ttnn.transpose` + LN | neighbour / lane | 200.1 / 102.6 | 147.2 / 81.1 |
+| mixer fc1 + GELU in the epilogue (`LIN_ACT`, not bit-identical) vs linear + GELU program | [40960, 64] @ [64, 64] / [20480, 128] @ [128, 128] | 107.8 / 106.4 | 95.8 / 108.7 |
+
+### Findings (measured)
+
+- **The attention was DRAM-bound, not compute-bound.** The fused attention costs about what the scale / mask /
+  softmax kernel alone cost (51.0 vs 50.6 µs self-attention), although it adds Q·Kᵀ and P·V: the time of the five
+  stock programs was the round trips of the `[1, 8, Sq, Sk]` fp32 scores and probabilities (4-8 MB each) and the
+  head split / merge copies. With the scores in L1 the per-core critical path is the stock softmax arithmetic on one
+  tile row; the decoder attention went from ~255 to ~114 µs per block.
+- **Stock LLK sequences compose bit-identically inside one kernel.** The `kernel_lib` softmax calls (reduce,
+  bcast-sub + exp, reduce + precise reciprocal, bcast-mul) work on plain CB ids in a `generic_op` and give the stock
+  softmax bit for bit; `matmul_block` with the in1 transpose flag reproduces the stock Q·Kᵀ (K = 32 in one block) and
+  a key-ordered DEST accumulation reproduces the stock P·V (the whole K in one block); the score value never needs to
+  leave DEST between Q·Kᵀ and the scale / mask SFPU ops. The global `UnpackToDestEn` (set when any CB unpacks to DEST)
+  did not perturb the FPU phases.
+- **The fp32 SFPU multiply by an immediate equals the multiply by a filled tile** (`mul_unary_tile` vs
+  `mul_binary_tile`, fp32 DEST): bit-identical, but only 0.03 ms of the plan; kept as the fused attention's form.
+- **Producers can emit the consumer's operand.** Writing `[x_hi | x_hi | x_lo | 1]` from the split LN / fused attention
+  (the `kcat_operand` LLK calls on the packed fp32 tile) costs 1-2 µs per producer and removes a 6.8 µs program plus
+  its 0.6 µs gap: −0.95 ms over 198 programs.
+- **The stock fp32 `ttnn.transpose` is exact** (unpack-to-dest + `transpose_dest`), so the mixer's transposes fold
+  into the LN kernels with `transpose_tile` bit for bit; the scattered tile writes cost nothing measurable (the LN is
+  DRAM-bound either way): −1.08 ms.
+- **A more accurate GELU still moves the plan.** `LIN_ACT` applies the GELU before the bf16 rounding (per-linear rms
+  error 2.06e-3 -> 1.66e-3 vs fp64), yet the 99-scene ego mean rises 0.149 -> 0.174 m: the chaotic scene sensitivity
+  of PORT_LOG known issues again; at −0.13 ms it is not worth that.
+- **The encoder K-concatenation pays off more than estimated** (−1.70 vs −1 to −2 ms): beyond the 3-pass glue, the
+  pre-projection matmuls on padded `[1, E, T, K]` rows read the operand once instead of three times.
+
+## Round 4
+
+Date 2026-10-10, plan `OPT_PLAN.md` §0.3 (item 1, the K = 1024 split operand traffic, then what the measurements of
+that step pointed at: the decoder's DRAM round trips). The device was shared with METEOR and other bundles' jobs;
+no job of this round hung, timed out or ran on a faulted chip (each job checks the FAULT marker after taking the
+lock). New kernel modes went through the hang protocol (smallest shapes under `TT_METAL_WATCHER=2`, eager twice,
+then the real shapes with a trace replay: `logs/diffusion-planner/opt_r4/*_tiny_watcher.log`, `*_real.log`).
+
+Method as in round 3 (`logs/diffusion-planner/opt_r4/`, scripts in `scripts/`): a device micro-check per kernel
+mode (`kl1_check.py`, `kl1_sweep.py`, `ko_check.py`, `fl1_check.py`, `dl1_check.py`: bit comparison with the shipped
+path, eager repeat, trace replay, µs per call over a 20-call trace), then `ab.sh <tag> <KNOB> <off> <on>` (device
+suite with the gates read-only and the per-scene report, 3 alternating `bench.py` pairs with raw readback dumps,
+`cmp.py`). Each A/B ran on the previous step's defaults, every job from a snapshot of the code.
+
+| A/B (same session, 3 pairs, kashiwanoha_dense) | off: b2b / trace / e2e ms | on: b2b / trace / e2e ms | Δ b2b | outputs | kept |
+|---|---|---|---:|---|---|
+| `KCAT_L1` 0 / 1 (round-3 defaults) | 30.16 / 30.22 / 40.78 | 28.53 / 28.59 / 39.30 | −1.63 | bit-identical | yes |
+| `KCAT_ACT_ONCE` 0 / 1 | 28.53 / 28.62 / 44.4 | 28.59 / 28.66 / 43.6 | +0.06 | bit-identical | no |
+| `ATTN_L1` 0 / 1 | 28.53 / 28.59 / 39.33 | 28.12 / 28.18 / 38.93 | −0.41 | bit-identical | yes |
+| `DEC_L1` 0 / 1 | 28.12 / 28.18 / 38.83 | 27.11 / 27.16 / 37.92 | −1.01 | bit-identical | yes |
+| `ENC_L1` 0 / 1 | 27.11 / 27.16 / 37.70 | 28.77 / 28.84 / 39.49 | +1.67 | final_x0 rel 8e-3, ego max 0.346 -> 0.379 m | no |
+| `DEC_L1` 1 / 2 | 27.11 / 27.17 / 37.91 | 27.02 / 27.07 / 37.90 | −0.09 | bit-identical | yes |
+| `FUS_L1` 0 / 1 | 27.02 / 27.07 / 37.76 | 27.40 / 27.47 / 38.16 | +0.39 | final_x0 rel 5e-4, ego max 0.346 -> 0.364 m | no |
+| `ATTN_L1` 1 / 2 | 27.02 / 27.08 / 37.82 | 26.62 / 26.69 / 37.71 | −0.39 | bit-identical | yes |
+
+AICLK 1350 MHz (min 1343) in every run; the device rows repeat to ±0.01 ms within a session (the `KCAT_ACT_ONCE`
+e2e column ran on a busier host). Final check of the shipped defaults (`93428bb`, no knob set;
+`logs/diffusion-planner/opt_r4/final/`): 44 / 44, readback and the 99 per-scene numbers identical to round 3; bench
+b2b 26.62 ms, trace p50 26.68 ms, e2e p50 37.36 ms.
+
+Kernel micro-benchmarks (µs per call, 20-call trace; every row bit-identical to the shipped path):
+
+| check | shape | DRAM (shipped) | L1 |
+|---|---|---:|---:|
+| K = 1024 operand build + matmul (`KCAT_L1`; best of 7 / 6 sharded configs) | `[352, 1024]` -> X' `[352, 3328]` @ `[3328, 256]`, 2d M2 N1 kb13 | 62.6 (23.4 + 41.4) | 41.3 (19.0 + 24.2) |
+| same | final p4 N = 324: K' 104 -> 99 tiles (11 slices), 2d M2 N1 kb9 | 62.3 | 41.8 |
+| K = 256 operand L1-sharded (`kl1_sweep.py`: K' padded to 26-32 to match the N blocks) | N = 256 / 768 / 1024, matmul only | 15.3 / 20.4 / 24.2 | 14.2 / 29.7 / 34.7 (best) |
+| operand build, GELU once (`KCAT_ACT_ONCE`) | K = 1024 / 512, gelu_tanh | 23.2 / 14.8 | 23.5 / 13.2 |
+| fused attention, inputs (+ mask) in L1 (`ATTN_L1`) | self 8 x 352 x 352 / cross 8 x 352 x 576 / fusion 8 x 576 x 576 | 51.0 / 61.0 / 106.2 | 35.1 / 54.2 / 88.4 |
+| split-row LN + residual + gate + emitted operand, inputs and outputs in L1 (`DEC_L1`) | `[352, 256]` | 24.1 | 20.6 |
+| K = 256 K-concatenated matmuls, in0 and output in L1 interleaved (`DEC_L1`) | N = 256 / 768 / 1024 | 14.9 / 19.7 / 24.1 | 12.7 / 15.6 / 20.0 |
+| fused self attention writing the operand to L1 (`DEC_L1`) | 8 x 352 x 352 | 37.6 | 36.0 |
+
+### Findings (measured)
+
+- **The decoder was bound by DRAM round trips of small tensors, not by compute.** The K = 1024 operand (4.7 MB)
+  costs as much to write and re-read as the matmul costs to compute: written to L1 shards and read in place, the
+  matmul alone drops 41.4 -> 24.2 µs. The same holds, at smaller scale, for every decoder intermediate (0.4-1.4 MB):
+  L1-interleaved inputs and outputs take 2-4 µs off each split LN, matmul and attention program.
+- **K / V re-reads.** The fused attention's 11 query-row cores of a head each read all K / V tiles; from L1 banks
+  instead of DRAM the self attention drops 51 -> 35 µs. No multicast kernel was needed for most of the estimated gain.
+- **L1 is not faster for large tensors.** The mixer's 5-10 MB intermediates in L1 interleaved made the trunk +1.67 ms
+  slower (120 cores reading every bank: NoC contention instead of 8 DRAM channels).
+- **Stock ops change program (and numerics) with an L1 output.** `ttnn.linear` with an L1 output ran the fusion kv
+  projection on 6 cores with a separate bias add (54 vs 16 µs; bit-identical here), and the mixer / fusion stock
+  matmuls, unary GELU and `ttnn.layer_norm` with L1 outputs were not bit-identical. L1 residency is safe and exact
+  for the bundle's own `generic_op` kernels and for matmuls with explicit program configs; stock ops need a profile
+  after the change (the first round-4 profile caught the 6-core linear).
+- **The operand build is write-bound at K = 1024**: running the GELU once instead of twice saves 1.6 µs at K = 512
+  but nothing at K = 1024 (not kept).
+- **Uneven block shards work**: X' `[352, ·]` over 6 shard rows of 64 (the last one half full) is read correctly by
+  the sharded-in0 matmul and written correctly through the TensorAccessor of the generic_op writer.
+
+## Round 5
+
+Date 2026-10-10, plan `OPT_PLAN.md` §0.4 (item 1, exact compaction, decoder buckets first; item 2, the host items).
+The device was shared with METEOR and other bundles' jobs (several of this round's jobs waited for the lock behind
+a CenterPoint A/B); no job of this round hung, timed out or ran on a faulted chip. One gate job failed fast with a
+host-side `TT_FATAL` (a slice end past the trimmed `cs` width: `ab_trim/gates_fail_slice.log`), fixed before any
+A/B. The one new kernel shape (the split-row fp32 LN at W = 1024, Wt = 32 members per row, which the r32 / r64 / r96
+buckets use for the final-layer LN) went through the hang protocol first: `TT_METAL_WATCHER=2`, rows 32 and 96,
+eager twice and a 2-call trace, bit-identical to the one-core-per-row kernel (`opt_r5/ln1024_bringup.log`).
+
+Method (`logs/diffusion-planner/opt_r5/`, scripts in `scripts/`): `compact_check.py` (one process with the bucket
+traces: every scene through the full `plan` and through its bucket, final_x0 on the needed rows, logits and the 11 ego
+iterates compared bit for bit, then b2b of every variant), `host_equiv.py` (old and new host package side by side on
+134 inputs: the `prepare` dataclasses, `plan_inputs` and six `make_output` calls per input with random and
+force-stop trajectories and two parameter sets, compared byte for byte), then per step the device suite with the
+gates read-only and the per-scene report, and 3 alternating `bench.py` pairs on 3 scenes (`ab_*/`).
+
+| A/B (same session, 3 pairs; kashiwanoha_dense / straight_road / scene-0103_kf14) | off: b2b ms | on: b2b ms | off -> on e2e p50 (kashiwanoha) | outputs | kept |
+|---|---|---|---|---|---|
+| `COMPACT` 0 / 1 (round-4 defaults) | 26.62 / 26.62 / 26.62 | 23.25 (r96) / 21.85 (r32) / 22.91 (r64) | 38.28 -> 33.79 | bit-identical on the read rows | yes |
+| `COMPACT` 1 / 2 | 23.25 / 21.85 / 22.91 | 19.99 / 17.42 / 19.23 | 33.21 -> 30.14 | bit-identical on the read rows | yes |
+| `HOST_FAST` 0 / 1 | 19.99 / 17.42 / 19.23 | 19.99 / 17.42 / 19.23 | 30.15 -> 27.31 (straight_road 27.05 -> 23.87, scene-0103 29.20 -> 26.22) | bit-exact host arrays | yes |
+| `INPUT_TRIM` 0 / 1 | 19.99 / 17.42 / 19.23 | 20.00 / 17.44 / 19.25 | 27.22 -> 26.84 (host_in 1.95 -> 1.58, h2d 0.94 -> 0.76, pack 0.38 -> 0.52) | readback bit-identical | yes |
+
+AICLK 1350 MHz (min 1343) in every run; the device rows repeat to ±0.003 ms within a session. The last A/B's `on`
+runs are the shipped defaults (`e933166`): 44 / 44 (`ab_trim/gates.log`), the 99 per-scene numbers identical to
+rounds 3 / 4.
+
+Per-bucket device time (`compact2_check.log`, b2b, before `INPUT_TRIM`'s +0.014 ms):
+
+| variant | rows (decoder agents / neighbour entities) | b2b ms | 99 gate scenes in it |
+|---|---|---:|---:|
+| `plan_r32` | 32 | 17.42 | 53 |
+| `plan_r64` | 64 | 19.24 | 44 |
+| `plan_r96` | 96 | 19.99 | 2 |
+| `plan_r128` / `plan_r192` | 128 / 192 | 21.11 / 22.30 | 0 |
+| `plan` | 352 / 320 | 26.65 | 0 |
+
+### Findings (measured)
+
+- **Compaction is exact here and needs no new numerics.** Valid agents and valid neighbour tokens are a prefix in
+  all 99 scenes, but the bucket does not rely on it: it covers the last needed row. The rows past R
+  are masked self-attention keys (exp gives exact zeros, added in the same tile order) and invalid neighbour tokens
+  (multiplied by `token_valid` = 0), so dropping them changes no read value. Every derived matmul config keeps the
+  swept K blocking; only `per_core_M` is fitted, which leaves each output element's accumulation unchanged.
+- **The decoder is latency-bound, not row-bound.** From 352 to 96 rows the decoder blocks fall only 12.0 -> 9.5 ms
+  (the profile below): its 616 programs per plan run at 8-43 µs each almost regardless of the row count (split-row LN
+  16.7 µs on 24 cores, self attention 15.9 µs, cross attention 43.3 µs on 24 cores, K = 256 matmuls 7.8-14.6 µs).
+  Smaller buckets mostly remove the encoder's neighbour work (3.9 -> 1.5 ms for the trunk at r96) and the K = 1024
+  operand traffic. What remains is per-program latency: the MK-D megakernel's target.
+- **Host time halves with plain numpy.** The node-faithful loops (trajectory smoothing, per-tensor normalisation
+  with boolean-mask assignment, all 320 neighbours denormalised and converted) cost 6.8 ms; the same arithmetic
+  vectorised in the same float64 / float32 order costs 3.9 ms, byte for byte equal. `math.hypot` with three arguments
+  stays a Python loop over 80 points (`np.hypot` rounds differently in the last bit).
+- **Upload trimming is small.** The tilization of the 19 inputs is ~1.6 ms of `host_in`; the all-zero test of every
+  input costs 0.14 ms of `pack`, so the net gain is 0.37 ms e2e.
 
 ## Known hangs (all rounds)
 
 | when (UTC) | command | cause | status |
 |---|---|---|---|
-| – | – | none: no hang, timeout, reset or FAULT marker in any device job of the port, the baseline or the release docs | – |
+| – | – | none: no hang, timeout, reset or FAULT marker in any device job of the port, the baseline, the release docs or rounds 1-5 (the FAULT markers of 2026-10-10 came from METEOR jobs of another agent; this bundle's queued jobs waited and resumed after the orchestrator cleared them). Round 3: one profile job was cancelled from this session before it ran (a `pkill` of its wrapper script); no FAULT, no device effect | – |
 
 ## Rejected / not kept
 
-None yet. Configurations that were measured and are **not** the default, for accuracy (PORT_LOG decision 11,
+Round 4: `KCAT_ACT_ONCE` (bit-identical, +0.06 ms), `ENC_L1` (+1.67 ms, not bit-identical), `FUS_L1` (+0.39 ms, not
+bit-identical), `ATTN_L1=1` (the fusion q / kv linears with L1 outputs on 6 cores; superseded by `ATTN_L1=2`), the
+K = 256 split operands L1 block-sharded (`kl1_sweep.py`: −1.1 µs at N = 256 with K' padded 25 -> 32, slower at N =
+768 / 1024; not wired). All stay selectable (knobs) except the last.
+
+Round 3: `LIN_ACT` (the mixer fc1 GELU in the matmul epilogue: −0.13 ms, 99-scene ego mean 0.149 -> 0.174 m, max
+0.327 -> 0.392 m; default 0, selectable) and `ATTN_SMSM=2` (the scale by an immediate in the smsm kernel:
+bit-identical, −0.03 ms; superseded by `ATTN_FUSED`, which uses that form).
+
+Round 2: every A/B'd step was kept. Measured on the way and not used: the first `ATTN_SMASK` kernel (an exact fp32
+add of the mask) was not bit-identical, because it drops the TF32 truncation the stock fp32 + bf16 add applies
+(`smask_real.log`: every unmasked score differs); it was replaced by the truncating form before any A/B, and an
+exact-score variant remains a separate precision experiment. `LN_KERNEL=1` (per-tile unpacker / SFPU inits) is 2-4 %
+slower than `LN_KERNEL=2` at kernel level (`ln_real.json` vs `ln_lean.json`) and stays selectable.
+
+
+Round 1 (measured, gates green, not kept because they are not bit-identical and raise the 99-scene error):
+
+- `ATTN_FAST=3`, the attention scale as an SFPU pre-activation of the mask add (one program instead of two over the
+  `[1, 8, Sq, Sk]` scores; sweep 75 -> 42 µs per call): −2.5 ms over mode 1, ego max 0.385 m / mean 0.173 m;
+- `ATTN_FAST=2`, the scale on Q before Q·Kᵀ: −1.97 ms, ego max 0.379 m / mean 0.172 m, neighbours 0.098 m.
+
+Both stay selectable for timing. The round-2 fused attention kernel replaces these pieces anyway.
+
+Earlier, configurations that were measured and are **not** the default, for accuracy (PORT_LOG decision 11,
 `OPT_BASELINE.md` "What the numerics defaults cost"):
 
 - the first device round's defaults (split only for the mixer inputs and the decoder pre-projection, bf16 SDPA):
@@ -98,8 +519,131 @@ None yet. Configurations that were measured and are **not** the default, for acc
 - the fastest graph (fused LayerNorm, no split, bf16 SDPA): 43.7 ms, fails the encoder PCC gate (`enc.ego` 0.99856 /
   0.99261).
 
+## Republish (2026-10-11)
+
+Republish prep of the optimized release (no optimization step; evidence under `logs/diffusion-planner/republish/`):
+
+| item | result |
+|---|---|
+| ttaw | re-vendored 0.23.0 -> 0.23.2 (common `60b6dd7`: Resize2d tail and `reshape_patch_present`, neither used by this model), `c91ce91`; host suite 108 passed / 45 skipped |
+| tt-metal patches | the shared tree carries `patches/tt-metal-reshape-rm-sys1419.patch` since 2026-10-10 (every round-2..5 device number was measured with it); the bundle now ships it and the image build asserts its marker |
+| device suite (gates read-only) on `bc966b4` | 44 passed; again 44 passed under `TT_METAL_TRACE_ALLOC_TRACKING=1`; 99 per-scene numbers identical to round 5 (`opt_r5/e2e_compact2.json`, 0 of 99 differ) |
+| stage bench (100 iterations) | kashiwanoha_dense r96: b2b **20.00 ms**, one plan 20.05 ms, e2e p50 26.73 ms (p99 27.48), `.npz` path 32.27 ms; straight_road r32: 17.44 / 17.50 / 23.13 ms; scene-0103_kf14 r64: 19.25 / 19.30 / 25.65 ms; AICLK 1350 MHz (min 1343) |
+| **fix `7e003c4`: `dispatch="worker"` crashed at capture** | found by this re-measure: the swept 2-D multicast configs of rounds 1-5 assume the 12-column ETH grid (decoder qkv: 12 output blocks on x), so the 11×10 WORKER grid, the A/B opt-in and the fallback without the ETH patch, failed with `TT_FATAL: Num output blocks along x (12) must be smaller than or equal to the number of columns in compute grid (11)`. Now each config is checked against `compute_with_storage_grid_size()` and falls back to the auto config (`mcast_fits`, host test `test_grid_fit_host.py`). ETH: unchanged configs, device suite 44 passed, 99 numbers identical. WORKER: `test_e2e_device.py` 2 passed with the same 99 numbers as ETH; b2b 20.61 ms (r96) / 18.01 ms (r32), 3 % slower than ETH |
+| load / served | `from_pretrained` 352 s with an empty JIT cache (6 traces), 12.9 s warm; served `/predict` 37.2 ms median (`model()` 26.8 ms, JSON decode 8.6 ms), client round trip 41.6 ms, READY 13.0 s on the host; the bundle's `smoke_test.py` passes against the host server |
+| accuracy re-checks | ORT oracle: 92 nuScenes instants worst ego max 0.346 m / mean 0.154 m, turn 92 / 92; scene-0061 sequence 0.112 / 0.055 m, 33 / 33; open loop ADE / FDE 4.75 / 11.87 m (CPU 4.74 / 11.87), blinker 88 / 92; neighbours of the shipped sample: per-agent max 8.6 cm median, 0.23 m p90, 3.72 m worst |
+| media | the 13 renders of `media/` re-rendered from the optimized p150 outputs |
+
 ## Profile at the end (trace replay)
 
+**After round 5** (`e933166`, the shipped defaults, the bench scene's bucket `plan_r96`;
+`logs/diffusion-planner/opt_r5/profile_r5final/`: `analyze.md`, `breakdown.md`, `ops_perf_results.csv.zst`): **1,418
+programs**, kernel sum **19.64 ms**, op-to-op gaps 0.89 ms, span 20.54 ms at 1.35 GHz (b2b 20.00 ms; replay ==
+eager). Stage groups: decoder blocks 9.49 ms (11 evaluations at 96 rows; 12.01 at 352), mixer blocks 5.66 ms (8.12),
+fusion 1.36 ms, mixer pre-projections 1.28 ms (2.08), decoder pre-projection 0.66 ms (1.03), solver 0.31 ms, the rest
+0.9 ms.
+
+Largest remaining items (kernel time, share of the 20.5 ms span):
+
+| item | ms | share | what limits it / next step |
+|---|---:|---:|---|
+| decoder split-row LNs with residual (154 programs on 24 cores, 16.7 µs) + final LNs | 2.58 + 0.34 | 14.2 % | gather / fold / broadcast latency, row-count independent; MK-D |
+| map mixer trunks: lanes (140 entities) 2.41, line strings (60) 1.19, route (25) 0.61 | 4.21 | 20.5 % | full capacity in every scene (kashiwanoha: 123 / 60 / 17 valid; p50 scenes 100 / 30 / 6): map-entity compaction (est. −1.5 to −2 ms at p50 scenes, ~−0.4 ms on the bench scene) |
+| mixer LNs (36 programs, 120 cores, 59 µs) | 2.13 | 10.4 % | DRAM-bound; MK-E |
+| decoder cross attention (33 x 43.3 µs on 24 cores: 8 heads x 3 query tiles, 18 key tiles each) | 1.43 | 7.0 % | one core per (head, query tile): a key-split form with an ordered gather (bit-identical like the split LN) could use 4x the cores (est. −0.8 ms) |
+| decoder K = 1024 split linears: matmuls 1.45 + operand builds 0.83 | 2.28 | 11.1 % | latency of 24-33-core programs at 96 rows; MK-D |
+| decoder K = 256 split linears (qkv, out, mlp fc1, cross q / out) | 2.26 | 11.0 % | 7.8-14.6 µs per program; MK-D |
+| mixer fc1 GELU programs | 0.71 | 3.5 % | bit-identical fusion needs a matmul epilogue that rounds to bf16 before the GELU |
+| fusion attention (6 x 94 µs) | 0.57 | 2.8 % | – |
+
+**After round 4** (`93428bb`, the shipped defaults; `logs/diffusion-planner/opt_r4/profile_r4final/`: `analyze.md`,
+`breakdown.md`, `ops_perf_results.csv.zst`): **1,411 programs**, kernel sum **26.00 ms**, op-to-op gaps 0.89 ms,
+span 26.89 ms at 1.35 GHz (b2b 26.62 ms). Stage groups: decoder blocks 12.01 ms (11 evaluations; 15.61 after round
+3), mixer blocks 8.12 ms, mixer pre-projections 2.08 ms, fusion 1.35 ms, decoder pre-projection 1.03 ms, solver 0.39
+ms, the rest 1.1 ms. (The first round-4 profile, `profile_r4final_c3b6a5a/`, showed the fusion at 1.74 ms: the
+6-core fusion linears fixed by `ATTN_L1=2`.)
+
+Largest remaining items (kernel time, share of the 26.9 ms span):
+
+| item | ms | share | what limits it / next step |
+|---|---:|---:|---|
+| decoder split-row LNs (4 per block) + fused self attention (88-core GenericOp group) | 4.02 | 14.9 % | gather / ordered fold / broadcast latency (~17-20 µs per LN); MK-D |
+| mixer LNs (36 programs, 120 cores) | 3.11 | 11.6 % | DRAM-bound; L1 was slower (`ENC_L1`); MK-E |
+| decoder K = 256 K-concatenated matmuls (qkv, out, mlp fc1, cross q / out) | 2.82 | 10.5 % | 10-18 µs on 48-72 cores; fp32 HiFi4 compute |
+| decoder K = 1024 split linears: matmuls 1.61 + operand builds 1.27 | 2.88 | 10.7 % | the operand build is write-bound (3 tiles per input tile, gelu_tanh); in-kernel split matmul (MK-D phase) |
+| mixer channel / token matmuls | ~2.6 | 9.7 % | bf16 / fp32 linears at 60-77 µs on 107-117 cores; MK-E |
+| decoder cross attention | 1.11 | 4.1 % | 33.6 µs per call |
+| mixer fc1 GELU programs | 1.04 | 3.9 % | bit-identical fusion needs a custom matmul epilogue (round to bf16 before the GELU) |
+| 1024-wide final-layer LN (11 cores) | 0.66 | 2.4 % | a multi-tile split form (est. −0.3 ms: the 32-tile fold stays serial on the root) |
+| pre-projections (operand builds 0.52, island glue 0.28, matmuls) | 2.08 | 7.7 % | – |
+
+**After round 3** (`0b750e3`, the shipped defaults; `logs/diffusion-planner/opt_r3/profile_r3final/`: `analyze.md`,
+`breakdown.md`, `ops_perf_results.csv.zst`): **1,405 programs**, kernel sum **29.71 ms**, op-to-op gaps 0.88 ms,
+span 30.59 ms at 1.35 GHz (b2b 30.16 ms). Stage groups: decoder blocks 15.61 ms (11 evaluations), mixer blocks
+8.15 ms, mixer pre-projections 2.11 ms, fusion 1.41 ms, decoder pre-projection 1.05 ms, solver 0.39 ms, the rest
+1.0 ms. By op code: GenericOp 14.18 ms (438 programs), Matmul 11.90 ms (567), Unary 1.44 ms (114, mostly the mixer
+fc1 GELU), BinaryNg 1.08 ms (137), Transpose 0.21 ms (36, from 84).
+
+Largest remaining items (kernel time, share of the 30.6 ms span):
+
+| item | ms | share | what limits it / next step |
+|---|---:|---:|---|
+| decoder K = 1024 split linears (mlp fc2 x 2, final p4): matmul `[352, 3328] @ [3328, N]` on 48-66 cores + the operand build after the tanh GELU | 2.95 + 1.61 | 14.9 % | in0 is 4.7 MB of fp32 written and read per call; L1-sharded operand, bf16 `w_lo` (precision change), or the in-kernel split matmul (MK-D building block) |
+| decoder split-row LayerNorms (4 per block, 88 cores) + the 1024-wide final LN (11 cores) | 2.75 + 0.40 + 0.67 | 12.5 % | latency-bound (gather / ordered fold / broadcast, ~20 µs per LN); MK-D |
+| decoder fused attention (self 49, cross 59 µs) + fusion (103 µs) | 3.58 + 0.62 | 13.7 % | the stock softmax arithmetic on one tile row per core (88 of 120 cores busy); K / V re-read by 11 cores per head |
+| decoder K = 256 split linears (qkv, out, mlp fc1, cross q / out) | 3.50 | 11.4 % | small matmuls (12-22 µs on 48-72 cores, configs from the round-2 sweep) |
+| mixer LNs (24 programs, 120 cores) | 3.13 | 10.2 % | DRAM-bound (x, residual, h, y fp32); MK-E (entity block resident in L1) |
+| mixer channel / token matmuls | ~2.6 | 8.5 % | bf16 / fp32 linears at 60-77 µs on 107-117 cores |
+| mixer fc1 GELU programs | 1.05 | 3.4 % | a bit-identical fusion needs a custom matmul epilogue (round to bf16, then GELU); `LIN_ACT` (not bit-identical) rejected |
+| pre-projections (12 operand builds, island glue, matmuls) | 2.11 | 6.9 % | – |
+
+**After round 2** (`c42c540`, the shipped defaults; `logs/diffusion-planner/opt_r2/profile_r2final/`: `analyze.md`,
+`breakdown.md`, the per-program CSVs as `.zst`): **2,155 programs**, kernel sum **39.99 ms**, op-to-op gaps 1.47 ms,
+span 41.46 ms at 1.35 GHz (b2b 40.90 ms). Stage groups: decoder blocks 22.45 ms (11 evaluations), mixer blocks
+9.21 ms, mixer pre-projections 3.59 ms, fusion 2.31 ms, decoder pre-projection 1.05 ms, solver 0.38 ms, the rest
+1.0 ms.
+
+| op code | after round 1 ms | after round 2 ms (programs) |
+|---|---:|---:|
+| Matmul | 18.1 | 16.28 (759) |
+| GenericOp (the round-2 kernels) | – | 12.91 (612) |
+| Softmax | 3.5 | 3.48 (72) |
+| BinaryNg | 31.5 | 2.29 (209) |
+| Unary / Typecast | 2.8 / 2.4 | 1.44 / 0.42 |
+| Transpose (mixer token mixing) | 1.3 | 1.36 |
+| Reduce (LN means) | 3.7 | 0.09 |
+
+Largest remaining single items: the decoder softmax (2.97 ms, 66 x 45 µs on 88 cores), the mixer LN programs (2.88 ms
+on 120 cores, DRAM-bound), the K = 1024 K-concatenated matmuls (2.94 ms) and their operand builds (1.62 ms), the
+decoder LN programs (2.38 ms on 88 cores + 0.67 ms for the 1024-wide final-layer LN, which stays on 11 cores: 352
+rows x 32 tiles do not fit the split form), the score scale + mask passes (2.29 ms decoder + 0.37 ms fusion), the
+attention matmuls (3.40 ms), the encoder pre-projection split glue (1.23 ms BinaryNg) and the mixer transposes
+(1.16 ms).
+
+
+**After round 1** (`dd6f249`, one traced replay under the device profiler, `logs/diffusion-planner/opt_r1/
+scripts/profile.sh`; numbers from the device-side per-program report `profile_r1/cpp_device_perf_report.csv.zst`,
+because the Tracy post-processing ran out of disk, see below): **6,282 programs** (unchanged: round 1 changed program
+configs, not the op graph), **kernel sum 65.06 ms** (was 99.27), span 70.96 ms at 1.35 GHz (b2b 70.12 ms).
+
+| op code | baseline ms | after round 1 ms |
+|---|---:|---:|
+| BinaryNg (fp32 LN decomposition, split-matmul glue, attention scale / mask, residual adds, adaLN gates, solver) | 31.5 | **31.5** (48 % of the kernel sum) |
+| Matmul | 52.3 | **18.1** |
+| Reduce (LN means) | 3.7 | 3.7 |
+| Softmax | 3.5 | 3.5 |
+| Unary / Typecast | 2.8 / 2.4 | 2.8 / 2.4 |
+| Transpose (mixer token mixing) | – | 1.3 |
+
+The element-wise decompositions that the gates need are now the larger half of the plan: that is what the round-2
+L2 kernels (fused split matmul, fused fp32 LayerNorm, fp32 attention kernel; `OPT_PLAN.md` items 2-4) remove.
+
+Disk incident: the Tracy run of the 6,282-program replay wrote 17 GB of raw logs (`tracy_ops_times.csv` 9.0 GB,
+`profile_log_device.csv` 7.4 GB) and filled the root disk while copying them into the report folder (09:06 UTC, a
+few minutes at 100 %). The partial copy was deleted at once and the raw logs compressed with zstd (17 GB -> 1.4 GB,
+`generated/profiler/diffusion-planner_opt_r1/.logs/`). Future profiles of this model: check free space first (> 25
+GB), and profile the replay only (the eager signposted pass doubles the logs).
+
 The baseline profile (`OPT_BASELINE.md`): 6,282 programs, kernel sum 99.27 ms, span 103.66 ms; by op code Matmul
 52.3 ms (1,365 programs), BinaryNg 31.5 ms (2,948), Reduce 3.7 ms, Softmax 3.5 ms, Unary 2.8 ms, Typecast 2.4 ms. The
 ten slowest programs are the neighbour mixer's channel matmuls `[1, 320, 64, 128] @ [128, 128]`, ~970 µs each on 8
@@ -107,6 +651,11 @@ cores.
 
 ## Remaining backlog (gains on the 102.0 ms replay unless marked e2e; estimates, not measurements)
 
+Updated ranking on the 20.0 ms base (bucket r96) after round 5: `OPT_PLAN.md` §0.5. After round 4 it was §0.4. After round 3 it was §0.3 (the K = 1024 operand traffic, the decoder
+LN latency, then compaction, MK-E and MK-D; the table "After round 3" above). After round 2 it was §0.2 (all done in
+round 3 except the K = 1024 operand).
+The list below is the original one; items 1-5 are done or partly done (rounds 1-2).
+
 From `OPT_BASELINE.md` "Ranked optimization opportunities". The gains are not additive: each item shrinks the base of
 the next.
 

PYTHON.md

@@ -68,7 +68,7 @@ the three v5.0 ONNX files and `diffusion_planner.param.json`, with an offline fa
 file and the weights' `major_version == 5` are checked), opens the chip (ETH dispatch, 12×10, 1 CQ; the other open
 parameters are `DEVICE_DEFAULTS` in `tt_diffusion_planner/device.py`, overridable with `DIFFUSION_PLANNER_*`), reads the
 ONNX initializers as data, uploads the weights and constants (48.1 MB), builds the graph, then compiles and captures
-the metal trace. If ETH dispatch cannot open (tt-metal without the patch), it warns and falls back to WORKER dispatch
+the metal traces. If ETH dispatch cannot open (tt-metal without the patch), it warns and falls back to WORKER dispatch
 (`model.info["device"]["fallback"]` names it). Any other keyword argument is a `TypeError`.
 
 The numerics are not arguments: the published configuration is the default of the `DIFFUSION_PLANNER_LN_FP32`,
@@ -78,9 +78,9 @@ accuracy figures of the card until the gates are re-run.
 
 ## Warm-up
 
-`from_pretrained` returns a warm model: it builds the graph, runs the plan once eagerly (this first run compiles every kernel into the JIT cache), then captures the whole plan as one metal trace (`warmup_variants="default"`: the variant `plan`) with program-cache misses forbidden, so no later call compiles anything. `model.warmup()` is idempotent; `warmup_variants="none"` defers the capture to `model.warmup()`.
+`from_pretrained` returns a warm model: it builds the graph, runs the plan once eagerly (this first run compiles every kernel into the JIT cache), then captures the whole plan as metal traces (`warmup_variants="default"`: the full-capacity variant `plan` and one variant per agent bucket, `plan_r32`, `plan_r64`, `plan_r96`, `plan_r128`, `plan_r192`; each call replays the smallest that holds the scene, exact) with program-cache misses forbidden, so no later call compiles anything. `model.warmup()` is idempotent; `warmup_variants="none"` defers the capture to `model.warmup()`.
 
-Measured on the shipped sample (`model.info["warmup_ms"]`, 2026-10-09): the load takes 315 s with an empty JIT cache and 8.6 s with a warm one (build 0.53 s: the ONNX initializers read and 48.1 MB of weights and constants uploaded; warm-up and capture 3.7 s; the rest is the device open). The first call then takes 120 ms and the second 120 ms (the stage bench's steady state: 118 ms p50 for decoded arrays, 125 ms for an `.npz` path). The trace holds 74.6 MB of DRAM (`trace_region_size` 192 MiB).
+Measured on the shipped sample (`model.info["warmup_ms"]`, 2026-10-11): the load takes 352 s with an empty JIT cache and 12.9 s with a warm one (build 0.76 s: the ONNX initializers read and 118.9 MB of weights and constants uploaded; warm-up and capture of the 6 traces 7.7 s; the rest is the device open). The first call then takes 37 ms and the second 32 ms (an `.npz` path; the stage bench's steady state: 27 ms p50 for decoded arrays, 32 ms for an `.npz` path). The 6 traces hold 60.3 MB of DRAM (`trace_region_size` 192 MiB).
 
 ## Call: `model(...)`
 
@@ -143,7 +143,7 @@ plans, the caller keeps the same state (SERVING.md 3.5 has the details):
 
 ## Lifetime and information
 
-- `model.close()` releases the trace and the persistent device tensors and closes the chip if the model opened it;
+- `model.close()` releases the traces and the persistent device tensors and closes the chip if the model opened it;
   idempotent. `with` calls it for you; an unclosed model is closed when Python exits.
 - `model.info`: weights (repo, tag, revision, path), device (dispatch, grid, CQs, fallback), variant, warm variants,
   warm-up times, runtime parameter defaults, the input schema, the numerics options and precision policy in effect, the
@@ -152,28 +152,29 @@ plans, the caller keeps the same state (SERVING.md 3.5 has the details):
 
 ## Speed
 
-Warm calls, batch 1, ETH dispatch, 1 CQ, 12×10, the pinned numerics (`code/scripts/bench.py`, 100 iterations; the numbers of `OPT_BASELINE.md`, 2026-10-08, on a shared host; p50, with p99 in brackets):
+Warm calls, batch 1, ETH dispatch, 1 CQ, 12×10, the pinned numerics (`code/scripts/bench.py`, 100 iterations, 2026-10-11, the optimized release, on a shared host; p50, with p99 in brackets):
 
-| stage | shipped sample `kashiwanoha_dense` |
+| stage | shipped sample `kashiwanoha_dense` (88 neighbours: bucket r96) |
 |---|---:|
-| `.npz` decode + schema check (path inputs only) | 7.02 (18.20) ms |
-| host pre-processing (the node's normalization, the encoder's host features, masks, `x_T`) | 4.75 (12.55) ms |
-| packing the 19 persistent trace inputs · ttnn host tensors · H2D | 0.45 · 2.41 · 1.14 ms |
-| **device trace, one blocking plan** | **102.13** (104.81) ms |
-| D2H (one packed read, 460 KB) | 0.58 ms |
-| host post-processing (trajectory, predicted paths, turn decision) | 4.32 (11.54) ms |
-| **`model(inputs=arrays)` end to end** | **117.90** (134.56) ms |
-| `model(inputs=<.npz path>)` | 124.76 (147.74) ms |
-| back-to-back replays (device time per plan) | 102.04 ms = 9.80 plans/s |
-
-The device time does not depend on the scene: every plan computes the full capacities (re-checked on c0d84f9: kashiwanoha_dense 102.10, straight_road 102.11, a nuScenes instant 102.09 ms). Throughput above one plan per ~118 ms needs pipelining of the host work of neighbouring requests (not implemented); 2 CQs do not help a synchronous request (`OPT_BASELINE.md`). Where the time goes and what comes next: `OPT_REPORT.md`.
+| `.npz` decode + schema check (path inputs only) | 5.76 (7.33) ms |
+| host pre-processing (the node's normalization, the encoder's host features, masks, `x_T`) | 2.45 (3.16) ms |
+| packing the persistent trace inputs · ttnn host tensors · H2D | 0.52 · 1.57 · 0.76 ms |
+| **device trace, one blocking plan** | **20.05** (20.08) ms |
+| D2H (one packed read) | 0.22 ms |
+| host post-processing (trajectory, predicted paths, turn decision) | 1.39 (1.84) ms |
+| **`model(inputs=arrays)` end to end** | **26.73** (27.48) ms |
+| `model(inputs=<.npz path>)` | 32.27 (33.61) ms |
+| back-to-back replays (device time per plan) | 20.00 ms = 50.0 plans/s |
+
+The device time depends on the agent bucket the scene needs (exact compaction: the smallest of 32 / 64 / 96 / 128 / 192 decoder rows that holds the ego and every valid neighbour row, else the full capacity). Back to back: r32 (≤ 31 neighbours, e.g. `straight_road`) 17.44 ms, r64 (a nuScenes instant with 42 neighbours) 19.25 ms, r96 20.00 ms, r128 21.11 ms, r192 22.30 ms, full capacity (> 191) 26.65 ms (the last three: `OPT_REPORT.md` round 5). The map entities are computed at full capacity in every plan. The first release took 102.04 ms for every scene (`OPT_BASELINE.md`). 2 CQs were not re-measured on this release (at the first release they did not help a synchronous request: `OPT_BASELINE.md`). Where the time goes and what comes next: `OPT_REPORT.md`.
 
 ## Limits
 
 - Batch 1 on the chip; one model per process.
 - Fixed shapes of the v5.0 export (320 neighbours, 140 lanes, 25 route lanes, 10 polygons, 60 line strings, 31 history
-  and 80 future steps) and 10 DPM-Solver steps (11 decoder evaluations) are compiled into the trace; every plan computes
-  the full capacities, so the device time does not depend on the scene.
+  and 80 future steps) and 10 DPM-Solver steps (11 decoder evaluations) are compiled into the traces; the neighbour trunk
+  and the decoder run on the smallest agent bucket that holds the scene (exact), so the device time depends on the
+  neighbour count (17.4-26.7 ms); the map entities are computed at full capacity.
 - The node's guidance services (start / stop / centerline guidance) are not available (the node's default is off).
 - Accuracy is agreement with the fp32 CPU reference of the same network (README "Demo & Performances"); the planner's
   driving quality is the weights' (trained by TIER IV on data that is not public).

README.md

@@ -1,22 +1,22 @@
 ---
 tags:
-- autonomous-driving
-- autoware
 - blackhole
-- diffusion
-- diffusion-planner
-- motion-prediction
 - p150
-- planning
-- tenstorrent
-- trajectory-generation
 - tt-dit-server
-- tt-metal
 - tt-model-cache
-- tt-model-catalog
 - tt-model-container
-- tt-nn
+- tenstorrent
 - ttnn
+- tt-metal
+- tt-nn
+- autoware
+- autonomous-driving
+- planning
+- trajectory-generation
+- motion-prediction
+- diffusion
+- diffusion-planner
+- tt-model-catalog
 pipeline_tag: robotics
 license: apache-2.0
 license_link: https://huggingface.co/AutowareFoundation/diffusion_planner
@@ -26,16 +26,16 @@ base_model:
 
 # diffusion-planner-p150
 
-Diffusion Planner v5.0 (Autoware diffusion_planner): the network Autoware deploys in `autoware_diffusion_planner`, ported to one Tenstorrent Blackhole p150 with tt-nn. The whole plan runs on the chip as one metal trace: the scene encoder, the 11 DiT decoder evaluations of the DPM-Solver++(2M) loop with their solver updates, and the turn-indicator head. The Autoware planner tensors in (ego and neighbour histories, lanes, route, polygons, line strings, goal, ego shape, turn-indicator history); an 8 s ego trajectory, the predicted 8 s paths of the neighbours and a turn-indicator command out, with the node's exact pre- and post-processing.
+Diffusion Planner v5.0 (Autoware diffusion_planner): the network Autoware deploys in `autoware_diffusion_planner`, ported to one Tenstorrent Blackhole p150 with tt-nn. The whole plan runs on the chip as one metal trace (one per agent bucket, picked per request): the scene encoder, the 11 DiT decoder evaluations of the DPM-Solver++(2M) loop with their solver updates, and the turn-indicator head. The Autoware planner tensors in (ego and neighbour histories, lanes, route, polygons, line strings, goal, ego shape, turn-indicator history); an 8 s ego trajectory, the predicted 8 s paths of the neighbours and a turn-indicator command out, with the node's exact pre- and post-processing.
 Weights: [AutowareFoundation/diffusion_planner `v5.0`](https://huggingface.co/AutowareFoundation/diffusion_planner/tree/423efde67f5414734da43a7ad856c17ceb8b51aa) · Paper: [arXiv:2501.15564](https://arxiv.org/abs/2501.15564) · Autoware package: [autoware_diffusion_planner](https://github.com/autowarefoundation/autoware_universe/tree/9ceaccf026c31ffc5319bc9eeb4bd7bede0af3fd/planning/autoware_diffusion_planner) · Training code: [tier4/Diffusion-Planner (TIER IV training fork)](https://github.com/tier4/Diffusion-Planner) · Port: [`code/`](https://huggingface.co/changh95/diffusion-planner-p150/tree/main/code)
 
-Runs on **p150** (mesh `P150`). Configuration: dispatch on the ETH cores, 1 command queue, 12×10 compute grid. Numerics (the default and the `serve.env` pins): fp32 residual streams and solver state, HiFi4 with fp32 accumulation, split hi / lo matmuls for the mixer inputs and every decoder linear, fp32 LayerNorm in the mixers and the decoder, fp32 matmul attention in the fusion encoder and the decoder: the configuration the end-to-end accuracy gates need (Caveats). All numbers on this card were measured in this configuration.
+Runs on **p150** (mesh `P150`). Configuration: dispatch on the ETH cores, 1 command queue, 12×10 compute grid. Numerics (the default and the `serve.env` pins): fp32 residual streams and solver state, HiFi4 with fp32 accumulation, split hi / lo matmuls for the mixer inputs and every decoder linear (each as one K-concatenated fp32 matmul), fp32 LayerNorm in the mixers and the decoder and fp32 matmul attention in the fusion encoder and the decoder (each as one fused custom kernel): the configuration the end-to-end accuracy gates need (Caveats). Optimized release: 5.1x faster on the device than the first release (102.0 -> 20.0 ms per plan on the shipped sample), same gates. All numbers on this card were measured in this configuration.
 
 Packaged and published with [tt-model-manager](https://github.com/tenstorrent/tt-model-manager) 0.1.0 (manifest schema 5.1).
 
 ## Quickstart (Python)
 
-Prerequisite: a tt-metal / ttnn environment at tt-metal [`44d66500520`](https://github.com/tenstorrent/tt-metal/commit/44d66500520fda9f2c7060c0f6b41ec48f7ab37e) with [`patches/tt-metal-eth-dispatch.patch`](patches/tt-metal-eth-dispatch.patch) applied. ttnn is not on PyPI.
+Prerequisite: a tt-metal / ttnn environment at tt-metal [`44d66500520`](https://github.com/tenstorrent/tt-metal/commit/44d66500520fda9f2c7060c0f6b41ec48f7ab37e) with [`patches/tt-metal-eth-dispatch.patch`](patches/tt-metal-eth-dispatch.patch) applied. The image also carries [`patches/tt-metal-reshape-rm-sys1419.patch`](patches/tt-metal-reshape-rm-sys1419.patch), a guard against an intermittent hang of row-major reshapes under ETH dispatch on Blackhole; the numbers on this card were measured with both. ttnn is not on PyPI.
 
 ```bash
 hf download changh95/diffusion-planner-p150 --exclude "image/*" --local-dir diffusion-planner-p150 && cd diffusion-planner-p150
@@ -56,8 +56,8 @@ print(out.poses[:5])
 print(out.turn_indicator["command_name"], out.predicted_agents.shape)
 ```
 
-- `from_pretrained` downloads the three v5.0 ONNX files and `diffusion_planner.param.json` (58.9 MB) of [`AutowareFoundation/diffusion_planner`](https://huggingface.co/AutowareFoundation/diffusion_planner) at the pinned commit `423efde67f5` (tag `v5.0`) to your HF cache (no token needed), checks their sha256, opens the chip, builds the graph and captures the metal trace. The first load compiles the kernels (315 s with an empty JIT cache, firmware and every kernel compiled); later loads take about 9 s (device open, reading and uploading the weights, warm-up and trace capture).
-- The trace is captured during the load, so the first call is as fast as the later calls and no call compiles anything.
+- `from_pretrained` downloads the three v5.0 ONNX files and `diffusion_planner.param.json` (58.9 MB) of [`AutowareFoundation/diffusion_planner`](https://huggingface.co/AutowareFoundation/diffusion_planner) at the pinned commit `423efde67f5` (tag `v5.0`) to your HF cache (no token needed), checks their sha256, opens the chip, builds the graph and captures the metal traces (6: the full capacity and 5 agent buckets). The first load compiles the kernels (352 s with an empty JIT cache, firmware and every kernel compiled); later loads take about 13 s (device open, reading and uploading the weights, warm-up and the capture of the 6 traces).
+- The traces are captured during the load, so no call compiles anything (first call 37 ms, later calls 32 ms with an `.npz` path).
 - The `with` block releases the trace and closes the chip. Without `with`, call `model.close()`.
 
 | | |
@@ -70,6 +70,7 @@ print(out.turn_indicator["command_name"], out.predicted_agents.shape)
 - The API gives the same output as the HTTP server `/predict`: both share the decoders, the device trace and the host post-processing (checked on the device by `test_api_equals_server`).
 - The API is stateless: one call is one plan. What the Autoware node keeps between plans stays with the caller: the turn-indicator hold window (1.0 s; the node's manager ships as `tt_diffusion_planner.host.postprocess.TurnIndicatorManager`), the initial solver state `sampled_trajectories` (zeros = the node's default temperature 0; noise for a temperature > 0; the previous plan for the RTC prefix) and the agent and ego histories. Details: [`code/PYTHON.md`](code/PYTHON.md) "What the caller keeps".
 - One model uses one chip; calls from several threads are serialised.
+- Each call replays the trace of the smallest agent bucket (32 / 64 / 96 / 128 / 192 rows, else the full 321) that holds the ego and every valid neighbour row (the node fills the rows from the front); the result is exact (bit-identical to the full-capacity trace), only the device time depends on the scene.
 - Full reference: [`code/PYTHON.md`](code/PYTHON.md). Runnable example: [`examples/quickstart.py`](examples/quickstart.py) (also writes `quickstart_bev.png`, the input tensors and the plan from above).
 
 ## Serving (HTTP)
@@ -95,11 +96,11 @@ The response for the shipped sample (served on the p150; trajectory cut to 3 of
  "model": "diffusion-planner-p150",
  "frame_id": "base_link",
  "meta": {"predicted_agent_columns": ["x", "y", "yaw", "cos", "sin"], "predicted_agent_rows": [0, 1, 2, "... 85 more"], "force_stop": false, "time_from_start_s": [0.1, 0.2, 0.3, "... 77 more"], "valid_counts": {"ego": 1, "neighbor": 88, "static": 0, "lane": 123, "route": 17, "polygon": 0, "line_string": 60, "goal": 1, "ego_shape": 1, "turn": 1}},
- "timing_ms": {"preprocess": 3.8, "device": 105.3, "postprocess": 3.2, "total": 122.8, "decode": 8.5, "model_call": 112.5},
+ "timing_ms": {"preprocess": 2.4, "device": 22.8, "postprocess": 1.2, "total": 36.9, "decode": 8.4, "model_call": 26.7},
  "num_poses": 80,
  "columns": ["x", "y", "yaw", "cos", "sin", "velocity", "acceleration"],
- "trajectory": [[0.365, 0.01, -0.017, 0.9962, -0.0169, 3.8059, -0.0431], [0.7631, -0.0022, -0.0414, 0.9966, -0.0413, 3.8016, -0.5485], [1.1524, -0.0282, -0.0675, 0.9942, -0.0674, 3.7468, -0.4708], "... 77 more rows"],
- "turn_indicator": {"command": 1, "command_name": "DISABLE", "keep_selected": true, "held": false, "logits": [-15.960084915161133, -5.125061511993408, -5.230083465576172, -0.7658059597015381, 5.504596710205078], "probabilities": [1.651601411190029e-09, 8.38487030705437e-05, 7.548934809165075e-05, 0.006556871347129345, 0.993184506893158]},
+ "trajectory": [[0.3644, 0.01, -0.0169, 0.9964, -0.0169, 3.8033, -0.0392], [0.7626, -0.0021, -0.0412, 0.9968, -0.0412, 3.7994, -0.5496], [1.1517, -0.0281, -0.0674, 0.9942, -0.0672, 3.7444, -0.4683], "... 77 more rows"],
+ "turn_indicator": {"command": 1, "command_name": "DISABLE", "keep_selected": true, "held": false, "logits": [-15.960197448730469, -5.119836807250977, -5.23490047454834, -0.7636525630950928, 5.502235412597656], "probabilities": [1.655269699085693e-09, 8.448460721410811e-05, 7.530191942350939e-05, 0.00658634165301919, 0.9931545257568359]},
  "predicted_agents": {"format": "npz", "key": "predicted_agents", "dtype": "float32", "shape": [88, 80, 5], "data": "<base64 npz>"}
 }
 ```
@@ -132,40 +133,41 @@ nuScenes renders: rendered from the nuScenes dataset (v1.0-mini, CAN bus expansi
 
 ## Demo & Performances
 
-Warm, batch 1; the stage bench 2026-10-08, the served rows, load times and the re-check 2026-10-09. Latency: the stage bench of [`OPT_BASELINE.md`](OPT_BASELINE.md) (`code/scripts/bench.py`, 100 iterations per stage) on the shipped sample `kashiwanoha_dense.npz` (88 neighbours, 123 lanes, 17 route lanes, 60 line strings; the device time is the same for every scene: every plan computes the full capacities); the served rows from uvicorn on the host (the app the container runs, with its serve pins) and a loopback client, 50 requests of the shipped sample. The host is shared with other jobs, so the host stages move with its load by several ms; the device rows repeat to ±0.05 ms. Accuracy: the p150 output against the fp32 CPU reference of the same network on 99 scenes: the 2 shipped samples, 5 research scenes and 92 nuScenes v1.0-mini planning instants (the frozen end-to-end gates), and against the research pipeline's ONNX Runtime outputs as an independent oracle.
+Warm, batch 1, measured 2026-10-11 on the optimized release (bundle `bc966b4`; the packaged code adds only the WORKER-grid fallback `7e003c4`, bit-identical on ETH, and the version string). Latency: the stage bench of [`OPT_REPORT.md`](OPT_REPORT.md) (`code/scripts/bench.py`, 100 iterations per stage) on the shipped sample `kashiwanoha_dense.npz` (88 neighbours, 123 lanes, 17 route lanes, 60 line strings: agent bucket r96). The device time depends on the agent bucket the scene needs (exact compaction, see Caveats): 17.44 ms for ≤ 31 neighbours up to 26.65 ms at the full 320. The served rows come from uvicorn on the host (the app the container runs, with its serve pins) and a loopback client, 50 requests of the shipped sample. The host is shared with other jobs, so the host stages move with its load by a few ms; the device rows repeat to ±0.02 ms. Accuracy: the p150 output against the fp32 CPU reference of the same network on 99 scenes: the 2 shipped samples, 5 research scenes and 92 nuScenes v1.0-mini planning instants (the frozen end-to-end gates), and against the research pipeline's ONNX Runtime outputs as an independent oracle. The first release's numbers are in [`OPT_BASELINE.md`](OPT_BASELINE.md).
 
 | Metric | Performance |
 |---|---:|
-| Agreement with the fp32 CPU reference, shipped sample `kashiwanoha_dense` (8 s ego plan, 88 neighbours) | ego max **2.3 cm** / mean 1.1 cm; turn command identical; neighbours: median per-agent max 8.6 cm |
-| Agreement, shipped sample `straight_road` | ego max 4.9 cm / mean 1.8 cm; turn command identical; neighbours 6.3 cm |
-| Agreement, all 99 gated scenes (2 samples, 5 research scenes, 92 nuScenes-mini instants) | worst ego max **0.313 m** / mean **0.143 m** (gates 1.0 / 0.3 m; nuscenes/scene-0103_kf14); ego mean: median 1.3 cm, 95th percentile 6.2 cm; turn command identical **99 / 99**; neighbours: median per-agent max ≤ 0.086 m (gate 1.5 m) |
-| Agreement with an independent oracle: the research pipeline's ONNX Runtime outputs (raw x0), 92 nuScenes instants + the 33-plan scene-0061 sequence | worst ego max 0.313 m / mean 0.143 m, turn command identical 92 / 92; sequence: worst ego max 0.118 m / mean 0.049 m, turn 33 / 33; every instant within the gates |
-| Module PCC vs the fp32 reference (encoder categories, encoding, teacher-forced decoder evaluation; gate 0.999) | ≥ 0.999952 / 0.999978 / ≥ 0.9999995 |
+| Agreement with the fp32 CPU reference, shipped sample `kashiwanoha_dense` (8 s ego plan, 88 neighbours) | ego max **2.0 cm** / mean 0.8 cm; turn command identical; neighbours: median per-agent max 8.6 cm |
+| Agreement, shipped sample `straight_road` | ego max 2.3 cm / mean 0.9 cm; turn command identical; neighbours 6.0 cm |
+| Agreement, all 99 gated scenes (2 samples, 5 research scenes, 92 nuScenes-mini instants) | worst ego max **0.346 m** / mean **0.154 m** (gates 1.0 / 0.3 m; nuscenes/scene-0103_kf14); ego mean: median 1.3 cm, 95th percentile 6.3 cm; turn command identical **99 / 99**; neighbours: median per-agent max ≤ 0.094 m (gate 1.5 m). First release: 0.313 / 0.143 m (two gated precision changes since, Caveats) |
+| Agreement with an independent oracle: the research pipeline's ONNX Runtime outputs (raw x0), 92 nuScenes instants + the 33-plan scene-0061 sequence | worst ego max 0.346 m / mean 0.154 m, turn command identical 92 / 92; sequence: worst ego max 0.112 m / mean 0.055 m, turn 33 / 33; every instant within the gates |
+| Module PCC vs the fp32 reference (encoder categories, encoding, teacher-forced decoder evaluation; gate 0.999) | ≥ 0.999948 / 0.999976 / ≥ 0.9999995 |
 | Open-loop vs the nuScenes log, 92 instants (a sanity check against one human driver, not a planning metric) | p150 ADE / FDE at 8 s 4.75 / 11.87 m (CPU reference 4.74 / 11.87; constant velocity 4.99 / 13.06); turn command = logged blinker 88 / 92 (CPU 88 / 92) |
-| Python `model()` call, shipped sample (host pre-processing, H2D, trace, D2H, host post-processing) | **117.9 ms p50** (p99 134.6) · 8.5 plans/s |
-| Served `/predict` `timing_ms.total` (uvicorn on the host, the shipped sample) | **123.1 ms median** (min 122.5; of which decode 8.6) |
-| Served client round trip, loopback (base64 `.npz` request, 0.15 MB) | 127.5 ms median |
-| Device trace, one blocking plan (encoder + 11 DiT evaluations + 10 solver updates + turn head) | **102.13 ms** |
-| Back-to-back trace replays | **102.04 ms per plan** · 9.80 plans/s |
-| Host pre-processing · pack · host tensors · H2D · D2H · host post-processing | 4.75 · 0.45 · 2.41 · 1.14 · 0.58 · 4.32 ms |
-| `from_pretrained` load: empty JIT cache / warm cache | 315 s / 8.6 s |
+| Python `model()` call, shipped sample (host pre-processing, H2D, trace, D2H, host post-processing) | **26.7 ms p50** (p99 27.5) · 37.4 plans/s (first release 117.9 ms) |
+| Served `/predict` `timing_ms.total` (uvicorn on the host, the shipped sample) | **37.2 ms median** (min 36.7; of which JSON decode 8.6, `model()` 26.8; first release 123.1 ms) |
+| Served client round trip, loopback (base64 `.npz` request, 0.15 MB) | 41.6 ms median (first release 127.5 ms) |
+| Device trace, one blocking plan (encoder + 11 DiT evaluations + 10 solver updates + turn head) | **20.05 ms** (bucket r96; straight_road r32 17.50, nuScenes scene-0103_kf14 r64 19.30; first release 102.13 ms for every scene) |
+| Back-to-back trace replays, by agent bucket (neighbours) | **20.00 ms per plan** · 50.0 plans/s on the shipped sample (r96, 64-95); r32 (≤ 31) 17.44 · r64 19.25 · r128 21.11 · r192 22.30 · full (> 191) 26.65 ms; first release 102.04 ms for every scene |
+| Host pre-processing · pack · host tensors · H2D · D2H · host post-processing | 2.45 · 0.52 · 1.57 · 0.76 · 0.22 · 1.39 ms |
+| `from_pretrained` load: empty JIT cache / warm cache (6 traces) | 352 s / 12.9 s |
 
-All numbers in this table were measured with dispatch on the ETH cores, 1 command queue and a 12×10 compute grid on one p150, with the pinned numerics of `serve.env`. Accuracy is agreement with the fp32 CPU reference of the same Autoware network (same weights, same pre- and post-processing); no dataset-level accuracy is claimed: the paper's benchmark is nuPlan closed loop (an account-gated dataset and simulator, not run here), and the deployed v5.0 weights were trained by TIER IV on data that is not public, so no public benchmark is in-domain. Details: [`VERIFICATION_2026-10-08.md`](VERIFICATION_2026-10-08.md), [`OPT_BASELINE.md`](OPT_BASELINE.md), [`OPT_REPORT.md`](OPT_REPORT.md).
+All numbers in this table were measured with dispatch on the ETH cores, 1 command queue and a 12×10 compute grid on one p150, with the pinned numerics of `serve.env`. Accuracy is agreement with the fp32 CPU reference of the same Autoware network (same weights, same pre- and post-processing); no dataset-level accuracy is claimed: the paper's benchmark is nuPlan closed loop (an account-gated dataset and simulator, not run here), and the deployed v5.0 weights were trained by TIER IV on data that is not public, so no public benchmark is in-domain. The r128, r192 and full-capacity rows are the same-session values of [`OPT_REPORT.md`](OPT_REPORT.md) round 5, reproduced by [`VERIFICATION_OPT_2026-10-11.md`](VERIFICATION_OPT_2026-10-11.md). Details: [`VERIFICATION_OPT_2026-10-11.md`](VERIFICATION_OPT_2026-10-11.md) (the optimization rounds), [`VERIFICATION_2026-10-08.md`](VERIFICATION_2026-10-08.md) (the first release), [`OPT_REPORT.md`](OPT_REPORT.md), [`OPT_BASELINE.md`](OPT_BASELINE.md).
 
 No GPU comparison: no GPU was available on the host where this port was built and measured, so this card makes no GPU speed claim. The reference rows are the port's own fp32 CPU reference on the same host (a correctness baseline, not a speed target). Autoware's CHANGELOG quotes 5.13 ms mean (300 runs) for an older single-step engine of this planner on an RTX PRO 6000 Blackwell with TensorRT (precision not stated); it is not like-for-like with this v5.0 multi-step port. p150 power was not measured, so no efficiency comparison is made.
 
 ## Caveats
 
-- First release: **baseline port, optimization pending.** The plan is one metal trace and kernel-bound (6,282 programs; op-to-op gaps 4.4 ms of a 103.7 ms span). The **first optimization target is the precision cost**: the numerics defaults that the end-to-end gates need (split hi / lo matmuls, fp32 LayerNorm, fp32 matmul attention) cost 33.4 ms per plan (102.0 ms vs 68.6 ms with the first device round's defaults, which fail the gates); fused kernels are to recover it without dropping precision. A second sink of similar size is unrelated to precision: the encoder's channel-MLP and pre-projection matmuls run on 4-8 cores (25.8 ms). [`OPT_REPORT.md`](OPT_REPORT.md) ranks what comes next.
+- Optimized release (5 optimization rounds, [`OPT_REPORT.md`](OPT_REPORT.md)): 102.0 -> 20.0 ms per plan on the device for the shipped sample. What changed: explicit multi-core program configs for the encoder and decoder matmuls; custom `generic_op` kernels for the fp32 LayerNorm, the K-concatenated split matmul operands, the whole fp32 attention (scores in L1) and the score scale / mask / softmax; L1-resident intermediates in the decoder; exact compaction by agent buckets (one trace per bucket); vectorised host pre- and post-processing. The plan is still a sequence of ~1,400 device programs (1,418 at r96, kernel-bound: kernel 19.64 ms, op-to-op gaps 0.89 ms); no megakernel ships in this release (a persistent mixer / DiT megakernel is work in progress and not measured yet). What remains and what comes next: [`OPT_REPORT.md`](OPT_REPORT.md).
+- Two of the optimizations change the numerics slightly, with the frozen gates green: the decoder split linears (round 2) and the encoder split linears (round 3) each accumulate their hi / lo products in one fp32 matmul instead of three. The 99-scene worst case moved from 0.313 / 0.143 m (ego max / mean) to 0.346 / 0.154 m, the gates are 1.0 / 0.3 m; scene by scene the error moves both ways. A held-out check on 50 further scenes (the 33-plan scene-0061 sequence, full-capacity and noisy-x_T scenes, every agent-bucket boundary) passed every gate ([`VERIFICATION_OPT_2026-10-11.md`](VERIFICATION_OPT_2026-10-11.md)); its worst ego max is 0.507 m, on a synthetic scene with 192 neighbours. Every other step is bit-identical to its predecessor.
 - Deployment status in Autoware: `autoware_diffusion_planner` is an alternative to the default rule-based planning stack, selected with `planning_setting:=diffusion_planner` (package README); it is aimed at Autoware's proposed new planning framework. This bundle is not a ROS 2 node (Python API and HTTP) and not a certified Autoware component; do not use it for safety-critical driving decisions or closed-loop vehicle control.
 - Stateless API: the node's state between plans (the turn-indicator hold window, the RTC prefix and temperature of the initial solver state, the agent buffers and the ego history) is the client's ("What the caller keeps" in [`code/PYTHON.md`](code/PYTHON.md)). The node's guidance services (start / stop / centerline guidance) are off, as in the node's default.
-- Precision policy of this release: fp32 residual streams and solver state; HiFi4 with fp32 accumulation for every matmul; the ego / neighbour pre-projection as a pad-relative fp32 island; split hi / lo matmuls (bf16 hi + fp32 lo parts, ~1e-5 relative, because a device fp32 matmul rounds its operands like TF32) for the mixer inputs and every decoder linear; an fp32 LayerNorm decomposition in the mixers and the decoder (the fused `ttnn.layer_norm` loses the per-entity signal on the mixers' offset-dominated rows); fp32 matmul attention in the fusion encoder and the decoder; the turn head in fp32; the other weights and the hidden MLP activations of the mixers and the fusion encoder in bf16. The plan's sensitivity to these choices is chaotic per scene: with the first device round's numerics, one nuScenes instant (scene-0103_kf14) moved 0.35 m on average while the module PCCs differed only in the 5th decimal. So every numerics change is re-checked on all 99 scenes.
+- Precision policy of this release: fp32 residual streams and solver state; HiFi4 with fp32 accumulation for every matmul; the ego / neighbour pre-projection as a pad-relative fp32 island; split hi / lo matmuls (bf16 hi + fp32 lo parts, ~1e-5 relative, because a device fp32 matmul rounds its operands like TF32; computed as one fp32 matmul over the K-concatenated operand `[x_hi | x_hi | x_lo | 1]`) for the mixer inputs and every decoder linear; an fp32 LayerNorm (a fused custom kernel with fp32 statistics) in the mixers and the decoder (the fused `ttnn.layer_norm` loses the per-entity signal on the mixers' offset-dominated rows); fp32 matmul attention in the fusion encoder and the decoder (a fused kernel with the scores in L1); the turn head in fp32; the other weights and the hidden MLP activations of the mixers and the fusion encoder in bf16. The plan's sensitivity to these choices is chaotic per scene: with the first device round's numerics, one nuScenes instant (scene-0103_kf14) moved 0.35 m on average while the module PCCs differed only in the 5th decimal. So every numerics change is re-checked on all 99 scenes.
 - Validation scope: the p150 output agrees with the fp32 CPU reference on 99 scenes (above). The two shipped samples and the five research scenes are synthetic-but-faithful scenes built with a Python port of the node's tensor construction; the 92 nuScenes instants are converted from nuScenes v1.0-mini, a domain the model never saw (Singapore and Boston, right-hand traffic in Boston, oracle tracks from 2 Hz annotations, no traffic-light states, no speed limits, stop areas instead of stop lines). On them the plans are plausible but conservative (moving plans about 17 % shorter than the logged drive); the open-loop numbers above are a sanity check against one human driver, not a planning metric.
-- Neighbour predictions: the gate is the median over agents of each agent's max displacement, so single agents can differ more. On the shipped sample `kashiwanoha_dense` (container smoke, served p150 output vs the stored CPU reference) the per-agent max displacement is 8.6 cm median and 0.26 m at the 90th percentile; the worst of the 88 agents differs by 4.31 m. The ego plan is gated on its max and mean; neighbour paths only on that median.
+- Neighbour predictions: the gate is the median over agents of each agent's max displacement, so single agents can differ more. On the shipped sample `kashiwanoha_dense` (p150 output vs the stored CPU reference, 2026-10-11) the per-agent max displacement is 8.6 cm median and 0.23 m at the 90th percentile; the worst of the 88 agents differs by 3.72 m. The ego plan is gated on its max and mean; neighbour paths only on that median.
 - Documented deviations from the node: the output stays in `base_link` (the node transforms it to `map` with the ego pose); `yaw` is what `tf2::getYaw` reads from the node's quaternion of the unnormalised cos / sin rotation (it differs from atan2(sin, cos) when |(cos, sin)| ≠ 1, as in the node); the speed masks follow the node's TensorRT path (`> FLT_EPSILON`); the turn indicator is decided without the hold window; `delay` is accepted and ignored (the node's multi-step mode never reads it).
 - Does not scale to multiple p150 in a mesh configuration. The build uses a 12×10 compute grid of Tensix cores: the dispatch functions move from one Tensix column to the ETH cores (`patches/tt-metal-eth-dispatch.patch`), so this build assumes that you do not need chip-to-chip ethernet communication.
-- `dispatch="worker"` (server: `DIFFUSION_PLANNER_DISPATCH=worker`) is an A/B opt-in. On a p150 it gives an 11×10 grid (101.87 ms per replay, the same as ETH within 0.2 %); if ETH dispatch is not available (tt-metal without the patch), the model falls back to it with a warning. The numbers on this card do not apply to that mode.
-- Batch 1, one plan per request; requests are serialised on the chip. The shapes of the v5.0 export (320 neighbours, 140 lanes, 25 route lanes, 10 polygons, 60 line strings, 31 history and 80 future steps) and 10 DPM-Solver steps are compiled into the trace, and every plan computes them in full, so the device time does not depend on the scene.
+- `dispatch="worker"` (server: `DIFFUSION_PLANNER_DISPATCH=worker`) is an A/B opt-in. On a p150 it gives an 11×10 grid (20.61 ms per replay on the shipped sample, 3 % slower than ETH; the 99 gated scenes give the same agreement numbers as ETH); if ETH dispatch is not available (tt-metal without the patch), the model falls back to it with a warning. The numbers on this card do not apply to that mode.
+- Batch 1, one plan per request; requests are serialised on the chip. The shapes of the v5.0 export (320 neighbours, 140 lanes, 25 route lanes, 10 polygons, 60 line strings, 31 history and 80 future steps) and 10 DPM-Solver steps are compiled into the traces. The neighbour trunk of the encoder and the decoder run on the smallest agent bucket that holds the scene (one trace per bucket, exact: the outputs equal the full-capacity trace bit for bit), so the device time depends on the neighbour count (17.4-26.7 ms); the map entities (lanes, route, polygons, line strings) are still computed at full capacity in every plan.
 - Not an OpenAI-compatible API; `GET /v1/models` is a stub so the tt-model ready card does not 404.
 - p150 power was not measured, so no efficiency comparison is made.
 
@@ -173,7 +175,7 @@ No GPU comparison: no GPU was available on the host where this port was built an
 
 - Weights: [AutowareFoundation/diffusion_planner](https://huggingface.co/AutowareFoundation/diffusion_planner) at tag `v5.0` (commit `423efde67f5414734da43a7ad856c17ceb8b51aa`), Apache-2.0 per its model card. Not redistributed here: the package only points to them. The upstream card states that TIER IV trained the models on TIER IV synthetic and real driving data; the dataset composition is not publicly documented.
 - Pre- and post-processing ported from autoware_universe [`planning/autoware_diffusion_planner`](https://github.com/autowarefoundation/autoware_universe/tree/9ceaccf026c31ffc5319bc9eeb4bd7bede0af3fd/planning/autoware_diffusion_planner) (Apache-2.0).
-- Port and serving code (`code/`): Apache-2.0. `patches/tt-metal-eth-dispatch.patch` modifies tt-metal (Apache-2.0).
+- Port and serving code (`code/`): Apache-2.0. `patches/tt-metal-eth-dispatch.patch` and `patches/tt-metal-reshape-rm-sys1419.patch` modify tt-metal (Apache-2.0).
 - Sample data: `code/tt_diffusion_planner/samples/kashiwanoha_dense.npz` is derived from the Apache-2.0 Lanelet2 map [AutowareFoundation/map-carla-kashiwanoha](https://huggingface.co/datasets/AutowareFoundation/map-carla-kashiwanoha) (0.2.0) with a scripted ego, route and agents; `straight_road.npz` is a procedural scene generated by this repo. Both Apache-2.0, each with its stored CPU-reference output (`*.reference.json`). Only these redistributable samples ship; the nuScenes-derived planning instants of the accuracy tables are not in this repository.
 - Demo media (`media/`, sources and changes in [`media/ATTRIBUTION.md`](media/ATTRIBUTION.md)):
   - nuScenes renders (`media/*_NC.*`), **non-commercial, CC BY-NC-SA 4.0**: Rendered from the nuScenes dataset (v1.0-mini, CAN bus expansion and map expansion v1.3), © Motional AD Inc., CC BY-NC-SA 4.0 and the nuScenes Terms of Use (https://www.nuscenes.org/terms-of-use). Non-commercial use only; adaptations under the same license. Motional does not endorse this work. Cite: H. Caesar et al., *nuScenes: A Multimodal Dataset for Autonomous Driving*, CVPR 2020.
@@ -185,11 +187,11 @@ These are the exact sources the container image was built from:
 
 | component | built from |
 | --- | --- |
-| tt-metal | [`44d66500520fda9f2c7060c0f6b41ec48f7ab37e`](https://github.com/tenstorrent/tt-metal/commit/44d66500520fda9f2c7060c0f6b41ec48f7ab37e) + [`patches/tt-metal-eth-dispatch.patch`](patches/tt-metal-eth-dispatch.patch) (sha256 `08d0ddf6…45cc`, 4 files; dirty tree: the image includes the patch) |
+| tt-metal | [`44d66500520fda9f2c7060c0f6b41ec48f7ab37e`](https://github.com/tenstorrent/tt-metal/commit/44d66500520fda9f2c7060c0f6b41ec48f7ab37e) + [`patches/tt-metal-eth-dispatch.patch`](patches/tt-metal-eth-dispatch.patch) (sha256 `08d0ddf6…45cc`, 4 files) + [`patches/tt-metal-reshape-rm-sys1419.patch`](patches/tt-metal-reshape-rm-sys1419.patch) (sha256 `6be07ce5…9cfa`, 1 file) (dirty tree: the image includes both patches; the image build asserts both markers) |
 | weights | [`AutowareFoundation/diffusion_planner@423efde67f5414734da43a7ad856c17ceb8b51aa`](https://huggingface.co/AutowareFoundation/diffusion_planner/tree/423efde67f5414734da43a7ad856c17ceb8b51aa) (tag `v5.0`), files `diffusion_planner_encoder.onnx, diffusion_planner_decoder.onnx, diffusion_planner_turn_indicator.onnx, diffusion_planner.param.json` (sha256 `2856886a…ca49`, `eb30c0c0…57ca`, `07acfb58…a732`, `ee3145b6…a268`, checked at load) |
 | Autoware reference | autoware_universe [`9ceaccf026c31ffc5319bc9eeb4bd7bede0af3fd`](https://github.com/autowarefoundation/autoware_universe/tree/9ceaccf026c31ffc5319bc9eeb4bd7bede0af3fd) (`planning/autoware_diffusion_planner`, package 0.53.0, `multi_step` mode) |
-| shared package | `ttaw` 0.20.0, vendored as `code/tt_diffusion_planner/ttaw` from the Autoware ports' shared `common` repository at commit `89dec49` (`code/tt_diffusion_planner/ttaw/VENDORED.json`: version, commit and per-file sha256) |
-| `code/` digest (image) | `c0e7abb7888098a9` (sha256, first 16 hex digits; `built.code_sha256` of `tt_kernel_manifest.json`) |
-| image | `tt-model/diffusion-planner-p150:3b96d8ea7190` (`sha256:3b96d8ea71902fe6a00f1792dd41290839b7758070431464cd613cde3f6bf909`) |
+| shared package | `ttaw` 0.23.2, vendored as `code/tt_diffusion_planner/ttaw` from the Autoware ports' shared `common` repository at commit `60b6dd7` (`code/tt_diffusion_planner/ttaw/VENDORED.json`: version, commit and per-file sha256) |
+| `code/` digest (image) | `5da22b97bbf89b03` (sha256, first 16 hex digits; `built.code_sha256` of `tt_kernel_manifest.json`) |
+| image | `tt-model/diffusion-planner-p150:dde78ac2f0be` (`sha256:dde78ac2f0be5b7e637ddceba1a7c30fd832c2a50dd3e728acbf187f86354cf4`) |
 | base images | build stage `ghcr.io/tenstorrent/tt-metal/tt-metalium/ubuntu-22.04-dev-amd64:latest` @ `sha256:df9d279c7f85c17c6fad982d196802682d669cca1b7ced9cbaad8181339cd5fc`; runtime stage `docker.io/library/ubuntu:22.04` @ `sha256:5ec03bb3441e8b0bf3b4f9cd4629a1ae763010dc3035bb8da3ae6cf026486401` (tt-model's `FROM` tags float; these are the digests this build resolved, see [`build_info.json`](build_info.json)) |
-| built | 2026-10-09T04:37:19+00:00 by tt-model 0.1.0 |
+| built | 2026-10-11T03:53:34+00:00 by tt-model 0.1.0 |

SERVING.md

@@ -9,13 +9,13 @@ image. The device open, the decoders, the Python-API contract and the HTTP contr
 
 | item | value |
 |---|---|
-| tt-metal tree | `/home/ubuntu/experiments/tt-models/tt-metal` (main `44d66500520`, `v0.80.0-dev20261006-78-g44d6650052`, + `patches/tt-metal-eth-dispatch.patch`; torch 2.11.0+cpu) |
+| tt-metal tree | `/home/ubuntu/experiments/tt-models/tt-metal` (main `44d66500520`, `v0.80.0-dev20261006-78-g44d6650052`, + `patches/tt-metal-eth-dispatch.patch` + `patches/tt-metal-reshape-rm-sys1419.patch`; torch 2.11.0+cpu) |
 | weights | `AutowareFoundation/diffusion_planner` @ `423efde67f5414734da43a7ad856c17ceb8b51aa` (tag `v5.0`; files `diffusion_planner_encoder.onnx`, `diffusion_planner_decoder.onnx`, `diffusion_planner_turn_indicator.onnx`, `diffusion_planner.param.json`; 58.9 MB; Apache-2.0; public, ungated; sha256 of every file checked at load) |
 | Autoware reference | `planning/autoware_diffusion_planner` @ autoware_universe `9ceaccf026c31ffc5319bc9eeb4bd7bede0af3fd` (package 0.53.0), the node's `multi_step` mode |
-| shared package | `ttaw` 0.20.0 (`common` @ `89dec49`) |
+| shared package | `ttaw` 0.23.2 (`common` @ `60b6dd7`) |
 | app | `tt_diffusion_planner.server.app:app` (uvicorn; the ASGI lifespan does weights -> device -> graph -> trace capture) |
-| device recipe | `ttnn.open_device(device_id, dispatch_core_config=DispatchCoreConfig(ETH), l1_small_size=32768, trace_region_size=192 MiB, num_command_queues=1)` (`DEVICE_DEFAULTS` in `code/tt_diffusion_planner/device.py`; the plan's trace holds 74.6 MB) |
-| graph | ONE metal trace per plan: the encoder (6 MLP-Mixer trunks, small encoders, 6 fusion blocks), the cross-attention K / V hoisted once, 11 DiT evaluations with the 10 fp32 DPM-Solver++(2M) updates and the prefix constraint, the turn-indicator head, one packed readback; 6,282 device programs |
+| device recipe | `ttnn.open_device(device_id, dispatch_core_config=DispatchCoreConfig(ETH), l1_small_size=32768, trace_region_size=192 MiB, num_command_queues=1)` (`DEVICE_DEFAULTS` in `code/tt_diffusion_planner/device.py`; the 6 traces hold 60.3 MB) |
+| graph | ONE metal trace per plan: the encoder (6 MLP-Mixer trunks, small encoders, 6 fusion blocks), the cross-attention K / V hoisted once, 11 DiT evaluations with the 10 fp32 DPM-Solver++(2M) updates and the prefix constraint, the turn-indicator head, one packed readback. Six traces are captured at startup, the full capacity and one per agent bucket (32 / 64 / 96 / 128 / 192 decoder rows); each request replays the smallest that holds its valid neighbour rows (exact). 1,418 device programs at r96 (6,282 in the first release); custom `generic_op` kernels in `code/tt_diffusion_planner/tt/kernels/` |
 | hardware | one Blackhole p150 (`hardware: p150`, `mesh_device: P150`, `TT_MESH_SHAPE=1x1`), 12×10 compute grid |
 
 ## 1. Run on the HOST (hardware validation, no Docker)
@@ -44,10 +44,10 @@ $ROOT/bin/devrun -t 1800 -- bash -c 'env DIFFUSION_PLANNER_DISPATCH=eth DIFFUSIO
 ```
 
 Boot log landmarks (they drive `tt-model serve`'s progress view, `boot_progress.py` TT_DIT_PHASES):
-`Loading weights` -> `Opening device` -> `Warming up: capturing trace ...` -> `Warmup complete` -> uvicorn
-`Application startup complete`. A cold JIT cache compiles every kernel first (315 s measured for
-`from_pretrained`); later boots take about 10 s to READY (measured on the host, warm JIT cache). Startup failures raise and uvicorn exits non-zero (no CPU
-fallback). SIGTERM: the lifespan releases the trace and closes the chip (within `tt-model stop`'s 120 s budget).
+`Loading weights` -> `Opening device` -> `Warming up: capturing trace ...` (once per trace variant) -> `Warmup complete` -> uvicorn
+`Application startup complete`. A cold JIT cache compiles every kernel first (352 s measured for
+`from_pretrained`); later boots take about 13 s to READY (measured on the host, warm JIT cache). The packaged container measured the same on 2026-10-11: `tt-model serve` READY after 6 min 0 s on its first start (server boot 358 s, empty JIT cache) and after 20.9 s on the next one (server boot 19.0 s). Startup failures raise and uvicorn exits non-zero (no CPU
+fallback). SIGTERM: the lifespan releases the traces and closes the chip (within `tt-model stop`'s 120 s budget).
 
 Offline override (no Hub): `DIFFUSION_PLANNER_WEIGHTS_DIR=<dir holding diffusion_planner_encoder.onnx, diffusion_planner_decoder.onnx, diffusion_planner_turn_indicator.onnx, diffusion_planner.param.json>`.
 
@@ -144,11 +144,11 @@ the predicted-agents array to its header):
  "model": "diffusion-planner-p150",
  "frame_id": "base_link",
  "meta": {"predicted_agent_columns": ["x", "y", "yaw", "cos", "sin"], "predicted_agent_rows": [0, 1, 2, "... 85 more"], "force_stop": false, "time_from_start_s": [0.1, 0.2, 0.3, "... 77 more"], "valid_counts": {"ego": 1, "neighbor": 88, "static": 0, "lane": 123, "route": 17, "polygon": 0, "line_string": 60, "goal": 1, "ego_shape": 1, "turn": 1}},
- "timing_ms": {"preprocess": 3.8, "device": 105.3, "postprocess": 3.2, "total": 122.8, "decode": 8.5, "model_call": 112.5},
+ "timing_ms": {"preprocess": 2.4, "device": 22.8, "postprocess": 1.2, "total": 36.9, "decode": 8.4, "model_call": 26.7},
  "num_poses": 80,
  "columns": ["x", "y", "yaw", "cos", "sin", "velocity", "acceleration"],
- "trajectory": [[0.365, 0.01, -0.017, 0.9962, -0.0169, 3.8059, -0.0431], [0.7631, -0.0022, -0.0414, 0.9966, -0.0413, 3.8016, -0.5485], [1.1524, -0.0282, -0.0675, 0.9942, -0.0674, 3.7468, -0.4708], "... 77 more rows"],
- "turn_indicator": {"command": 1, "command_name": "DISABLE", "keep_selected": true, "held": false, "logits": [-15.960084915161133, -5.125061511993408, -5.230083465576172, -0.7658059597015381, 5.504596710205078], "probabilities": [1.651601411190029e-09, 8.38487030705437e-05, 7.548934809165075e-05, 0.006556871347129345, 0.993184506893158]},
+ "trajectory": [[0.3644, 0.01, -0.0169, 0.9964, -0.0169, 3.8033, -0.0392], [0.7626, -0.0021, -0.0412, 0.9968, -0.0412, 3.7994, -0.5496], [1.1517, -0.0281, -0.0674, 0.9942, -0.0672, 3.7444, -0.4683], "... 77 more rows"],
+ "turn_indicator": {"command": 1, "command_name": "DISABLE", "keep_selected": true, "held": false, "logits": [-15.960197448730469, -5.119836807250977, -5.23490047454834, -0.7636525630950928, 5.502235412597656], "probabilities": [1.655269699085693e-09, 8.448460721410811e-05, 7.530191942350939e-05, 0.00658634165301919, 0.9931545257568359]},
  "predicted_agents": {"format": "npz", "key": "predicted_agents", "dtype": "float32", "shape": [88, 80, 5], "data": "<base64 npz>"}
 }
 ```
@@ -168,7 +168,7 @@ the predicted-agents array to its header):
   `~/debug/denoising_steps`).
 
 `timing_ms`: `decode` (base64 + npz parsing + the schema check), `preprocess` (the node's normalization and the encoder's
-host features), `device` (host tensors + H2D + one trace replay + D2H), `postprocess`, `model_call`, `total` (server
+host features), `device` (host tensors + H2D + one trace replay of the request's agent bucket + D2H), `postprocess`, `model_call`, `total` (server
 side, after the body arrived).
 
 ### 3.3 Errors
@@ -187,9 +187,10 @@ when the boot failed), **500** `inference failed: <ExceptionType>: <message>` if
 | `TT_MESH_SHAPE` | launcher (`runtime.mesh_shape_env`) | `1x1`; anything else -> RuntimeError at startup |
 | `TT_DEVICE_ID` | you | chip to open, default 0 |
 | `DIFFUSION_PLANNER_DISPATCH` | `serve.env` | `eth` (default) \| `worker` (A/B only, 11×10) \| `auto` (ETH if the patch is present) |
-| `DIFFUSION_PLANNER_NUM_CQS` | `serve.env` | `1` (2 CQs replay ~10 % slower on ETH, OPT_BASELINE.md) |
+| `DIFFUSION_PLANNER_NUM_CQS` | `serve.env` | `1` (at the first release 2 CQs replayed ~10 % slower on ETH, OPT_BASELINE.md; not re-measured since) |
 | `DIFFUSION_PLANNER_VARIANT` | `serve.env` | `default` (the only v5.0 graph) |
 | `DIFFUSION_PLANNER_LN_FP32`, `_SPLIT_MATMUL`, `_ATTN_MATMUL`, `_HIDDEN_FP32`, `_ATTN_FP32_ACC` | `serve.env` | the numerics knobs, pinned at their gated defaults (`enc.mixer.*,dec.*`; `enc.island.*,enc.pre.*,dec.*`; `enc.fusion.attn,dec.*`; empty; empty: `tt/config.py` `KNOBS`); changing one changes the numerics and needs the gates re-run |
+| `DIFFUSION_PLANNER_ENC_CH2D`, `_ATTN_FAST`, `_DEC_MMCFG`, `_LN_KERNEL`, `_LN_RESID`, `_SPLIT_KCAT`, `_ATTN_SMASK`, `_ATTN_SMSM`, `_ATTN_FUSED`, `_KCAT_EMIT`, `_LN_TR`, `_KCAT_L1`, `_ATTN_L1`, `_DEC_L1`, `_ENC_KCAT`, `_KCAT_ACT`, `_LN_SPLIT`, `_LN_SFPU_BCAST`, `_COMPACT`, `_AGENT_BUCKETS`, `_HOST_FAST`, `_INPUT_TRIM` (and the rejected `_KCAT_ACT_ONCE`, `_ENC_L1`, `_FUS_L1`, `_LIN_ACT` = 0) | `serve.env` | the optimization knobs, pinned at the shipped defaults (each one an A/B of `OPT_REPORT.md`; `tt/config.py` `KNOBS`); the gates were run with exactly these values |
 | `DIFFUSION_PLANNER_PRECISION` | `serve.env` | empty: the precision policy `DEFAULT_PRECISION` of `tt/config.py` (HiFi4 + fp32 accumulation everywhere); extra rules such as `dec.*=HiFi2+fp32` are for experiments |
 | `DIFFUSION_PLANNER_WARMUP` | you | JSON list of warm-up variants, `default` or `none` |
 | `DIFFUSION_PLANNER_TRACE_REGION`, `DIFFUSION_PLANNER_L1_SMALL`, `DIFFUSION_PLANNER_WORKER_L1_SIZE` | you | device-open overrides (validated values: `DEVICE_DEFAULTS` in `device.py`) |
@@ -197,8 +198,8 @@ when the boot failed), **500** `inference failed: <ExceptionType>: <message>` if
 | `TT_METAL_VISIBLE_DEVICES`, `MESH_DEVICE`, `HF_HUB_DISABLE_IMPLICIT_TOKEN` | `serve.env` / launcher | `0`, `P150`, `1` (the weights are public: no token is ever sent) |
 
 Server-side pipeline: JSON -> `tt_diffusion_planner.io` decoders (schema check) -> `DiffusionPlanner.__call__` (the node's
-host pre-processing ported from `planning/autoware_diffusion_planner` -> H2D into 19 persistent device inputs ->
-`execute_trace` of the whole plan -> one D2H -> the node's host post-processing) under one lock -> `Output.to_dict()`.
+host pre-processing ported from `planning/autoware_diffusion_planner` -> H2D into the persistent device inputs ->
+`execute_trace` of the plan's agent-bucket trace -> one D2H -> the node's host post-processing) under one lock -> `Output.to_dict()`.
 
 ### 3.5 What the client keeps (the API is stateless)
 
@@ -228,11 +229,12 @@ sequence of plans keeps the same state and puts it into the request or applies i
 
 - Batch 1, one plan per request; concurrent clients queue on the lock.
 - Fixed shapes of the v5.0 export (320 neighbours, 140 lanes, 25 route lanes, 10 polygons, 60 line strings, 31 history
-  and 80 future steps) and 10 DPM-Solver steps (11 decoder evaluations) are compiled into the trace; every plan computes
-  the full capacities whatever the scene holds, so the device time does not depend on the scene.
+  and 80 future steps) and 10 DPM-Solver steps (11 decoder evaluations) are compiled into the traces. The neighbour trunk
+  and the decoder run on the smallest agent bucket that holds the scene (exact; one trace per bucket), so the device
+  time depends on the neighbour count (17.4-26.7 ms per plan); the map entities are computed at full capacity.
 - The node's guidance services (start / stop / centerline guidance) are off (the node's default too); a guided mode
   would need the host in the solver loop.
-- Warm-up captures the trace inside the lifespan, so READY means warm; the first cold boot pays the ttnn JIT.
+- Warm-up captures the traces inside the lifespan, so READY means warm; the first cold boot pays the ttnn JIT.
 - Weights are pinned by sha in `weights.revision`, exported by the launcher and repeated in `serve.env`;
   `snapshot_download(..., revision=<sha>)` is a cache hit after `serve`'s pre-download and falls back to
   `local_files_only=True` if the Hub is unreachable. The sha256 of the four files and `major_version == 5` are checked

VERIFICATION_OPT_2026-10-11.md

@@ -0,0 +1,157 @@
+# diffusion-planner-p150: independent verification of the optimization rounds 1-5, 2026-10-11
+
+## Verdict
+
+**PASS** (3 low findings, none blocking). Every claim in `OPT_REPORT.md` that I re-measured reproduced on the device
+from a clean checkout of `d9a0436` (the OPT round-5 HEAD), in one session with alternating runs against the published
+baseline `b8114ec`:
+
+- **Device suite (gates read-only):** 44 / 44 passed, run twice (`TTAW_GATES_READONLY=1`). All 99 per-scene numbers
+  are identical to the round-5 record (`logs/diffusion-planner/opt_r5/e2e_compact2.json`, 0 of 99 differ): worst ego
+  max 0.346 m, worst mean 0.154 m (gates 1.0 / 0.3 m), turn command 99 / 99, neighbours ≤ 0.094 m (gate 1.5 m).
+- **No gate loosened:** `test_e2e_device.gates.json`, `test_pcc_device.gates.json`, `test_e2e_device.py` and
+  `test_pcc_device.py` are byte-identical between the published commit `b8114ec` and `d9a0436`.
+- **Speed:** the published baseline gave 102.03 ms back-to-back and 113.1 ms e2e. HEAD gives **20.00 ms** b2b and
+  26.8-27.2 ms e2e on the bench scene. Both match the claims (20.00 ms and 26.84 ms) to within 0.01 ms on the device
+  and within host noise on e2e. The round-end commits r1 to r4 also reproduce to within 0.01 ms.
+- **Held-out (50 scenes outside the 99-scene gate set):** all 50 pass the four gates on HEAD. The set covers the
+  33-instant scene-0061 sequence, 4 adversarial scenes and 13 agent-bucket boundary scenes.
+- **Compaction is exact:** HEAD (`COMPACT=2`) and HEAD with `COMPACT=0` give bit-identical public outputs on all 50
+  held-out scenes. This includes the r128, r192 and full buckets, which no earlier check had covered.
+- **Soak:** 1,200 `model()` calls cycled over 6 scenes covering the r32, r64, r96, r128, r192 and full plans. There
+  were 0 output mismatches against each scene's first output, no hang and no fault. The bench runs also replayed the
+  traces back to back more than 1,350 times on HEAD.
+- **Host suite (fake ttnn):** 108 passed, 45 skipped (the 44 device tests and the ONNX Runtime module).
+
+## Setup
+
+| item | value |
+|---|---|
+| Verified commit | `d9a0436` (OPT round-5 HEAD), checked out as a clean `git worktree` (`dirty=0` printed by every job); the bundle's main working tree carries another agent's uncommitted round-6 work, which was **not** tested |
+| Before | `b8114ec` (the published release, `PUSH_READY`); round-end commits `dd6f249` (r1), `c42c540` (r2), `0b750e3` (r3), `93428bb` (r4) |
+| Configuration | shipped defaults: every `DIFFUSION_PLANNER_*` variable unset (the code defaults equal the `serve.env` pins, `test_serve_env_pins_the_numerics`). ETH dispatch, 1 CQ, 12×10, AICLK 1350 MHz (min 1343) in every run |
+| Device hygiene | every device step ran through `bin/devrun` with explicit test files. `TT_METAL_DISPATCH_TIMEOUT_COMMAND_TO_EXECUTE`, `TT_METAL_OPERATION_TIMEOUT_SECONDS` and `TT_METAL_WATCHER` were unset; pytest used `-o faulthandler_timeout=900`. Device work began only after the orchestrator cleared the CenterPoint FAULT marker of 2026-10-10 20:31 (cleared 00:53 UTC) |
+| Window | 2026-10-11 00:54-01:45 UTC, one session; host load 3.7-5.3 (shared host) |
+| Logs / scripts | `logs/diffusion-planner/verify_opt_r1/` (workspace): `job.sh`, `driver.sh`, `gen_heldout.py`, `dev_heldout.py`, `cmp_heldout.py`, `sum_bench.py`, every log / JSON |
+
+## Speed: same-session re-measurement
+
+`code/scripts/bench.py --iters 30 --warmup 5 --b2b-iters 50 --b2b-rounds 3` was run on 3 scenes:
+`kashiwanoha_dense` (88 neighbours, bucket r96), `straight_road` (r32) and nuScenes `scene-0103_kf14` (r64). The runs
+alternated base / HEAD, 3 pairs, then r4 / HEAD-with-`COMPACT=0`, 2 pairs. Each value below is the median b2b; the
+spread is the range over the pairs.
+
+| build | kashiwanoha b2b | straight_road b2b | scene-0103 b2b | one plan (trace) kashiwanoha | e2e p50 kashiwanoha | claimed in `OPT_REPORT.md` |
+|---|---:|---:|---:|---:|---:|---|
+| published `b8114ec` (3 runs) | 102.03 (±0.00) | 102.03 | 102.03 | 102.09 | 113.11-113.24 | 102.03 / 102.09 / 113.1 (same-session baseline) |
+| r1 `dd6f249` | 70.12 | 70.12 | 70.12 | 70.17 | 80.93 | 70.12 / 70.18 / 81.2 |
+| r2 `c42c540` | 40.89 | 40.89 | 40.89 | 40.95 | 51.77 | 40.89 / 40.95 / 51.98 |
+| r3 `0b750e3` | 30.16 | 30.16 | 30.16 | 30.22 | 41.01 | 30.16 / 30.22 / 40.84 |
+| r4 `93428bb` (2 runs) | 26.62 | 26.62 | 26.62 | 26.68-26.70 | 37.25-37.39 | 26.62 / 26.68 / 37.36 |
+| HEAD `COMPACT=0` (2 runs) | 26.64 | 26.64 | 26.64 | 26.70 | 33.58-33.63 | +0.014 ms over r4 (INPUT_TRIM concat); host items −3.5 ms |
+| **HEAD default (3 runs)** | **20.00** (r96) | **17.44** (r32) | **19.25** (r64) | 20.05-20.06 | **26.82-27.17** | 20.00 / 17.44 / 19.25; 20.06; 26.84 |
+
+- Every device number reproduces the report to within ±0.02 ms.
+- Speed-up on the bench scene: **5.10×** b2b (102.03 -> 20.00 ms) and **4.2×** e2e (113.1 -> 26.9 ms).
+- `staged_equals_model` was true in every run: the staged path is bit-identical to `model()`.
+- Load cost: HEAD's warm `from_pretrained` takes 11.3-11.5 s, against 8.3-8.8 s for the base. HEAD captures 6
+  traces with 60.3 MB of trace buffers and 480 program-cache entries; the base has 74.6 MB and 309 entries.
+- `timing_ms.device` over the 99 gate scenes (from my suite run): p50 20.41 ms, range 20.19-26.77 ms. The report gives
+  p50 20.83 ms. This value includes the upload, so it moves with host load.
+
+## Accuracy
+
+### Frozen gates (99 scenes + per-module PCC)
+
+- 44 passed in 28.3 s (second run; the first run in the same window also passed).
+- The per-scene report (`suite_head.e2e.json`) matches `opt_r5/e2e_compact2.json` on all 99 scenes and all
+  fields except device time.
+- The round-3 precision change (`ENC_KCAT`, 0.327 -> 0.346 m worst ego max) is therefore the current state, as
+  reported.
+
+### Held-out set: 50 scenes outside the 99-scene gate set
+
+The CPU references come from the bundle's fp32 `ReferencePlanner` (CPU only, generated by `gen_heldout.py`). The
+device outputs come from the public API, `model(inputs=raw)`, so the bucket is chosen per request. The gates are the
+same as the e2e gates: ego max ≤ 1.0 m, ego mean ≤ 0.3 m, turn command equal, and neighbour median-of-max ≤ 1.5 m.
+
+| set | n | HEAD: pass / worst ego max / worst mean / worst nb median | published `b8114ec`: worst ego max / mean / nb | HEAD vs HEAD `COMPACT=0` |
+|---|---:|---|---|---|
+| scene-0061 sequence (2 Hz), not in the gate set | 33 | 33 / 33; 0.112 m; 0.055 m; 0.027 m; turn 33 / 33 | 0.118; 0.049; 0.027 | 33 / 33 bit-identical |
+| adversarial: full capacity (320 neighbours, 140 lanes, 25 route, 10 polygons), ego only, x_T = 0.5 N(0,1), x_T = 1.0 N(0,1) | 4 | 4 / 4; 0.058 m; 0.030 m; 0.101 m | 0.062; 0.031; 0.099 | 4 / 4 bit-identical |
+| bucket boundaries: 31, 32, 63, 64, 95, 96, 127, 128, 191, 192 neighbours (variants r32, r64, r64, r96, r96, r128, r128, r192, r192, full), one far slot (100 -> r128, 319 -> full), 31 neighbours + x_T noise | 13 | 13 / 13; **0.507 m** (n192); 0.174 m; 0.076 m | 0.476 (n191); 0.164; 0.076 | 13 / 13 bit-identical |
+
+- The precision changes of r2.3 (`SPLIT_KCAT`) and r3.2 (`ENC_KCAT`) hold on the held-out scenes.
+- Error moves both ways scene by scene, as the report describes. The worst case got worse on dense synthetic scenes:
+  n192 went from 0.432 to 0.507 m and n128 from 0.340 to 0.361 m, still at about half the 1.0 m gate.
+- The worst neighbour max, which is not gated, is 6.3 m on HEAD against 5.8 m on the base, both on the dense
+  synthetic scenes.
+
+### Soak (default path)
+
+- 1,200 `model()` calls in 33.9 s, cycled over 6 scenes: ego only (r32), scene-0061_kf14 (r64), temp05_kashi (r96),
+  n127 (r128), n191 (r192) and full capacity (plan).
+- Each call's poses, predicted agents and turn command were compared byte for byte with that scene's first call:
+  0 mismatches.
+- `timing_ms.device` p50 24.0 ms, max 30.4 ms (full plan).
+- No hang, no fault marker, every devrun job rc = 0.
+
+## Audit
+
+| check | result |
+|---|---|
+| Gates loosened | none (the gate JSONs and the device test files are unchanged since the publish) |
+| Defaults without an A/B row | none. All 26 knobs added by rounds 1-5 have an `OPT_REPORT.md` row. The four rejected ones (`LIN_ACT`, `KCAT_ACT_ONCE`, `ENC_L1`, `FUS_L1`) default to off, and the `serve.env` pins equal the code defaults (host test) |
+| Custom kernels (P3) | every `generic_op` (fattn, ln32 / ln32s, kcat, smask, smsm) passes the per-core runtime-arg count (`N_RT`) as a compile-time arg; the per-core values are runtime args, the buffer addresses common runtime args |
+| Bit-identity claims | r5.1 / r5.2 compaction re-checked on 50 new scenes (above); r1-r4 b2b numbers reproduce; the 99-scene numbers are identical to round 3, as claimed |
+| Hangs / flakiness | none observed (2 suites, 3 held-out runs, 13 benches, 1,200-call soak) |
+| Precision changes vs held-out | r2.3 and r3.2 pass the held-out set (above); not recorded by the optimizer (finding 2) |
+
+## Findings
+
+1. **(low, coverage) The device gate set does not exercise the r128, r192 and full-capacity traces.**
+   - The 99 gate scenes fall in r32, r64 and r96 only (53 / 44 / 2), and `compact_check.py` covered the same three
+     buckets.
+   - Three of the six shipped traces therefore had no gate or bit-identity check before this verification.
+   - They are correct today: they are bit-identical to `COMPACT=0` and within the gates on my boundary scenes.
+   - A future change to a bucket config could still break them without any gate noticing.
+   - Recommendation: add boundary scenes for every bucket to the e2e test as an additional case set. This adds
+     coverage and does not loosen any threshold. Alternatively, run a `COMPACT=0` vs default bit-identity test over
+     scenes that reach every bucket.
+2. **(low, process) No held-out check is recorded for the two kept precision changes (r2.3 `SPLIT_KCAT`, r3.2
+   `ENC_KCAT`).**
+   - `OPT_PLAN.md` asks for one ("plus a 2-scene held-out check"), and the task rules ask for held-out frames on any
+     precision change.
+   - My 50-scene held-out run passes, so the changes stand.
+   - The worst-case ego error on dense synthetic scenes grew from 0.48 to 0.51 m.
+3. **(info) The bundle's main working tree is dirty.** It holds another agent's uncommitted round-6 work (new
+   `fattn_ks_*` kernels and edits to `config.py`, `decoder.py`, `encoder.py` and the ln32s kernels). Anyone who runs
+   from that directory runs code that this verification did not cover. Everything here ran on clean worktrees of the
+   commits named above.
+
+## Reproduce
+
+```bash
+S=<scratch>; R=/home/ubuntu/experiments/tt-models
+# clean worktrees under $S/wt/<name>/bundles/diffusion-planner-p150 with research/, assets/, common/ symlinked
+# next to bundles/ (the tests resolve the goldens as PKG.parents[3]/research)
+git -C $R/bundles/diffusion-planner-p150 worktree add --detach $S/wt/head/bundles/diffusion-planner-p150 d9a0436
+tools/research-venv/bin/python gen_heldout.py $S/wt/head/bundles/diffusion-planner-p150/code $S/heldout   # CPU
+bash driver.sh      # every device step via bin/devrun -t <s> -- bash job.sh <wt> suite|heldout|bench|soak <log> [ENV=..]
+python cmp_heldout.py; python sum_bench.py $S
+```
+
+## Addendum: republish prep (2026-10-11, 01:50-02:25 UTC)
+
+Re-checked on the code that is packaged for the optimized release, after this verification:
+
+| step | commit | result | evidence (`logs/diffusion-planner/republish/`) |
+|---|---|---|---|
+| ttaw re-vendor 0.23.0 -> 0.23.2 (common `60b6dd7`) + the reshape-guard patch shipped | `c91ce91` | host suite 108 passed / 45 skipped; `vendor.py --check` 0 differences | `host_suite_v0232.log` |
+| device suite, gates read-only, alloc tracking, quickstart, stage bench (3 scenes, 100 iterations) | `bc966b4` | 44 passed (twice); 99 per-scene numbers identical to round 5; b2b 20.00 / 17.44 / 19.25 ms, e2e p50 26.73 / 23.13 / 25.65 ms | `job1.out`, `device_suite.log`, `alloc_tracking.log`, `bench_recheck.json` |
+| load times, 132 demo / agreement scenes, served bench | `bc966b4` | 352 s cold / 12.9 s warm; ORT oracle 92 / 92 and 33 / 33 within the gates; served 37.2 ms median, `smoke_test.py` PASS | `job2.out`, `load_*.json`, `gates_*_tt_vs_ort.json`, `served_bench.json` |
+| **new finding (fixed)**: `dispatch="worker"` failed at capture (`TT_FATAL: Num output blocks along x (12) must be smaller than or equal to the number of columns in compute grid (11)`): swept matmul configs assumed the 12×10 ETH grid | `7e003c4` | ETH: 44 passed, 99 numbers identical (configs unchanged on 12×10). WORKER (11×10): `test_e2e_device.py` 2 passed, 99 numbers identical to ETH, b2b 20.61 ms (r96) / 18.01 ms (r32). Host suite 108 passed (+2 new host tests) | `job3.out`, `device_suite_fix.log`, `device_e2e_worker.log`, `bench_worker.log`, `host_suite_fix.log` |
+
+The WORKER failure was outside the scope of the audit above (it covered the shipped ETH configuration only) and
+affected only the A/B opt-in and the fallback for a tt-metal without the ETH patch; the published ETH numbers are
+unchanged by the fix.

build_info.json

@@ -1,13 +1,13 @@
 {
  "schema": "ttaw-build-info/1",
  "bundle": "diffusion-planner-p150",
- "recorded_at": "2026-10-09T04:48:03Z",
+ "recorded_at": "2026-10-11T04:02:55Z",
  "image": {
-  "tag": "tt-model/diffusion-planner-p150:3b96d8ea7190",
-  "digest": "sha256:3b96d8ea71902fe6a00f1792dd41290839b7758070431464cd613cde3f6bf909",
-  "built_at": "2026-10-09T04:37:19+00:00",
-  "code_sha256": "c0e7abb7888098a9319ab5c66a10c4fd4009fc1542326b26507554c04f05a481",
-  "size_bytes": 4069025450,
+  "tag": "tt-model/diffusion-planner-p150:dde78ac2f0be",
+  "digest": "sha256:dde78ac2f0be5b7e637ddceba1a7c30fd832c2a50dd3e728acbf187f86354cf4",
+  "built_at": "2026-10-11T03:53:34+00:00",
+  "code_sha256": "5da22b97bbf89b03f83133089a6a3d01a8862a7a1601437774063d8118d9cae3",
+  "size_bytes": 4069615514,
   "layers": 23
  },
  "base_images": {

code/PYTHON.md

@@ -68,7 +68,7 @@ the three v5.0 ONNX files and `diffusion_planner.param.json`, with an offline fa
 file and the weights' `major_version == 5` are checked), opens the chip (ETH dispatch, 12×10, 1 CQ; the other open
 parameters are `DEVICE_DEFAULTS` in `tt_diffusion_planner/device.py`, overridable with `DIFFUSION_PLANNER_*`), reads the
 ONNX initializers as data, uploads the weights and constants (48.1 MB), builds the graph, then compiles and captures
-the metal trace. If ETH dispatch cannot open (tt-metal without the patch), it warns and falls back to WORKER dispatch
+the metal traces. If ETH dispatch cannot open (tt-metal without the patch), it warns and falls back to WORKER dispatch
 (`model.info["device"]["fallback"]` names it). Any other keyword argument is a `TypeError`.
 
 The numerics are not arguments: the published configuration is the default of the `DIFFUSION_PLANNER_LN_FP32`,
@@ -78,9 +78,9 @@ accuracy figures of the card until the gates are re-run.
 
 ## Warm-up
 
-`from_pretrained` returns a warm model: it builds the graph, runs the plan once eagerly (this first run compiles every kernel into the JIT cache), then captures the whole plan as one metal trace (`warmup_variants="default"`: the variant `plan`) with program-cache misses forbidden, so no later call compiles anything. `model.warmup()` is idempotent; `warmup_variants="none"` defers the capture to `model.warmup()`.
+`from_pretrained` returns a warm model: it builds the graph, runs the plan once eagerly (this first run compiles every kernel into the JIT cache), then captures the whole plan as metal traces (`warmup_variants="default"`: the full-capacity variant `plan` and one variant per agent bucket, `plan_r32`, `plan_r64`, `plan_r96`, `plan_r128`, `plan_r192`; each call replays the smallest that holds the scene, exact) with program-cache misses forbidden, so no later call compiles anything. `model.warmup()` is idempotent; `warmup_variants="none"` defers the capture to `model.warmup()`.
 
-Measured on the shipped sample (`model.info["warmup_ms"]`, 2026-10-09): the load takes 315 s with an empty JIT cache and 8.6 s with a warm one (build 0.53 s: the ONNX initializers read and 48.1 MB of weights and constants uploaded; warm-up and capture 3.7 s; the rest is the device open). The first call then takes 120 ms and the second 120 ms (the stage bench's steady state: 118 ms p50 for decoded arrays, 125 ms for an `.npz` path). The trace holds 74.6 MB of DRAM (`trace_region_size` 192 MiB).
+Measured on the shipped sample (`model.info["warmup_ms"]`, 2026-10-11): the load takes 352 s with an empty JIT cache and 12.9 s with a warm one (build 0.76 s: the ONNX initializers read and 118.9 MB of weights and constants uploaded; warm-up and capture of the 6 traces 7.7 s; the rest is the device open). The first call then takes 37 ms and the second 32 ms (an `.npz` path; the stage bench's steady state: 27 ms p50 for decoded arrays, 32 ms for an `.npz` path). The 6 traces hold 60.3 MB of DRAM (`trace_region_size` 192 MiB).
 
 ## Call: `model(...)`
 
@@ -143,7 +143,7 @@ plans, the caller keeps the same state (SERVING.md 3.5 has the details):
 
 ## Lifetime and information
 
-- `model.close()` releases the trace and the persistent device tensors and closes the chip if the model opened it;
+- `model.close()` releases the traces and the persistent device tensors and closes the chip if the model opened it;
   idempotent. `with` calls it for you; an unclosed model is closed when Python exits.
 - `model.info`: weights (repo, tag, revision, path), device (dispatch, grid, CQs, fallback), variant, warm variants,
   warm-up times, runtime parameter defaults, the input schema, the numerics options and precision policy in effect, the
@@ -152,28 +152,29 @@ plans, the caller keeps the same state (SERVING.md 3.5 has the details):
 
 ## Speed
 
-Warm calls, batch 1, ETH dispatch, 1 CQ, 12×10, the pinned numerics (`code/scripts/bench.py`, 100 iterations; the numbers of `OPT_BASELINE.md`, 2026-10-08, on a shared host; p50, with p99 in brackets):
+Warm calls, batch 1, ETH dispatch, 1 CQ, 12×10, the pinned numerics (`code/scripts/bench.py`, 100 iterations, 2026-10-11, the optimized release, on a shared host; p50, with p99 in brackets):
 
-| stage | shipped sample `kashiwanoha_dense` |
+| stage | shipped sample `kashiwanoha_dense` (88 neighbours: bucket r96) |
 |---|---:|
-| `.npz` decode + schema check (path inputs only) | 7.02 (18.20) ms |
-| host pre-processing (the node's normalization, the encoder's host features, masks, `x_T`) | 4.75 (12.55) ms |
-| packing the 19 persistent trace inputs · ttnn host tensors · H2D | 0.45 · 2.41 · 1.14 ms |
-| **device trace, one blocking plan** | **102.13** (104.81) ms |
-| D2H (one packed read, 460 KB) | 0.58 ms |
-| host post-processing (trajectory, predicted paths, turn decision) | 4.32 (11.54) ms |
-| **`model(inputs=arrays)` end to end** | **117.90** (134.56) ms |
-| `model(inputs=<.npz path>)` | 124.76 (147.74) ms |
-| back-to-back replays (device time per plan) | 102.04 ms = 9.80 plans/s |
-
-The device time does not depend on the scene: every plan computes the full capacities (re-checked on c0d84f9: kashiwanoha_dense 102.10, straight_road 102.11, a nuScenes instant 102.09 ms). Throughput above one plan per ~118 ms needs pipelining of the host work of neighbouring requests (not implemented); 2 CQs do not help a synchronous request (`OPT_BASELINE.md`). Where the time goes and what comes next: `OPT_REPORT.md`.
+| `.npz` decode + schema check (path inputs only) | 5.76 (7.33) ms |
+| host pre-processing (the node's normalization, the encoder's host features, masks, `x_T`) | 2.45 (3.16) ms |
+| packing the persistent trace inputs · ttnn host tensors · H2D | 0.52 · 1.57 · 0.76 ms |
+| **device trace, one blocking plan** | **20.05** (20.08) ms |
+| D2H (one packed read) | 0.22 ms |
+| host post-processing (trajectory, predicted paths, turn decision) | 1.39 (1.84) ms |
+| **`model(inputs=arrays)` end to end** | **26.73** (27.48) ms |
+| `model(inputs=<.npz path>)` | 32.27 (33.61) ms |
+| back-to-back replays (device time per plan) | 20.00 ms = 50.0 plans/s |
+
+The device time depends on the agent bucket the scene needs (exact compaction: the smallest of 32 / 64 / 96 / 128 / 192 decoder rows that holds the ego and every valid neighbour row, else the full capacity). Back to back: r32 (≤ 31 neighbours, e.g. `straight_road`) 17.44 ms, r64 (a nuScenes instant with 42 neighbours) 19.25 ms, r96 20.00 ms, r128 21.11 ms, r192 22.30 ms, full capacity (> 191) 26.65 ms (the last three: `OPT_REPORT.md` round 5). The map entities are computed at full capacity in every plan. The first release took 102.04 ms for every scene (`OPT_BASELINE.md`). 2 CQs were not re-measured on this release (at the first release they did not help a synchronous request: `OPT_BASELINE.md`). Where the time goes and what comes next: `OPT_REPORT.md`.
 
 ## Limits
 
 - Batch 1 on the chip; one model per process.
 - Fixed shapes of the v5.0 export (320 neighbours, 140 lanes, 25 route lanes, 10 polygons, 60 line strings, 31 history
-  and 80 future steps) and 10 DPM-Solver steps (11 decoder evaluations) are compiled into the trace; every plan computes
-  the full capacities, so the device time does not depend on the scene.
+  and 80 future steps) and 10 DPM-Solver steps (11 decoder evaluations) are compiled into the traces; the neighbour trunk
+  and the decoder run on the smallest agent bucket that holds the scene (exact), so the device time depends on the
+  neighbour count (17.4-26.7 ms); the map entities are computed at full capacity.
 - The node's guidance services (start / stop / centerline guidance) are not available (the node's default is off).
 - Accuracy is agreement with the fp32 CPU reference of the same network (README "Demo & Performances"); the planner's
   driving quality is the weights' (trained by TIER IV on data that is not public).

code/scripts/bench.py

@@ -74,12 +74,12 @@ def counts(arrays: Dict[str, np.ndarray]) -> Dict[str, int]:
             "line_strings": rows(arrays["line_strings"][0], (1, 2))}
 
 
-def unpack(out: Dict[str, np.ndarray]) -> Dict[str, Any]:
+def unpack(out: Dict[str, np.ndarray], tt: Any) -> Dict[str, Any]:
     """The packed readback -> the raw outputs of ``TtDiffusionPlanner.forward`` (same reshapes)."""
     from tt_diffusion_planner.reference import config as C
     from tt_diffusion_planner.tt import config as T
 
-    final = out["final_x0"].reshape(T.AGENTS, T.STATE_COLS)[:C.MAX_NUM_AGENTS]
+    final = tt.unpack(out)
     steps = out["ego_steps"].reshape(-1, T.STATE_COLS)
     return {"final_x0": final.reshape(C.MAX_NUM_AGENTS, C.OUTPUT_T + 1, C.POSE_DIM).astype(np.float32),
             "logit": out["logit"].reshape(-1)[:C.TURN_INDICATOR_OUTPUT_DIM].astype(np.float32),
@@ -101,6 +101,7 @@ def bench_scene(model, scene: Dict[str, Any], a: argparse.Namespace) -> Dict[str
     obs = model.normalization.observation
     bench = StageBench(f"diffusion-planner {scene['name']}")
     timing: Dict[str, list] = {}
+    variant = model.tt.variant_for(hp.prepare(load_named_arrays(arrays, C.INPUT_SCHEMA), obs))   # COMPACT bucket
     sync = lambda: ttnn.synchronize_device(dev)                     # noqa: E731
     for _ in range(a.warmup):
         ref = model(inputs=arrays)
@@ -120,7 +121,7 @@ def bench_scene(model, scene: Dict[str, Any], a: argparse.Namespace) -> Dict[str
             with bench.stage("host_pre"):
                 prep = hp.prepare(raw, obs)
             with bench.stage("pack"):
-                packed = I.plan_inputs(prep)
+                packed = model.tt.filter_inputs(I.plan_inputs(prep))     # INPUT_TRIM (as model() does)
             with bench.stage("host_in"):
                 host = {k: to_host_tensor(v, slots[k].dtype, slots[k].layout, shape=slots[k].shape)
                         for k, v in packed.items()}
@@ -128,25 +129,37 @@ def bench_scene(model, scene: Dict[str, Any], a: argparse.Namespace) -> Dict[str
                 runner.upload(host)
                 sync()
             with bench.stage("trace"):
-                runner.replay("plan")
+                runner.replay(variant)
                 sync()
             with bench.stage("d2h"):
-                out = runner.read("plan")
+                out = runner.read(variant)
             with bench.stage("host_post"):
-                res = model._postprocess(unpack(out), prep, params)
+                res = model._postprocess(unpack(out, model.tt), prep, params)
         if scene["schema_npz"]:
             for _ in range(a.iters):
                 with bench.stage("e2e_path"):
                     model(inputs=scene["path"])
-        rounds = [time_b2b(lambda: runner.replay("plan"), sync, n=a.b2b_iters, warmup=3)
+        rounds = [time_b2b(lambda: runner.replay(variant), sync, n=a.b2b_iters, warmup=3)
                   for _ in range(a.b2b_rounds)]
         for r in rounds:
             bench.add("b2b", r)
+    if getattr(a, "dump", None):
+        Path(a.dump).parent.mkdir(parents=True, exist_ok=True)
+        flat: Dict[str, np.ndarray] = {}
+
+        def walk(prefix, v):
+            if isinstance(v, dict):
+                for k, w in v.items():
+                    walk(f"{prefix}{k}/", w)
+            else:
+                flat[prefix.rstrip("/") or "out"] = np.asarray(v)
+        walk("", out)
+        np.savez(f"{a.dump}.{scene['name'].replace('/', '_')}.npz", **flat)
     same = bool(np.array_equal(res.poses, ref.poses) and np.array_equal(res.predicted_agents, ref.predicted_agents)
                 and res.turn_indicator["command"] == ref.turn_indicator["command"])
     summary = bench.summary()
     b2b = statistics.median(rounds)
-    return {"name": scene["name"], "path": scene["path"], "valid": counts(arrays), "stages_ms": summary,
+    return {"name": scene["name"], "path": scene["path"], "valid": counts(arrays), "variant": variant, "stages_ms": summary,
             "timing_ms": {k: {"p50": float(np.percentile(v, 50)), "p99": float(np.percentile(v, 99)),
                               "min": float(min(v))} for k, v in timing.items()},
             "b2b_rounds_ms": rounds, "plans_per_s_b2b": 1000.0 / b2b, "aiclk": clk.summary(),
@@ -166,6 +179,8 @@ def main() -> None:
     ap.add_argument("--chip", type=int, default=0, help="sysfs chip index for the AICLK sampler")
     ap.add_argument("--tag", default="")
     ap.add_argument("--json")
+    ap.add_argument("--dump", help="prefix: save each scene's raw packed readback as <prefix>.<scene>.npz "
+                    "(bit-identity checks of structural rewrites)")
     a = ap.parse_args()
     from tt_diffusion_planner.reference.weights import find_weights_dir
 
@@ -191,7 +206,7 @@ def main() -> None:
             res["scenes"][r["name"]] = r
             cfg = model.device_info
             print(f"\n## {r['name']}  [{cfg.get('dispatch')} {res['num_command_queues']}CQ {cfg.get('grid')}] "
-                  f"valid {r['valid']}")
+                  f"valid {r['valid']} variant {r['variant']}")
             print(r["table"])
             print("timing_ms p50:", {k: round(v["p50"], 3) for k, v in r["timing_ms"].items()},
                   f"| b2b rounds {[round(x, 3) for x in r['b2b_rounds_ms']]} ms -> {r['plans_per_s_b2b']:.2f} plans/s",

code/scripts/profile_ops.py

@@ -209,7 +209,10 @@ def main() -> None:
         tt, runner, dev = model.tt, model.runner, model.device
         print("loaded in %.1f s:" % (time.perf_counter() - t0), json.dumps(model.device_info), flush=True)
         raw = load_named_arrays(a.input, C.INPUT_SCHEMA)
-        inputs = I.plan_inputs(hp.prepare(raw, model.normalization.observation))
+        prep = hp.prepare(raw, model.normalization.observation)
+        inputs = I.plan_inputs(prep)
+        variant = tt.variant_for(prep)                 # COMPACT: the served bucket trace
+        print("variant", variant, flush=True)
         served = model(inputs=raw)                     # one served plan: upload + replay + readback
         ttnn.synchronize_device(dev)
         read_device_profiler(dev)                      # warm-up / capture / first plan out of the buffer
@@ -219,7 +222,7 @@ def main() -> None:
             try:
                 t1 = time.perf_counter()
                 with signposted("eager"):
-                    eager = runner.run_eager("plan", inputs=inputs)
+                    eager = runner.run_eager(variant, inputs=inputs)
                     ttnn.synchronize_device(dev)
                 timings["eager_ms"] = (time.perf_counter() - t1) * 1e3
             finally:
@@ -232,11 +235,11 @@ def main() -> None:
         read_device_profiler(dev)
         t1 = time.perf_counter()
         with signposted("trace"):
-            runner.replay("plan", n=a.replays)
+            runner.replay(variant, n=a.replays)
             ttnn.synchronize_device(dev)
         timings["trace_ms"] = (time.perf_counter() - t1) * 1e3 / a.replays
         read_device_profiler(dev)
-        out = runner.read("plan")
+        out = runner.read(variant)
         same = bool(np.array_equal(out["final_x0"], eager["final_x0"])) if a.eager else None
         desc = {"device": model.device_info, "timings_ms": timings, "replays": a.replays,
                 "replay_equals_eager": same, "turn_command": int(served.turn_indicator["command"]),

code/tt_diffusion_planner/__init__.py

@@ -16,7 +16,7 @@ never edit it here). ``tt_diffusion_planner.reference`` is the fp32 CPU referenc
 ``tt_diffusion_planner.host`` the node's pre- and post-processing (numpy).
 """
 
-__version__ = "0.1.0"
+__version__ = "0.2.0"
 __all__ = ["DiffusionPlanner", "Output", "open_device", "load_inputs", "INPUT_SCHEMA", "__version__"]
 
 _LAZY = {

code/tt_diffusion_planner/host/features.py

@@ -115,6 +115,18 @@ class DecoderMasks:
     current_states: np.ndarray  # [321, 4] normalised (prefix constraint)
 
 
+
+def _host_fast() -> bool:
+    """``DIFFUSION_PLANNER_HOST_FAST`` (OPT round 5 item 3), read once: the vectorised host functions (bit-exact)
+    instead of the first port's ``*_ref`` versions."""
+    from ..tt.config import KNOBS
+
+    return bool(KNOBS.read().HOST_FAST)
+
+
+HOST_FAST = _host_fast()
+
+
 def _any_nonzero(a: np.ndarray, axes) -> np.ndarray:
     return np.any(a != 0, axis=axes)
 
@@ -129,16 +141,31 @@ def encoder_features(norm: Mapping[str, np.ndarray], masks: Mapping[str, np.ndar
 
     # neighbours: keep the 6 newest rows (encoder.py:176-181, 441-451)
     nb_raw = f32("neighbor_agents_past")
-    nb = np.zeros_like(nb_raw)
-    nb[:, C.NEIGHBOR_HISTORY_KEEP] = nb_raw[:, C.NEIGHBOR_HISTORY_KEEP]
-    nb_type = nb[:, -1, 8:11].copy()
-    x8 = nb[..., :8]
-    step_valid = _any_nonzero(x8, -1)                      # [320, 31]
-    nb_valid = step_valid.any(axis=-1)                     # [320]
-    feat = np.concatenate([x8, step_valid[..., None].astype(np.float32)], axis=-1)
-    feat[..., 4:6] = 0.0                                   # velocities zeroed after the validity test
-    feat[~nb_valid] = 0.0
-    pos["neighbor"] = _onehot_pos(x8[:, -1, :4], C.POS_CLASS["neighbor"])
+    if HOST_FAST:
+        K = C.NEIGHBOR_HISTORY_KEEP                            # the other rows are zero: computed on the kept ones
+        x8k = nb_raw[:, K, :8]
+        svk = _any_nonzero(x8k, -1)                            # [320, 6] valid steps
+        nb_valid = svk.any(axis=-1)                            # [320]
+        nb_type = nb_raw[:, -1, 8:11].copy()                   # row 30 is kept
+        feat = np.zeros(nb_raw.shape[:2] + (9,), np.float32)
+        fk = feat[:, K]
+        fk[..., :8] = x8k
+        fk[..., 8] = svk
+        fk[..., 4:6] = 0.0                                     # velocities zeroed after the validity test
+        fk[~nb_valid] = 0.0
+        feat[:, K] = fk
+        pos["neighbor"] = _onehot_pos(nb_raw[:, -1, :4], C.POS_CLASS["neighbor"])
+    else:
+        nb = np.zeros_like(nb_raw)
+        nb[:, C.NEIGHBOR_HISTORY_KEEP] = nb_raw[:, C.NEIGHBOR_HISTORY_KEEP]
+        nb_type = nb[:, -1, 8:11].copy()
+        x8 = nb[..., :8]
+        step_valid = _any_nonzero(x8, -1)                      # [320, 31]
+        nb_valid = step_valid.any(axis=-1)                     # [320]
+        feat = np.concatenate([x8, step_valid[..., None].astype(np.float32)], axis=-1)
+        feat[..., 4:6] = 0.0                                   # velocities zeroed after the validity test
+        feat[~nb_valid] = 0.0
+        pos["neighbor"] = _onehot_pos(x8[:, -1, :4], C.POS_CLASS["neighbor"])
 
     static = f32("static_objects")
     st_valid = _any_nonzero(static[..., :10], -1)

code/tt_diffusion_planner/host/normalize.py

@@ -19,7 +19,18 @@ __all__ = ["FLT_EPSILON", "normalize_inputs", "normalize_array", "speed_masks"]
 FLT_EPSILON = np.float32(np.finfo(np.float32).eps)  # std::numeric_limits<float>::epsilon()
 
 
-def normalize_array(value: np.ndarray, mean: np.ndarray, std: np.ndarray) -> np.ndarray:
+def _host_fast() -> bool:
+    """``DIFFUSION_PLANNER_HOST_FAST`` (OPT round 5 item 3), read once: the vectorised host functions (bit-exact)
+    instead of the first port's ``*_ref`` versions."""
+    from ..tt.config import KNOBS
+
+    return bool(KNOBS.read().HOST_FAST)
+
+
+HOST_FAST = _host_fast()
+
+
+def normalize_array_ref(value: np.ndarray, mean: np.ndarray, std: np.ndarray) -> np.ndarray:
     """``normalize_vector`` of ``preprocessing_utils.cpp:36-63`` on one tensor (a new float32 array).
 
     Rows have ``std.size`` columns; a single-value mean / std broadcasts over the columns (the C++ ``mean.size() == 1``
@@ -43,6 +54,57 @@ def normalize_array(value: np.ndarray, mean: np.ndarray, std: np.ndarray) -> np.
     return rows.reshape(data.shape)
 
 
+_NORMALIZERS: Dict[Tuple[int, int], tuple] = {}
+
+
+def _normalizer(mean: np.ndarray, std: np.ndarray):
+    """``(cols, m, s)`` of a normalizer, validated once per (mean, std) object pair (the weights' constants)."""
+    key = (id(mean), id(std))
+    hit = _NORMALIZERS.get(key)
+    if hit is not None and hit[0] is mean and hit[1] is std:
+        return hit[2]
+    m32 = np.asarray(mean, np.float32).reshape(-1)
+    s32 = np.asarray(std, np.float32).reshape(-1)
+    if m32.size != s32.size:
+        raise ValueError("Mean and std must be same size")
+    cols = s32.size
+    if cols and np.any(np.abs(s32) < FLT_EPSILON):
+        raise ValueError("Standard deviation is zero, cannot normalize data")
+    m = np.ascontiguousarray(np.broadcast_to(m32 if m32.size > 1 else m32[:1], (cols,)))
+    s = np.ascontiguousarray(np.broadcast_to(s32 if s32.size > 1 else s32[:1], (cols,)))
+    out = (cols, m, s, np.ones(cols, np.uint8) if cols < 256 else None)
+    if len(_NORMALIZERS) > 256:
+        _NORMALIZERS.clear()
+    _NORMALIZERS[key] = (mean, std, out)
+    return out
+
+
+def normalize_array(value: np.ndarray, mean: np.ndarray, std: np.ndarray) -> np.ndarray:
+    """``normalize_vector`` of ``preprocessing_utils.cpp:36-63`` on one tensor (a new float32 array).
+
+    Rows have ``std.size`` columns; a single-value mean / std broadcasts over the columns (the C++ ``mean.size() == 1``
+    branch); a zero standard deviation raises like the C++. Only the rows that are not all-small are computed
+    (``(v - mean) / std`` in float32, element by element); the others are copied unchanged."""
+    if not HOST_FAST:
+        return normalize_array_ref(value, mean, std)
+    data = np.asarray(value, dtype=np.float32)   # not modified: the result is a new array
+    cols, m, s, ones = _normalizer(mean, std)
+    if cols == 0 or data.size % cols:
+        raise ValueError(f"data size {data.size} is not divisible by the normalizer size {cols}")
+    rows = data.reshape(-1, cols)
+    big = ~(np.abs(rows) < FLT_EPSILON)             # NaN counts as not small, as in the C++ test
+    if ones is not None:                            # per-row count of the non-small values: a uint8 dot product
+        keep = (big.view(np.uint8) @ ones) != 0     # (faster than a short-axis any())
+    else:
+        keep = big.any(axis=1)
+    idx = np.flatnonzero(keep)
+    if 2 * idx.size > keep.size:                    # mostly valid rows: compute all, keep the small rows
+        return np.where(keep[:, None], (rows - m) / s, rows).reshape(data.shape)
+    out = rows.copy()
+    out[idx] = (rows[idx] - m) / s
+    return out.reshape(data.shape)
+
+
 def normalize_inputs(raw: Mapping[str, np.ndarray],
                      observation: Mapping[str, Tuple[np.ndarray, np.ndarray]]) -> Dict[str, np.ndarray]:
     """``normalize_input_data(input_data_map, normalization_map)``: every key except the four skipped ones must have

code/tt_diffusion_planner/host/pipeline.py

@@ -19,6 +19,7 @@ from ..ttaw.io import InputError, encode_array
 from ..ttaw.outputs import Trajectory
 from .features import DecoderMasks, EncoderFeatures, decoder_masks, encoder_features
 from .normalize import normalize_inputs, speed_masks
+from .postprocess import HOST_FAST
 from .postprocess import (TurnIndicatorManager, denoising_steps_ego, denormalize, predicted_paths,
                           trajectory_from_poses)
 
@@ -83,15 +84,21 @@ def make_output(final_x0: np.ndarray, logit: np.ndarray, prepared: Prepared, nor
     ``RUNTIME_PARAMS``; ``denoising_steps``: the 11 iterates ``[321, 81, 4]`` when ``return_denoising_steps``."""
     x0 = np.asarray(final_x0, np.float32).reshape(C.MAX_NUM_AGENTS, C.OUTPUT_T + 1, C.POSE_DIM)
     mean, std = normalization.state()
-    denorm = denormalize(x0, mean, std)  # [321, 80, 4]
+    rows = prepared.neighbor_rows
+    if HOST_FAST:   # only the rows read below: the ego and the emitted neighbours (element-wise: the same values)
+        idx = np.concatenate([[0], np.asarray(rows, np.int64) + 1])
+        pick = lambda v: v[idx] if v.shape[0] == C.MAX_NUM_AGENTS else v  # noqa: E731
+        denorm = (x0[idx, 1:] * pick(std) + pick(mean)).astype(np.float32)  # [1 + n, 80, 4]
+    else:
+        denorm = denormalize(x0, mean, std)  # [321, 80, 4]
     traj = trajectory_from_poses(denorm[0], (0.0, 0.0, 0.0),
                                  velocity_smoothing_window=int(params["velocity_smoothing_window"]),
                                  enable_force_stop=prepared.enable_force_stop,
                                  stopping_threshold=float(params["stopping_threshold"]))
     manager = TurnIndicatorManager(keep_offset=float(params["turn_indicator_keep_offset"]))
     decision = manager.evaluate(np.asarray(logit, np.float32).reshape(-1), 0.0, prepared.prev_report)
-    rows = prepared.neighbor_rows
-    paths = predicted_paths(denorm, C.MAX_NUM_NEIGHBORS)[rows] if rows.size else np.zeros(
+    paths = (predicted_paths(denorm, rows.size) if HOST_FAST else predicted_paths(denorm, C.MAX_NUM_NEIGHBORS)[rows]) \
+        if rows.size else np.zeros(
         (0, C.OUTPUT_T, len(PREDICTED_AGENT_COLUMNS)), np.float32)
     out_meta: Dict[str, Any] = {"predicted_agent_columns": list(PREDICTED_AGENT_COLUMNS),
                                 "predicted_agent_rows": [int(r) for r in rows],

code/tt_diffusion_planner/host/postprocess.py

@@ -102,7 +102,19 @@ class EgoTrajectory:
                          self.velocity, self.acceleration], axis=-1).astype(np.float32)
 
 
-def trajectory_from_poses(poses_xycs: np.ndarray, base_position: Tuple[float, float, float] = (0.0, 0.0, 0.0), *,
+
+def _host_fast() -> bool:
+    """``DIFFUSION_PLANNER_HOST_FAST`` (OPT round 5 item 3), read once: the vectorised host functions (bit-exact)
+    instead of the first port's ``*_ref`` versions."""
+    from ..tt.config import KNOBS
+
+    return bool(KNOBS.read().HOST_FAST)
+
+
+HOST_FAST = _host_fast()
+
+
+def trajectory_from_poses_ref(poses_xycs: np.ndarray, base_position: Tuple[float, float, float] = (0.0, 0.0, 0.0), *,
                           velocity_smoothing_window: int = C.VELOCITY_SMOOTHING_WINDOW,
                           enable_force_stop: bool = True,
                           stopping_threshold: float = C.STOPPING_THRESHOLD) -> EgoTrajectory:
@@ -151,6 +163,61 @@ def trajectory_from_poses(poses_xycs: np.ndarray, base_position: Tuple[float, fl
     return EgoTrajectory(pos, cs, quat, tf2_yaw(quat), vel, accel, tfs, force_stop)
 
 
+def trajectory_from_poses(poses_xycs: np.ndarray, base_position: Tuple[float, float, float] = (0.0, 0.0, 0.0), *,
+                          velocity_smoothing_window: int = C.VELOCITY_SMOOTHING_WINDOW,
+                          enable_force_stop: bool = True,
+                          stopping_threshold: float = C.STOPPING_THRESHOLD) -> EgoTrajectory:
+    """``get_trajectory_from_poses`` (``postprocessing_utils.cpp:362-455``) for poses ``[N, 4]`` = denormalised
+    (x, y, cos, sin) float32 of one agent, positions at z = ``base_position[2]`` (identity transform)."""
+    if not HOST_FAST:
+        return trajectory_from_poses_ref(poses_xycs, base_position, velocity_smoothing_window=velocity_smoothing_window,
+                                         enable_force_stop=enable_force_stop, stopping_threshold=stopping_threshold)
+    p = np.asarray(poses_xycs, np.float32)
+    n = p.shape[0]
+    dt = C.TRAJECTORY_DT
+    pos = np.zeros((n, 3), np.float64)
+    pos[:, 0] = p[:, 0].astype(np.float64)
+    pos[:, 1] = p[:, 1].astype(np.float64)
+    pos[:, 2] = float(base_position[2])
+    cs = p[:, 2:4].astype(np.float64).copy()
+    quat = quaternion_from_cos_sin(cs[:, 0], cs[:, 1])
+    # the node's loops, vectorised with the same float64 / float32 operations in the same order (bit-exact)
+    prev = np.asarray(base_position, np.float64)
+    px, py, pz = float(prev[0]), float(prev[1]), float(prev[2])
+    hyp = math.hypot
+    dist = []
+    for x, y, z in pos.tolist():            # math.hypot (3 args) kept: np.hypot rounds differently
+        dist.append(hyp(x - px, y - py, z - pz))
+        px, py, pz = x, y, z
+    vel = (np.asarray(dist, np.float64) / dt).astype(np.float32)
+    w = int(velocity_smoothing_window)
+    if n <= w:
+        raise ValueError("velocity_smoothing_window must be smaller than number of points")
+    thr = np.float32(stopping_threshold)
+    m = n - w + 1
+    v64 = vel.astype(np.float64)
+    acc = np.zeros(m, np.float64)
+    for k in range(w):                      # acc = ((0 + v[i]) + v[i+1]) + ... per window, as the C++ loop
+        acc = acc + v64[k:k + m]
+    sm = (acc / float(w)).astype(np.float32)  # window i reads only original values (it writes vel[i] last)
+    vel[:m] = sm
+    force_stop = False
+    stop = -1
+    if enable_force_stop and m > 1:
+        hit = np.flatnonzero((np.abs(sm[:-1]) > thr) & (np.abs(sm[1:]) < thr))
+        if hit.size:
+            stop, force_stop = int(hit[0]) + 1, True
+    if force_stop:                          # from the first stop on: zero velocity, the pose frozen
+        vel[stop:] = np.float32(0.0)
+        pos[stop:], cs[stop:], quat[stop:] = pos[stop - 1], cs[stop - 1], quat[stop - 1]
+    else:
+        vel[m:] = vel[m - 1]
+    accel = np.zeros(n, np.float32)
+    accel[:n - 1] = ((vel[1:].astype(np.float64) - vel[:-1].astype(np.float64)) / dt).astype(np.float32)
+    tfs = dt * (np.arange(n, dtype=np.float64) + 1.0)
+    return EgoTrajectory(pos, cs, quat, tf2_yaw(quat), vel, accel, tfs, force_stop)
+
+
 def predicted_paths(denorm: np.ndarray, n_neighbors: int) -> np.ndarray:
     """``create_predicted_objects`` poses for the first ``n_neighbors`` neighbours: ``[n, 80, 5]`` float32 =
     x, y, yaw (``tf2::getYaw``), cos, sin. ``denorm`` is ``[321, 80, 4]`` (agent 0 = ego)."""

code/tt_diffusion_planner/tests/test_bundle_host.py

@@ -133,7 +133,11 @@ def test_serve_env_pins_the_numerics(manifest):
     env = manifest["serve"]["env"]
     pins = KNOBS.serve_env()
     assert set(pins) == {f"DIFFUSION_PLANNER_{k}" for k in
-                         ("LN_FP32", "HIDDEN_FP32", "SPLIT_MATMUL", "ATTN_FP32_ACC", "ATTN_MATMUL")}
+                         ("LN_FP32", "HIDDEN_FP32", "SPLIT_MATMUL", "ATTN_FP32_ACC", "ATTN_MATMUL",
+                          "ENC_CH2D", "ATTN_FAST",
+                          "DEC_MMCFG", "LN_KERNEL", "LN_RESID", "SPLIT_KCAT", "LN_SFPU_BCAST", "LN_SPLIT", "KCAT_ACT", "ATTN_SMASK",
+                          "ATTN_SMSM", "ENC_KCAT", "LIN_ACT", "ATTN_FUSED", "KCAT_EMIT", "LN_TR", "KCAT_L1", "KCAT_ACT_ONCE", "ATTN_L1", "DEC_L1", "ENC_L1", "FUS_L1",
+                          "COMPACT", "AGENT_BUCKETS", "HOST_FAST", "INPUT_TRIM")}
     assert {k: env.get(k) for k in pins} == pins
     assert env.get("DIFFUSION_PLANNER_PRECISION") == ""
     for profile in manifest.get("serve_profiles") or []:

code/tt_diffusion_planner/tests/test_grid_fit_host.py

@@ -0,0 +1,23 @@
+# SPDX-License-Identifier: Apache-2.0
+"""The swept 2-D multicast configs fit the 12x10 ETH grid and fall back to the auto config on the 11x10 WORKER grid
+(republish 2026-10-11: ``dispatch="worker"`` failed at capture with "Num output blocks along x (12) must be smaller
+than or equal to the number of columns in compute grid (11)")."""
+from types import SimpleNamespace
+
+from tt_diffusion_planner.tt.layers import DEC_MM_CONFIGS, DEC_ROWS, KCAT_CONFIGS, fit_pcm, mcast_fits
+
+ETH, WORKER = SimpleNamespace(x=12, y=10), SimpleNamespace(x=11, y=10)
+
+
+def test_every_swept_config_fits_the_eth_grid():
+    for rows in (DEC_ROWS, 32, 64, 96, 128, 192):
+        for (k, n, _ps), (tr, pcm, pcn, _kb, _sw) in DEC_MM_CONFIGS.items():
+            assert mcast_fits(tr, fit_pcm(tr, pcm, rows, ETH), pcn, rows, n, ETH), (rows, k, n)
+    for (rows, _ktp, n), (tr, pcm, pcn, _kb, _sw) in KCAT_CONFIGS.items():
+        assert mcast_fits(tr, pcm, pcn, rows, n, ETH), (rows, n)
+
+
+def test_the_qkv_config_does_not_fit_the_worker_grid():
+    tr, pcm, pcn, _kb, _sw = KCAT_CONFIGS[(352, 25, 768)]
+    assert not mcast_fits(tr, pcm, pcn, 352, 768, WORKER)        # 24 N tiles / 2 = 12 blocks on x > 11
+    assert mcast_fits(False, 2, 1, 352, 256, WORKER)

code/tt_diffusion_planner/tests/test_tt_params_host.py

@@ -156,3 +156,31 @@ def test_plan_inputs_layout(prepared):
     np.testing.assert_array_equal(y0[:C.MAX_NUM_AGENTS, 4:], prepared.x_T.reshape(C.MAX_NUM_AGENTS, -1)[:, 4:])
     warm = I.warmup_inputs()
     assert all(warm[k].shape == v for k, v in I.INPUT_SPECS.items())
+
+
+def test_compact_bucket_covers_every_needed_row():
+    """``COMPACT``: the bucket holds the ego, every valid self-attention key and every emitted neighbour row (the
+    rows past it are masked keys, never read back), and falls back to the full 352 rows."""
+    from types import SimpleNamespace
+
+    from tt_diffusion_planner.tt.model import bucket_rows, needed_rows
+
+    def prep(valid_idx, emitted):
+        v = np.zeros(C.MAX_NUM_AGENTS, bool)
+        v[0] = True
+        v[list(valid_idx)] = True
+        nb = np.zeros(C.MAX_NUM_NEIGHBORS, bool)
+        nb[[i - 1 for i in valid_idx]] = True
+        return SimpleNamespace(decoder=SimpleNamespace(agent_valid=v), neighbor_rows=np.asarray(emitted, int),
+                               features=SimpleNamespace(valid={"neighbor": nb}))
+
+    b = (32, 64, 96, 128, 192)
+    assert needed_rows(prep([], [])) == 1 and bucket_rows(prep([], []), b) == 32
+    assert needed_rows(prep([1, 2, 31], [0, 1, 30])) == 32 and bucket_rows(prep([31], []), b) == 32
+    assert bucket_rows(prep([32], []), b) == 64                      # the 33rd row is a valid key
+    assert bucket_rows(prep([5], [70]), b) == 96                     # emitted neighbour 70 = decoder row 71
+    assert bucket_rows(prep([200], []), b) == T.AGENTS
+    assert bucket_rows(prep([5], []), ()) == T.AGENTS                # COMPACT off
+    p = prep([3], [])
+    p.features.valid["neighbor"][40] = True                          # a valid neighbour token (row 41)
+    assert needed_rows(p) == 42 and bucket_rows(p, b) == 64

code/tt_diffusion_planner/tt/attention.py

@@ -0,0 +1,96 @@
+# SPDX-License-Identifier: Apache-2.0
+"""The fp32 matmul attention of the fusion transformer and the decoder (C20 ``attention_matmul`` math) with stock-op
+program choices tuned for the planner's shapes (``ATTN_FAST``, OPT round 1 item 4a).
+
+``softmax(scale * Q K^T + mask) V`` as in ``ttaw.ops.attention.attention_matmul``, with:
+
+- (``mode >= 1``) ``P V`` on a ``MatmulMultiCoreReuseProgramConfig`` with one output tile per core (8 heads x
+  Sq / 32 = 88 cores at Sq = 352) instead of the auto config's 11-18 cores; the K-sum in one block
+  (``in0_block_w`` = Sk / 32): bit-identical to the auto config in the device sweep (73 -> 21 us self, 71 -> 32 us
+  cross, 81 -> 42 us fusion; ``logs/diffusion-planner/opt_r1/attn_sweep.log``);
+- (``mode == 2``) the scale applied to Q (``[1, 8, Sq, 32]``, 11x smaller than the scores) before ``Q K^T``
+  instead of to the scores: one small multiply instead of the large one; a precision change (rounding order);
+- (``mode == 3``) the scale as an SFPU pre-activation of the mask add (``input_tensor_a_activations``): one
+  program instead of two over the scores (75 -> 42 us), but NOT bit-identical (final_x0 rel 4e-3; e2e ego mean
+  0.143 -> 0.173 m over the 99 scenes): measured and rejected (OPT_REPORT round 1).
+- ``Q K^T`` keeps the auto config (the reuse configs of the sweep were slower).
+"""
+from __future__ import annotations
+
+from typing import Any, Optional
+
+__all__ = ["attention_matmul", "pv_config"]
+
+TILE = 32
+# per-core L1 budget for the P block of one core (fp32, double-buffered): 18 tiles x 4 KiB x 2 = 144 KiB
+MAX_KB_TILES = 18
+
+
+def pv_config(device: Any, heads: int, sq: int, sk: int):
+    """``MatmulMultiCoreReuseProgramConfig`` for ``P [1, H, Sq, Sk] @ V [1, H, Sk, 32]``: one ``[32, 32]`` output
+    tile per core (``per_core_M = 1`` when ``H * Sq / 32`` fits the grid, else 2)."""
+    import ttnn
+
+    grid = device.compute_with_storage_grid_size()
+    mt, kt = sq // TILE, sk // TILE
+    pcm = 1 if heads * mt <= grid.x * grid.y else 2
+    kb = max(d for d in range(1, min(kt, MAX_KB_TILES) + 1) if kt % d == 0)
+    return ttnn.MatmulMultiCoreReuseProgramConfig(compute_with_storage_grid_size=grid, in0_block_w=kb,
+                                                  out_subblock_h=1, out_subblock_w=1, per_core_M=pcm, per_core_N=1)
+
+
+def attention_matmul(q: Any, k: Any, v: Any, *, scale: float, attn_mask: Any = None,
+                     compute_kernel_config: Any = None, mode: int = 1, pv_pc: Optional[Any] = None,
+                     smask: bool = False, smsm: int = 0):
+    """``q`` ``[1, H, Sq, D]``, ``k`` / ``v`` ``[1, H, Sk, D]`` (tile-aligned ``Sk``) -> ``[1, H, Sq, D]``.
+    ``mode=0``: exactly ``ttaw.ops.attention.attention_matmul``. ``smask`` (``ATTN_SMASK``, mode 1 with a mask):
+    the scale multiply and the mask add as one bit-identical generic_op (``tt/smask_kernel.py``); ``smsm``
+    (``ATTN_SMSM`` 1 / 2, with ``smask``): the scale, the mask and the softmax as one generic_op
+    (``tt/smsm_kernel.py``; 2 = the scale multiply by an immediate)."""
+    import ttnn
+
+    from ..ttaw.ops import attention as A
+    from ..ttaw.precision import compute_kernel_config as ckc
+
+    cfg = compute_kernel_config if compute_kernel_config is not None else ckc("HiFi4", fp32_acc=True)
+    dev = q.device() if callable(getattr(q, "device", None)) else None
+    if mode and (dev is None or not hasattr(dev, "compute_with_storage_grid_size")
+                 or not hasattr(getattr(ttnn, "UnaryOpType", None), "MUL_UNARY_SFPU")):
+        mode = 0                           # the host fake ttnn: the stock path (same math)
+    if not mode:
+        return A.attention_matmul(q, k, v, scale=scale, attn_mask=attn_mask, compute_kernel_config=cfg)
+    sk = int(k.shape[-2])
+    if sk % TILE:
+        raise ValueError(f"attention_matmul needs a tile-aligned Sk (got {sk})")
+    if mode == 2:
+        q = ttnn.multiply(q, float(scale))
+    scores = ttnn.matmul(q, k, transpose_b=True, compute_kernel_config=cfg)
+    fused = False
+    if smsm and smask and mode == 1 and attn_mask is not None:
+        from .smsm_kernel import scale_mask_softmax
+        from .smsm_kernel import supported as smsm_supported
+
+        if smsm_supported(scores, attn_mask):
+            probs = scale_mask_softmax(scores, scale, attn_mask, scale_mode=int(smsm) - 1)
+            if pv_pc is None:
+                pv_pc = pv_config(dev, int(q.shape[1]), int(q.shape[-2]), sk)
+            return ttnn.matmul(probs, v, compute_kernel_config=cfg, program_config=pv_pc)
+    if smask and mode == 1 and attn_mask is not None:
+        from .smask_kernel import scale_mask, supported
+
+        if supported(scores, attn_mask):
+            scores, fused = scale_mask(scores, scale, attn_mask), True
+    if fused:
+        pass
+    elif mode == 3 and attn_mask is not None:
+        act = [ttnn.UnaryWithParam(ttnn.UnaryOpType.MUL_UNARY_SFPU, float(scale))]
+        scores = ttnn.add(scores, attn_mask, input_tensor_a_activations=act)
+    else:
+        if mode != 2:
+            scores = ttnn.multiply(scores, float(scale))
+        if attn_mask is not None:
+            scores = ttnn.add(scores, attn_mask)
+    probs = ttnn.softmax(scores, dim=-1, numeric_stable=True, compute_kernel_config=cfg)
+    if pv_pc is None:
+        pv_pc = pv_config(dev, int(q.shape[1]), int(q.shape[-2]), sk)
+    return ttnn.matmul(probs, v, compute_kernel_config=cfg, program_config=pv_pc)

code/tt_diffusion_planner/tt/config.py

@@ -97,6 +97,108 @@ def _knobs():
         Knob("ATTN_MATMUL", "enc.fusion.attn,dec.*", "comma-separated module globs (enc.fusion.attn, "
              "dec.self_attn, dec.cross_attn) whose attention runs as fp32 matmuls + softmax (C20 attention_matmul) "
              "instead of the bf16 SDPA kernel; 'none' = empty"),
+        Knob("ENC_CH2D", True, "the encoder linears on batched [1, E, T, C] activations run as one 2-D matmul "
+             "over all rows (free [1, 1, E*T, C] view for T = 64, else a 1-D program config with fuse_batch) "
+             "instead of the stock per-element tiling on 4-8 cores (OPT round 1 item 1); 0 = the stock call"),
+        Knob("ATTN_FAST", 1, "fp32 matmul attention variant (tt/attention.py, OPT round 1 item 4a): 0 = "
+             "ttaw.ops.attention.attention_matmul; 1 = P.V on one output tile per core (bit-identical); 2 = 1 + "
+             "the scale on Q; 3 = 1 + the scale fused into the mask add (not bit-identical, rejected)",
+             choices=(0, 1, 2, 3)),
+        Knob("DEC_MMCFG", True, "explicit 2-D multicast program configs for the decoder's 352-row matmuls "
+             "(tt/layers.py DEC_MM_CONFIGS, the fastest bit-identical config per shape and split pass of the device "
+             "sweep, OPT round 1 item 2a); 0 = the auto config"),
+        Knob("LN_KERNEL", 2, "the fp32 LayerNorm decomposition (LN_FP32 modules) as one fused generic_op "
+             "program (tt/ln_kernel.py, the same SFPU LLK sequence, bit-identical: OPT round 2 item 3) instead of "
+             "7-9 stock programs: 1 = per-tile unpacker / SFPU inits, 2 = one init per phase; 0 = the stock "
+             "decomposition", choices=(0, 1, 2)),
+        Knob("LN_RESID", True, "with LN_KERNEL: the residual add in front of a fused fp32 LayerNorm (mixer blocks: "
+             "x + y, x + ch2; decoder: h + gate * a, h + m2b, ...) in the same program, which writes the new stream "
+             "and the normalised rows (OPT round 2 item 5); 0 = separate add / multiply programs"),
+        Knob("SPLIT_KCAT", 2, "decoder split matmuls (SPLIT_MATMUL dec.*, tile-aligned K) as one fp32 matmul over "
+             "the concatenated K axis [x_hi|x_hi|x_lo|1] @ [w_hi;w_lo;w_hi;b] (OPT round 2 item 2): 0 = three "
+             "passes + adds (the published numerics); 1 = stock-op operand build (typecast / subtract / concat); 2 = "
+             "the operand built by one generic_op (tt/kcat_kernel.py); not bit-identical to 0 (one fp32 accumulation, "
+             "exact bias), gates green",
+             choices=(0, 1, 2)),
+        Knob("ATTN_SMASK", True, "fp32 matmul attention (ATTN_FAST=1) with a mask: the score scale multiply and the "
+             "mask add as one generic_op (tt/smask_kernel.py, the same SFPU LLK calls: bit-identical; OPT round 2 "
+             "item 4b) instead of two binary_ng programs; 0 = the two programs"),
+        Knob("ATTN_SMSM", 1, "with ATTN_SMASK: the score scale, the mask add and the softmax as one generic_op "
+             "(tt/smsm_kernel.py: the smask SFPU sequence, then the stock softmax kernel_lib calls on the row in L1; "
+             "OPT round 3 item 1, bit-identical) instead of the smask program + ttnn.softmax: 1 = the scale as a "
+             "multiply by a scale tile (mul_binary_tile, as smask), 2 = the same fp32 SFPU multiply by an immediate "
+             "(mul_unary_tile, no scale tile copy); 0 = the two programs", choices=(0, 1, 2)),
+        Knob("ATTN_FUSED", True, "with ATTN_SMSM: each fp32 matmul attention (decoder self / cross, fusion) as one "
+             "generic_op (tt/fattn_kernel.py: Q K^T into DEST, the smsm phases, P V accumulated in DEST, Q / K / V "
+             "read in place and the heads merged on write; OPT round 3 item 4) instead of head split, Q K^T, smsm, "
+             "P V and head merge; 0 = those programs"),
+        Knob("KCAT_EMIT", True, "with SPLIT_KCAT=2: the producers of the decoder's K-concatenated split linears' "
+             "inputs (the split-row LayerNorms -> qkv / mlp fc1 / cross q, the fused attention -> attn out / cross "
+             "out) write the split operand [x_hi | x_hi | x_lo | 1] themselves (the kcat LLK calls on their output "
+             "tile; OPT round 3 item 5) instead of x + a kcat operand program; 0 = the operand programs"),
+        Knob("LN_TR", True, "with LN_KERNEL and LN_RESID: the two transposes around each mixer block's token-mixing "
+             "MLP done inside the fused LayerNorm programs (n1 writes its output per-entity transposed, n2 reads the "
+             "token-mixing output transposed; the stock transpose LLK, exact; OPT round 3 item 6) instead of two "
+             "ttnn.transpose programs (mixer trunks on the one-core-per-row LN kernel); 0 = the transpose programs"),
+        Knob("ENC_KCAT", True, "with SPLIT_KCAT=2: the encoder's split linears (SPLIT_MATMUL enc.pre.*, "
+             "enc.island.*; any K, blocks padded to whole tiles) as one K-concatenated fp32 matmul too (OPT round 3 "
+             "item 2; one fp32 accumulation, exact bias: a precision change of the SPLIT_KCAT kind); 0 = three "
+             "passes + adds"),
+        Knob("LIN_ACT", False, "the GELU of the encoder mixer linears (token / channel MLP fc1, bf16 output) in the "
+             "matmul epilogue (an explicit copy of the stock auto config + fused_activation, OPT round 3 item 3) "
+             "instead of a unary program on the rounded bf16 output; a precision change (GELU before the bf16 "
+             "rounding); 0 = the unary program"),
+        Knob("KCAT_ACT", True, "with SPLIT_KCAT=2: the GELU of a decoder linear that feeds another K-concatenated "
+             "split linear (preproj.fc1 -> fc2, mlp fc1 -> fc2) applied inside the next one's operand build (the "
+             "same SFPU LLK: bit-identical) instead of its own unary program; 0 = the unary program"),
+        Knob("LN_SPLIT", True, "with LN_KERNEL: a fused fp32 LayerNorm over few tile rows (the decoder's 11) spread "
+             "over Wt cores per row (tt/ln_kernel.py layer_norm_fp32_split: the root core folds the gathered tiles "
+             "in order and broadcasts the statistics, bit-identical) instead of one core per row; 0 = one core per "
+             "row"),
+        Knob("KCAT_L1", True, "with SPLIT_KCAT=2: the decoder's K = 1024 split operands (mlp fc2, final p4) written "
+             "L1 block-sharded by the operand build and read in place by a 2-D multicast matmul with in0 sharded "
+             "(tt/layers.py KCAT_L1_CONFIGS; OPT round 4 item 1, bit-identical) instead of a DRAM round trip of "
+             "the 4.7 MB operand; 0 = DRAM interleaved"),
+        Knob("KCAT_ACT_ONCE", False, "with KCAT_ACT: the operand build applies the deferred GELU to one DEST copy of "
+             "each tile and copies it to the second (copy_dest_values; OPT round 4 item 2, bit-identical) instead "
+             "of running the GELU on both copies (measured, not kept: +0.06 ms); 0 = twice"),
+        Knob("ATTN_L1", 2, "with ATTN_FUSED: the fused attention's inputs in L1 (interleaved) instead of DRAM: the "
+             "decoder qkv / cross q projections write L1, the hoisted cross K / V heads and the masks are L1-resident "
+             "(OPT round 4 item 3, bit-identical; the 11 query-row cores of a head re-read the same K / V tiles); "
+             "1 = also the fusion q / kv projections (the stock linear then runs on 6 cores with a separate bias "
+             "add: +0.3 ms, item 8); 2 = not those; 0 = DRAM", choices=(0, 1, 2)),
+        Knob("DEC_L1", 2, "the decoder blocks' intermediates (the split-row LN outputs and stream, the fused "
+             "attention outputs, the out / mlp / cross-out linear outputs) in L1 (interleaved) instead of DRAM (OPT "
+             "round 4 item 4, bit-identical); 2 = also the pre-projection and the final layer (item 6); 0 = DRAM",
+             choices=(0, 1, 2)),
+        Knob("ENC_L1", False, "the encoder mixer blocks' intermediates (the fused LN outputs and stream, the token / "
+             "channel MLP outputs) in L1 (interleaved) instead of DRAM (OPT round 4 item 5, bit-identical; measured, "
+             "not kept: +1.67 ms); 0 = DRAM"),
+        Knob("FUS_L1", False, "the encoder fusion blocks' intermediates (LN outputs, stream adds, attention / out / "
+             "MLP outputs) in L1 (interleaved) instead of DRAM (OPT round 4 item 7; measured, not kept: +0.39 ms and "
+             "not bit-identical, the stock ops pick other programs for L1 outputs); 0 = DRAM"),
+        Knob("COMPACT", 2, "exact compaction (OPT round 5): besides the full-capacity plan, one trace per agent "
+             "bucket (AGENT_BUCKETS), picked per request from the host masks. 1 = the decoder runs on the first R "
+             "rows only (the needed rows: ego, every valid self-attention key, every emitted neighbour, every valid "
+             "neighbour token; the rows past R are masked keys and never read; item 1); 2 = also the encoder's "
+             "neighbour trunk and head on the first R neighbours, the other neighbour tokens zero (invalid: "
+             "token_valid zeroes them anyway; item 2). The matmul configs keep their K blocking, so the needed rows "
+             "are the same values; 0 = the full 352 rows / 320 neighbours always", choices=(0, 1, 2)),
+        Knob("AGENT_BUCKETS", "32,64,96,128,192", "with COMPACT: the decoder row buckets (multiples of 32 below "
+             "352), one captured trace each"),
+        Knob("HOST_FAST", True, "the host pre- and post-processing vectorised (OPT round 5 item 3: the normalisation "
+             "computes only the rows that are not all-small, the neighbour features only the 6 kept history rows, "
+             "the trajectory's velocity smoothing / force stop / acceleration as array operations in the same "
+             "float64 / float32 order, the denormalisation and predicted paths on the emitted rows only); bit-exact "
+             "(the host suite compares the arrays); 0 = the first port's *_ref functions"),
+        Knob("INPUT_TRIM", True, "per-request uploads trimmed (OPT round 5 item 4, H5): the solver's current states "
+             "``cs`` uploaded as one tile column [352, 32] (the 4 t = 0 columns + zeros) and widened to 324 columns in "
+             "the trace by a tile-aligned concat with a zero block (the same values), and an input whose array is all "
+             "+0.0 (``y0`` at temperature 0, ``static_x``) not uploaded again while its device buffer already holds "
+             "zeros from this model's last upload; 0 = every input every plan, ``cs`` at full width"),
+        Knob("LN_SFPU_BCAST", True, "with LN_KERNEL: the row statistics written to every column by the SFPU row "
+             "reduce itself (kernels/ln32_sfpu.h, the stock reduce's arithmetic with extra stores: bit-identical) "
+             "instead of a RISC-V column fill of the packed column-0 tile; 0 = the fill"),
     ])
 
 

code/tt_diffusion_planner/tt/decoder.py

@@ -32,7 +32,8 @@ from ..reference import config as C
 from ..ttaw.ops import attention as A
 from . import config as T
 from . import params as P
-from .layers import ATTN, Build, Const, LayerNorm, make_linear
+from .attention import attention_matmul
+from .layers import KcatOperand, operand_ktp, ATTN, Build, Const, LayerNorm, chain2, make_linear
 
 __all__ = ["TtDecoder", "TtTurnHead", "STEP_KEYS"]
 
@@ -42,7 +43,10 @@ STEP_KEYS = ("n1_g", "n1_b", "gate_msa", "n2_g", "n2_b", "gate_mlp")
 class TtDecoder:
     """DiT weights, the per-step tables as device rows, and the evaluation / solver graph."""
 
-    def __init__(self, build: Build, p: Mapping[str, np.ndarray], tables: P.StepTables):
+    def __init__(self, build: Build, p: Mapping[str, np.ndarray], tables: P.StepTables,
+                 rows: Sequence[int] = (T.AGENTS,)):
+        """``rows``: the decoder row counts the graph runs on (``T.AGENTS`` and, with ``COMPACT``, the agent
+        buckets): the agent embedding rows and the cross-attention mask are uploaded once per count."""
         self.build = build
         self.tables = tables
         D = "dec.block"
@@ -53,7 +57,8 @@ class TtDecoder:
                                 out=build.hidden("dec.preproj.fc1"), activation="gelu")
         self.pre2 = make_linear(build, P.linear(p, "decoder.dit.preproj.fc2", bias=False), "dec.preproj.fc2",
                                 out=self.stream)
-        self.agent_rows = Const(build, P.agent_rows(p, T.AGENTS), self.stream)
+        self.rows = tuple(sorted({int(r) for r in rows} | {T.AGENTS}))
+        self.agent_rows = {r: Const(build, P.agent_rows(p, r), self.stream) for r in self.rows}
         self.blocks = []
         for i in range(C.DIT_DEPTH):
             B = f"decoder.dit.blocks.{i}"
@@ -92,7 +97,21 @@ class TtDecoder:
         self.attn_mm = {"self": build.attn_matmul("dec.self_attn"), "cross": build.attn_matmul("dec.cross_attn")}
         if self.attn_mm["cross"]:   # fp32 cross-attention over the 576 encoder rows, the 12 pad tokens masked
             row = np.where(np.arange(T.TOKENS) < T.TOKENS_REAL, 0.0, -np.inf).astype(np.float32)
-            self.cross_mask = Const(build, np.broadcast_to(row, (T.AGENTS, T.TOKENS)), "bfloat16")
+            self.cross_mask = {r: Const(build, np.broadcast_to(row, (r, T.TOKENS)), "bfloat16",
+                                        memory_config=build.attn_mem()) for r in self.rows}
+        if build.dec_mem() is not None:         # DEC_L1: the blocks' intermediates in L1 (interleaved)
+            for blk in self.blocks:
+                for key in ("out", "m1a", "m1b", "cout", "m2a", "m2b"):
+                    blk[key].out_mem = build.dec_mem()
+                for key in ("ln", "n3", "n4"):
+                    blk[key].mem = build.dec_mem()
+        self.edge_mem = build.dec_mem() if build.dec_l1 >= 2 else None
+        if self.edge_mem is not None:           # DEC_L1=2: the pre-projection and the final layer too
+            self.pre1.out_mem = self.pre2.out_mem = self.p1.out_mem = self.edge_mem
+            self.fin_ln.mem = self.p0.mem = self.p3.mem = self.edge_mem
+        if build.attn_mem() is not None:        # ATTN_L1: the fused attention reads Q / K / V from L1
+            for blk in self.blocks:
+                blk["qkv"].out_mem = blk["cq"].out_mem = build.attn_mem()
 
     def upload_step(self, tables: P.StepTables, k: int) -> Dict[str, Any]:
         b = self.build
@@ -110,20 +129,76 @@ class TtDecoder:
 
         if not self.attn_mm["cross"] and int(enc.shape[-2]) != T.TOKENS_REAL:
             enc = ttnn.slice(enc, [0, 0, 0, 0], [1, 1, T.TOKENS_REAL, C.HIDDEN_DIM])
-        return [(A.split_heads(k(enc), C.NUM_HEADS), A.split_heads(v(enc), C.NUM_HEADS))
-                for k, v in zip(self.k_lin, self.v_lin)]
+        kv = [(A.split_heads(k(enc), C.NUM_HEADS), A.split_heads(v(enc), C.NUM_HEADS))
+              for k, v in zip(self.k_lin, self.v_lin)]
+        mem = self.build.attn_mem()
+        if mem is not None and self.attn_mm["cross"]:    # ATTN_L1: the hoisted heads L1-resident for 11 evaluations
+            kv = [(ttnn.to_memory_config(k, mem), ttnn.to_memory_config(v, mem)) for k, v in kv]
+        return kv
 
     def _attend(self, kind: str, q, k, v, mask):
         """Attention of ``kind`` ("self" / "cross") -> concatenated heads ``[1, 1, 352, 256]``."""
         if self.attn_mm[kind]:
-            return A.merge_heads(A.attention_matmul(q, k, v, scale=self.scale, attn_mask=mask))
+            return A.merge_heads(attention_matmul(q, k, v, scale=self.scale, attn_mask=mask,
+                                                 mode=self.build.attn_fast, smask=self.build.attn_smask,
+                                                 smsm=self.build.attn_smsm))
         return A.sdpa(q, k, v, scale=self.scale, attn_mask=mask, concat_heads=True, fp32_acc=self.attn_fp32[kind])
 
+    def _fused(self, kind: str, q, k, v, mask, flat_heads, kcat: int = 0):
+        """``ATTN_FUSED``: the attention of ``kind`` as one program (``tt/fattn_kernel.py``) reading Q / K / V in
+        place, returning the merged heads ``[1, 1, 352, 256]``; ``flat_heads`` = the head-block column offsets (in
+        heads) of Q, K, V inside one projection output (self), None for the cross-attention (Q flat, K / V the
+        hoisted heads). None when the fused program does not apply (then the stock chain runs)."""
+        if not (self.build.attn_fused and self.attn_mm[kind] and mask is not None and self.build.attn_smask
+                and self.build.attn_smsm and self.build.attn_fast == 1):
+            return None
+        from .fattn_kernel import flat_at, fused_attention, heads_at, supported
+
+        if not supported((q, k, v), mask):
+            return None
+        H = C.NUM_HEADS
+        if flat_heads is not None:
+            ats = [flat_at(t, o * H) for t, o in zip((q, k, v), flat_heads)]
+        else:
+            ats = [flat_at(q, 0), heads_at(k), heads_at(v)]
+        out = fused_attention(q, k, v, mask, self.scale, H, *ats, kcat_ktp=kcat, memory_config=self.build.dec_mem())
+        return KcatOperand(out, kcat) if kcat else out
+
     def evaluate(self, x, rows: Mapping[str, Any], kv: Sequence[Tuple[Any, Any]], self_mask):
-        """One decoder evaluation: ``x`` ``[1, 1, 352, 324]`` fp32 (prefix-constrained) -> ``m * mask0`` fp32."""
+        """One decoder evaluation: ``x`` ``[1, 1, R, 324]`` fp32 (prefix-constrained; R = 352 or an agent bucket)
+        -> ``m * mask0`` fp32."""
         import ttnn
 
-        h = ttnn.add(self.pre2(self.pre1(x)), self.agent_rows())                       # [1, 1, 352, 256]
+        R = int(tuple(x.shape)[-2])
+
+        h = ttnn.add(chain2(self.pre1, self.pre2, x, self.build.kcat_act), self.agent_rows[R](),   # [1, 1, R, 256]
+                     **({} if self.edge_mem is None else {"memory_config": self.edge_mem}))
+        if self.build.ln_resid:                                       # LN_RESID: the stream adds fused into the LNs
+            pending = None
+            for b, r, (kc, vc) in zip(self.blocks, rows["blocks"], kv):
+                ek = {key: operand_ktp(self.build, b[key]) for key in ("qkv", "out", "m1a", "cq", "cout", "m2a")}
+                if pending is None:
+                    n1 = b["ln"](h, r["n1_g"], r["n1_b"], kcat=ek["qkv"])
+                else:
+                    h, n1 = b["ln"].residual(h, pending, None, r["n1_g"], r["n1_b"], kcat=ek["qkv"])
+                qkv = b["qkv"](n1)
+                a = self._fused("self", qkv, qkv, qkv, self_mask, (0, 1, 2), ek["out"])
+                if a is None:
+                    q, k, v = A.split_qkv(qkv, C.NUM_HEADS)
+                    a = self._attend("self", q, k, v, self_mask)
+                a = b["out"](a)
+                h, n2 = b["ln"].residual(h, a, r["gate_msa"], r["n2_g"], r["n2_b"], kcat=ek["m1a"])
+                m = chain2(b["m1a"], b["m1b"], n2, self.build.kcat_act)
+                h, n3 = b["n3"].residual(h, m, r["gate_mlp"], kcat=ek["cq"])
+                cq = b["cq"](n3)
+                cmask = self.cross_mask[R]() if self.attn_mm["cross"] else None
+                c = self._fused("cross", cq, kc, vc, cmask, None, ek["cout"])
+                if c is None:
+                    c = self._attend("cross", A.split_heads(cq, C.NUM_HEADS), kc, vc, cmask)
+                h, n4 = b["n4"].residual(h, b["cout"](c), kcat=ek["m2a"])
+                pending = chain2(b["m2a"], b["m2b"], n4, self.build.kcat_act)
+            _, f = self.fin_ln.residual(h, pending, None, rows["g"], rows["b"], write_h=False)
+            return self.p4(self.p3(self.p1(self.p0(f))))
         for b, r, (kc, vc) in zip(self.blocks, rows["blocks"], kv):
             q, k, v = A.split_qkv(b["qkv"](b["ln"](h, r["n1_g"], r["n1_b"])), C.NUM_HEADS)
             a = b["out"](self._attend("self", q, k, v, self_mask))
@@ -131,7 +206,7 @@ class TtDecoder:
             m = b["m1b"](b["m1a"](b["ln"](h, r["n2_g"], r["n2_b"])))
             h = ttnn.add(h, ttnn.multiply(m, r["gate_mlp"]))
             qc = A.split_heads(b["cq"](b["n3"](h)), C.NUM_HEADS)
-            cmask = self.cross_mask() if self.attn_mm["cross"] else None
+            cmask = self.cross_mask[R]() if self.attn_mm["cross"] else None
             h = ttnn.add(h, b["cout"](self._attend("cross", qc, kc, vc, cmask)))
             h = ttnn.add(h, b["m2b"](b["m2a"](b["n4"](h))))
         f = self.fin_ln(h, rows["g"], rows["b"])

code/tt_diffusion_planner/tt/encoder.py

@@ -29,6 +29,7 @@ from ..reference import config as C
 from ..ttaw.ops import attention as A
 from . import config as T
 from . import params as P
+from .attention import attention_matmul
 from .layers import ATTN, Build, Const, LayerNorm, make_linear
 
 __all__ = ["MixerTrunk", "TtEncoder", "MIXER_CATS"]
@@ -48,6 +49,7 @@ class MixerTrunk:
 
     def __init__(self, build: Build, p: Mapping[str, np.ndarray], cat: str):
         self.cat, self.E, self.T = cat, ENTITIES[cat], T.MIXER_T[cat]
+        self.build = build
         self.island = cat in T.ISLAND_ROWS
         mods = P.mixer_module(p, cat)
         mix = f"enc.mixer.{cat}"
@@ -83,18 +85,23 @@ class MixerTrunk:
                 "n2": LayerNorm(build, blk["n2"], mix),
                 "ch1": make_linear(build, blk["ch1"], mix, out=build.hidden(mix), activation="gelu"),
                 "ch2": make_linear(build, blk["ch2"], mix, out=self.stream)})
+        if build.enc_mem() is not None:         # ENC_L1: the mixer blocks' intermediates in L1 (interleaved)
+            for b in self.blocks:
+                b["n1"].mem = b["n2"].mem = build.enc_mem()
+                for key in ("tk1", "tk2", "ch1", "ch2"):
+                    b[key].out_mem = build.enc_mem()
 
     # ------------------------------------------------------------------------------------------------------------
     def _tokens_view(self, x):
         """``[1, E, 128, 64]`` <-> ``[1, 1, E*128, 64]`` (free: 128 rows are whole tiles)."""
         import ttnn
 
-        return ttnn.reshape(x, (1, 1, self.E * C.MIXER_CHANNELS, x.shape[-1]))
+        return ttnn.reshape(x, (1, 1, int(x.shape[1]) * C.MIXER_CHANNELS, x.shape[-1]))
 
     def _entities_view(self, x):
         import ttnn
 
-        return ttnn.reshape(x, (1, self.E, C.MIXER_CHANNELS, x.shape[-1]))
+        return ttnn.reshape(x, (1, int(x.shape[2]) // C.MIXER_CHANNELS, C.MIXER_CHANNELS, x.shape[-1]))
 
     def pre(self, x):
         """``[1, E, T, C]`` -> ``x0`` ``[1, E, 64, 128]`` (the ``enc.<cat>.pre`` tap)."""
@@ -129,6 +136,31 @@ class MixerTrunk:
         """The 6 MixerBlocks on ``[1, E, 64, 128]`` (the ``enc.<cat>.mixer`` tap)."""
         import ttnn
 
+        if self.build.ln_resid and self.blocks[0]["n1"].can_transpose(x):
+            # LN_TR: the transposes around the token-mixing MLP inside the LN programs (n1 writes LN(h)^T, n2 reads
+            # the token-mixing output transposed): the same values, two programs less per block
+            pending = None
+            for b in self.blocks:
+                if pending is None:
+                    n1t = b["n1"].transposed(x)
+                else:
+                    x, n1t = b["n1"].transposed(x, res=pending)
+                t = self._entities_view(b["tk2"](b["tk1"](self._tokens_view(n1t))))   # [1, E, 128, 64]
+                x, n2 = b["n2"].residual_t(x, t)
+                pending = b["ch2"](b["ch1"](n2))
+            return ttnn.add(x, pending, memory_config=ttnn.DRAM_MEMORY_CONFIG)
+        if self.build.ln_resid:                                               # LN_RESID: adds fused into the LNs
+            pending = None
+            for b in self.blocks:
+                if pending is None:
+                    n1 = b["n1"](x)
+                else:
+                    x, n1 = b["n1"].residual(x, pending)
+                y = self._tokens_view(ttnn.transpose(n1, -2, -1))
+                y = ttnn.transpose(self._entities_view(b["tk2"](b["tk1"](y))), -2, -1)
+                x, n2 = b["n2"].residual(x, y)
+                pending = b["ch2"](b["ch1"](n2))
+            return ttnn.add(x, pending)
         for b in self.blocks:
             y = self._tokens_view(ttnn.transpose(b["n1"](x), -2, -1))       # LN over the 128 channels, then T
             y = ttnn.transpose(self._entities_view(b["tk2"](b["tk1"](y))), -2, -1)
@@ -141,7 +173,7 @@ class MixerTrunk:
         import ttnn
 
         m = ttnn.mean(x, dim=2, keepdim=True)                                # [1, E, 1, 128]
-        return ttnn.reshape(m, (1, 1, self.E, C.MIXER_CHANNELS))
+        return ttnn.reshape(m, (1, 1, int(x.shape[1]), C.MIXER_CHANNELS))
 
 
 class _Head:
@@ -194,9 +226,14 @@ class TtEncoder:
               "polygon_x", "line_string_x", "goal_x", "ego_shape_x", "turn_x", "token_valid", "pos_aug",
               "fusion_key_row")
 
-    def __init__(self, build: Build, p: Mapping[str, np.ndarray]):
+    def __init__(self, build: Build, p: Mapping[str, np.ndarray], nb_rows=()):
+        """``nb_rows`` (``COMPACT``): the neighbour entity counts the trunk may run on (the agent buckets): a zero
+        token block for the rest is uploaded per count."""
         self.build = build
         self.fstream = build.stream("enc.fusion")
+        E = ENTITIES["neighbor"]
+        self.nb_zeros = {int(n): Const(build, np.zeros((E - int(n), C.HIDDEN_DIM), np.float32), self.fstream)
+                         for n in nb_rows if 0 < int(n) < E}
         self.trunks = {cat: MixerTrunk(build, p, cat) for cat in MIXER_CATS}
         self.heads = {cat: _Head(build, p, cat, self.fstream) for cat in MIXER_CATS}
         st = "enc.head.static"
@@ -222,11 +259,39 @@ class TtEncoder:
                 "fc1": make_linear(build, P.linear(p, f"{B}.mlp.fc1"), fu, out=build.hidden(fu),
                                    activation="gelu"),
                 "fc2": make_linear(build, P.linear(p, f"{B}.mlp.fc2"), fu, out=self.fstream)})
+        if self.attn_mm and build.attn_mem() is not None and build.attn_l1 == 1:   # ATTN_L1=1: Q, K | V to L1
+            for blk in self.blocks:
+                blk["q"].out_mem = blk["kv"].out_mem = build.attn_mem()
+        self.fus_mem = None
+        if build.fus_l1:                        # FUS_L1: the fusion blocks' intermediates in L1 (interleaved)
+            import ttnn
+
+            self.fus_mem = ttnn.L1_MEMORY_CONFIG
+            for blk in self.blocks:
+                blk["n1"].mem = blk["n2"].mem = self.fus_mem
+                for key in ("out", "fc1", "fc2"):
+                    blk[key].out_mem = self.fus_mem
         self.final_norm = LayerNorm(build, P.norm(p, "encoder.fusion.norm"), fu)
         self.scale = float(C.ATTN_SCALE)
         self.attn_fp32 = build.attn_fp32_acc("enc.fusion.attn")
 
-    def forward(self, ctx: Mapping[str, Any], taps: Optional[Dict[str, Any]] = None):
+    def _fused(self, qs, kv, mask):
+        """``ATTN_FUSED``: the fusion attention as one program (``tt/fattn_kernel.py``; Q from ``[1, 1, 576, 256]``,
+        K | V from ``[1, 1, 576, 512]`` in place) -> merged heads, or None (the stock chain runs)."""
+        b = self.build
+        if not (b.attn_fused and self.attn_mm and b.attn_smask and b.attn_smsm and b.attn_fast == 1):
+            return None
+        from .fattn_kernel import flat_at, fused_attention, supported
+
+        if not supported((qs, kv), mask):
+            return None
+        H = C.NUM_HEADS
+        return fused_attention(qs, kv, kv, mask, self.scale, H, flat_at(qs, 0), flat_at(kv, 0), flat_at(kv, H),
+                               memory_config=self.fus_mem)
+
+    def forward(self, ctx: Mapping[str, Any], taps: Optional[Dict[str, Any]] = None, nb: Optional[int] = None):
+        """``nb`` (``COMPACT``): run the neighbour trunk and head on the first ``nb`` entities only and fill the other
+        neighbour tokens with zeros: they are invalid entities, which ``token_valid`` zeroes anyway."""
         import ttnn
 
         def tap(name, t):
@@ -237,10 +302,17 @@ class TtEncoder:
         out: Dict[str, Any] = {}
         for cat in MIXER_CATS:
             tr = self.trunks[cat]
-            x0 = tap(f"enc.{cat}.pre", tr.pre(ctx[f"{cat}_x"]))
-            x = tap(f"enc.{cat}.mixer", tr.mix(x0))
+            xin = ctx[f"{cat}_x"]
             aux = ctx.get(f"{cat}_aux") if cat in ("neighbor", "lane", "route") else None
+            cut = cat == "neighbor" and nb in self.nb_zeros
+            if cut:                                          # COMPACT: the first nb neighbours
+                xin = ttnn.slice(xin, [0, 0, 0, 0], [1, nb] + list(xin.shape)[2:])
+                aux = ttnn.slice(aux, [0, 0, 0, 0], [1, 1, nb, int(aux.shape[-1])])
+            x0 = tap(f"enc.{cat}.pre", tr.pre(xin))
+            x = tap(f"enc.{cat}.mixer", tr.mix(x0))
             out[cat] = self.heads[cat](tr.pool(x), aux)
+            if cut:
+                out[cat] = ttnn.concat([out[cat], self.nb_zeros[nb]()], dim=2)
         out["static"] = self.static2(self.static1(ctx["static_x"]))
         for cat, enc in self.small.items():
             out[cat] = enc(ctx[f"{cat}_x"])
@@ -249,15 +321,24 @@ class TtEncoder:
         x = ttnn.concat([out[name] for name, _ in C.TOKEN_LAYOUT] + [self.pad_tokens()], dim=2)   # [1,1,576,256]
         x = ttnn.multiply(x, ctx["token_valid"])                              # invalid entities -> 0
         x = tap("enc.tokens", ttnn.add(x, self.pos(ctx["pos_aug"])))          # + valid * (pos W + b)
-        mask = A.expand_key_bias(ctx["fusion_key_row"], T.TOKENS)             # [1, 1, 576, 576], once per plan
+        mask = A.expand_key_bias(ctx["fusion_key_row"], T.TOKENS,             # [1, 1, 576, 576], once per plan
+                                 memory_config=self.build.attn_mem() if self.attn_mm else None)
         for i, b in enumerate(self.blocks):
             kv = b["kv"](x)                                                   # K | V from the un-normalised x
-            q, k, v = A.split_q_kv(b["q"](b["n1"](x)), kv, C.NUM_HEADS)
-            if self.attn_mm:
-                a = A.merge_heads(A.attention_matmul(q, k, v, scale=self.scale, attn_mask=mask))
+            qs = b["q"](b["n1"](x))
+            a = self._fused(qs, kv, mask)
+            if a is not None:
+                pass
+            elif self.attn_mm:
+                q, k, v = A.split_q_kv(qs, kv, C.NUM_HEADS)
+                a = A.merge_heads(attention_matmul(q, k, v, scale=self.scale, attn_mask=mask,
+                                                     mode=self.build.attn_fast, smask=self.build.attn_smask,
+                                                     smsm=self.build.attn_smsm))
             else:
+                q, k, v = A.split_q_kv(qs, kv, C.NUM_HEADS)
                 a = A.sdpa(q, k, v, scale=self.scale, attn_mask=mask, concat_heads=True, fp32_acc=self.attn_fp32)
-            x = ttnn.add(x, b["out"](a))
-            x = ttnn.add(x, b["fc2"](b["fc1"](b["n2"](x))))
+            fkw = {} if self.fus_mem is None else {"memory_config": self.fus_mem}
+            x = ttnn.add(x, b["out"](a), **fkw)
+            x = ttnn.add(x, b["fc2"](b["fc1"](b["n2"](x))), **fkw)
             tap(f"enc.fusion.{i}", x)
         return tap("enc.encoding", self.final_norm(x))

code/tt_diffusion_planner/tt/fattn_kernel.py

@@ -0,0 +1,140 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Fused fp32 matmul attention as one ``ttnn.generic_op`` (``ATTN_FUSED``, OPT round 3 item 4).
+
+``fused_attention(q, k, v, mask, scale, heads, q_at, k_at, v_at)`` -> the merged-heads output ``[1, 1, Sq, H * 32]``
+fp32: ``softmax(scale * Q K^T + mask) V`` per head with head dim 32 (one tile), i.e. the programs
+``nlp_create_qkv_heads`` (or ``split_heads``), ``Q K^T`` matmul, the scale + mask + softmax (``ATTN_SMSM``), the
+``P V`` matmul and ``nlp_concat_heads`` of ``tt/attention.py`` in one program (``kernels/fattn_*.cpp``). Each unit
+is one head x one query tile row: ``Q K^T`` into DEST two key tiles at a time, the smask SFPU sequence, the stock
+softmax ``kernel_lib`` calls on the row in L1, then ``P V`` accumulated in DEST in key order, so the scores and the
+probabilities never leave L1 and the math is that of the stock programs (meant to be bit-identical: the device
+check compares with them).
+
+Q / K / V are read in place from their producers: ``q_at`` / ``k_at`` / ``v_at`` = ``(base, row_stride,
+head_stride)`` in tiles, e.g. the self-attention ``qkv`` projection ``[1, 1, S, 768]`` gives Q at ``(0, 24, 1)``,
+K at ``(8, 24, 1)``, V at ``(16, 24, 1)``; a ``[1, H, S, 32]`` heads tensor is ``(0, 1, S / 32)``.
+"""
+from __future__ import annotations
+
+import os
+import struct
+from typing import Any, Sequence
+
+__all__ = ["fused_attention", "supported", "flat_at", "heads_at"]
+
+_KDIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "kernels")
+TB = 4096
+N_RT = 2
+NDST = 4
+
+
+def flat_at(t: Any, col_tile: int):
+    """``(base, row_stride, head_stride)`` of heads stored as consecutive 32-column tiles of ``[1, 1, S, W]``,
+    the first head at column tile ``col_tile``."""
+    return (int(col_tile), int(t.padded_shape[-1]) // 32, 1)
+
+
+def heads_at(t: Any):
+    """``(base, row_stride, head_stride)`` of a ``[1, H, S, 32]`` heads tensor."""
+    return (0, 1, int(t.padded_shape[-2]) // 32)
+
+
+def supported(tensors: Sequence[Any], mask: Any) -> bool:
+    import ttnn
+
+    try:
+        if not hasattr(ttnn, "generic_op") or mask is None:
+            return False
+        for t in tensors:
+            if t.dtype != ttnn.float32 or t.layout != ttnn.TILE_LAYOUT or t.is_sharded():
+                return False
+        return (mask.layout == ttnn.TILE_LAYOUT and mask.dtype in (ttnn.bfloat16, ttnn.float32)
+                and not mask.is_sharded() and int(mask.padded_shape[1]) == 1)
+    except Exception:  # noqa: BLE001 - the host fake ttnn
+        return False
+
+
+def fused_attention(q: Any, k: Any, v: Any, mask: Any, scale: float, heads: int, q_at, k_at, v_at,
+                    kcat_ktp: int = 0, memory_config: Any = None):
+    """``kcat_ktp`` (``KCAT_EMIT``): write the split operand ``[o_hi | o_hi | o_lo | 1 | 0..]`` of the next
+    K-concatenated linear (``[1, 1, Sq, 32 * kcat_ktp]``, the ``kcat_operand`` layout) instead of the output."""
+    import ttnn
+
+    dev = q.device()
+    Sq, Sk = int(mask.padded_shape[-2]), int(mask.padded_shape[-1])
+    Mt, Wt = Sq // 32, Sk // 32
+    units = heads * Mt
+    Wp = (Wt + 1) // 2
+    rb = -(-Wt // NDST) * NDST
+    p_pad = rb - Wt
+    ow = 32 * kcat_ktp if kcat_ktp else heads * 32
+    out = ttnn.allocate_tensor_on_device(ttnn.Shape([1, 1, Sq, ow]), ttnn.float32, ttnn.TILE_LAYOUT, dev,
+                                         memory_config or ttnn.DRAM_MEMORY_CONFIG)
+    g = dev.compute_with_storage_grid_size()
+    n = min(units, g.x * g.y)
+    cs = [(i % g.x, i // g.x) for i in range(n)]
+    full, rem = divmod(n, g.x)
+    rs = []
+    if full:
+        rs.append(ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(g.x - 1, full - 1)))
+    if rem:
+        rs.append(ttnn.CoreRange(ttnn.CoreCoord(0, full), ttnn.CoreCoord(rem - 1, full)))
+    crs = ttnn.CoreRangeSet(set(rs))
+    base, extra = divmod(units, n)
+    rd, wr, cp = ttnn.RuntimeArgs(), ttnn.RuntimeArgs(), ttnn.RuntimeArgs()
+    u0 = 0
+    for i, (cx, cy) in enumerate(cs):
+        kk = base + (1 if i < extra else 0)
+        rd[cx][cy] = [u0, kk]
+        wr[cx][cy] = [u0, kk]
+        cp[cx][cy] = [kk, 0]
+        u0 += kk
+    mbf = mask.dtype == ttnn.bfloat16
+    tbm = 2048 if mbf else 4096
+
+    def cb(idx, pages, dt, tb):
+        return ttnn.CBDescriptor(total_size=pages * tb, core_ranges=crs, format_descriptors=[
+            ttnn.CBFormatDescriptor(buffer_index=idx, data_format=dt, page_size=tb)])
+
+    f32 = ttnn.float32
+    cbs = [cb(0, 2, f32, TB), cb(1, 2 * Wp, mask.dtype, tbm), cb(2, 4, f32, TB), cb(3, 1, f32, TB),
+           cb(4, 1, f32, TB), cb(5, Wt, f32, TB), cb(16, 2, f32, TB), cb(24, Wt, f32, TB), cb(25, 1, f32, TB),
+           cb(26, Wt, f32, TB), cb(27, 1, f32, TB), cb(28, rb, f32, TB)]
+    if kcat_ktp:
+        cbs += [cb(17, 2, f32, TB), cb(18, 1, f32, TB), cb(29, 1, f32, TB)]
+        if kcat_ktp > 3 * heads + 1:
+            cbs.append(cb(19, 1, f32, TB))
+    um = [ttnn.UnpackToDestMode.Default] * 64
+    if kcat_ktp:
+        um[29] = ttnn.UnpackToDestMode.UnpackToDestFp32
+    if not mbf:
+        um[1] = ttnn.UnpackToDestMode.UnpackToDestFp32
+    ccfg = ttnn.ComputeConfigDescriptor(math_fidelity=ttnn.MathFidelity.HiFi4, fp32_dest_acc_en=True,
+                                        math_approx_mode=False)
+    ccfg.unpack_to_dest_mode = um
+    bits = int.from_bytes(struct.pack("<f", float(scale)), "little")
+    SRC = ttnn.KernelDescriptor.SourceType.FILE_PATH
+
+    def acc(t):
+        return list(ttnn.TensorAccessorArgs(t).get_compile_time_args())
+
+    def kd(name, ct, rt, common, config):
+        return ttnn.KernelDescriptor(kernel_source=os.path.join(_KDIR, name), source_type=SRC, core_ranges=crs,
+                                     compile_time_args=ct, defines=[], runtime_args=rt, common_runtime_args=common,
+                                     config=config)
+
+    ats = [int(x) for a in (q_at, k_at, v_at) for x in a]
+    ks = [kd("fattn_reader.cpp", [Wt, Mt, N_RT, tbm] + ats + acc(q) + acc(k) + acc(v) + acc(mask), rd,
+             [q.buffer_address(), k.buffer_address(), v.buffer_address(), mask.buffer_address()],
+             ttnn.ReaderConfigDescriptor()),
+          kd("fattn_writer.cpp", [Mt, heads, N_RT, int(bool(kcat_ktp)), int(kcat_ktp)] + acc(out), wr,
+             [out.buffer_address()],
+             ttnn.WriterConfigDescriptor()),
+          kd("fattn_compute.cpp", [Wt, N_RT, int(mbf), NDST, p_pad, bits, int(bool(kcat_ktp))], cp, [], ccfg)]
+    ins = []
+    for t in (q, k, v):
+        if all(t is not x for x in ins):
+            ins.append(t)
+    ins += [mask, out]
+    ttnn.generic_op(ins, ttnn.ProgramDescriptor(kernels=ks, semaphores=[], cbs=cbs))
+    return out

code/tt_diffusion_planner/tt/inputs.py

@@ -18,10 +18,13 @@ from ..ttaw.ops.attention import key_bias_row
 from . import config as T
 from .params import lane_aux_features, pad_rows
 
-__all__ = ["plan_inputs", "INPUT_SPECS", "warmup_inputs", "decoder_state", "BF16_INPUTS"]
+__all__ = ["plan_inputs", "INPUT_SPECS", "warmup_inputs", "decoder_state", "BF16_INPUTS", "CS_COLS"]
 
 f32 = np.float32
 
+# INPUT_TRIM: the current states as one tile column (widened to STATE_COLS in the trace)
+CS_COLS = T.TILE if T.KNOBS.read().INPUT_TRIM else T.STATE_COLS
+
 # name -> device shape (all fp32 TILE except the bf16 key-bias rows)
 INPUT_SPECS: Dict[str, tuple] = {
     "ego_x": (1, 1, T.MIXER_T["ego"], T.MIXER_CIN["ego"]),
@@ -41,7 +44,7 @@ INPUT_SPECS: Dict[str, tuple] = {
     "pos_aug": (1, 1, T.TOKENS, T.POS_AUG_DIM),
     "fusion_key_row": (1, 1, 1, T.TOKENS),
     "agent_key_row": (1, 1, 1, T.AGENTS),
-    "cs": (1, 1, T.AGENTS, T.STATE_COLS),
+    "cs": (1, 1, T.AGENTS, CS_COLS),
     "y0": (1, 1, T.AGENTS, T.STATE_COLS),
 }
 BF16_INPUTS = ("fusion_key_row", "agent_key_row")
@@ -84,7 +87,7 @@ def plan_inputs(prep: Prepared) -> Dict[str, np.ndarray]:
         "fusion_key_row": key_bias_row(pad_rows(np.asarray(f.key_valid, bool), T.TOKENS)),
         "agent_key_row": key_bias_row(pad_rows(np.asarray(prep.decoder.agent_valid, bool), T.AGENTS)),
         "cs": decoder_state(np.zeros((C.MAX_NUM_AGENTS, C.OUTPUT_T + 1, C.POSE_DIM), f32),
-                            prep.decoder.current_states),
+                            prep.decoder.current_states)[..., :CS_COLS],
         "y0": decoder_state(prep.x_T),
     }
     for name, arr in out.items():

code/tt_diffusion_planner/tt/kcat_kernel.py

@@ -0,0 +1,109 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Operand build of the K-concatenated split matmul (``SPLIT_KCAT=2``, OPT round 2 item 2) as one ``generic_op``.
+
+``kcat_operand(x)``: fp32 TILE ``[..., M, K]`` (K a multiple of 32) -> fp32 ``[..., M, 3K + 32]`` =
+``[x_hi | x_hi | x_lo | ones]`` with ``x_hi`` = bf16(x) (the ``ttnn.typecast`` LLK) and ``x_lo = x - x_hi``; the
+last tile has columns 0 and 1 = 1 (the exact bias rows of the concatenated weight). It replaces the 5 stock programs
+of ``SPLIT_KCAT=1`` (typecast, typecast, subtract, concat over 4 tensors). x tiles are split in contiguous ranges
+over the grid (``kernels/kcat_*.cpp``).
+"""
+from __future__ import annotations
+
+import os
+from typing import Any
+
+__all__ = ["kcat_operand", "kcat_tiles", "supported"]
+
+_KDIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "kernels")
+TB = 4096
+N_RT = 2
+
+
+def supported(x: Any) -> bool:
+    """True for a device fp32 TILE interleaved tensor with tile-aligned rows and columns (not the host fake)."""
+    import ttnn
+
+    try:
+        shp = list(x.padded_shape)
+        return (x.dtype == ttnn.float32 and x.layout == ttnn.TILE_LAYOUT and not x.is_sharded()
+                and shp[-1] % 32 == 0 and shp[-2] % 32 == 0 and hasattr(ttnn, "generic_op"))
+    except Exception:  # noqa: BLE001 - the host fake ttnn
+        return False
+
+
+def kcat_tiles(k: int, pad: int) -> int:
+    """Row tiles of X' for an input of K = ``k`` columns: 3 Kt + 1, rounded up to a multiple of ``pad``."""
+    return -(-(3 * (k // 32) + 1) // pad) * pad
+
+
+ACTS = {None: 0, "gelu": 1, "gelu_tanh": 2}
+
+
+def kcat_operand(x: Any, pad: int = 1, act: Any = None, memory_config: Any = None, act_once: bool = False):
+    """``pad``: X' row tiles rounded up to a multiple of ``pad`` with zero tiles (the matmul's K block must divide
+    them; 3 Kt + 1 is odd, e.g. 97 for K = 1024). ``act`` (``KCAT_ACT``): ``"gelu"`` / ``"gelu_tanh"`` applied to
+    x first (the previous linear's activation, the same LLK as the stock ``ttnn.gelu`` program). ``memory_config``:
+    of X' (default DRAM interleaved; ``KCAT_L1``: the L1 block-sharded layout the consumer matmul reads in place, the
+    writer's TensorAccessor resolves the shard of each tile). ``act_once`` (``KCAT_ACT_ONCE``): the activation on one
+    DEST copy of the tile, copied to the second by ``copy_dest_values`` (instead of on both copies)."""
+    import ttnn
+
+    dev = x.device()
+    shp = list(x.padded_shape)
+    K = shp[-1]
+    assert K % 32 == 0 and shp[-2] % 32 == 0 and x.dtype == ttnn.float32 and x.layout == ttnn.TILE_LAYOUT
+    Kt = K // 32
+    rows = 1
+    for d in shp[:-1]:
+        rows *= d
+    rows //= 32
+    total = rows * Kt
+    Ktp = kcat_tiles(K, pad)
+    oshape = list(x.shape)[:-1] + [32 * Ktp]
+    out = ttnn.allocate_tensor_on_device(ttnn.Shape(oshape), ttnn.float32, ttnn.TILE_LAYOUT, dev,
+                                         memory_config or ttnn.DRAM_MEMORY_CONFIG)
+    g = dev.compute_with_storage_grid_size()
+    n = min(total, g.x * g.y)
+    cs = [(i % g.x, i // g.x) for i in range(n)]
+    full, rem = divmod(n, g.x)
+    rs = []
+    if full:
+        rs.append(ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(g.x - 1, full - 1)))
+    if rem:
+        rs.append(ttnn.CoreRange(ttnn.CoreCoord(0, full), ttnn.CoreCoord(rem - 1, full)))
+    crs = ttnn.CoreRangeSet(set(rs))
+    base, extra = divmod(total, n)
+    rd, wr, cp = ttnn.RuntimeArgs(), ttnn.RuntimeArgs(), ttnn.RuntimeArgs()
+    t0 = 0
+    for i, (cx, cy) in enumerate(cs):
+        k = base + (1 if i < extra else 0)
+        rd[cx][cy] = [t0, k]
+        wr[cx][cy] = [t0, k]
+        cp[cx][cy] = [k, 0]
+        t0 += k
+
+    def cb(idx, pages):
+        return ttnn.CBDescriptor(total_size=pages * TB, core_ranges=crs, format_descriptors=[
+            ttnn.CBFormatDescriptor(buffer_index=idx, data_format=ttnn.float32, page_size=TB)])
+
+    cbs = [cb(0, 2), cb(16, 2), cb(17, 2), cb(18, 1)] + ([cb(19, 1)] if Ktp > 3 * Kt + 1 else [])
+    um = [ttnn.UnpackToDestMode.Default] * 64
+    um[0] = ttnn.UnpackToDestMode.UnpackToDestFp32
+    ccfg = ttnn.ComputeConfigDescriptor(math_fidelity=ttnn.MathFidelity.HiFi4, fp32_dest_acc_en=True,
+                                        math_approx_mode=False)
+    ccfg.unpack_to_dest_mode = um
+    SRC = ttnn.KernelDescriptor.SourceType.FILE_PATH
+
+    def acc(t):
+        return list(ttnn.TensorAccessorArgs(t).get_compile_time_args())
+
+    def kd(name, ct, rt, common, config):
+        return ttnn.KernelDescriptor(kernel_source=os.path.join(_KDIR, name), source_type=SRC, core_ranges=crs,
+                                     compile_time_args=ct, defines=[], runtime_args=rt, common_runtime_args=common,
+                                     config=config)
+
+    ks = [kd("kcat_reader.cpp", [N_RT] + acc(x), rd, [x.buffer_address()], ttnn.ReaderConfigDescriptor()),
+          kd("kcat_writer.cpp", [Kt, N_RT, Ktp] + acc(out), wr, [out.buffer_address()], ttnn.WriterConfigDescriptor()),
+          kd("kcat_compute.cpp", [N_RT, ACTS[act], int(bool(act_once))], cp, [], ccfg)]
+    ttnn.generic_op([x, out], ttnn.ProgramDescriptor(kernels=ks, semaphores=[], cbs=cbs))
+    return out

code/tt_diffusion_planner/tt/kernels/README.md

@@ -1,10 +1,32 @@
 # Custom kernels of diffusion-planner-p150
 
-None for the functional port: PLAN.md section 2.12 maps every op of the v5.0 graphs to stock ttnn (matmul / linear,
-layer_norm, gelu, SDPA through the shared C20 wrapper, eltwise). Fused kernels (entity-parallel MLP-Mixer, fusion
-block, DiT step + solver update) and the single-megakernel attempt (decision D19) are optimization-phase work; when
-they land, their `.cpp` sources go here (package data in the repo-root `pyproject.toml`, asserted by a `verify:` line of
-`tt-model.yaml`), each with a numpy / torch oracle test and the hang protocol of PLAN.md section 4.4.
-
-| file | op | used by | status |
-|---|---|---|---|
+The functional port maps every op to stock ttnn (PLAN.md section 2.12). The optimization rounds add `ttnn.generic_op`
+programs whose sources live here (package data in the repo-root `pyproject.toml`, asserted by a `verify:` line of
+`tt-model.yaml`). Each one has a device check against the stock graph it replaces (bit for bit where the math is
+the same) and went through the hang protocol of PLAN.md section 4.4 (smallest shape under `TT_METAL_WATCHER=2`,
+eager twice, a trace): `logs/diffusion-planner/opt_r2/scripts/ln_check.py`, `kcat_check.py`, `smask_check.py`,
+`logs/diffusion-planner/opt_r3/scripts/smsm_check.py`, `fattn_check.py`, `kemit_check.py`, `lntr_check.py`.
+
+Per-core runtime-argument counts are compile-time args (probe P3: the generic_op program hash ignores them).
+
+| files | op | used by | knob | status |
+|---|---|---|---|---|
+| `ln32_reader.cpp`, `ln32_compute.cpp`, `ln32_writer.cpp` | fp32 LayerNorm over the last dim (the 9-program `layer_norm_fp32` decomposition in one program, the same SFPU LLK calls in the same order), optionally with the residual add `h = x + r (* gate)` in front | `tt/ln_kernel.py`, `tt/layers.py` `LayerNorm` (mixers, decoder) | `LN_KERNEL`, `LN_RESID` | bit-identical to the stock decomposition (OPT round 2) |
+| `kcat_reader.cpp`, `kcat_compute.cpp`, `kcat_writer.cpp` | operand of the K-concatenated split matmul (optionally after the previous linear's GELU, `KCAT_ACT`): fp32 `x` -> `[bf16(x) \| bf16(x) \| x - bf16(x) \| 1 \| 0...]` | `tt/kcat_kernel.py`, `tt/layers.py` `SplitLinear` (decoder) | `SPLIT_KCAT=2` | bit-identical to the stock-op operand (`SPLIT_KCAT=1`) (OPT round 2) |
+| `smask_reader.cpp`, `smask_compute.cpp`, `smask_writer.cpp` | attention score scale + additive mask: `s * scale + mask` in one pass (the two stock binary_ng programs; with a bf16 mask the stock add runs on the FPU, which truncates the scores to TF32, reproduced with an SFPU AND) | `tt/smask_kernel.py`, `tt/attention.py` (decoder self / cross, fusion) | `ATTN_SMASK` | bit-identical (OPT round 2) |
+| `smsm_reader.cpp`, `smsm_compute.cpp`, `smsm_writer.cpp` | attention score scale + mask + softmax over one tile row (head, 32 queries) per pass: the smask SFPU sequence into an L1 row, then the stock `ttnn.softmax(numeric_stable=True)` compute (`kernel_lib` reduce / bcast / exp calls in the stock order: FPU row max, `exp(x - max)`, FPU row sum + precise fp32 reciprocal, bcast multiply); the scale either by a scale tile (`mul_binary_tile`) or an immediate (`mul_unary_tile`, the same fp32 SFPU multiply) | `tt/smsm_kernel.py`, `tt/attention.py` (decoder self / cross, fusion) | `ATTN_SMSM` | bit-identical to smask + `ttnn.softmax` (OPT round 3) |
+| `fattn_reader.cpp`, `fattn_compute.cpp`, `fattn_writer.cpp` | the whole fp32 matmul attention per (head, query tile row): `Q K^T` into DEST (`matmul_block`, in1 transposed), the smsm phases (scale by an immediate, TF32 truncation, mask, stock softmax `kernel_lib` calls) on the row in L1, `P V` accumulated in DEST in key order; Q / K / V read in place from the projection outputs (or the hoisted cross K / V heads), the output written as tile (r, h) of the merged heads | `tt/fattn_kernel.py`, `tt/decoder.py` / `tt/encoder.py` (decoder self / cross, fusion) | `ATTN_FUSED` | bit-identical to head split + `Q K^T` + smsm + `P V` + head merge (OPT round 3) |
+| `ln32s_reader.cpp`, `ln32s_compute.cpp`, `ln32s_writer.cpp` | the fused fp32 LayerNorm (+ residual) with each tile row spread over Wt cores: member j owns column tile j, the root gathers the tiles (semaphore 0), folds them in order and broadcasts mean / rstd (semaphore 1, monotonic counts) | `tt/ln_kernel.py` `layer_norm_fp32_split`, `LayerNorm` (decoder: 11 tile rows on 88 cores) | `LN_SPLIT` | bit-identical (OPT round 2) |
+| `ln32_sfpu.h` | SFPU row sum of one fp32 tile with the sum stored to every column: the stock `sfpu_reduce<SUM, Float32, REDUCE_ROW>` arithmetic (same loads, replayed adds, butterfly) with three extra stores, so the LN statistics come out already broadcast (no RISC-V column fill) | `ln32_compute.cpp`, `ln32s_compute.cpp` | `LN_SFPU_BCAST` | bit-identical (OPT round 2) |
+
+`ln32_sfpu.h` is included by the two LN compute kernels through `compiler_include_paths`; the JIT build cache keys a
+kernel on its own source, compile-time args and defines, so after editing the header, also touch the including
+`.cpp` files (or clear the kernel cache) before measuring.
+
+`KCAT_EMIT` (OPT round 3): `ln32s_*.cpp` and `fattn_*.cpp` can write their output tile as the split operand
+`[x_hi | x_hi | x_lo | 1 | 0..]` of the next K-concatenated decoder linear (the `kcat_compute.cpp` LLK calls on the
+packed fp32 tile, the `kcat_writer.cpp` layout) instead of the plain output: bit-identical to output + `kcat_operand`.
+
+`LN_TR` (OPT round 3): `ln32_*.cpp` can read the residual per-entity transposed (`[.., E, W, T]`, transposed back with
+`transpose_tile`, the stock `ttnn.transpose` LLK with the fp32 unpack-to-dest: exact) and write the output the same
+way (via c_18): the two transposes around the mixer's token-mixing MLP, bit-identical.

code/tt_diffusion_planner/tt/kernels/fattn_compute.cpp

@@ -0,0 +1,222 @@
+// SPDX-License-Identifier: Apache-2.0
+// Fused fp32 matmul attention (tt/fattn_kernel.py, ATTN_FUSED), compute. Per unit (head h, query tile row r; Wt key
+// tiles), the five stock programs of tt/attention.py (ATTN_FAST=1 + ATTN_SMSM) in one pass, the scores and the
+// probabilities kept in L1:
+//   phase Q: s_j = Q_r K_j^T in DEST (matmul_block, in1 transposed by the unpacker: the stock Q K^T matmul with its
+//            K = 32 in one block), two key tiles at a time;
+//   phase A: x_j = trunc_tf32(s_j * scale) + mask_j on the SFPU (kernels/smsm_compute.cpp, scale mode 1), packed to
+//            cb_x;
+//   phase B: the stock numeric-stable softmax of the row (kernels/smsm_compute.cpp phase B) -> cb_p;
+//   phase C: o = sum_j P_j V_j accumulated in DEST in key order (the stock P V matmul: the whole K in one block),
+//            packed to cb_o; the writer stores it as tile (r, h) of the merged-heads output.
+//   (kcat_out, KCAT_EMIT) the output tile as the operand of the next K-concatenated split linear: o_hi = bf16(o)
+//            and o_lo = o - o_hi (kernels/kcat_compute.cpp's LLK calls on the packed fp32 tile) -> cb_o / cb_lo.
+// CT args: [0] Wt, [1] per-core RT-arg count (P3), [2] trunc_tf32, [3] ndst, [4] P pad tiles per row (to a multiple
+//          of ndst), [5] scale bits (fp32), [6] kcat_out. Per-core RT args: [nunits, 0].
+#include <cstdint>
+
+#include "api/compute/common.h"
+#include "api/compute/compute_kernel_api.h"
+#include "api/compute/eltwise_binary.h"
+#include "api/compute/eltwise_binary_sfpu.h"
+#include "api/compute/eltwise_unary/binop_with_scalar.h"
+#include "api/compute/eltwise_unary/bitwise.h"
+#include "api/compute/eltwise_unary/eltwise_unary.h"
+#include "api/compute/eltwise_unary/typecast.h"
+#include "api/compute/tile_move_copy.h"
+#include "api/compute/matmul.h"
+#include "api/compute/bcast.h"
+#include "api/compute/softmax.h"
+#include "api/compute/reduce.h"
+#include "ttnn/cpp/ttnn/kernel_lib/reduce_helpers_compute.hpp"
+#include "ttnn/cpp/ttnn/kernel_lib/eltwise/api/chain.hpp"
+#include "ttnn/cpp/ttnn/kernel_lib/eltwise/api/convenience.hpp"
+#include "ttnn/cpp/ttnn/kernel_lib/eltwise/unary/math.hpp"
+#include "ttnn/cpp/ttnn/kernel_lib/eltwise/core/optional.hpp"
+
+namespace ckl = compute_kernel_lib;
+using namespace ckernel;
+
+constexpr uint32_t Wt = get_compile_time_arg_val(0);
+constexpr uint32_t trunc_tf32 = get_compile_time_arg_val(2);
+constexpr uint32_t ndst = get_compile_time_arg_val(3);
+constexpr uint32_t p_pad = get_compile_time_arg_val(4);
+constexpr uint32_t scale_bits = get_compile_time_arg_val(5);
+constexpr uint32_t kcat_out = get_compile_time_arg_val(6);
+constexpr uint32_t cb_q = 0, cb_m = 1, cb_k = 2, cb_max_scaler = 3, cb_sum_scaler = 4, cb_v = 5;
+constexpr uint32_t cb_o = 16, cb_lo = 17, cb_y = 29;
+constexpr uint32_t cb_x = 24, cb_max = 25, cb_exps = 26, cb_recip = 27, cb_p = 28;
+constexpr auto RNE = ckernel::DstRoundingMode::NearestEven;
+constexpr uint32_t Wp = (Wt + 1) / 2;
+constexpr uint32_t Wr = Wt + p_pad;
+
+template <std::uint32_t dfb_in, std::uint32_t dfb_max_scaler, std::uint32_t dfb_max, std::uint32_t dfb_out>
+void calc_numeric_stable(std::uint32_t W, std::uint32_t nd) {
+    compute_kernel_lib::reduce<
+        PoolType::MAX,
+        ReduceDim::REDUCE_ROW,
+        dfb_in,
+        dfb_max_scaler,
+        dfb_max,
+        compute_kernel_lib::ReduceInputPolicy::WaitUpfrontNoPop,
+        compute_kernel_lib::ReduceDataFormatReconfigMode::INPUT>(compute_kernel_lib::ReduceInputBlockShape::row(W));
+    ckl::eltwise_chain(
+        ckl::IterationShape::tiles(W).block_size(nd),
+        ckl::BinaryFpu<
+            ckl::BinaryFpuOp::Sub,
+            ckl::input(
+                dfb_in,
+                ckl::WaitPolicy::Upfront,
+                ckl::PopPolicy::AtEnd,
+                ckl::InputTileMapping::Block,
+                ckl::DataFormatReconfig::Disabled),
+            ckl::input(dfb_max, ckl::BroadcastDim::Col, ckl::WaitPolicy::Upfront, ckl::PopPolicy::AtEnd)>{},
+        ckl::Exp<ckl::Approx::Exact, ckl::Dst::D0>{},
+        ckl::PackTile<ckl::output(
+            dfb_out,
+            ckl::ReservePolicy::PerBlockSize,
+            ckl::PushPolicy::PerBlockSize,
+            ckl::DataFormatReconfig::Disabled)>{});
+    cb_wait_front(dfb_out, W);
+}
+
+void kernel_main() {
+    const uint32_t nunits = get_arg_val<uint32_t>(0);
+    if (nunits == 0) {
+        return;
+    }
+    compute_kernel_hw_startup(cb_q, cb_k, cb_x);
+    cb_wait_front(cb_max_scaler, 1);
+    cb_wait_front(cb_sum_scaler, 1);
+    for (uint32_t u = 0; u < nunits; ++u) {
+        // ---- phases Q + A: x_j = trunc(Q K_j^T * scale) + mask_j, two key tiles at a time ----
+        cb_wait_front(cb_q, 1);
+        cb_wait_front(cb_m, 2 * Wp);
+        pack_reconfig_data_format(cb_x);
+        for (uint32_t p = 0; p < Wp; ++p) {
+            const uint32_t j = 2 * p;
+            const bool two = (j + 1) < Wt;
+            cb_wait_front(cb_k, 2);
+            reconfig_data_format(cb_k, cb_q);
+            matmul_block_init(cb_q, cb_k, 1, 1, 1, 1);
+            tile_regs_acquire();
+            matmul_block(cb_q, cb_k, 0, 0, 0, 1, 1, 1, 1);
+            matmul_block(cb_q, cb_k, 0, 1, 1, 1, 1, 1, 1);
+            binop_with_scalar_tile_init();
+            mul_unary_tile(0, scale_bits);
+            mul_unary_tile(1, scale_bits);
+            if constexpr (trunc_tf32) {
+                bitwise_and_tile_init();
+                bitwise_and_tile<DataFormat::Int32>(0, 0xFFFFE000u);
+                bitwise_and_tile<DataFormat::Int32>(1, 0xFFFFE000u);
+            }
+            reconfig_data_format_srca(cb_k, cb_m);
+            copy_init(cb_m);
+            copy_tile(cb_m, j, 2);
+            copy_tile(cb_m, j + 1, 3);
+            add_binary_tile_init();
+            add_binary_tile<RNE>(0, 2, 0);
+            add_binary_tile<RNE>(1, 3, 1);
+            const uint32_t np = two ? 2 : 1;
+            cb_reserve_back(cb_x, np);
+            tile_regs_commit();
+            tile_regs_wait();
+            pack_tile(0, cb_x);
+            if (two) {
+                pack_tile(1, cb_x);
+            }
+            tile_regs_release();
+            cb_push_back(cb_x, np);
+            cb_pop_front(cb_k, 2);
+        }
+        cb_pop_front(cb_m, 2 * Wp);
+        cb_pop_front(cb_q, 1);
+
+        // ---- phase B: the stock numeric-stable softmax of the row -> cb_p ----
+        reconfig_data_format(cb_x, cb_x);
+        pack_reconfig_data_format(cb_exps);
+        copy_init(cb_x);
+        calc_numeric_stable<cb_x, cb_max_scaler, cb_max, cb_exps>(Wt, ndst);
+        reconfig_data_format(cb_exps, cb_sum_scaler);
+        compute_kernel_lib::reduce<
+            PoolType::SUM,
+            ReduceDim::REDUCE_ROW,
+            cb_exps,
+            cb_sum_scaler,
+            cb_recip,
+            compute_kernel_lib::ReduceInputPolicy::WaitUpfrontNoPop>(
+            compute_kernel_lib::ReduceInputBlockShape::row(Wt),
+            compute_kernel_lib::ReduceInputMemoryLayout::contiguous(),
+            compute_kernel_lib::NoAccumulation{},
+            [](std::uint32_t) {
+                if constexpr (DST_ACCUM_MODE) {
+                    recip_tile_init<ReciprocalDestAcc::FP32, ReciprocalApproxMode::Precise>();
+                    recip_tile<ReciprocalDestAcc::FP32, ReciprocalApproxMode::Precise>(0);
+                } else {
+                    recip_tile_init();
+                    recip_tile(0);
+                }
+            });
+        ckl::mul<
+            ckl::input(cb_exps, ckl::WaitPolicy::Upfront, ckl::PopPolicy::AtEnd, ckl::InputTileMapping::Block),
+            ckl::input(cb_recip, ckl::BroadcastDim::Col, ckl::WaitPolicy::Upfront, ckl::PopPolicy::AtEnd),
+            ckl::output(cb_p, ckl::ReservePolicy::PerBlockSize, ckl::PushPolicy::PerBlockSize)>(
+            ckl::IterationShape::tiles(Wt).block_size(ndst));
+        if constexpr (p_pad > 0) {
+            cb_reserve_back(cb_p, p_pad);
+            cb_push_back(cb_p, p_pad);
+        }
+
+        // ---- phase C: o = sum_j P_j V_j (key order, one DEST accumulation) ----
+        cb_wait_front(cb_p, Wr);
+        reconfig_data_format(cb_v, cb_p);
+        pack_reconfig_data_format(cb_o);
+        matmul_block_init(cb_p, cb_v, 0, 1, 1, 1);
+        tile_regs_acquire();
+        for (uint32_t j = 0; j < Wt; ++j) {
+            cb_wait_front(cb_v, 1);
+            matmul_block(cb_p, cb_v, j, 0, 0, 0, 1, 1, 1);
+            cb_pop_front(cb_v, 1);
+        }
+        if constexpr (kcat_out) {
+            cb_reserve_back(cb_y, 1);
+            tile_regs_commit();
+            tile_regs_wait();
+            pack_tile(0, cb_y);
+            tile_regs_release();
+            cb_push_back(cb_y, 1);
+            cb_pop_front(cb_p, Wr);
+            // ---- the split operand of the next linear (kernels/kcat_compute.cpp) ----
+            cb_wait_front(cb_y, 1);
+            reconfig_data_format_srca(cb_y);
+            tile_regs_acquire();
+            copy_init(cb_y);
+            copy_tile(cb_y, 0, 0);
+            copy_tile(cb_y, 0, 1);
+            typecast_tile_init<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>();
+            typecast_tile<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>(0);
+            sub_binary_tile_init();
+            sub_binary_tile<ckernel::DstRoundingMode::NearestEven>(1, 0, 1);
+            cb_reserve_back(cb_o, 1);
+            cb_reserve_back(cb_lo, 1);
+            tile_regs_commit();
+            tile_regs_wait();
+            pack_tile(0, cb_o);
+            pack_tile(1, cb_lo);
+            tile_regs_release();
+            cb_push_back(cb_o, 1);
+            cb_push_back(cb_lo, 1);
+            cb_pop_front(cb_y, 1);
+        } else {
+            cb_reserve_back(cb_o, 1);
+            tile_regs_commit();
+            tile_regs_wait();
+            pack_tile(0, cb_o);
+            tile_regs_release();
+            cb_push_back(cb_o, 1);
+            cb_pop_front(cb_p, Wr);
+        }
+    }
+    cb_pop_front(cb_max_scaler, 1);
+    cb_pop_front(cb_sum_scaler, 1);
+}

code/tt_diffusion_planner/tt/kernels/fattn_reader.cpp

@@ -0,0 +1,95 @@
+// SPDX-License-Identifier: Apache-2.0
+// Fused fp32 matmul attention (tt/fattn_kernel.py), reader (RISCV_0). Once: the two reduce scaler tiles of the stock
+// softmax (fp32 1.0 in row 0 of every face). Per unit u of this core's range (head h = u / Mt, query tile row
+// r = u % Mt): the Wt mask tiles of row r (+ 1 unused tile when Wt is odd), the Q tile (r, h), the K tiles (j, h) in
+// pairs (the last pair of an odd row carries one unused tile), then the V tiles (j, h) one by one.
+// Tile (row, head) of a source = base + row * row_stride + head * head_stride (tiles): a [1, 1, S, W] projection
+// output read in place (row_stride = W / 32, head_stride = 1, base = the column tile of the first head) or a
+// [1, H, S, 32] heads tensor (row_stride = 1, head_stride = S / 32).
+// CT args: [0] Wt, [1] Mt, [2] per-core RT-arg count (P3), [3] mask tile bytes, [4..12] (base, row_stride,
+//          head_stride) of Q, K, V, then the TensorAccessorArgs of q, k, v, mask.
+// Common RT args: [q_addr, k_addr, v_addr, m_addr]. Per-core RT args: [u0, n].
+#include <cstdint>
+
+#include "api/dataflow/dataflow_api.h"
+
+constexpr uint32_t Wt = get_compile_time_arg_val(0);
+constexpr uint32_t Mt = get_compile_time_arg_val(1);
+constexpr uint32_t TBM = get_compile_time_arg_val(3);
+constexpr uint32_t qb = get_compile_time_arg_val(4), qr = get_compile_time_arg_val(5), qh = get_compile_time_arg_val(6);
+constexpr uint32_t kb = get_compile_time_arg_val(7), kr = get_compile_time_arg_val(8), kh = get_compile_time_arg_val(9);
+constexpr uint32_t vb = get_compile_time_arg_val(10), vr = get_compile_time_arg_val(11),
+                   vh = get_compile_time_arg_val(12);
+constexpr auto q_args = TensorAccessorArgs<13>();
+constexpr auto k_args = TensorAccessorArgs<q_args.next_compile_time_args_offset()>();
+constexpr auto v_args = TensorAccessorArgs<k_args.next_compile_time_args_offset()>();
+constexpr auto m_args = TensorAccessorArgs<v_args.next_compile_time_args_offset()>();
+constexpr uint32_t cb_q = 0, cb_m = 1, cb_k = 2, cb_max_scaler = 3, cb_sum_scaler = 4, cb_v = 5;
+constexpr uint32_t TB = 4096;
+constexpr uint32_t Wp = (Wt + 1) / 2;
+
+inline void fill_scaler(uint32_t cb) {
+    cb_reserve_back(cb, 1);
+    auto* p = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_write_ptr(cb));
+    for (uint32_t i = 0; i < 1024; ++i) {
+        p[i] = 0;
+    }
+    for (uint32_t f = 0; f < 4; ++f) {
+        for (uint32_t c = 0; c < 16; ++c) {
+            p[f * 256 + c] = 0x3F800000u;
+        }
+    }
+    cb_push_back(cb, 1);
+}
+
+void kernel_main() {
+    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 m_addr = get_common_arg_val<uint32_t>(3);
+    const uint32_t u0 = get_arg_val<uint32_t>(0);
+    const uint32_t n = get_arg_val<uint32_t>(1);
+    if (n == 0) {
+        return;
+    }
+    const auto qa = TensorAccessor(q_args, q_addr, TB);
+    const auto ka = TensorAccessor(k_args, k_addr, TB);
+    const auto va = TensorAccessor(v_args, v_addr, TB);
+    const auto ma = TensorAccessor(m_args, m_addr, TBM);
+    fill_scaler(cb_max_scaler);
+    fill_scaler(cb_sum_scaler);
+    for (uint32_t u = u0; u < u0 + n; ++u) {
+        const uint32_t h = u / Mt;
+        const uint32_t r = u % Mt;
+        cb_reserve_back(cb_m, 2 * Wp);
+        {
+            uint32_t p = get_write_ptr(cb_m);
+            for (uint32_t j = 0; j < Wt; ++j) {
+                noc_async_read(ma.get_noc_addr(r * Wt + j), p, TBM);
+                p += TBM;
+            }
+        }
+        cb_reserve_back(cb_q, 1);
+        noc_async_read(qa.get_noc_addr(qb + r * qr + h * qh), get_write_ptr(cb_q), TB);
+        noc_async_read_barrier();
+        cb_push_back(cb_m, 2 * Wp);
+        cb_push_back(cb_q, 1);
+        for (uint32_t pp = 0; pp < Wp; ++pp) {
+            const uint32_t j = 2 * pp;
+            cb_reserve_back(cb_k, 2);
+            const uint32_t p = get_write_ptr(cb_k);
+            noc_async_read(ka.get_noc_addr(kb + j * kr + h * kh), p, TB);
+            if (j + 1 < Wt) {
+                noc_async_read(ka.get_noc_addr(kb + (j + 1) * kr + h * kh), p + TB, TB);
+            }
+            noc_async_read_barrier();
+            cb_push_back(cb_k, 2);
+        }
+        for (uint32_t j = 0; j < Wt; ++j) {
+            cb_reserve_back(cb_v, 1);
+            noc_async_read(va.get_noc_addr(vb + j * vr + h * vh), get_write_ptr(cb_v), TB);
+            noc_async_read_barrier();
+            cb_push_back(cb_v, 1);
+        }
+    }
+}

code/tt_diffusion_planner/tt/kernels/fattn_writer.cpp

@@ -0,0 +1,82 @@
+// SPDX-License-Identifier: Apache-2.0
+// Fused fp32 matmul attention (tt/fattn_kernel.py), writer (RISCV_1): per unit u (head h = u / Mt, query tile row
+// r = u % Mt) the output tile as tile (r, h) of the merged-heads [1, 1, Sq, H * 32] tensor (what nlp_concat_heads
+// makes of the [1, H, Sq, 32] P V output).
+// With kcat_out (KCAT_EMIT) the output is instead the split operand [o_hi | o_hi | o_lo | 1 | 0..] of the next
+// K-concatenated linear (row tiles Ktp, kernels/kcat_writer.cpp's layout): o_hi to (r, h) and (r, H + h), o_lo to
+// (r, 2 H + h); the unit of head 0 also writes the ones tile (r, 3 H) and the zero pad tiles.
+// CT args: [0] Mt, [1] H, [2] per-core RT-arg count (P3), [3] kcat_out, [4] Ktp, then the TensorAccessorArgs of out.
+// Common RT args: [out_addr]. Per-core RT args: [u0, n].
+#include <cstdint>
+
+#include "api/dataflow/dataflow_api.h"
+
+constexpr uint32_t Mt = get_compile_time_arg_val(0);
+constexpr uint32_t H = get_compile_time_arg_val(1);
+constexpr uint32_t kcat_out = get_compile_time_arg_val(3);
+constexpr uint32_t Ktp = get_compile_time_arg_val(4);
+constexpr auto o_args = TensorAccessorArgs<5>();
+constexpr uint32_t cb_o = 16, cb_lo = 17, cb_one = 18, cb_zero = 19;
+constexpr uint32_t TB = 4096;
+
+void kernel_main() {
+    const uint32_t o_addr = get_common_arg_val<uint32_t>(0);
+    const uint32_t u0 = get_arg_val<uint32_t>(0);
+    const uint32_t n = get_arg_val<uint32_t>(1);
+    const auto o = TensorAccessor(o_args, o_addr, TB);
+    if constexpr (!kcat_out) {
+        for (uint32_t u = u0; u < u0 + n; ++u) {
+            const uint32_t h = u / Mt;
+            const uint32_t r = u % Mt;
+            cb_wait_front(cb_o, 1);
+            noc_async_write(get_read_ptr(cb_o), o.get_noc_addr(r * H + h), TB);
+            noc_async_writes_flushed();
+            cb_pop_front(cb_o, 1);
+        }
+    } else {
+        // ones tile: element (i, c) at face (i / 16) * 2 + (c / 16), offset (i % 16) * 16 + c % 16
+        uint32_t one_ptr = 0, zero_ptr = 0;
+        if (u0 < Mt) {
+            cb_reserve_back(cb_one, 1);
+            one_ptr = get_write_ptr(cb_one);
+            auto* p = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(one_ptr);
+            for (uint32_t i = 0; i < 1024; ++i) {
+                p[i] = 0;
+            }
+            for (uint32_t i = 0; i < 32; ++i) {
+                const uint32_t base = (i / 16) * 512 + (i % 16) * 16;
+                p[base] = 0x3F800000u;
+                p[base + 1] = 0x3F800000u;
+            }
+            if constexpr (Ktp > 3 * H + 1) {
+                cb_reserve_back(cb_zero, 1);
+                zero_ptr = get_write_ptr(cb_zero);
+                auto* z = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(zero_ptr);
+                for (uint32_t i = 0; i < 1024; ++i) {
+                    z[i] = 0;
+                }
+            }
+        }
+        for (uint32_t u = u0; u < u0 + n; ++u) {
+            const uint32_t h = u / Mt;
+            const uint32_t r = u % Mt;
+            const uint32_t ob = r * Ktp;
+            cb_wait_front(cb_o, 1);
+            const uint32_t hp = get_read_ptr(cb_o);
+            noc_async_write(hp, o.get_noc_addr(ob + h), TB);
+            noc_async_write(hp, o.get_noc_addr(ob + H + h), TB);
+            cb_wait_front(cb_lo, 1);
+            noc_async_write(get_read_ptr(cb_lo), o.get_noc_addr(ob + 2 * H + h), TB);
+            if (h == 0) {
+                noc_async_write(one_ptr, o.get_noc_addr(ob + 3 * H), TB);
+                for (uint32_t z = 3 * H + 1; z < Ktp; ++z) {
+                    noc_async_write(zero_ptr, o.get_noc_addr(ob + z), TB);
+                }
+            }
+            noc_async_writes_flushed();
+            cb_pop_front(cb_o, 1);
+            cb_pop_front(cb_lo, 1);
+        }
+    }
+    noc_async_write_barrier();
+}

code/tt_diffusion_planner/tt/kernels/kcat_compute.cpp

@@ -0,0 +1,74 @@
+// SPDX-License-Identifier: Apache-2.0
+// Split-matmul operand build (tt/kcat_kernel.py), compute: per fp32 x tile, x_hi = bf16(x) (typecast_tile
+// Float32 -> Float16_b, the stock ttnn.typecast LLK) kept as fp32, and x_lo = x - x_hi (sub_binary_tile, exact).
+// With act (KCAT_ACT) the previous linear's activation is applied first, as the stock unary program does it
+// (math_approx_mode false): 1 = gelu_tile<0> (ttnn.gelu, accurate), 2 = gelu_tanh_tile (GeluVariant.Tanh).
+// CT args: [0] per-core RT-arg count (P3), [1] act, [2] once (KCAT_ACT_ONCE: the activation on one DEST copy, then
+// copy_dest_values to the second; the same values). Per-core RT args: [n, 0].
+#include <cstdint>
+
+#include "api/compute/common.h"
+#include "api/compute/compute_kernel_api.h"
+#include "api/compute/copy_dest_values.h"
+#include "api/compute/eltwise_binary_sfpu.h"
+#include "api/compute/eltwise_unary/eltwise_unary.h"
+#include "api/compute/eltwise_unary/gelu.h"
+#include "api/compute/eltwise_unary/typecast.h"
+#include "api/compute/tile_move_copy.h"
+
+using namespace ckernel;
+
+constexpr uint32_t act = get_compile_time_arg_val(1);
+constexpr uint32_t once = get_compile_time_arg_val(2);
+constexpr uint32_t cb_x = 0, cb_hi = 16, cb_lo = 17;
+
+void kernel_main() {
+    const uint32_t n = get_arg_val<uint32_t>(0);
+    if (n == 0) {
+        return;
+    }
+    compute_kernel_hw_startup(cb_x, cb_hi);
+    for (uint32_t t = 0; t < n; ++t) {
+        cb_wait_front(cb_x, 1);
+        tile_regs_acquire();
+        copy_init(cb_x);
+        copy_tile(cb_x, 0, 0);
+        if constexpr (once && act != 0) {
+            if constexpr (act == 1) {
+                gelu_tile_init<0u>();
+                gelu_tile<0u>(0);
+            } else {
+                gelu_tanh_tile_init();
+                gelu_tanh_tile(0);
+            }
+            copy_dest_values_init();
+            copy_dest_values<DataFormat::Float32>(0, 1);
+        } else {
+            copy_tile(cb_x, 0, 1);
+        }
+        if constexpr (once && act != 0) {
+        } else if constexpr (act == 1) {
+            gelu_tile_init<0u>();
+            gelu_tile<0u>(0);
+            gelu_tile<0u>(1);
+        } else if constexpr (act == 2) {
+            gelu_tanh_tile_init();
+            gelu_tanh_tile(0);
+            gelu_tanh_tile(1);
+        }
+        typecast_tile_init<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>();
+        typecast_tile<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>(0);
+        sub_binary_tile_init();
+        sub_binary_tile<ckernel::DstRoundingMode::NearestEven>(1, 0, 1);
+        cb_reserve_back(cb_hi, 1);
+        cb_reserve_back(cb_lo, 1);
+        tile_regs_commit();
+        tile_regs_wait();
+        pack_tile(0, cb_hi);
+        pack_tile(1, cb_lo);
+        tile_regs_release();
+        cb_push_back(cb_hi, 1);
+        cb_push_back(cb_lo, 1);
+        cb_pop_front(cb_x, 1);
+    }
+}

code/tt_diffusion_planner/tt/kernels/kcat_reader.cpp

@@ -0,0 +1,24 @@
+// SPDX-License-Identifier: Apache-2.0
+// Split-matmul operand build (tt/kcat_kernel.py), reader (RISCV_0): this core's contiguous range of x tiles.
+// CT args: [0] per-core RT-arg count (P3), then the TensorAccessorArgs of x.
+// Common RT args: [x_addr]. Per-core RT args: [t0, n] (flattened x tile indices).
+#include <cstdint>
+
+#include "api/dataflow/dataflow_api.h"
+
+constexpr auto x_args = TensorAccessorArgs<1>();
+constexpr uint32_t cb_x = 0;
+constexpr uint32_t TB = 4096;
+
+void kernel_main() {
+    const uint32_t x_addr = get_common_arg_val<uint32_t>(0);
+    const uint32_t t0 = get_arg_val<uint32_t>(0);
+    const uint32_t n = get_arg_val<uint32_t>(1);
+    const auto x = TensorAccessor(x_args, x_addr, TB);
+    for (uint32_t t = t0; t < t0 + n; ++t) {
+        cb_reserve_back(cb_x, 1);
+        noc_async_read(x.get_noc_addr(t), get_write_ptr(cb_x), TB);
+        noc_async_read_barrier();
+        cb_push_back(cb_x, 1);
+    }
+}

code/tt_diffusion_planner/tt/kernels/kcat_writer.cpp

@@ -0,0 +1,66 @@
+// SPDX-License-Identifier: Apache-2.0
+// Split-matmul operand build (tt/kcat_kernel.py), writer (RISCV_1): x tile (r, j) of [.., M, K] goes to
+// X' (r, j) and (r, Kt + j) as x_hi and to (r, 2 Kt + j) as x_lo; the core holding (r, 0) also writes the ones tile
+// (r, 3 Kt) (columns 0 and 1 = 1.0: the bias rows of the concatenated weight) and the zero tiles
+// (r, 3 Kt + 1 .. Ktp - 1) that pad K' to a multiple of the matmul's K block.
+// CT args: [0] Kt, [1] per-core RT-arg count (P3), [2] Ktp (row tiles of X'), then the TensorAccessorArgs of X'.
+// Common RT args: [out_addr]. Per-core RT args: [t0, n].
+#include <cstdint>
+
+#include "api/dataflow/dataflow_api.h"
+
+constexpr uint32_t Kt = get_compile_time_arg_val(0);
+constexpr uint32_t OKt = get_compile_time_arg_val(2);
+constexpr auto o_args = TensorAccessorArgs<3>();
+constexpr uint32_t cb_hi = 16, cb_lo = 17, cb_one = 18, cb_zero = 19;
+constexpr uint32_t TB = 4096;
+
+void kernel_main() {
+    const uint32_t o_addr = get_common_arg_val<uint32_t>(0);
+    const uint32_t t0 = get_arg_val<uint32_t>(0);
+    const uint32_t n = get_arg_val<uint32_t>(1);
+    const auto o = TensorAccessor(o_args, o_addr, TB);
+    // ones tile: element (i, c) at face (i / 16) * 2 + (c / 16), offset (i % 16) * 16 + c % 16
+    cb_reserve_back(cb_one, 1);
+    const uint32_t one_ptr = get_write_ptr(cb_one);
+    {
+        auto* p = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(one_ptr);
+        for (uint32_t i = 0; i < 1024; ++i) {
+            p[i] = 0;
+        }
+        for (uint32_t i = 0; i < 32; ++i) {
+            const uint32_t base = (i / 16) * 512 + (i % 16) * 16;
+            p[base] = 0x3F800000u;
+            p[base + 1] = 0x3F800000u;
+        }
+    }
+    uint32_t zero_ptr = 0;
+    if constexpr (OKt > 3 * Kt + 1) {
+        cb_reserve_back(cb_zero, 1);
+        zero_ptr = get_write_ptr(cb_zero);
+        auto* z = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(zero_ptr);
+        for (uint32_t i = 0; i < 1024; ++i) {
+            z[i] = 0;
+        }
+    }
+    for (uint32_t t = t0; t < t0 + n; ++t) {
+        const uint32_t r = t / Kt, j = t - r * Kt;
+        const uint32_t ob = r * OKt;
+        cb_wait_front(cb_hi, 1);
+        const uint32_t hp = get_read_ptr(cb_hi);
+        noc_async_write(hp, o.get_noc_addr(ob + j), TB);
+        noc_async_write(hp, o.get_noc_addr(ob + Kt + j), TB);
+        cb_wait_front(cb_lo, 1);
+        noc_async_write(get_read_ptr(cb_lo), o.get_noc_addr(ob + 2 * Kt + j), TB);
+        if (j == 0) {
+            noc_async_write(one_ptr, o.get_noc_addr(ob + 3 * Kt), TB);
+            for (uint32_t z = 3 * Kt + 1; z < OKt; ++z) {
+                noc_async_write(zero_ptr, o.get_noc_addr(ob + z), TB);
+            }
+        }
+        noc_async_writes_flushed();
+        cb_pop_front(cb_hi, 1);
+        cb_pop_front(cb_lo, 1);
+    }
+    noc_async_write_barrier();
+}

code/tt_diffusion_planner/tt/kernels/ln32_compute.cpp

@@ -0,0 +1,282 @@
+// SPDX-License-Identifier: Apache-2.0
+// Fused fp32 LayerNorm (tt/ln_kernel.py), compute: the exact op sequence of tt/layers.py layer_norm_fp32 (stock
+// ttnn programs: mean -> sub -> mul -> mean -> add eps -> rsqrt -> mul -> mul gamma -> add beta), with the same SFPU
+// LLK calls in the same order, so the result is meant to be bit-identical to the 9-program decomposition:
+//   - ttnn.mean (accurate fp32 path, reduce_op.cpp use_sfpu_fp32_mean): copy tile 0, add_binary_tile fold of tiles
+//     1..Wt-1 in order, sfpu_reduce<SUM, Float32, REDUCE_ROW>, mul_unary_tile(1/W) (the AVG post-mul);
+//   - binary_ng fp32 ops: sub / add as sub_binary_tile / add_binary_tile<NearestEven>, mul as mul_binary_tile,
+//     the broadcast operand built by the dataflow (column / row fill), lhs in DST 0;
+//   - ttnn.rsqrt(fast_and_approximate_mode=False): rsqrt_tile<RsqrtMode::Default> (math_approx_mode false).
+// Everything stays fp32 (UnpackToDestFp32 copies, fp32 DST, fp32 CBs): every intermediate the stock graph packs to
+// an fp32 DRAM tensor round-trips losslessly.
+//
+// With has_res (LN_RESID, the residual add of the stream fused in front): h = x + res (* rgate), the stock
+// `ttnn.add(x, ttnn.multiply(res, rgate))` / `ttnn.add(x, res)` as mul_binary_tile / add_binary_tile, packed to
+// c_9 (the LayerNorm input) and to c_17 (written out as the new stream when write_h).
+//
+// CT args: [0] Wt, [1] has_gamma, [2] has_beta, [3] 1/W bits, [4] per-core RT-arg count (P3), [5] has_res,
+//          [6] has_rgate, [7] write_h, [9] sfpu_bcast (the statistics broadcast over the columns by the SFPU
+//          reduce itself, ln32_sfpu.h, and packed straight to c_5; else column 0 packed to c_4 and broadcast by the
+//          writer), [8] lean (one copy_init / SFPU-binary init per phase instead of one per
+//          tile: every CB copied from is an fp32 UnpackToDestFp32 CB, and add / sub / mul share one init).
+//          [10] res_t (LN_TR: the residual tiles arrive transposed and are transposed back with transpose_tile, the
+//          stock ttnn.transpose LLK, exact with the fp32 unpack-to-dest; no rgate), [11] out_t (the y tiles are
+//          transposed the same way, via c_18, before they are packed for the writer).
+// Per-core RT args: [n_rows].
+#include <cstdint>
+
+#include "api/compute/common.h"
+#include "api/compute/compute_kernel_api.h"
+#include "api/compute/eltwise_binary_sfpu.h"
+#include "api/compute/eltwise_unary/binop_with_scalar.h"
+#include "api/compute/eltwise_unary/eltwise_unary.h"
+#include "api/compute/eltwise_unary/rsqrt.h"
+#include "api/compute/tile_move_copy.h"
+#include "api/compute/transpose.h"
+#include "ln32_sfpu.h"
+
+using namespace ckernel;
+
+constexpr uint32_t Wt = get_compile_time_arg_val(0);
+constexpr uint32_t has_gamma = get_compile_time_arg_val(1);
+constexpr uint32_t has_beta = get_compile_time_arg_val(2);
+constexpr uint32_t inv_w_bits = get_compile_time_arg_val(3);
+constexpr uint32_t has_res = get_compile_time_arg_val(5);
+constexpr uint32_t has_rgate = get_compile_time_arg_val(6);
+constexpr uint32_t write_h = get_compile_time_arg_val(7);
+constexpr bool lean = get_compile_time_arg_val(8) != 0;
+constexpr bool sfpu_bcast = get_compile_time_arg_val(9) != 0;
+constexpr uint32_t res_t = get_compile_time_arg_val(10);
+constexpr uint32_t out_t = get_compile_time_arg_val(11);
+
+// per-tile (re-)inits, skipped in lean mode (done once at the start of the phase instead)
+ALWI void ci(uint32_t cb) {
+    if constexpr (!lean) {
+        copy_init(cb);
+    }
+}
+#define BIN_INIT(op)                  \
+    do {                              \
+        if constexpr (!lean) {        \
+            op##_binary_tile_init();  \
+        }                             \
+    } while (0)
+
+constexpr uint32_t cb_x = 0, cb_g = 1, cb_b = 2, cb_eps = 3, cb_stat = 4, cb_bc = 5, cb_xc = 6, cb_r = 7, cb_rg = 8,
+                   cb_h = 9, cb_out = 16, cb_hout = 17, cb_yt = 18;
+constexpr uint32_t cb_src = has_res ? cb_h : cb_x;
+constexpr auto RNE = ckernel::DstRoundingMode::NearestEven;
+
+// row mean of the Wt tiles of `cb` (already waited) into DST 0; `square` folds xc * xc instead of x
+template <bool square>
+ALWI void row_sum_to_dst0(uint32_t cb) {
+    copy_init(cb);
+    if constexpr (lean) {
+        add_binary_tile_init();
+    }
+    if constexpr (square) {
+        copy_tile(cb, 0, 0);
+        copy_tile(cb, 0, 1);
+        BIN_INIT(mul);
+        mul_binary_tile(0, 1, 0);
+    } else {
+        copy_tile(cb, 0, 0);
+    }
+    for (uint32_t w = 1; w < Wt; ++w) {
+        ci(cb);
+        copy_tile(cb, w, 1);
+        if constexpr (square) {
+            copy_tile(cb, w, 2);
+            BIN_INIT(mul);
+            mul_binary_tile(1, 2, 1);
+        }
+        BIN_INIT(add);
+        add_binary_tile(0, 1, 0);
+    }
+    sfpu_reduce_init<PoolType::SUM, DataFormat::Float32>();
+    if constexpr (sfpu_bcast) {
+        ln_row_sum_bcast_tile(0);
+    } else {
+        sfpu_reduce<PoolType::SUM, DataFormat::Float32, ReduceDim::REDUCE_ROW>(0, 1, 1);
+    }
+    binop_with_scalar_tile_init();
+    mul_unary_tile(0, inv_w_bits);
+}
+
+ALWI void pack_stat() {
+    constexpr uint32_t cb = sfpu_bcast ? cb_bc : cb_stat;
+    cb_reserve_back(cb, 1);
+    tile_regs_commit();
+    tile_regs_wait();
+    pack_tile(0, cb);
+    tile_regs_release();
+    cb_push_back(cb, 1);
+}
+
+void kernel_main() {
+    const uint32_t n_rows = get_arg_val<uint32_t>(0);
+    if (n_rows == 0) {
+        return;
+    }
+    compute_kernel_hw_startup(cb_x, cb_out);
+    // the constant tiles are waited for where they are first used (the reader sends x row 0 first)
+    for (uint32_t r = 0; r < n_rows; ++r) {
+        cb_wait_front(cb_x, Wt);
+        if constexpr (has_res) {
+            // h = x + res (* rgate)
+            cb_wait_front(cb_r, Wt);
+            if constexpr (has_rgate) {
+                cb_wait_front(cb_rg, Wt);
+            }
+            if constexpr (lean) {
+                copy_init(cb_r);
+                add_binary_tile_init();
+            }
+            for (uint32_t w = 0; w < Wt; ++w) {
+                tile_regs_acquire();
+                if constexpr (res_t) {
+                    transpose_init(cb_r);
+                    transpose_tile(cb_r, w, 0);
+                    copy_init(cb_x);
+                    add_binary_tile_init();
+                } else {
+                    ci(cb_r);
+                    copy_tile(cb_r, w, 0);
+                }
+                if constexpr (has_rgate) {
+                    ci(cb_rg);
+                    copy_tile(cb_rg, w, 1);
+                    BIN_INIT(mul);
+                    mul_binary_tile(0, 1, 0);
+                }
+                ci(cb_x);
+                copy_tile(cb_x, w, 1);
+                BIN_INIT(add);
+                add_binary_tile<RNE>(1, 0, 1);
+                cb_reserve_back(cb_h, 1);
+                if constexpr (write_h) {
+                    cb_reserve_back(cb_hout, 1);
+                }
+                tile_regs_commit();
+                tile_regs_wait();
+                pack_tile(1, cb_h);
+                if constexpr (write_h) {
+                    pack_tile(1, cb_hout);
+                }
+                tile_regs_release();
+                cb_push_back(cb_h, 1);
+                if constexpr (write_h) {
+                    cb_push_back(cb_hout, 1);
+                }
+            }
+            cb_pop_front(cb_r, Wt);
+            cb_pop_front(cb_x, Wt);
+            cb_wait_front(cb_h, Wt);
+        }
+        // mean
+        tile_regs_acquire();
+        row_sum_to_dst0<false>(cb_src);
+        pack_stat();
+
+        // xc = x - mean
+        cb_wait_front(cb_bc, 1);
+        if constexpr (lean) {
+            copy_init(cb_src);
+            sub_binary_tile_init();
+        }
+        for (uint32_t w = 0; w < Wt; ++w) {
+            tile_regs_acquire();
+            ci(cb_src);
+            copy_tile(cb_src, w, 0);
+            ci(cb_bc);
+            copy_tile(cb_bc, 0, 1);
+            BIN_INIT(sub);
+            sub_binary_tile<RNE>(0, 1, 0);
+            cb_reserve_back(cb_xc, 1);
+            tile_regs_commit();
+            tile_regs_wait();
+            pack_tile(0, cb_xc);
+            tile_regs_release();
+            cb_push_back(cb_xc, 1);
+        }
+        cb_pop_front(cb_bc, 1);
+        cb_pop_front(cb_src, Wt);
+
+        // rstd = rsqrt(mean(xc * xc) + eps)
+        cb_wait_front(cb_xc, Wt);
+        tile_regs_acquire();
+        row_sum_to_dst0<true>(cb_xc);
+        cb_wait_front(cb_eps, 1);
+        copy_init(cb_eps);
+        copy_tile(cb_eps, 0, 1);
+        add_binary_tile_init();
+        add_binary_tile<RNE>(0, 1, 0);
+        rsqrt_tile_init();
+        rsqrt_tile<RsqrtMode::Default>(0);
+        pack_stat();
+
+        // y = xc * rstd (* gamma) (+ beta)
+        cb_wait_front(cb_bc, 1);
+        if constexpr (has_gamma) {
+            cb_wait_front(cb_g, Wt);
+        }
+        if constexpr (has_beta) {
+            cb_wait_front(cb_b, Wt);
+        }
+        if constexpr (lean) {
+            copy_init(cb_xc);
+            mul_binary_tile_init();
+        }
+        for (uint32_t w = 0; w < Wt; ++w) {
+            tile_regs_acquire();
+            ci(cb_xc);
+            copy_tile(cb_xc, w, 0);
+            ci(cb_bc);
+            copy_tile(cb_bc, 0, 1);
+            BIN_INIT(mul);
+            mul_binary_tile(0, 1, 0);
+            if constexpr (has_gamma) {
+                ci(cb_g);
+                copy_tile(cb_g, w, 1);
+                BIN_INIT(mul);
+                mul_binary_tile(0, 1, 0);
+            }
+            if constexpr (has_beta) {
+                ci(cb_b);
+                copy_tile(cb_b, w, 1);
+                BIN_INIT(add);
+                add_binary_tile<RNE>(0, 1, 0);
+            }
+            if constexpr (out_t) {
+                cb_reserve_back(cb_yt, 1);
+                tile_regs_commit();
+                tile_regs_wait();
+                pack_tile(0, cb_yt);
+                tile_regs_release();
+                cb_push_back(cb_yt, 1);
+                cb_wait_front(cb_yt, 1);
+                tile_regs_acquire();
+                transpose_init(cb_yt);
+                transpose_tile(cb_yt, 0, 0);
+                cb_reserve_back(cb_out, 1);
+                tile_regs_commit();
+                tile_regs_wait();
+                pack_tile(0, cb_out);
+                tile_regs_release();
+                cb_push_back(cb_out, 1);
+                cb_pop_front(cb_yt, 1);
+                copy_init(cb_xc);
+                mul_binary_tile_init();
+            } else {
+                cb_reserve_back(cb_out, 1);
+                tile_regs_commit();
+                tile_regs_wait();
+                pack_tile(0, cb_out);
+                tile_regs_release();
+                cb_push_back(cb_out, 1);
+            }
+        }
+        cb_pop_front(cb_bc, 1);
+        cb_pop_front(cb_xc, Wt);
+    }
+}

code/tt_diffusion_planner/tt/kernels/ln32_reader.cpp

@@ -0,0 +1,123 @@
+// SPDX-License-Identifier: Apache-2.0
+// Fused fp32 LayerNorm (tt/ln_kernel.py), reader (RISCV_0): the affine rows once, then this core's tile rows of x.
+//
+// CT args: [0] Wt, [1] has_gamma, [2] has_beta, [3] eps bits (fp32), [4] per-core RT-arg count (P3: part of the
+//          program hash), [5] has_res, [6] has_rgate, [7] res_t (LN_TR: res is stored transposed per entity, i.e.
+//          [.., E, W, T]: tile (r, w) of res is read from tile (e, w, tr) of that layout, e = r / Tt, tr = r % Tt;
+//          the compute transposes it back), [8] Tt (tile rows per entity), then the TensorAccessorArgs of x, gamma,
+//          beta, res, rgate (x's when absent).
+// Common RT args: [x_addr, gamma_addr, beta_addr, res_addr, rgate_addr]. Per-core RT args: [row0, n_rows].
+// CBs: c_0 x (Wt tiles per row, double-buffered), c_1 gamma rows (Wt tiles, row 0 replicated over the tile),
+//      c_2 beta rows (same), c_3 the eps tile (every element eps), c_7 res (as c_0), c_8 the residual gate rows.
+#include <cstdint>
+
+#include "api/dataflow/dataflow_api.h"
+
+constexpr uint32_t Wt = get_compile_time_arg_val(0);
+constexpr uint32_t has_gamma = get_compile_time_arg_val(1);
+constexpr uint32_t has_beta = get_compile_time_arg_val(2);
+constexpr uint32_t eps_bits = get_compile_time_arg_val(3);
+constexpr uint32_t has_res = get_compile_time_arg_val(5);
+constexpr uint32_t has_rgate = get_compile_time_arg_val(6);
+constexpr uint32_t res_t = get_compile_time_arg_val(7);
+constexpr uint32_t Tt = get_compile_time_arg_val(8);
+constexpr auto x_args = TensorAccessorArgs<9>();
+constexpr auto g_args = TensorAccessorArgs<x_args.next_compile_time_args_offset()>();
+constexpr auto b_args = TensorAccessorArgs<g_args.next_compile_time_args_offset()>();
+constexpr auto r_args = TensorAccessorArgs<b_args.next_compile_time_args_offset()>();
+constexpr auto rg_args = TensorAccessorArgs<r_args.next_compile_time_args_offset()>();
+
+constexpr uint32_t cb_x = 0, cb_g = 1, cb_b = 2, cb_eps = 3, cb_r = 7, cb_rg = 8;
+constexpr uint32_t TB = 4096;  // fp32 tile bytes
+
+// Replicate row 0 of each of the n fp32 tiles at p over its 32 rows (binary_ng fill_tile_with_first_row: local
+// NoC doubling), all tiles per doubling step under one barrier.
+FORCE_INLINE void fill_first_rows(uint32_t p, uint32_t n_tiles) {
+    constexpr uint32_t kFace = 1024, kRow = 64;
+    for (uint32_t n = kRow; n < kFace; n <<= 1) {
+        for (uint32_t t = 0; t < n_tiles; ++t) {
+            const uint32_t f0 = p + t * TB, f1 = f0 + kFace;
+            noc_async_read(get_noc_addr(f0), f0 + n, n);
+            noc_async_read(get_noc_addr(f1), f1 + n, n);
+        }
+        noc_async_read_barrier();
+    }
+    for (uint32_t t = 0; t < n_tiles; ++t) {
+        const uint32_t f0 = p + t * TB;
+        noc_async_read(get_noc_addr(f0), f0 + 2 * kFace, 2 * kFace);
+    }
+    noc_async_read_barrier();
+}
+
+template <typename A>
+FORCE_INLINE void read_rows(uint32_t cb, const A& acc) {
+    cb_reserve_back(cb, Wt);
+    uint32_t p = get_write_ptr(cb);
+    for (uint32_t w = 0; w < Wt; ++w) {
+        noc_async_read(acc.get_noc_addr(w), p + w * TB, TB);
+    }
+    noc_async_read_barrier();
+    fill_first_rows(p, Wt);
+    cb_push_back(cb, Wt);
+}
+
+template <typename A, typename R>
+FORCE_INLINE void read_x_row(uint32_t r, const A& x, const R& res) {
+    cb_reserve_back(cb_x, Wt);
+    uint32_t p = get_write_ptr(cb_x);
+    for (uint32_t w = 0; w < Wt; ++w) {
+        noc_async_read(x.get_noc_addr(r * Wt + w), p + w * TB, TB);
+    }
+    if constexpr (has_res) {
+        cb_reserve_back(cb_r, Wt);
+        uint32_t q = get_write_ptr(cb_r);
+        for (uint32_t w = 0; w < Wt; ++w) {
+            const uint32_t ri = res_t ? (r / Tt) * (Wt * Tt) + w * Tt + (r % Tt) : r * Wt + w;
+            noc_async_read(res.get_noc_addr(ri), q + w * TB, TB);
+        }
+        noc_async_read_barrier();
+        cb_push_back(cb_r, Wt);
+    } else {
+        noc_async_read_barrier();
+    }
+    cb_push_back(cb_x, Wt);
+}
+
+void kernel_main() {
+    const uint32_t x_addr = get_common_arg_val<uint32_t>(0);
+    const uint32_t g_addr = get_common_arg_val<uint32_t>(1);
+    const uint32_t b_addr = get_common_arg_val<uint32_t>(2);
+    const uint32_t r_addr = get_common_arg_val<uint32_t>(3);
+    const uint32_t rg_addr = get_common_arg_val<uint32_t>(4);
+    const uint32_t row0 = get_arg_val<uint32_t>(0);
+    const uint32_t n_rows = get_arg_val<uint32_t>(1);
+    if (n_rows == 0) {
+        return;
+    }
+    const auto x = TensorAccessor(x_args, x_addr, TB);
+    const auto res = TensorAccessor(r_args, r_addr, TB);
+
+    // the first x row goes first (the compute starts on it); the affine rows are needed only by its last phase
+    read_x_row(row0, x, res);
+    // eps tile
+    cb_reserve_back(cb_eps, 1);
+    {
+        auto* e = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_write_ptr(cb_eps));
+        for (uint32_t i = 0; i < 1024; ++i) {
+            e[i] = eps_bits;
+        }
+    }
+    cb_push_back(cb_eps, 1);
+    if constexpr (has_rgate) {
+        read_rows(cb_rg, TensorAccessor(rg_args, rg_addr, TB));
+    }
+    if constexpr (has_gamma) {
+        read_rows(cb_g, TensorAccessor(g_args, g_addr, TB));
+    }
+    if constexpr (has_beta) {
+        read_rows(cb_b, TensorAccessor(b_args, b_addr, TB));
+    }
+    for (uint32_t r = row0 + 1; r < row0 + n_rows; ++r) {
+        read_x_row(r, x, res);
+    }
+}

code/tt_diffusion_planner/tt/kernels/ln32_sfpu.h

@@ -0,0 +1,54 @@
+// SPDX-License-Identifier: Apache-2.0
+// SFPU helper of the fused fp32 LayerNorm kernels (ln32_compute.cpp, ln32s_compute.cpp): the row sum of one fp32
+// tile with the result in EVERY column.
+//
+// ln_row_sum_bcast is perform_reduce_row_sum_tile of tt-metal's blackhole ckernel_sfpu_reduce.h (the REDUCE_ROW SUM
+// Float32 single-tile path behind sfpu_reduce<SUM, Float32, REDUCE_ROW>, used by the accurate ttnn.mean): the same
+// loads, the same replayed vertical adds (the replay recorded by sfpu_reduce_init<SUM, Float32>) and the same
+// horizontal butterfly, which leaves each row's full sum in all 8 lanes of its sub-vector. The stock kernel stores
+// that register to the even columns of the left face only (column 0 is what the packer keeps); this one stores it to
+// the four column groups (even / odd x left / right face), so the tile is the column-0 result broadcast over the
+// row: bit for bit the values the stock graph broadcasts with binary_ng's column fill, without a RISC-V fill loop.
+#pragma once
+
+#ifdef TRISC_MATH
+namespace ckernel::sfpu {
+
+template <bool is_fp32_dest_acc_en>
+inline void ln_row_sum_bcast() {
+    constexpr InstrModLoadStore MODE = GetSfpLoadStoreInstrMod<DataFormat::Float32, is_fp32_dest_acc_en>();
+    for (std::uint32_t face_pair = 0; face_pair < 2; face_pair++) {
+        const std::uint32_t face_pair_base = face_pair * 2 * ROWS_PER_FACE;
+        for (std::uint32_t row_group = 0; row_group < 2; row_group++) {
+            const std::uint32_t a = face_pair_base + row_group * 8;
+            const std::uint32_t b = a + 4;
+            load_and_clear_high_bits<false>(p_sfpu::LREG0, MODE, ADDR_MOD_7, a);
+            load_and_clear_high_bits<false>(p_sfpu::LREG1, MODE, ADDR_MOD_7, a + 2);
+            load_and_clear_high_bits<false>(p_sfpu::LREG2, MODE, ADDR_MOD_7, a + ROWS_PER_FACE);
+            load_and_clear_high_bits<false>(p_sfpu::LREG3, MODE, ADDR_MOD_7, a + ROWS_PER_FACE + 2);
+            load_and_clear_high_bits<false>(p_sfpu::LREG4, MODE, ADDR_MOD_7, b);
+            load_and_clear_high_bits<false>(p_sfpu::LREG5, MODE, ADDR_MOD_7, b + 2);
+            load_and_clear_high_bits<false>(p_sfpu::LREG6, MODE, ADDR_MOD_7, b + ROWS_PER_FACE);
+            load_and_clear_high_bits<false>(p_sfpu::LREG7, MODE, ADDR_MOD_7, b + ROWS_PER_FACE + 2);
+            lltt::replay(0, 6);
+            horizontal_reduce<false>();
+            TT_SFPSTORE(p_sfpu::LREG0, MODE, ADDR_MOD_7, a);
+            TT_SFPSTORE(p_sfpu::LREG4, MODE, ADDR_MOD_7, b);
+            TT_SFPSTORE(p_sfpu::LREG0, MODE, ADDR_MOD_7, a + 2);
+            TT_SFPSTORE(p_sfpu::LREG4, MODE, ADDR_MOD_7, b + 2);
+            TT_SFPSTORE(p_sfpu::LREG0, MODE, ADDR_MOD_7, a + ROWS_PER_FACE);
+            TT_SFPSTORE(p_sfpu::LREG4, MODE, ADDR_MOD_7, b + ROWS_PER_FACE);
+            TT_SFPSTORE(p_sfpu::LREG0, MODE, ADDR_MOD_7, a + ROWS_PER_FACE + 2);
+            TT_SFPSTORE(p_sfpu::LREG4, MODE, ADDR_MOD_7, b + ROWS_PER_FACE + 2);
+        }
+    }
+}
+
+}  // namespace ckernel::sfpu
+#endif
+
+// row sums of the fp32 tile in DST idst, broadcast over all columns (sfpu_reduce_init<SUM, Float32> first)
+ALWI void ln_row_sum_bcast_tile(uint32_t idst) {
+    MATH(SFPU_UNARY_CALL(DST_SYNC_MODE, DST_ACCUM_MODE, ln_row_sum_bcast, (DST_ACCUM_MODE), idst,
+                         VectorMode::RC_custom));
+}

code/tt_diffusion_planner/tt/kernels/ln32_writer.cpp

@@ -0,0 +1,84 @@
+// SPDX-License-Identifier: Apache-2.0
+// Fused fp32 LayerNorm (tt/ln_kernel.py), writer (RISCV_1): per tile row, broadcasts the two statistic tiles the
+// compute packs (column 0 = the row's mean, then its rsqrt(var + eps)) over all 32 columns, then drains the Wt
+// output tiles to DRAM.
+//
+// With write_h (LN_RESID) the Wt tiles of the new stream h = x + res (* rgate) (c_17) go first, to h_addr.
+//
+// CT args: [0] Wt, [1] per-core RT-arg count (P3), [2] write_h, [3] sfpu_bcast (no fill here), [4] out_t (LN_TR:
+//          the output tiles, transposed by the compute, go to the per-entity transposed layout [.., E, W, T]: tile
+//          (r, w) -> (e, w, tr), e = r / Tt, tr = r % Tt), [5] Tt, then the TensorAccessorArgs of out and h (out's
+//          when absent).
+// Common RT args: [out_addr, h_addr]. Per-core RT args: [row0, n_rows].
+// CBs: c_4 statistic (from compute, column 0 valid), c_5 broadcast statistic (to compute), c_16 out, c_17 h.
+#include <cstdint>
+
+#include "api/dataflow/dataflow_api.h"
+
+constexpr uint32_t Wt = get_compile_time_arg_val(0);
+constexpr uint32_t write_h = get_compile_time_arg_val(2);
+constexpr uint32_t sfpu_bcast = get_compile_time_arg_val(3);   // the compute broadcasts the statistics itself
+constexpr uint32_t out_t = get_compile_time_arg_val(4);
+constexpr uint32_t Tt = get_compile_time_arg_val(5);
+constexpr auto out_args = TensorAccessorArgs<6>();
+constexpr auto h_args = TensorAccessorArgs<out_args.next_compile_time_args_offset()>();
+
+constexpr uint32_t cb_stat = 4, cb_bc = 5, cb_out = 16, cb_hout = 17;
+constexpr uint32_t TB = 4096;
+
+// dst tile = src tile's column 0 replicated over the 32 columns (binary_ng fill_tile_with_first_column, fp32;
+// faces 0/1 take face 0's column 0, faces 2/3 face 2's).
+FORCE_INLINE void bcast_col(uint32_t src, uint32_t dst) {
+    auto* s = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(src);
+    auto* d = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(dst);
+    for (uint32_t fo = 0; fo < 1024; fo += 512) {
+        for (uint32_t ro = 0; ro < 256; ro += 16) {
+            const uint32_t v = s[fo + ro];
+            volatile tt_l1_ptr uint32_t* l = d + fo + ro;
+            volatile tt_l1_ptr uint32_t* r = l + 256;
+            for (uint32_t c = 0; c < 16; ++c) {
+                l[c] = v;
+                r[c] = v;
+            }
+        }
+    }
+}
+
+FORCE_INLINE void stat() {
+    cb_wait_front(cb_stat, 1);
+    cb_reserve_back(cb_bc, 1);
+    bcast_col(get_read_ptr(cb_stat), get_write_ptr(cb_bc));
+    cb_push_back(cb_bc, 1);
+    cb_pop_front(cb_stat, 1);
+}
+
+void kernel_main() {
+    const uint32_t out_addr = get_common_arg_val<uint32_t>(0);
+    const uint32_t h_addr = get_common_arg_val<uint32_t>(1);
+    const auto hacc = TensorAccessor(h_args, h_addr, TB);
+    const uint32_t row0 = get_arg_val<uint32_t>(0);
+    const uint32_t n_rows = get_arg_val<uint32_t>(1);
+    const auto out = TensorAccessor(out_args, out_addr, TB);
+    for (uint32_t r = row0; r < row0 + n_rows; ++r) {
+        if constexpr (write_h) {
+            for (uint32_t w = 0; w < Wt; ++w) {
+                cb_wait_front(cb_hout, 1);
+                noc_async_write(get_read_ptr(cb_hout), hacc.get_noc_addr(r * Wt + w), TB);
+                noc_async_writes_flushed();
+                cb_pop_front(cb_hout, 1);
+            }
+        }
+        if constexpr (!sfpu_bcast) {
+            stat();  // mean
+            stat();  // rsqrt(var + eps)
+        }
+        for (uint32_t w = 0; w < Wt; ++w) {
+            cb_wait_front(cb_out, 1);
+            const uint32_t oi = out_t ? (r / Tt) * (Wt * Tt) + w * Tt + (r % Tt) : r * Wt + w;
+            noc_async_write(get_read_ptr(cb_out), out.get_noc_addr(oi), TB);
+            noc_async_writes_flushed();
+            cb_pop_front(cb_out, 1);
+        }
+    }
+    noc_async_write_barrier();
+}

code/tt_diffusion_planner/tt/kernels/ln32s_compute.cpp

@@ -0,0 +1,214 @@
+// SPDX-License-Identifier: Apache-2.0
+// Split-row fused fp32 LayerNorm (tt/ln_kernel.py, LN_SPLIT), compute: member j of a tile row owns column tile j.
+// The same LLK sequence as ln32_compute.cpp (and the stock 9-program decomposition): the folds over the Wt tiles run
+// on the root, in tile order, over the gathered tiles (c_10), so every statistic is bit for bit the stock one.
+//   h  = x (+ res (* rgate))                       -> c_11 (to the root's gather) [+ c_17 for write_h]
+//   root: mean = fold(c_10) / W                    -> c_4 (the writer broadcasts it into every member's c_5 page 0)
+//   xc = h - mean; sq = xc * xc                    -> c_6 (kept), sq -> c_11
+//   root: rstd = rsqrt(fold(c_10) / W + eps)       -> c_4 (broadcast into c_5 page 1)
+//   y  = xc * rstd (* gamma) (+ beta)              -> c_16
+// CT args: [0] Wt, [1] has_gamma, [2] has_beta, [3] 1/W bits, [4] per-core RT-arg count (P3), [5] has_res,
+//          [6] has_rgate, [7] write_h, [8] sfpu_bcast (the root's statistics broadcast over the columns by the
+//          SFPU reduce, ln32_sfpu.h; else column 0 and the writer fills), [9] kcat_out (KCAT_EMIT: y is emitted
+//          as the split operand of the next K-concatenated linear, y_hi = bf16(y) -> c_16 and y_lo = y - y_hi ->
+//          c_19, the LLK calls of kernels/kcat_compute.cpp on the packed fp32 y tile c_18).
+//          Per-core RT args: [n_rows, is_root].
+#include <cstdint>
+
+#include "api/compute/common.h"
+#include "api/compute/compute_kernel_api.h"
+#include "api/compute/eltwise_binary_sfpu.h"
+#include "api/compute/eltwise_unary/binop_with_scalar.h"
+#include "api/compute/eltwise_unary/eltwise_unary.h"
+#include "api/compute/eltwise_unary/rsqrt.h"
+#include "api/compute/eltwise_unary/typecast.h"
+#include "api/compute/tile_move_copy.h"
+#include "ln32_sfpu.h"
+
+using namespace ckernel;
+
+constexpr uint32_t Wt = get_compile_time_arg_val(0);
+constexpr uint32_t has_gamma = get_compile_time_arg_val(1);
+constexpr uint32_t has_beta = get_compile_time_arg_val(2);
+constexpr uint32_t inv_w_bits = get_compile_time_arg_val(3);
+constexpr uint32_t has_res = get_compile_time_arg_val(5);
+constexpr uint32_t has_rgate = get_compile_time_arg_val(6);
+constexpr uint32_t write_h = get_compile_time_arg_val(7);
+constexpr bool sfpu_bcast = get_compile_time_arg_val(8) != 0;
+constexpr uint32_t kcat_out = get_compile_time_arg_val(9);
+
+constexpr uint32_t cb_x = 0, cb_g = 1, cb_b = 2, cb_eps = 3, cb_stat = 4, cb_bc = 5, cb_xc = 6, cb_r = 7, cb_rg = 8,
+                   cb_h = 9, cb_gat = 10, cb_snd = 11, cb_out = 16, cb_hout = 17, cb_y = 18, cb_lo = 19;
+constexpr uint32_t cb_src = has_res ? cb_h : cb_x;
+constexpr auto RNE = ckernel::DstRoundingMode::NearestEven;
+
+// fold of the Wt gathered tiles (in order) -> row sum * 1/W in DST 0
+ALWI void fold_gathered() {
+    cb_wait_front(cb_gat, Wt);
+    copy_init(cb_gat);
+    add_binary_tile_init();
+    copy_tile(cb_gat, 0, 0);
+    for (uint32_t w = 1; w < Wt; ++w) {
+        copy_tile(cb_gat, w, 1);
+        add_binary_tile(0, 1, 0);
+    }
+    sfpu_reduce_init<PoolType::SUM, DataFormat::Float32>();
+    if constexpr (sfpu_bcast) {
+        ln_row_sum_bcast_tile(0);
+    } else {
+        sfpu_reduce<PoolType::SUM, DataFormat::Float32, ReduceDim::REDUCE_ROW>(0, 1, 1);
+    }
+    binop_with_scalar_tile_init();
+    mul_unary_tile(0, inv_w_bits);
+}
+
+ALWI void pack_one(uint32_t dst, uint32_t cb) {
+    cb_reserve_back(cb, 1);
+    tile_regs_commit();
+    tile_regs_wait();
+    pack_tile(dst, cb);
+    tile_regs_release();
+    cb_push_back(cb, 1);
+}
+
+void kernel_main() {
+    const uint32_t n_rows = get_arg_val<uint32_t>(0);
+    const bool root = get_arg_val<uint32_t>(1) != 0;
+    if (n_rows == 0) {
+        return;
+    }
+    compute_kernel_hw_startup(cb_x, cb_out);
+    for (uint32_t r = 0; r < n_rows; ++r) {
+        // h (this member's tile) -> c_11 for the gather (and c_9 / c_17)
+        cb_wait_front(cb_x, 1);
+        tile_regs_acquire();
+        if constexpr (has_res) {
+            cb_wait_front(cb_r, 1);
+            if constexpr (has_rgate) {
+                cb_wait_front(cb_rg, 1);
+            }
+            copy_init(cb_r);
+            copy_tile(cb_r, 0, 0);
+            if constexpr (has_rgate) {
+                copy_tile(cb_rg, 0, 1);
+                mul_binary_tile_init();
+                mul_binary_tile(0, 1, 0);
+            }
+            copy_tile(cb_x, 0, 1);
+            add_binary_tile_init();
+            add_binary_tile<RNE>(1, 0, 1);
+            cb_reserve_back(cb_h, 1);
+            cb_reserve_back(cb_snd, 1);
+            if constexpr (write_h) {
+                cb_reserve_back(cb_hout, 1);
+            }
+            tile_regs_commit();
+            tile_regs_wait();
+            pack_tile(1, cb_h);
+            pack_tile(1, cb_snd);
+            if constexpr (write_h) {
+                pack_tile(1, cb_hout);
+            }
+            tile_regs_release();
+            cb_push_back(cb_h, 1);
+            cb_push_back(cb_snd, 1);
+            if constexpr (write_h) {
+                cb_push_back(cb_hout, 1);
+            }
+            cb_pop_front(cb_r, 1);
+            cb_pop_front(cb_x, 1);
+            cb_wait_front(cb_h, 1);
+        } else {
+            copy_init(cb_x);
+            copy_tile(cb_x, 0, 0);
+            pack_one(0, cb_snd);
+        }
+        // root: mean
+        if (root) {
+            tile_regs_acquire();
+            fold_gathered();
+            pack_one(0, cb_stat);
+            cb_pop_front(cb_gat, Wt);
+        }
+        // xc = h - mean -> c_6; sq = xc * xc -> c_11
+        cb_wait_front(cb_bc, 1);
+        tile_regs_acquire();
+        copy_init(cb_src);
+        copy_tile(cb_src, 0, 0);
+        copy_tile(cb_bc, 0, 1);
+        sub_binary_tile_init();
+        sub_binary_tile<RNE>(0, 1, 0);
+        pack_one(0, cb_xc);
+        cb_pop_front(cb_bc, 1);
+        cb_pop_front(cb_src, 1);
+        cb_wait_front(cb_xc, 1);
+        tile_regs_acquire();
+        copy_init(cb_xc);
+        copy_tile(cb_xc, 0, 0);
+        copy_tile(cb_xc, 0, 1);
+        mul_binary_tile_init();
+        mul_binary_tile(0, 1, 0);
+        pack_one(0, cb_snd);
+        // root: rstd
+        if (root) {
+            tile_regs_acquire();
+            fold_gathered();
+            cb_wait_front(cb_eps, 1);
+            copy_init(cb_eps);
+            copy_tile(cb_eps, 0, 1);
+            add_binary_tile_init();
+            add_binary_tile<RNE>(0, 1, 0);
+            rsqrt_tile_init();
+            rsqrt_tile<RsqrtMode::Default>(0);
+            pack_one(0, cb_stat);
+            cb_pop_front(cb_gat, Wt);
+        }
+        // y = xc * rstd (* gamma) (+ beta)
+        cb_wait_front(cb_bc, 1);
+        tile_regs_acquire();
+        copy_init(cb_xc);
+        copy_tile(cb_xc, 0, 0);
+        copy_tile(cb_bc, 0, 1);
+        mul_binary_tile_init();
+        mul_binary_tile(0, 1, 0);
+        if constexpr (has_gamma) {
+            cb_wait_front(cb_g, 1);
+            copy_tile(cb_g, 0, 1);
+            mul_binary_tile(0, 1, 0);
+        }
+        if constexpr (has_beta) {
+            cb_wait_front(cb_b, 1);
+            copy_tile(cb_b, 0, 1);
+            add_binary_tile_init();
+            add_binary_tile<RNE>(0, 1, 0);
+        }
+        if constexpr (kcat_out) {
+            pack_one(0, cb_y);
+            cb_pop_front(cb_bc, 1);
+            cb_pop_front(cb_xc, 1);
+            cb_wait_front(cb_y, 1);
+            tile_regs_acquire();
+            copy_init(cb_y);
+            copy_tile(cb_y, 0, 0);
+            copy_tile(cb_y, 0, 1);
+            typecast_tile_init<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>();
+            typecast_tile<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>(0);
+            sub_binary_tile_init();
+            sub_binary_tile<RNE>(1, 0, 1);
+            cb_reserve_back(cb_out, 1);
+            cb_reserve_back(cb_lo, 1);
+            tile_regs_commit();
+            tile_regs_wait();
+            pack_tile(0, cb_out);
+            pack_tile(1, cb_lo);
+            tile_regs_release();
+            cb_push_back(cb_out, 1);
+            cb_push_back(cb_lo, 1);
+            cb_pop_front(cb_y, 1);
+        } else {
+            pack_one(0, cb_out);
+            cb_pop_front(cb_bc, 1);
+            cb_pop_front(cb_xc, 1);
+        }
+    }
+}

code/tt_diffusion_planner/tt/kernels/ln32s_reader.cpp

@@ -0,0 +1,130 @@
+// SPDX-License-Identifier: Apache-2.0
+// Split-row fused fp32 LayerNorm (tt/ln_kernel.py, LN_SPLIT), reader (RISCV_0). One tile row is spread over Wt
+// cores (member j owns column tile j); member 0 (the root) gathers the Wt tiles that the stock reduction folds in
+// order and broadcasts the statistics back, so the arithmetic stays the stock decomposition's (bit-identical).
+// Per row this kernel: reads x (r, j) (+ res (r, j)); on the root, waits for the Wt gathered tiles (semaphore 0,
+// monotonic count) and hands them to the compute (c_10); then waits for the two statistic broadcasts (semaphore 1,
+// monotonic count) and hands each to the compute (c_5, 2 pages: mean, rstd).
+// CT args: [0] Wt, [1] has_gamma, [2] has_beta, [3] eps bits, [4] per-core RT-arg count (P3), [5] has_res,
+//          [6] has_rgate, then the TensorAccessorArgs of x, gamma, beta, res, rgate.
+// Common RT args: [x_addr, gamma_addr, beta_addr, res_addr, rgate_addr].
+// Per-core RT args: [row0, n_rows, row_stride, j, root_x, root_y, (member x, y) x Wt].
+#include <cstdint>
+
+#include "api/dataflow/dataflow_api.h"
+
+constexpr uint32_t Wt = get_compile_time_arg_val(0);
+constexpr uint32_t has_gamma = get_compile_time_arg_val(1);
+constexpr uint32_t has_beta = get_compile_time_arg_val(2);
+constexpr uint32_t eps_bits = get_compile_time_arg_val(3);
+constexpr uint32_t has_res = get_compile_time_arg_val(5);
+constexpr uint32_t has_rgate = get_compile_time_arg_val(6);
+constexpr auto x_args = TensorAccessorArgs<7>();
+constexpr auto g_args = TensorAccessorArgs<x_args.next_compile_time_args_offset()>();
+constexpr auto b_args = TensorAccessorArgs<g_args.next_compile_time_args_offset()>();
+constexpr auto r_args = TensorAccessorArgs<b_args.next_compile_time_args_offset()>();
+constexpr auto rg_args = TensorAccessorArgs<r_args.next_compile_time_args_offset()>();
+
+constexpr uint32_t cb_x = 0, cb_g = 1, cb_b = 2, cb_eps = 3, cb_bc = 5, cb_r = 7, cb_rg = 8, cb_gat = 10;
+constexpr uint32_t TB = 4096;
+
+FORCE_INLINE void fill_first_row(uint32_t p) {
+    constexpr uint32_t kFace = 1024, kRow = 64;
+    const uint32_t f0 = p, f1 = p + kFace;
+    for (uint32_t n = kRow; n < kFace; n <<= 1) {
+        noc_async_read(get_noc_addr(f0), f0 + n, n);
+        noc_async_read(get_noc_addr(f1), f1 + n, n);
+        noc_async_read_barrier();
+    }
+    noc_async_read(get_noc_addr(f0), f0 + 2 * kFace, 2 * kFace);
+    noc_async_read_barrier();
+}
+
+template <typename A>
+FORCE_INLINE void read_row_tile(uint32_t cb, const A& acc, uint32_t j) {
+    cb_reserve_back(cb, 1);
+    const uint32_t p = get_write_ptr(cb);
+    noc_async_read(acc.get_noc_addr(j), p, TB);
+    noc_async_read_barrier();
+    fill_first_row(p);
+    cb_push_back(cb, 1);
+}
+
+void kernel_main() {
+    const uint32_t x_addr = get_common_arg_val<uint32_t>(0);
+    const uint32_t g_addr = get_common_arg_val<uint32_t>(1);
+    const uint32_t b_addr = get_common_arg_val<uint32_t>(2);
+    const uint32_t r_addr = get_common_arg_val<uint32_t>(3);
+    const uint32_t rg_addr = get_common_arg_val<uint32_t>(4);
+    const uint32_t row0 = get_arg_val<uint32_t>(0);
+    const uint32_t n_rows = get_arg_val<uint32_t>(1);
+    const uint32_t stride = get_arg_val<uint32_t>(2);
+    const uint32_t j = get_arg_val<uint32_t>(3);
+    if (n_rows == 0) {
+        return;
+    }
+    const bool root = j == 0;
+    const auto x = TensorAccessor(x_args, x_addr, TB);
+    const auto res = TensorAccessor(r_args, r_addr, TB);
+    volatile tt_l1_ptr uint32_t* sem_g = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_semaphore(0));
+    volatile tt_l1_ptr uint32_t* sem_b = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_semaphore(1));
+    uint32_t n_gat = 0, n_bc = 0;
+
+    auto read_x = [&](uint32_t r) {
+        cb_reserve_back(cb_x, 1);
+        noc_async_read(x.get_noc_addr(r * Wt + j), get_write_ptr(cb_x), TB);
+        if constexpr (has_res) {
+            cb_reserve_back(cb_r, 1);
+            noc_async_read(res.get_noc_addr(r * Wt + j), get_write_ptr(cb_r), TB);
+            noc_async_read_barrier();
+            cb_push_back(cb_r, 1);
+        } else {
+            noc_async_read_barrier();
+        }
+        cb_push_back(cb_x, 1);
+    };
+    auto gather = [&]() {
+        cb_reserve_back(cb_gat, Wt);
+        n_gat += Wt;
+        noc_semaphore_wait_min(sem_g, n_gat);
+        cb_push_back(cb_gat, Wt);
+    };
+    auto bcast_in = [&]() {
+        cb_reserve_back(cb_bc, 1);
+        n_bc += 1;
+        noc_semaphore_wait_min(sem_b, n_bc);
+        cb_push_back(cb_bc, 1);
+    };
+
+    read_x(row0);
+    if constexpr (has_rgate) {
+        read_row_tile(cb_rg, TensorAccessor(rg_args, rg_addr, TB), j);
+    }
+    cb_reserve_back(cb_eps, 1);
+    {
+        auto* e = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_write_ptr(cb_eps));
+        for (uint32_t i = 0; i < 1024; ++i) {
+            e[i] = eps_bits;
+        }
+    }
+    cb_push_back(cb_eps, 1);
+    if constexpr (has_gamma) {
+        read_row_tile(cb_g, TensorAccessor(g_args, g_addr, TB), j);
+    }
+    if constexpr (has_beta) {
+        read_row_tile(cb_b, TensorAccessor(b_args, b_addr, TB), j);
+    }
+    for (uint32_t i = 0; i < n_rows; ++i) {
+        if (i > 0) {
+            read_x(row0 + i * stride);
+        }
+        if (root) {
+            gather();
+        }
+        bcast_in();
+        if (root) {
+            gather();
+        }
+        bcast_in();
+    }
+}

code/tt_diffusion_planner/tt/kernels/ln32s_writer.cpp

@@ -0,0 +1,161 @@
+// SPDX-License-Identifier: Apache-2.0
+// Split-row fused fp32 LayerNorm (tt/ln_kernel.py, LN_SPLIT), writer (RISCV_1). Per row: (write_h) the new stream
+// tile; sends this member's tile for the mean fold (c_11, from the compute) to page j of the root's gather CB c_10
+// and raises the root's semaphore 0; on the root, broadcasts the packed mean (c_4, column 0) over the 32 columns and
+// sends it to page 0 of every member's c_5 (semaphore 1 raised); the same for the squared tile and rstd (page 1);
+// then writes the output tile.
+// With kcat_out (KCAT_EMIT) out is the split operand [y_hi | y_hi | y_lo | 1 | 0..] of the next K-concatenated
+// linear (row tiles Ktp, the kcat_writer layout): y_hi to (r, j) and (r, Wt + j), y_lo to (r, 2 Wt + j); the root
+// also writes the ones tile (r, 3 Wt) and the zero pad tiles.
+// CT args: [0] Wt, [1] per-core RT-arg count (P3), [2] write_h, [3] sfpu_bcast, [4] kcat_out, [5] Ktp, then the
+//          TensorAccessorArgs of out and h.
+// Common RT args: [out_addr, h_addr]. Per-core RT args: as the reader's.
+#include <cstdint>
+
+#include "api/dataflow/dataflow_api.h"
+
+constexpr uint32_t Wt = get_compile_time_arg_val(0);
+constexpr uint32_t write_h = get_compile_time_arg_val(2);
+constexpr uint32_t sfpu_bcast = get_compile_time_arg_val(3);   // the root's compute packs a broadcast tile
+constexpr uint32_t kcat_out = get_compile_time_arg_val(4);
+constexpr uint32_t Ktp = get_compile_time_arg_val(5);
+constexpr auto out_args = TensorAccessorArgs<6>();
+constexpr auto h_args = TensorAccessorArgs<out_args.next_compile_time_args_offset()>();
+
+constexpr uint32_t cb_stat = 4, cb_bc = 5, cb_gat = 10, cb_snd = 11, cb_scr = 12, cb_out = 16, cb_hout = 17,
+                   cb_lo = 19, cb_one = 20, cb_zero = 21;
+constexpr uint32_t TB = 4096;
+
+FORCE_INLINE void bcast_col(uint32_t src, uint32_t dst) {
+    auto* s = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(src);
+    auto* d = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(dst);
+    for (uint32_t fo = 0; fo < 1024; fo += 512) {
+        for (uint32_t ro = 0; ro < 256; ro += 16) {
+            const uint32_t v = s[fo + ro];
+            volatile tt_l1_ptr uint32_t* l = d + fo + ro;
+            volatile tt_l1_ptr uint32_t* r = l + 256;
+            for (uint32_t c = 0; c < 16; ++c) {
+                l[c] = v;
+                r[c] = v;
+            }
+        }
+    }
+}
+
+void kernel_main() {
+    const uint32_t out_addr = get_common_arg_val<uint32_t>(0);
+    const uint32_t h_addr = get_common_arg_val<uint32_t>(1);
+    const uint32_t row0 = get_arg_val<uint32_t>(0);
+    const uint32_t n_rows = get_arg_val<uint32_t>(1);
+    const uint32_t stride = get_arg_val<uint32_t>(2);
+    const uint32_t j = get_arg_val<uint32_t>(3);
+    const uint32_t root_x = get_arg_val<uint32_t>(4);
+    const uint32_t root_y = get_arg_val<uint32_t>(5);
+    if (n_rows == 0) {
+        return;
+    }
+    const bool root = j == 0;
+    const auto out = TensorAccessor(out_args, out_addr, TB);
+    const auto hacc = TensorAccessor(h_args, h_addr, TB);
+    const uint32_t gat_base = get_write_ptr(cb_gat);   // the same L1 address on every core of the program
+    const uint32_t bc_base = get_write_ptr(cb_bc);
+    const uint64_t root_sem = get_noc_addr(root_x, root_y, get_semaphore(0));
+    const uint32_t sem_b_addr = get_semaphore(1);
+    cb_reserve_back(cb_scr, 1);
+    const uint32_t scr = get_write_ptr(cb_scr);
+
+    auto send_to_root = [&]() {
+        cb_wait_front(cb_snd, 1);
+        noc_async_write(get_read_ptr(cb_snd), get_noc_addr(root_x, root_y, gat_base + j * TB), TB);
+        noc_async_write_barrier();
+        noc_semaphore_inc(root_sem, 1);
+        cb_pop_front(cb_snd, 1);
+    };
+    auto bcast_out = [&](uint32_t page) {
+        cb_wait_front(cb_stat, 1);
+        uint32_t src = get_read_ptr(cb_stat);
+        if constexpr (!sfpu_bcast) {
+            bcast_col(src, scr);
+            src = scr;
+        }
+        for (uint32_t m = 0; m < Wt; ++m) {
+            const uint32_t mx = get_arg_val<uint32_t>(6 + 2 * m), my = get_arg_val<uint32_t>(7 + 2 * m);
+            noc_async_write(src, get_noc_addr(mx, my, bc_base + page * TB), TB);
+        }
+        noc_async_write_barrier();
+        cb_pop_front(cb_stat, 1);
+        for (uint32_t m = 0; m < Wt; ++m) {
+            const uint32_t mx = get_arg_val<uint32_t>(6 + 2 * m), my = get_arg_val<uint32_t>(7 + 2 * m);
+            noc_semaphore_inc(get_noc_addr(mx, my, sem_b_addr), 1);
+        }
+    };
+
+    uint32_t one_ptr = 0, zero_ptr = 0;
+    if constexpr (kcat_out) {
+        if (root) {
+            // ones tile: element (i, c) at face (i / 16) * 2 + (c / 16), offset (i % 16) * 16 + c % 16
+            cb_reserve_back(cb_one, 1);
+            one_ptr = get_write_ptr(cb_one);
+            auto* p = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(one_ptr);
+            for (uint32_t i = 0; i < 1024; ++i) {
+                p[i] = 0;
+            }
+            for (uint32_t i = 0; i < 32; ++i) {
+                const uint32_t base = (i / 16) * 512 + (i % 16) * 16;
+                p[base] = 0x3F800000u;
+                p[base + 1] = 0x3F800000u;
+            }
+            if constexpr (Ktp > 3 * Wt + 1) {
+                cb_reserve_back(cb_zero, 1);
+                zero_ptr = get_write_ptr(cb_zero);
+                auto* z = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(zero_ptr);
+                for (uint32_t i = 0; i < 1024; ++i) {
+                    z[i] = 0;
+                }
+            }
+        }
+    }
+
+    for (uint32_t i = 0; i < n_rows; ++i) {
+        const uint32_t r = row0 + i * stride;
+        if constexpr (write_h) {
+            cb_wait_front(cb_hout, 1);
+            noc_async_write(get_read_ptr(cb_hout), hacc.get_noc_addr(r * Wt + j), TB);
+            noc_async_writes_flushed();
+            cb_pop_front(cb_hout, 1);
+        }
+        send_to_root();
+        if (root) {
+            bcast_out(0);
+        }
+        send_to_root();
+        if (root) {
+            bcast_out(1);
+        }
+        if constexpr (kcat_out) {
+            const uint32_t ob = r * Ktp;
+            cb_wait_front(cb_out, 1);
+            const uint32_t hp = get_read_ptr(cb_out);
+            noc_async_write(hp, out.get_noc_addr(ob + j), TB);
+            noc_async_write(hp, out.get_noc_addr(ob + Wt + j), TB);
+            cb_wait_front(cb_lo, 1);
+            noc_async_write(get_read_ptr(cb_lo), out.get_noc_addr(ob + 2 * Wt + j), TB);
+            if (root) {
+                noc_async_write(one_ptr, out.get_noc_addr(ob + 3 * Wt), TB);
+                for (uint32_t z = 3 * Wt + 1; z < Ktp; ++z) {
+                    noc_async_write(zero_ptr, out.get_noc_addr(ob + z), TB);
+                }
+            }
+            noc_async_writes_flushed();
+            cb_pop_front(cb_out, 1);
+            cb_pop_front(cb_lo, 1);
+        } else {
+            cb_wait_front(cb_out, 1);
+            noc_async_write(get_read_ptr(cb_out), out.get_noc_addr(r * Wt + j), TB);
+            noc_async_writes_flushed();
+            cb_pop_front(cb_out, 1);
+        }
+    }
+    noc_async_write_barrier();
+    noc_async_atomic_barrier();
+}

code/tt_diffusion_planner/tt/kernels/smask_compute.cpp

@@ -0,0 +1,69 @@
+// SPDX-License-Identifier: Apache-2.0
+// Attention score scale + mask (tt/smask_kernel.py), compute: out = s * scale + mask, the two stock binary_ng fp32
+// programs `ttnn.multiply(s, scale)` (SFPU mul_binary_tile against a tile filled with the fp32 scale) and
+// `ttnn.add(., mask)`, in one pass. With an fp32 mask the stock add is the SFPU add_binary_tile<NearestEven>; with a
+// bf16 mask (trunc_tf32 = 1) binary_ng takes its FPU path, whose SrcA truncates the fp32 scores to TF32 (measured:
+// logs/diffusion-planner/opt_r2/smask_diag.log, the stock result == (exact & 0xFFFFE000) for every finite element),
+// so the scaled scores are truncated the same way (SFPU bitwise AND on the raw bits) before the exact add of the
+// 0 / -inf mask. Bit-identical to the two programs in both cases.
+// CT args: [0] H (even), [1] per-core RT-arg count (P3), [2] trunc_tf32. Per-core RT args: [n, 0].
+#include <cstdint>
+
+#include "api/compute/common.h"
+#include "api/compute/compute_kernel_api.h"
+#include "api/compute/eltwise_binary_sfpu.h"
+#include "api/compute/eltwise_unary/bitwise.h"
+#include "api/compute/eltwise_unary/eltwise_unary.h"
+#include "api/compute/tile_move_copy.h"
+
+using namespace ckernel;
+
+constexpr uint32_t H = get_compile_time_arg_val(0);
+constexpr uint32_t trunc_tf32 = get_compile_time_arg_val(2);
+constexpr uint32_t cb_s = 0, cb_m = 1, cb_c = 2, cb_out = 16;
+constexpr auto RNE = ckernel::DstRoundingMode::NearestEven;
+
+void kernel_main() {
+    const uint32_t n = get_arg_val<uint32_t>(0);
+    if (n == 0) {
+        return;
+    }
+    compute_kernel_hw_startup(cb_s, cb_out);
+    cb_wait_front(cb_c, 1);
+    for (uint32_t t = 0; t < n; ++t) {
+        cb_wait_front(cb_m, 1);
+        cb_wait_front(cb_s, H);
+        for (uint32_t h = 0; h < H; h += 2) {
+            tile_regs_acquire();
+            copy_init(cb_s);
+            copy_tile(cb_s, h, 0);
+            copy_tile(cb_s, h + 1, 1);
+            copy_init(cb_c);
+            copy_tile(cb_c, 0, 2);
+            mul_binary_tile_init();
+            mul_binary_tile(0, 2, 0);
+            mul_binary_tile(1, 2, 1);
+            if constexpr (trunc_tf32) {
+                bitwise_and_tile_init();
+                bitwise_and_tile<DataFormat::Int32>(0, 0xFFFFE000u);
+                bitwise_and_tile<DataFormat::Int32>(1, 0xFFFFE000u);
+            }
+            reconfig_data_format_srca(cb_c, cb_m);
+            copy_init(cb_m);
+            copy_tile(cb_m, 0, 3);
+            reconfig_data_format_srca(cb_m, cb_s);
+            add_binary_tile_init();
+            add_binary_tile<RNE>(0, 3, 0);
+            add_binary_tile<RNE>(1, 3, 1);
+            cb_reserve_back(cb_out, 2);
+            tile_regs_commit();
+            tile_regs_wait();
+            pack_tile(0, cb_out);
+            pack_tile(1, cb_out);
+            tile_regs_release();
+            cb_push_back(cb_out, 2);
+        }
+        cb_pop_front(cb_s, H);
+        cb_pop_front(cb_m, 1);
+    }
+}

code/tt_diffusion_planner/tt/kernels/smask_reader.cpp

@@ -0,0 +1,50 @@
+// SPDX-License-Identifier: Apache-2.0
+// Attention score scale + mask (tt/smask_kernel.py), reader (RISCV_0): the scale tile once, then per mask tile m of
+// this core's range the bf16 mask tile and the H fp32 score tiles (h, m) (the mask broadcasts over the heads).
+// CT args: [0] H, [1] mask tiles per head plane (Mt * Kt), [2] scale bits (fp32), [3] per-core RT-arg count (P3),
+//          [4] mask tile bytes (bf16 2048 / fp32 4096), then the TensorAccessorArgs of scores and mask.
+// Common RT args: [s_addr, m_addr]. Per-core RT args: [m0, n].
+#include <cstdint>
+
+#include "api/dataflow/dataflow_api.h"
+
+constexpr uint32_t H = get_compile_time_arg_val(0);
+constexpr uint32_t plane = get_compile_time_arg_val(1);
+constexpr uint32_t scale_bits = get_compile_time_arg_val(2);
+constexpr uint32_t TBM = get_compile_time_arg_val(4);
+constexpr auto s_args = TensorAccessorArgs<5>();
+constexpr auto m_args = TensorAccessorArgs<s_args.next_compile_time_args_offset()>();
+constexpr uint32_t cb_s = 0, cb_m = 1, cb_c = 2;
+constexpr uint32_t TB = 4096;
+
+void kernel_main() {
+    const uint32_t s_addr = get_common_arg_val<uint32_t>(0);
+    const uint32_t m_addr = get_common_arg_val<uint32_t>(1);
+    const uint32_t m0 = get_arg_val<uint32_t>(0);
+    const uint32_t n = get_arg_val<uint32_t>(1);
+    if (n == 0) {
+        return;
+    }
+    const auto s = TensorAccessor(s_args, s_addr, TB);
+    const auto mk = TensorAccessor(m_args, m_addr, TBM);
+    cb_reserve_back(cb_c, 1);
+    {
+        auto* p = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_write_ptr(cb_c));
+        for (uint32_t i = 0; i < 1024; ++i) {
+            p[i] = scale_bits;
+        }
+    }
+    cb_push_back(cb_c, 1);
+    for (uint32_t m = m0; m < m0 + n; ++m) {
+        cb_reserve_back(cb_m, 1);
+        noc_async_read(mk.get_noc_addr(m), get_write_ptr(cb_m), TBM);
+        cb_reserve_back(cb_s, H);
+        const uint32_t p = get_write_ptr(cb_s);
+        for (uint32_t h = 0; h < H; ++h) {
+            noc_async_read(s.get_noc_addr(h * plane + m), p + h * TB, TB);
+        }
+        noc_async_read_barrier();
+        cb_push_back(cb_m, 1);
+        cb_push_back(cb_s, H);
+    }
+}

code/tt_diffusion_planner/tt/kernels/smask_writer.cpp

@@ -0,0 +1,31 @@
+// SPDX-License-Identifier: Apache-2.0
+// Attention score scale + mask (tt/smask_kernel.py), writer (RISCV_1): the H output tiles (h, m) per mask tile m.
+// CT args: [0] H, [1] plane (Mt * Kt), [2] per-core RT-arg count (P3), then the TensorAccessorArgs of out.
+// Common RT args: [out_addr]. Per-core RT args: [m0, n].
+#include <cstdint>
+
+#include "api/dataflow/dataflow_api.h"
+
+constexpr uint32_t H = get_compile_time_arg_val(0);
+constexpr uint32_t plane = get_compile_time_arg_val(1);
+constexpr auto o_args = TensorAccessorArgs<3>();
+constexpr uint32_t cb_out = 16;
+constexpr uint32_t TB = 4096;
+
+void kernel_main() {
+    const uint32_t o_addr = get_common_arg_val<uint32_t>(0);
+    const uint32_t m0 = get_arg_val<uint32_t>(0);
+    const uint32_t n = get_arg_val<uint32_t>(1);
+    const auto o = TensorAccessor(o_args, o_addr, TB);
+    for (uint32_t m = m0; m < m0 + n; ++m) {
+        for (uint32_t h = 0; h < H; h += 2) {
+            cb_wait_front(cb_out, 2);
+            const uint32_t p = get_read_ptr(cb_out);
+            noc_async_write(p, o.get_noc_addr(h * plane + m), TB);
+            noc_async_write(p + TB, o.get_noc_addr((h + 1) * plane + m), TB);
+            noc_async_writes_flushed();
+            cb_pop_front(cb_out, 2);
+        }
+    }
+    noc_async_write_barrier();
+}

code/tt_diffusion_planner/tt/kernels/smsm_compute.cpp

@@ -0,0 +1,179 @@
+// SPDX-License-Identifier: Apache-2.0
+// Attention score scale + mask + softmax (tt/smsm_kernel.py, ATTN_SMSM), compute. Per tile row (one head, 32 query
+// rows, Wt key tiles):
+//   phase A: x = trunc_tf32(s * scale) + mask, the smask kernel's SFPU sequence (kernels/smask_compute.cpp: the two
+//            stock binary_ng programs bit for bit), packed to cb_x (L1) instead of DRAM;
+//   phase B: the stock ttnn.softmax(numeric_stable=True) of that row, i.e. the no-mask path of
+//            ttnn/.../softmax/device/kernels/attention/compute/softmax.cpp with the same kernel_lib calls: row max
+//            (FPU reduce), exp(x - max) (FPU bcast sub + SFPU exp), row sum (FPU reduce) + precise fp32
+//            reciprocal, x * 1/sum (FPU bcast mul).
+// The stock softmax unpacks its fp32 input to SrcA (TF32); x holds TF32 values already (phase A truncates them and
+// the mask is 0 / -inf), so the stock program saw exactly these operands.
+// CT args: [0] Wt, [1] per-core RT-arg count (P3), [2] trunc_tf32, [3] ndst (block), [4] out pad tiles per row,
+//          [5] scale mode: 0 = multiply by a tile filled with the scale (mul_binary_tile, as the stock binary_ng
+//          program), 1 = mul_unary_tile with the scale bits (the same SFPU fp32 multiply, no scale tile copy).
+// Per-core RT args: [nrows, 0].
+#include <cstdint>
+
+#include "api/compute/common.h"
+#include "api/compute/compute_kernel_api.h"
+#include "api/compute/eltwise_binary.h"
+#include "api/compute/eltwise_binary_sfpu.h"
+#include "api/compute/eltwise_unary/binop_with_scalar.h"
+#include "api/compute/eltwise_unary/bitwise.h"
+#include "api/compute/eltwise_unary/eltwise_unary.h"
+#include "api/compute/tile_move_copy.h"
+#include "api/compute/bcast.h"
+#include "api/compute/softmax.h"
+#include "api/compute/reduce.h"
+#include "ttnn/cpp/ttnn/kernel_lib/reduce_helpers_compute.hpp"
+#include "ttnn/cpp/ttnn/kernel_lib/eltwise/api/chain.hpp"
+#include "ttnn/cpp/ttnn/kernel_lib/eltwise/api/convenience.hpp"
+#include "ttnn/cpp/ttnn/kernel_lib/eltwise/unary/math.hpp"
+#include "ttnn/cpp/ttnn/kernel_lib/eltwise/core/optional.hpp"
+
+namespace ckl = compute_kernel_lib;
+using namespace ckernel;
+
+constexpr uint32_t Wt = get_compile_time_arg_val(0);
+constexpr uint32_t trunc_tf32 = get_compile_time_arg_val(2);
+constexpr uint32_t ndst = get_compile_time_arg_val(3);
+constexpr uint32_t out_pad = get_compile_time_arg_val(4);
+constexpr uint32_t scale_mode = get_compile_time_arg_val(5);
+constexpr uint32_t scale_bits = get_compile_time_arg_val(6);
+constexpr uint32_t cb_s = 0, cb_m = 1, cb_c = 2, cb_max_scaler = 3, cb_sum_scaler = 4;
+constexpr uint32_t cb_out = 16;
+constexpr uint32_t cb_x = 24, cb_max = 25, cb_exps = 26, cb_recip = 27;
+constexpr auto RNE = ckernel::DstRoundingMode::NearestEven;
+constexpr uint32_t Wp = (Wt + 1) / 2;  // tile pairs per row (the last one half-used when Wt is odd)
+
+// the stock calc_numeric_stable (softmax.cpp), CB ids instead of DFB handles
+template <std::uint32_t dfb_in, std::uint32_t dfb_max_scaler, std::uint32_t dfb_max, std::uint32_t dfb_out>
+void calc_numeric_stable(std::uint32_t W, std::uint32_t nd) {
+    compute_kernel_lib::reduce<
+        PoolType::MAX,
+        ReduceDim::REDUCE_ROW,
+        dfb_in,
+        dfb_max_scaler,
+        dfb_max,
+        compute_kernel_lib::ReduceInputPolicy::WaitUpfrontNoPop,
+        compute_kernel_lib::ReduceDataFormatReconfigMode::INPUT>(compute_kernel_lib::ReduceInputBlockShape::row(W));
+    ckl::eltwise_chain(
+        ckl::IterationShape::tiles(W).block_size(nd),
+        ckl::BinaryFpu<
+            ckl::BinaryFpuOp::Sub,
+            ckl::input(
+                dfb_in,
+                ckl::WaitPolicy::Upfront,
+                ckl::PopPolicy::AtEnd,
+                ckl::InputTileMapping::Block,
+                ckl::DataFormatReconfig::Disabled),
+            ckl::input(dfb_max, ckl::BroadcastDim::Col, ckl::WaitPolicy::Upfront, ckl::PopPolicy::AtEnd)>{},
+        ckl::Exp<ckl::Approx::Exact, ckl::Dst::D0>{},
+        ckl::PackTile<ckl::output(
+            dfb_out,
+            ckl::ReservePolicy::PerBlockSize,
+            ckl::PushPolicy::PerBlockSize,
+            ckl::DataFormatReconfig::Disabled)>{});
+    cb_wait_front(dfb_out, W);
+}
+
+void kernel_main() {
+    const uint32_t nrows = get_arg_val<uint32_t>(0);
+    if (nrows == 0) {
+        return;
+    }
+    compute_kernel_hw_startup(cb_s, cb_c, cb_x);
+    cb_wait_front(cb_c, 1);
+    cb_wait_front(cb_max_scaler, 1);
+    cb_wait_front(cb_sum_scaler, 1);
+    for (uint32_t r = 0; r < nrows; ++r) {
+        // ---- phase A: x = trunc(s * scale) + mask (pairs of tiles) ----
+        reconfig_data_format(cb_s, cb_c);
+        pack_reconfig_data_format(cb_x);
+        cb_wait_front(cb_m, 2 * Wp);
+        for (uint32_t p = 0; p < Wp; ++p) {
+            const uint32_t j = 2 * p;
+            const bool two = (j + 1) < Wt;
+            cb_wait_front(cb_s, 2);
+            tile_regs_acquire();
+            copy_init(cb_s);
+            copy_tile(cb_s, 0, 0);
+            copy_tile(cb_s, 1, 1);
+            if constexpr (scale_mode == 0) {
+                copy_init(cb_c);
+                copy_tile(cb_c, 0, 2);
+                mul_binary_tile_init();
+                mul_binary_tile(0, 2, 0);
+                mul_binary_tile(1, 2, 1);
+            } else {
+                binop_with_scalar_tile_init();
+                mul_unary_tile(0, scale_bits);
+                mul_unary_tile(1, scale_bits);
+            }
+            if constexpr (trunc_tf32) {
+                bitwise_and_tile_init();
+                bitwise_and_tile<DataFormat::Int32>(0, 0xFFFFE000u);
+                bitwise_and_tile<DataFormat::Int32>(1, 0xFFFFE000u);
+            }
+            reconfig_data_format_srca(scale_mode == 0 ? cb_c : cb_s, cb_m);
+            copy_init(cb_m);
+            copy_tile(cb_m, j, 2);
+            copy_tile(cb_m, j + 1, 3);
+            reconfig_data_format_srca(cb_m, cb_s);
+            add_binary_tile_init();
+            add_binary_tile<RNE>(0, 2, 0);
+            add_binary_tile<RNE>(1, 3, 1);
+            const uint32_t np = two ? 2 : 1;
+            cb_reserve_back(cb_x, np);
+            tile_regs_commit();
+            tile_regs_wait();
+            pack_tile(0, cb_x);
+            if (two) {
+                pack_tile(1, cb_x);
+            }
+            tile_regs_release();
+            cb_push_back(cb_x, np);
+            cb_pop_front(cb_s, 2);
+        }
+        cb_pop_front(cb_m, 2 * Wp);
+
+        // ---- phase B: the stock numeric-stable softmax of the row ----
+        reconfig_data_format(cb_x, cb_x);
+        pack_reconfig_data_format(cb_exps);
+        copy_init(cb_x);
+        calc_numeric_stable<cb_x, cb_max_scaler, cb_max, cb_exps>(Wt, ndst);
+        reconfig_data_format(cb_exps, cb_sum_scaler);
+        compute_kernel_lib::reduce<
+            PoolType::SUM,
+            ReduceDim::REDUCE_ROW,
+            cb_exps,
+            cb_sum_scaler,
+            cb_recip,
+            compute_kernel_lib::ReduceInputPolicy::WaitUpfrontNoPop>(
+            compute_kernel_lib::ReduceInputBlockShape::row(Wt),
+            compute_kernel_lib::ReduceInputMemoryLayout::contiguous(),
+            compute_kernel_lib::NoAccumulation{},
+            [](std::uint32_t) {
+                if constexpr (DST_ACCUM_MODE) {
+                    recip_tile_init<ReciprocalDestAcc::FP32, ReciprocalApproxMode::Precise>();
+                    recip_tile<ReciprocalDestAcc::FP32, ReciprocalApproxMode::Precise>(0);
+                } else {
+                    recip_tile_init();
+                    recip_tile(0);
+                }
+            });
+        ckl::mul<
+            ckl::input(cb_exps, ckl::WaitPolicy::Upfront, ckl::PopPolicy::AtEnd, ckl::InputTileMapping::Block),
+            ckl::input(cb_recip, ckl::BroadcastDim::Col, ckl::WaitPolicy::Upfront, ckl::PopPolicy::AtEnd),
+            ckl::output(cb_out, ckl::ReservePolicy::PerBlockSize, ckl::PushPolicy::PerBlockSize)>(
+            ckl::IterationShape::tiles(Wt).block_size(ndst));
+        if constexpr (out_pad > 0) {
+            cb_reserve_back(cb_out, out_pad);
+            cb_push_back(cb_out, out_pad);
+        }
+    }
+    cb_pop_front(cb_c, 1);
+    cb_pop_front(cb_max_scaler, 1);
+    cb_pop_front(cb_sum_scaler, 1);
+}

code/tt_diffusion_planner/tt/kernels/smsm_reader.cpp

@@ -0,0 +1,83 @@
+// SPDX-License-Identifier: Apache-2.0
+// Attention score scale + mask + softmax (tt/smsm_kernel.py), reader (RISCV_0). Once: the scale tile (fp32 filled
+// with the scale) and the two reduce scaler tiles of the stock softmax reader (fp32 1.0 in row 0 of every face, the
+// rest 0: dataflow_kernel_lib::calculate_and_prepare_reduce_scaler for MAX / SUM, REDUCE_ROW). Per tile row r of
+// this core's range (head h = r / Mt, query tile row mr = r % Mt): the Wt mask tiles of row mr (+ 1 unused tile when
+// Wt is odd, so the compute can take the mask in pairs), then the Wt fp32 score tiles in pairs (the last pair of an
+// odd row carries one unused tile).
+// CT args: [0] Wt, [1] Mt, [2] scale bits (fp32), [3] per-core RT-arg count (P3), [4] mask tile bytes, then the
+//          TensorAccessorArgs of scores and mask. Common RT args: [s_addr, m_addr]. Per-core RT args: [r0, n].
+#include <cstdint>
+
+#include "api/dataflow/dataflow_api.h"
+
+constexpr uint32_t Wt = get_compile_time_arg_val(0);
+constexpr uint32_t Mt = get_compile_time_arg_val(1);
+constexpr uint32_t scale_bits = get_compile_time_arg_val(2);
+constexpr uint32_t TBM = get_compile_time_arg_val(4);
+constexpr auto s_args = TensorAccessorArgs<5>();
+constexpr auto m_args = TensorAccessorArgs<s_args.next_compile_time_args_offset()>();
+constexpr uint32_t cb_s = 0, cb_m = 1, cb_c = 2, cb_max_scaler = 3, cb_sum_scaler = 4;
+constexpr uint32_t TB = 4096;
+constexpr uint32_t Wp = (Wt + 1) / 2;
+
+inline void fill_scaler(uint32_t cb) {
+    cb_reserve_back(cb, 1);
+    auto* p = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_write_ptr(cb));
+    for (uint32_t i = 0; i < 1024; ++i) {
+        p[i] = 0;
+    }
+    for (uint32_t f = 0; f < 4; ++f) {
+        for (uint32_t c = 0; c < 16; ++c) {
+            p[f * 256 + c] = 0x3F800000u;
+        }
+    }
+    cb_push_back(cb, 1);
+}
+
+void kernel_main() {
+    const uint32_t s_addr = get_common_arg_val<uint32_t>(0);
+    const uint32_t m_addr = get_common_arg_val<uint32_t>(1);
+    const uint32_t r0 = get_arg_val<uint32_t>(0);
+    const uint32_t n = get_arg_val<uint32_t>(1);
+    if (n == 0) {
+        return;
+    }
+    const auto s = TensorAccessor(s_args, s_addr, TB);
+    const auto mk = TensorAccessor(m_args, m_addr, TBM);
+    cb_reserve_back(cb_c, 1);
+    {
+        auto* p = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_write_ptr(cb_c));
+        for (uint32_t i = 0; i < 1024; ++i) {
+            p[i] = scale_bits;
+        }
+    }
+    cb_push_back(cb_c, 1);
+    fill_scaler(cb_max_scaler);
+    fill_scaler(cb_sum_scaler);
+    for (uint32_t r = r0; r < r0 + n; ++r) {
+        const uint32_t mr = r % Mt;
+        cb_reserve_back(cb_m, 2 * Wp);
+        {
+            uint32_t p = get_write_ptr(cb_m);
+            for (uint32_t j = 0; j < Wt; ++j) {
+                noc_async_read(mk.get_noc_addr(mr * Wt + j), p, TBM);
+                p += TBM;
+            }
+        }
+        noc_async_read_barrier();
+        cb_push_back(cb_m, 2 * Wp);
+        uint32_t t = r * Wt;
+        for (uint32_t pp = 0; pp < Wp; ++pp) {
+            cb_reserve_back(cb_s, 2);
+            const uint32_t p = get_write_ptr(cb_s);
+            noc_async_read(s.get_noc_addr(t), p, TB);
+            if (2 * pp + 1 < Wt) {
+                noc_async_read(s.get_noc_addr(t + 1), p + TB, TB);
+            }
+            noc_async_read_barrier();
+            cb_push_back(cb_s, 2);
+            t += 2;
+        }
+    }
+}

code/tt_diffusion_planner/tt/kernels/smsm_writer.cpp

@@ -0,0 +1,43 @@
+// SPDX-License-Identifier: Apache-2.0
+// Attention score scale + mask + softmax (tt/smsm_kernel.py), writer (RISCV_1): per tile row r of this core's range
+// the Wt probability tiles in the compute's blocks of ndst (the last one clamped), then the row's pad tiles (pushed
+// by the compute to keep cb_out aligned; not written).
+// CT args: [0] Wt, [1] ndst, [2] per-core RT-arg count (P3), [3] out pad tiles per row, then the TensorAccessorArgs of
+//          out. Common RT args: [out_addr]. Per-core RT args: [r0, n].
+#include <cstdint>
+
+#include "api/dataflow/dataflow_api.h"
+
+constexpr uint32_t Wt = get_compile_time_arg_val(0);
+constexpr uint32_t ndst = get_compile_time_arg_val(1);
+constexpr uint32_t out_pad = get_compile_time_arg_val(3);
+constexpr auto o_args = TensorAccessorArgs<4>();
+constexpr uint32_t cb_out = 16;
+constexpr uint32_t TB = 4096;
+
+void kernel_main() {
+    const uint32_t o_addr = get_common_arg_val<uint32_t>(0);
+    const uint32_t r0 = get_arg_val<uint32_t>(0);
+    const uint32_t n = get_arg_val<uint32_t>(1);
+    const auto o = TensorAccessor(o_args, o_addr, TB);
+    for (uint32_t r = r0; r < r0 + n; ++r) {
+        uint32_t t = r * Wt;
+        for (uint32_t j = 0; j < Wt; j += ndst) {
+            const uint32_t b = (j + ndst > Wt) ? (Wt - j) : ndst;
+            cb_wait_front(cb_out, b);
+            uint32_t p = get_read_ptr(cb_out);
+            for (uint32_t i = 0; i < b; ++i) {
+                noc_async_write(p, o.get_noc_addr(t + i), TB);
+                p += TB;
+            }
+            noc_async_writes_flushed();
+            cb_pop_front(cb_out, b);
+            t += b;
+        }
+        if constexpr (out_pad > 0) {
+            cb_wait_front(cb_out, out_pad);
+            cb_pop_front(cb_out, out_pad);
+        }
+    }
+    noc_async_write_barrier();
+}

code/tt_diffusion_planner/tt/layers.py

@@ -19,8 +19,8 @@ from ..ttaw.tensors import to_device, ttnn_dtype
 from . import config as T
 from .params import Lin, Norm
 
-__all__ = ["Build", "Linear", "SplitLinear", "make_linear", "LayerNorm", "Const", "ATTN", "policy",
-           "layer_norm_fp32"]
+__all__ = ["Build", "Linear", "SplitLinear", "make_linear", "chain2", "LayerNorm", "Const", "ATTN", "policy",
+           "layer_norm_fp32", "fold_batch"]
 
 ATTN = "bfloat16"  # dtype of the Q / K / V projections (SDPA takes bf16)
 
@@ -50,6 +50,30 @@ class Build:
         self.split = tuple(T.globs(knobs.SPLIT_MATMUL) if split is None else split)
         self.attn_fp32 = tuple(T.globs(knobs.ATTN_FP32_ACC) if attn_fp32_acc is None else attn_fp32_acc)
         self.attn_mm = tuple(T.globs(knobs.ATTN_MATMUL) if attn_matmul is None else attn_matmul)
+        self.ch2d = bool(knobs.ENC_CH2D)
+        self.attn_fast = int(knobs.ATTN_FAST)
+        self.dec_mmcfg = bool(knobs.DEC_MMCFG) and hasattr(device, "compute_with_storage_grid_size")
+        on_device = hasattr(device, "compute_with_storage_grid_size")
+        self.ln_kernel = int(knobs.LN_KERNEL) if on_device else 0
+        self.ln_resid = bool(knobs.LN_RESID)
+        self.split_kcat = int(knobs.SPLIT_KCAT)
+        self.attn_smask = bool(knobs.ATTN_SMASK)
+        self.attn_smsm = int(knobs.ATTN_SMSM)
+        self.attn_fused = bool(knobs.ATTN_FUSED)
+        self.kcat_emit = bool(knobs.KCAT_EMIT)
+        self.ln_tr = bool(knobs.LN_TR)
+        self.enc_kcat = bool(knobs.ENC_KCAT)
+        self.lin_act = bool(knobs.LIN_ACT)
+        self.kcat_act = bool(knobs.KCAT_ACT)
+        self.ln_split = bool(knobs.LN_SPLIT)
+        self.ln_sfpu_bcast = bool(knobs.LN_SFPU_BCAST)
+        self.kcat_l1 = bool(knobs.KCAT_L1)
+        self.kcat_act_once = bool(knobs.KCAT_ACT_ONCE)
+        self.attn_l1 = int(knobs.ATTN_L1) if on_device else 0
+        self.dec_l1 = int(knobs.DEC_L1) if on_device else 0
+        self.enc_l1 = bool(knobs.ENC_L1) and on_device
+        self.fus_l1 = bool(knobs.FUS_L1) and on_device
+        self.compact_enc = int(knobs.COMPACT) >= 2
         self.uploaded_bytes = 0
 
     @staticmethod
@@ -77,7 +101,16 @@ class Build:
 
     def options(self) -> Dict[str, Any]:
         return {"ln_fp32": list(self.ln_fp32), "hidden_fp32": list(self.hidden_fp32), "split": list(self.split),
-                "attn_fp32_acc": list(self.attn_fp32), "attn_matmul": list(self.attn_mm)}
+                "attn_fp32_acc": list(self.attn_fp32), "attn_matmul": list(self.attn_mm), "enc_ch2d": self.ch2d,
+                "attn_fast": self.attn_fast,
+                "dec_mmcfg": self.dec_mmcfg, "ln_kernel": self.ln_kernel, "ln_resid": self.ln_resid,
+                "split_kcat": self.split_kcat, "attn_smask": self.attn_smask, "attn_smsm": self.attn_smsm, "attn_fused": self.attn_fused, "kcat_emit": self.kcat_emit, "ln_tr": self.ln_tr,
+                "enc_kcat": self.enc_kcat, "lin_act": self.lin_act,
+                "kcat_act": self.kcat_act, "ln_split": self.ln_split,
+                "ln_sfpu_bcast": self.ln_sfpu_bcast, "kcat_l1": self.kcat_l1,
+                "kcat_act_once": self.kcat_act_once, "attn_l1": self.attn_l1,
+                "dec_l1": self.dec_l1, "enc_l1": self.enc_l1,
+                "fus_l1": self.fus_l1}
 
     def prec(self, module: str) -> Precision:
         return self.policy.resolve(module)
@@ -89,6 +122,33 @@ class Build:
         """The residual-stream dtype of ``module``."""
         return self.prec(module).activations
 
+    def attn_mem(self):
+        """``ATTN_L1``: the memory config of the fused attention's inputs (Q / K / V projections, hoisted cross K / V,
+        masks): L1 interleaved, else None (DRAM)."""
+        if not self.attn_l1:
+            return None
+        import ttnn
+
+        return ttnn.L1_MEMORY_CONFIG
+
+    def dec_mem(self):
+        """``DEC_L1``: the memory config of the decoder blocks' intermediates (stream, LN operands, linear outputs,
+        attention outputs): L1 interleaved, else None (DRAM)."""
+        if not self.dec_l1:
+            return None
+        import ttnn
+
+        return ttnn.L1_MEMORY_CONFIG
+
+    def enc_mem(self):
+        """``ENC_L1``: the memory config of the mixer blocks' intermediates (LN outputs and stream, token / channel
+        MLP outputs): L1 interleaved, else None (DRAM)."""
+        if not self.enc_l1:
+            return None
+        import ttnn
+
+        return ttnn.L1_MEMORY_CONFIG
+
     def upload(self, array: np.ndarray, dtype: str, *, shape4: bool = True):
         """Host float array -> DRAM TILE device tensor (rank padded to 4 with leading 1s)."""
         a = np.ascontiguousarray(np.asarray(array, np.float32))
@@ -98,6 +158,251 @@ class Build:
         return to_device(a, self.device, dtype)
 
 
+# OPT round 1 item 2a: explicit 2-D multicast program configs for the decoder's 352-row matmuls, the fastest
+# bit-identical candidate per (K, N, split pass) of the device sweep (logs/diffusion-planner/opt_r1/mm_sweep.json;
+# "hh" = x_hi @ w_hi, "hl" = x_hi @ w_lo, "lh" = x_lo @ w_hi + bias; a plain Linear uses "hh").
+# value: (transpose_mcast, per_core_M, per_core_N, in0_block_w, out_subblock_w); every candidate was bit-identical to
+# the auto config (fp32 DEST accumulation over K either way).
+DEC_ROWS = 352
+DEC_MM_CONFIGS = {
+    (324, 512, "hh"): (True, 1, 3, 11, 1), (324, 512, "hl"): (True, 2, 2, 11, 2), (324, 512, "lh"): (True, 1, 3, 11, 1),
+    (512, 256, "hh"): (False, 2, 1, 8, 1), (512, 256, "hl"): (False, 2, 1, 8, 1), (512, 256, "lh"): (False, 2, 1, 8, 1),
+    (256, 768, "hh"): (False, 2, 3, 8, 1), (256, 768, "hl"): (False, 2, 2, 8, 2), (256, 768, "lh"): (True, 1, 3, 8, 1),
+    (256, 256, "hh"): (False, 2, 2, 8, 2), (256, 256, "hl"): (False, 2, 2, 8, 2), (256, 256, "lh"): (False, 2, 1, 8, 1),
+    (256, 1024, "hh"): (False, 2, 3, 8, 1), (256, 1024, "hl"): (False, 2, 3, 8, 1),
+    (256, 1024, "lh"): (False, 2, 3, 8, 1),
+    (1024, 256, "hh"): (True, 1, 1, 8, 1), (1024, 256, "hl"): (True, 1, 1, 8, 1), (1024, 256, "lh"): (False, 2, 1, 8, 1),
+    (1024, 324, "hh"): (False, 2, 1, 8, 1), (1024, 324, "hl"): (False, 2, 1, 8, 1),
+    (1024, 324, "lh"): (False, 2, 1, 8, 1),
+}
+
+
+def compact_rows(rows: int) -> bool:
+    """``COMPACT`` (OPT round 5): a decoder row count other than :data:`DEC_ROWS` that the derived configs serve (a
+    tile-aligned agent bucket below the full capacity)."""
+    return rows % 32 == 0 and 32 <= rows < DEC_ROWS
+
+
+def fit_pcm(tr: bool, pcm: int, rows: int, grid) -> int:
+    """``per_core_M`` of a 2-D multicast config for ``rows`` rows: the swept value at :data:`DEC_ROWS`, else the
+    smallest value whose M blocks fit the grid axis that carries them (the in0 block width, ``per_core_N`` and the
+    subblocks stay as swept: the K accumulation order of every output element is unchanged)."""
+    if rows == DEC_ROWS:
+        return pcm
+    mt = -(-rows // 32)
+    lim = grid.x if tr else grid.y
+    p = 1
+    while -(-mt // p) > lim:
+        p += 1
+    return p
+
+
+def mcast_fits(tr: bool, pcm: int, pcn: int, rows: int, n: int, grid) -> bool:
+    """True when a 2-D multicast config's output blocks fit the grid: the swept configs assume the 12x10 ETH-dispatch
+    grid (N / per_core_N up to 12 blocks on x); on a smaller grid (WORKER dispatch, 11x10) the callers fall back to
+    the auto config instead of failing at capture (never hard-code the grid: CLAUDE.md)."""
+    mb, nb = -(-(-(-rows // 32)) // pcm), -(-(-(-n // 32)) // pcn)
+    gx, gy = (mb, nb) if tr else (nb, mb)
+    return gx <= grid.x and gy <= grid.y
+
+
+class DecConfigs:
+    """``DEC_MMCFG`` configs of one decoder linear: ``get(rows)`` -> ``{"hh" | "hl" | "lh": program config}`` for
+    :data:`DEC_ROWS` or (``COMPACT``) a smaller agent bucket, else None."""
+
+    def __init__(self, build: "Build", k: int, n: int):
+        self.build, self.k, self.n, self._cache = build, k, n, {}
+
+    def get(self, rows: int) -> Optional[Dict[str, Any]]:
+        rows = int(rows)
+        if rows != DEC_ROWS and not compact_rows(rows):
+            return None
+        if rows not in self._cache:
+            import ttnn
+
+            grid = self.build.device.compute_with_storage_grid_size()
+            out = {}
+            for ps in ("hh", "hl", "lh"):
+                tr, pcm, pcn, kb, sw = DEC_MM_CONFIGS[(self.k, self.n, ps)]
+                if not mcast_fits(tr, fit_pcm(tr, pcm, rows, grid), pcn, rows, self.n, grid):
+                    out = None                   # a grid smaller than the swept one: the auto configs
+                    break
+                out[ps] = ttnn.MatmulMultiCoreReuseMultiCastProgramConfig(
+                    compute_with_storage_grid_size=grid, in0_block_w=kb, out_subblock_h=1, out_subblock_w=sw,
+                    per_core_M=fit_pcm(tr, pcm, rows, grid), per_core_N=pcn, transpose_mcast=tr,
+                    fused_activation=None, fuse_batch=True)
+            self._cache[rows] = out
+        return self._cache[rows]
+
+
+def dec_configs(build: "Build", module: str, shape) -> Optional[DecConfigs]:
+    """The :class:`DecConfigs` of a decoder linear of weight ``shape`` (``DEC_MMCFG``), else None."""
+    if not (build.dec_mmcfg and module.startswith("dec.")):
+        return None
+    k, n = int(shape[0]), int(shape[1])
+    if (k, n, "hh") not in DEC_MM_CONFIGS:
+        return None
+    import ttnn
+
+    if not hasattr(ttnn, "MatmulMultiCoreReuseMultiCastProgramConfig"):
+        return None                      # the host fake ttnn
+    return DecConfigs(build, k, n)
+
+
+# OPT round 2 item 2 (SPLIT_KCAT): X' row tiles 3 Kt + 1 rounded up to a multiple of kcat_pad(K) (the K block of
+# the program config must divide them), and the fastest 2-D multicast config per (rows, K' tiles, N) of the device
+# sweep (logs/diffusion-planner/opt_r2/kcat_sweep.json): value (transpose_mcast, per_core_M, per_core_N, in0_block_w,
+# out_subblock_w); shapes not listed use the auto config.
+def kcat_pad(k: int) -> int:
+    """Sweep: K = 1024 (3 Kt + 1 = 97, prime) pads to 104 = 8 x 13; the other K keep 3 Kt + 1 (25 = 5 x 5, 49)."""
+    return 8 if k >= 1024 else 1
+
+
+KCAT_CONFIGS: Dict[tuple, tuple] = {
+    (352, 49, 256): (False, 2, 1, 7, 1), (352, 25, 768): (False, 2, 2, 5, 2), (352, 25, 256): (False, 2, 1, 5, 1),
+    (352, 25, 1024): (False, 2, 3, 5, 1), (352, 104, 256): (False, 2, 1, 13, 1),
+    (352, 104, 324): (False, 2, 1, 13, 1), (576, 25, 256): (False, 2, 1, 5, 1),
+}
+
+
+def kcat_config(build: "Build", ktp: int, n: int, rows: int):
+    """The program config of the K-concatenated matmul ``[rows, 32 ktp] @ [32 ktp, n]``: the swept entry, or
+    (``COMPACT``) the :data:`DEC_ROWS` entry with ``per_core_M`` fitted to a smaller agent bucket; None = auto."""
+    import ttnn
+
+    if not hasattr(ttnn, "MatmulMultiCoreReuseMultiCastProgramConfig") or not hasattr(
+            build.device, "compute_with_storage_grid_size"):
+        return None
+    grid = build.device.compute_with_storage_grid_size()
+    ent = KCAT_CONFIGS.get((rows, ktp, n))
+    if ent is None and compact_rows(rows):
+        ent = KCAT_CONFIGS.get((DEC_ROWS, ktp, n))
+    if ent is None:
+        return None
+    tr, pcm, pcn, kb, sw = ent
+    pcm = fit_pcm(tr, pcm, rows, grid) if (rows, ktp, n) not in KCAT_CONFIGS else pcm
+    if not mcast_fits(tr, pcm, pcn, rows, n, grid):
+        return None                              # a grid smaller than the swept one: the auto config
+    return ttnn.MatmulMultiCoreReuseMultiCastProgramConfig(
+        compute_with_storage_grid_size=grid, in0_block_w=kb, out_subblock_h=1, out_subblock_w=sw,
+        per_core_M=pcm, per_core_N=pcn,
+        transpose_mcast=tr, fused_activation=None, fuse_batch=True)
+
+
+# OPT round 4 item 1 (KCAT_L1): the K = 1024 split operands X' [352, 32 Ktp] written L1 block-sharded by the operand
+# build and read in place by a 2-D multicast matmul with in0 sharded (no 4.7 MB DRAM round trip). The K slices of
+# X' are the matmul's N blocks, so Ktp is padded to a multiple of them (zero tiles, exact). Value per (rows, K, N):
+# (transpose_mcast, per_core_M, per_core_N, pad, in0_block_w), the fastest bit-identical candidate of the device
+# check (logs/diffusion-planner/opt_r4/kl1_real.log: operand + matmul 62.6 -> 41.3 us for N = 256, 62.3 -> 41.8 us
+# for N = 324).
+KCAT_L1_CONFIGS: Dict[tuple, tuple] = {
+    (352, 1024, 256): (False, 2, 1, 8, 13),
+    (352, 1024, 324): (False, 2, 1, 11, 9),
+}
+
+
+def kcat_l1_layout(build: "Build", rows: int, k: int, n: int, ktp: int):
+    """``(memory_config, program_config)`` of the L1-sharded operand path for ``[rows, 32 ktp] @ [32 ktp, n]``."""
+    import ttnn
+
+    tr, pcm, pcn, _pad, kb = KCAT_L1_CONFIGS[(DEC_ROWS, k, n)]
+    grid = build.device.compute_with_storage_grid_size()
+    pcm = fit_pcm(tr, pcm, rows, grid)          # COMPACT: a smaller agent bucket (the K blocking unchanged)
+    mt, nt = -(-rows // 32), -(-n // 32)
+    nb, mb = -(-nt // pcn), -(-mt // pcm)
+    assert ktp % nb == 0 and (ktp // nb) % kb == 0, (ktp, nb, kb)
+    gx, gy = (mb, nb) if tr else (nb, mb)
+    if gx > grid.x or gy > grid.y:
+        return None                              # a grid smaller than the swept one: the DRAM operand path
+    crs = ttnn.CoreRangeSet({ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(gx - 1, gy - 1))})
+    spec = ttnn.ShardSpec(crs, [32 * pcm, 32 * (ktp // nb)],
+                          ttnn.ShardOrientation.COL_MAJOR if tr else ttnn.ShardOrientation.ROW_MAJOR)
+    mem = ttnn.MemoryConfig(ttnn.TensorMemoryLayout.BLOCK_SHARDED, ttnn.BufferType.L1, spec)
+    pc = ttnn.MatmulMultiCoreReuseMultiCastProgramConfig(
+        compute_with_storage_grid_size=grid, in0_block_w=kb, out_subblock_h=1,
+        out_subblock_w=max(d for d in (1, 2, 4) if pcn % d == 0), per_core_M=pcm, per_core_N=pcn,
+        transpose_mcast=tr, fused_activation=None, fuse_batch=True)
+    return mem, pc
+
+
+def _activate(y, act):
+    """The stock activation programs of the linears (``None`` / ``"gelu"`` / ``"gelu_tanh"``)."""
+    import ttnn
+
+    if act == "gelu":
+        return ttnn.gelu(y, fast_and_approximate_mode=False)
+    if act == "gelu_tanh":
+        return ttnn.gelu(y, variant=ttnn.GeluVariant.Tanh)
+    return y
+
+
+def _identity(t):
+    return t
+
+
+def fold_batch(x, n_out: int, enabled: bool):
+    """Run a linear on a batched activation ``[B0, B1, T, K] @ [K, N]`` as one 2-D matmul over all ``B0*B1*T`` rows
+    (``ENC_CH2D``, OPT round 1 item 1). The stock auto-config tiles M per batch element (``fuse_batch=False``), which
+    puts the encoder's ``[1, E, T, C]`` channel MLPs on 4-8 cores (970 us at E = 320). Returns ``(x2, back, pc)``:
+
+    - ``T % 32 == 0`` (the MixerBlock channel MLPs, T = 64): the free TILE view ``[1, 1, B*T, K]`` (auto config,
+      2-D multicast over the grid), ``back`` reshapes the output to ``[B0, B1, T, N]``;
+    - otherwise (the pre-projections, T = 6 / 20 / 40 padded per element): an explicit 1-D in1-multicast program config
+      with ``fuse_batch=True`` (rows of tiles split over the grid, the full N per core, ``in0_block_w = Kt``);
+      ``back`` is the identity.
+
+    The products and their fp32 DEST accumulation over K are the same; only the core assignment changes."""
+    import ttnn
+
+    shape = tuple(x.shape)
+    if not enabled or len(shape) != 4 or shape[0] * shape[1] == 1:
+        return x, _identity, None
+    b0, b1, t, k = shape
+    if t % 32 == 0:
+        x2 = ttnn.reshape(x, (1, 1, b0 * b1 * t, k))
+        return x2, (lambda y: ttnn.reshape(y, (b0, b1, t, y.shape[-1]))), None
+    cfg_cls = getattr(ttnn, "MatmulMultiCoreReuseMultiCast1DProgramConfig", None)
+    dev = x.device() if callable(getattr(x, "device", None)) else None
+    if cfg_cls is None or dev is None or not hasattr(dev, "compute_with_storage_grid_size"):
+        return x, _identity, None      # the host fake: numerics are the same
+    grid = dev.compute_with_storage_grid_size()
+    cores = grid.x * grid.y
+    m_tiles = b0 * b1 * (-(-t // 32))
+    kt, nt = -(-k // 32), -(-n_out // 32)
+    per_core_m = -(-m_tiles // cores)
+    sub_w = max(d for d in (1, 2, 4) if nt % d == 0)     # fp32 DEST: subblock h * w <= 4
+    pc = cfg_cls(compute_with_storage_grid_size=grid, in0_block_w=kt, out_subblock_h=1, out_subblock_w=sub_w,
+                 per_core_M=per_core_m, per_core_N=nt, fuse_batch=True, fused_activation=None, mcast_in0=False)
+    return x, _identity, pc
+
+
+def tall_config(x, wshape, fused_activation=None):
+    """The stock auto config of a tall 2-D matmul ``[1, 1, M, K] @ [K, N]`` (M / 32 >= the grid, K, N <= 128): the
+    1-D in1-multicast config the profile shows for the mixer linears (``MatmulMultiCoreReuseMultiCast1DProgramConfig``,
+    ``per_core_M = ceil(Mt / cores)``, ``in0_block_w = min(Kt, 2)``, the full N per core, ``mcast_in0 = False``), so a
+    fused activation can be added without changing the matmul (``LIN_ACT``). None for any other shape."""
+    import ttnn
+
+    shp = tuple(x.shape)
+    dev = x.device() if callable(getattr(x, "device", None)) else None
+    cls = getattr(ttnn, "MatmulMultiCoreReuseMultiCast1DProgramConfig", None)
+    if cls is None or dev is None or not hasattr(dev, "compute_with_storage_grid_size") or len(shp) != 4:
+        return None
+    if shp[0] * shp[1] != 1 or shp[-2] % 32 or wshape[0] % 32 or wshape[1] % 32:
+        return None
+    grid = dev.compute_with_storage_grid_size()
+    cores = grid.x * grid.y
+    mt, kt, nt = shp[-2] // 32, wshape[0] // 32, wshape[1] // 32
+    if mt < cores or nt > 4 or kt > 4:
+        return None
+    pcm = -(-mt // cores)
+    sw = max(d for d in (1, 2, 4) if nt % d == 0 and d <= 4)
+    sh = max(d for d in (1, 2, 4) if pcm % d == 0 and d * sw <= 4)
+    return cls(compute_with_storage_grid_size=grid, in0_block_w=min(kt, 2), out_subblock_h=sh, out_subblock_w=sw,
+               out_block_h=pcm, out_block_w=nt, per_core_M=pcm, per_core_N=nt, fuse_batch=False,
+               fused_activation=fused_activation, mcast_in0=False)
+
+
 class Linear:
     """``y = x @ w (+ b)`` with fused ``activation`` ("gelu" = exact erf, "gelu_tanh"); output dtype ``out``."""
 
@@ -110,12 +415,28 @@ class Linear:
         self.b = None if lin.b is None else build.upload(np.asarray(lin.b).reshape(1, -1), wd)
         self.cfg = build.cfg(module)
         self.shape = tuple(lin.w.shape)
+        self.fold = build.ch2d and module.startswith("enc.")
+        self.dec_pc = dec_configs(build, module, self.shape)
+        self.act_fuse = build.lin_act and module.startswith("enc.mixer.")
+        self.out_mem = None                      # ATTN_L1: output memory config (None = DRAM)
 
     def __call__(self, x):
         import ttnn
 
-        return ttnn.linear(x, self.w, bias=self.b, activation=self.activation, dtype=self.out,
-                           compute_kernel_config=self.cfg)
+        x, back, pc = fold_batch(x, self.shape[1], self.fold)
+        dpc = None if self.dec_pc is None else self.dec_pc.get(tuple(x.shape)[-2])
+        if dpc is not None and x.dtype == ttnn.bfloat16:
+            pc = dpc["hh"]
+        act = self.activation
+        if self.act_fuse and act == "gelu" and pc is None:
+            pc = tall_config(x, self.shape, ttnn.UnaryWithParam(ttnn.UnaryOpType.GELU, 0.0))
+            if pc is not None:                 # LIN_ACT: GELU in the matmul epilogue (fp32 DEST, before rounding)
+                act = None
+        kw = {} if pc is None else {"program_config": pc}
+        if self.out_mem is not None:
+            kw["memory_config"] = self.out_mem
+        return back(ttnn.linear(x, self.w, bias=self.b, activation=act, dtype=self.out,
+                                compute_kernel_config=self.cfg, **kw))
 
 
 def layer_norm_fp32(x, gamma=None, beta=None, *, eps: float = C.LN_EPS):
@@ -136,6 +457,26 @@ def layer_norm_fp32(x, gamma=None, beta=None, *, eps: float = C.LN_EPS):
     return y
 
 
+class KcatOperand:
+    """A split operand ``[x_hi | x_hi | x_lo | 1 | 0..]`` written by its producer (``KCAT_EMIT``: the split
+    LayerNorm, the fused attention) for a K-concatenated :class:`SplitLinear` (instead of x + a ``kcat_operand``
+    program). ``t``: the device tensor ``[..., M, 32 * ktp]``."""
+
+    def __init__(self, t, ktp: int):
+        self.t, self.ktp = t, int(ktp)
+
+    @property
+    def shape(self):
+        return self.t.shape
+
+
+def operand_ktp(build: "Build", consumer) -> int:
+    """Row tiles of the split operand ``consumer`` takes when its producer may emit it (``KCAT_EMIT``), else 0."""
+    if not getattr(build, "kcat_emit", False) or not isinstance(consumer, SplitLinear):
+        return 0
+    return consumer.operand_ktp()
+
+
 class SplitLinear:
     """``y = x @ w (+ b)`` to ~1e-5 relative from three device matmuls on split operands (the "bf16x3" scheme):
     ``x_hi @ w_hi + x_hi @ w_lo + x_lo @ w_hi`` with ``*_hi`` the bf16 roundings (bf16 x bf16 products are exact
@@ -158,27 +499,193 @@ class SplitLinear:
         self.activation = activation
         self.out = ttnn_dtype(out)
         self.shape = tuple(w.shape)
+        self.fold = build.ch2d and module.startswith("enc.")
+        self.dec_pc = dec_configs(build, module, self.shape)
+        # OPT round 2 item 2 (SPLIT_KCAT): the three passes as ONE fp32 matmul over the concatenated K axis,
+        # [x_hi | x_hi | x_lo | ones] @ [w_hi; w_lo; w_hi; (b_hi, b_lo, 0...)]: the same exact / TF32-truncated
+        # products, accumulated in one fp32 DEST sum (the bias exact as two K rows instead of the packer epilogue).
+        # Decoder modules with a tile-aligned K only (the 352-row linears); others keep the 3-pass form.
+        # OPT round 3 item 2 (ENC_KCAT): the encoder's split linears (pre-projections, the ego / neighbour island)
+        # in the same form, any K: each weight block padded to whole tiles with zero rows (the operand keeps the
+        # input's tile padding, which meets those zero rows as in the 3-pass form).
+        self.kcat = None
+        self.out_mem = None                      # ATTN_L1: output memory config of the K-concatenated matmul
+        k = self.shape[0]
+        enc_kcat = build.enc_kcat and (module.startswith("enc.pre.") or module.startswith("enc.island."))
+        if build.split_kcat and ((module.startswith("dec.") and k % 32 == 0) or enc_kcat):
+            rows = np.zeros((32, self.shape[1]), np.float32)
+            if lin.b is not None:
+                b = np.asarray(lin.b, np.float32).reshape(-1)
+                b_hi = round_to_bf16(b)
+                rows[0], rows[1] = b_hi, b - b_hi
+            from .kcat_kernel import kcat_tiles
+
+            kp = -(-k // 32) * 32
+            self.kcat_kp = kp
+
+            def padk(a):
+                return np.concatenate([a, np.zeros((kp - k, self.shape[1]), np.float32)], axis=0)
+
+            l1_key = (DEC_ROWS, kp, self.shape[1])
+            use_l1 = (build.kcat_l1 and build.split_kcat == 2 and module.startswith("dec.")
+                      and l1_key in KCAT_L1_CONFIGS and hasattr(build.device, "compute_with_storage_grid_size"))
+            self.kcat_padv = KCAT_L1_CONFIGS[l1_key][3] if use_l1 else kcat_pad(kp)
+            self.kcat_ktp = kcat_tiles(kp, self.kcat_padv)
+            # KCAT_L1: (memory config, program config) of the L1-sharded operand path for DEC_ROWS rows
+            self.kcat_l1 = {} if use_l1 else None        # {rows: layout}, filled per row count on first use
+            pad_rows = np.zeros((32 * self.kcat_ktp - 3 * kp - 32, self.shape[1]), np.float32)
+            wk = np.concatenate([padk(w_hi), padk((w - w_hi).astype(np.float32)), padk(w_hi), rows, pad_rows],
+                                axis=0)
+            self.kcat = build.upload(wk, "float32")
+            self.kcat_pc = {}                         # {rows: program config or None}
+            self.kcat_mode = build.split_kcat
+            self._act_once = build.kcat_act_once
+            self._ones = {}
+            self._build = build
+
+    def _ones_tile(self, x):
+        """``[.., M, 32 (Kt' - 3 Kt)]`` fp32 with columns 0 and 1 = 1 (the bias rows of the concatenated weight), the
+        rest 0 (the K padding)."""
+        key = tuple(tuple(x.shape)[:-1])
+        if key not in self._ones:
+            o = np.zeros(key + (32 * self.kcat_ktp - 3 * self.kcat_kp,), np.float32)
+            o[..., :2] = 1.0
+            self._ones[key] = self._build.upload(o, "float32")
+        return self._ones[key]
+
+    def _call_kcat(self, x, pre_act=None):
+        import ttnn
 
-    def __call__(self, x):
+        f32 = ttnn.float32
+        if isinstance(x, KcatOperand):                # KCAT_EMIT: the producer wrote the operand
+            assert pre_act is None and x.ktp == self.kcat_ktp
+            return self._kcat_mm(x.t, x.t)
+        if x.dtype != f32:
+            x = ttnn.typecast(x, f32)
+        from .kcat_kernel import kcat_operand, supported
+
+        if self.kcat_mode == 2 and supported(x):    # else (the host fake ttnn) the stock build: the same values
+            rows = tuple(x.shape)[-2]
+            if self.kcat_l1 is not None and (rows == DEC_ROWS or compact_rows(rows)) and not self.fold:
+                if rows not in self.kcat_l1:
+                    self.kcat_l1[rows] = kcat_l1_layout(self._build, rows, self.kcat_kp, self.shape[1],
+                                                        self.kcat_ktp)
+            if self.kcat_l1 is not None and self.kcat_l1.get(rows) is not None and not self.fold:
+                mem, pc = self.kcat_l1[rows]               # KCAT_L1: X' in L1, read in place by the sharded-in0 matmul
+                xk = kcat_operand(x, self.kcat_padv, pre_act, memory_config=mem, act_once=self._act_once)
+                y = ttnn.matmul(xk, self.kcat, dtype=f32, compute_kernel_config=self.cfg, program_config=pc,
+                                memory_config=self.out_mem or ttnn.DRAM_MEMORY_CONFIG)
+                xk.deallocate()
+                return y
+            return self._kcat_mm(kcat_operand(x, self.kcat_padv, pre_act, act_once=self._act_once), x)
+        x = _activate(x, pre_act)
+        x_hi = ttnn.typecast(ttnn.typecast(x, ttnn.bfloat16), f32)
+        x_lo = ttnn.subtract(x, x_hi)
+        xk = ttnn.concat([x_hi, x_hi, x_lo, self._ones_tile(x)], dim=-1)
+        return self._kcat_mm(xk, x)
+
+    def _kcat_mm(self, xk, x):
+        import ttnn
+
+        if self.fold:                        # encoder [1, E, T, K'] operands: one 2-D matmul (ENC_CH2D)
+            xk, back, pc = fold_batch(xk, self.shape[1], True)
+            kw = {} if pc is None else {"program_config": pc}
+            if self.out_mem is not None:
+                kw["memory_config"] = self.out_mem
+            return back(ttnn.matmul(xk, self.kcat, dtype=ttnn.float32, compute_kernel_config=self.cfg, **kw))
+        kw = self._kcat_kw(x)
+        if self.out_mem is not None:
+            kw["memory_config"] = self.out_mem
+        return ttnn.matmul(xk, self.kcat, dtype=ttnn.float32, compute_kernel_config=self.cfg, **kw)
+
+    def _kcat_ok(self, x) -> bool:
+        """The K-concatenated form applies: tile-aligned K (any operand build), or the generic_op operand build
+        (which keeps the input's tile padding for an unaligned K)."""
+        if self.kcat is None:
+            return False
+        if isinstance(x, KcatOperand) or self.kcat_kp == self.shape[0]:
+            return True
+        from .kcat_kernel import supported
+
+        return self.kcat_mode == 2 and supported(x)
+
+    def _kcat_kw(self, x):
+        rows = int(tuple(x.shape)[-2])
+        if rows not in self.kcat_pc:
+            self.kcat_pc[rows] = kcat_config(self._build, self.kcat_ktp, self.shape[1], rows)
+        pc = self.kcat_pc[rows]
+        return {} if pc is None else {"program_config": pc}
+
+    def operand_ktp(self) -> int:
+        """Row tiles of the operand a producer may write for this linear (``KCAT_EMIT``): the generic_op form with a
+        tile-aligned K and no K-block padding of a different layout; else 0."""
+        if self.kcat is None or self.kcat_mode != 2 or self.kcat_kp != self.shape[0]:
+            return 0
+        return self.kcat_ktp
+
+    def fuses_input_act(self) -> bool:
+        """True when this linear can take its input before the previous linear's activation (``KCAT_ACT``)."""
+        return self.kcat is not None and self.kcat_mode == 2 and self.kcat_kp == self.shape[0]
+
+    def __call__(self, x, *, pre_act=None, defer_act: bool = False):
+        """``pre_act``: apply ``"gelu"`` / ``"gelu_tanh"`` to ``x`` first (fused into the operand build when
+        :meth:`fuses_input_act`); ``defer_act``: return the pre-activation fp32 output (the next linear applies
+        :attr:`activation` through ``pre_act``)."""
         import ttnn
 
         f32 = ttnn.float32
+        if self._kcat_ok(x):
+            y = self._call_kcat(x, pre_act)
+            if defer_act:
+                return y
+            y = _activate(y, self.activation)
+            if self.out != f32:
+                y = ttnn.typecast(y, self.out)
+            return y
+        x = _activate(x, pre_act)
+        if defer_act:
+            assert self.out == f32
+            act, self.activation = self.activation, None
+            try:
+                return self(x)
+            finally:
+                self.activation = act
+        x, back, pc = fold_batch(x, self.shape[1], self.fold)
+        pcs = {"hh": pc, "hl": pc, "lh": pc}
+        dpc = None if self.dec_pc is None else self.dec_pc.get(tuple(x.shape)[-2])
+        if dpc is not None:
+            pcs = dpc
+
+        def kw(ps):
+            p = pcs[ps]
+            return {"compute_kernel_config": self.cfg, **({} if p is None else {"program_config": p})}
         if x.dtype != f32:          # a bf16 input is its own hi part: two matmuls
-            y = ttnn.matmul(x, self.w_hi, dtype=f32, compute_kernel_config=self.cfg)
-            y = ttnn.add(y, ttnn.linear(x, self.w_lo, bias=self.b, dtype=f32, compute_kernel_config=self.cfg))
+            y = ttnn.matmul(x, self.w_hi, dtype=f32, **kw("hh"))
+            y = ttnn.add(y, ttnn.linear(x, self.w_lo, bias=self.b, dtype=f32, **kw("hl")))
         else:
             x_hi = ttnn.typecast(x, ttnn.bfloat16)
             x_lo = ttnn.subtract(x, ttnn.typecast(x_hi, f32))
-            y = ttnn.matmul(x_hi, self.w_hi, dtype=f32, compute_kernel_config=self.cfg)
-            y = ttnn.add(y, ttnn.matmul(x_hi, self.w_lo, dtype=f32, compute_kernel_config=self.cfg))
-            y = ttnn.add(y, ttnn.linear(x_lo, self.w_hi, bias=self.b, dtype=f32, compute_kernel_config=self.cfg))
+            y = ttnn.matmul(x_hi, self.w_hi, dtype=f32, **kw("hh"))
+            y = ttnn.add(y, ttnn.matmul(x_hi, self.w_lo, dtype=f32, **kw("hl")))
+            y = ttnn.add(y, ttnn.linear(x_lo, self.w_hi, bias=self.b, dtype=f32, **kw("lh")))
         if self.activation == "gelu":
             y = ttnn.gelu(y, fast_and_approximate_mode=False)
         elif self.activation == "gelu_tanh":
             y = ttnn.gelu(y, variant=ttnn.GeluVariant.Tanh)
         if self.out != f32:
             y = ttnn.typecast(y, self.out)
-        return y
+        return back(y)
+
+
+def chain2(first, second, x, enabled: bool):
+    """``second(first(x))``; with ``enabled`` (``KCAT_ACT``) and a K-concatenated ``second``, ``first``'s activation
+    runs inside ``second``'s operand build instead of as its own program (same LLK, same values)."""
+    import ttnn
+
+    if (enabled and isinstance(first, SplitLinear) and isinstance(second, SplitLinear) and first.activation
+            and first.out == ttnn.float32 and second.fuses_input_act()):
+        return second(first(x, defer_act=True), pre_act=first.activation)
+    return second(first(x))
 
 
 def make_linear(build: Build, lin: Lin, module: str, *, out: str, activation: Optional[str] = None):
@@ -196,26 +703,101 @@ class LayerNorm:
     def __init__(self, build: Build, nrm: Optional[Norm], module: str):
         self.cfg = build.cfg(module)
         self.mode = build.ln_mode(module)
+        self.fused = build.ln_kernel
+        self.resid = build.ln_resid
+        self.split = build.ln_split
+        self.sbc = build.ln_sfpu_bcast
+        self.tr = build.ln_tr
+        self.mem = None                          # DEC_L1: memory config of the split-row LN's outputs
         self.g = self.b = None
         if nrm is not None:
             self.g = build.upload(np.asarray(nrm.gamma).reshape(1, -1), "float32")
             self.b = build.upload(np.asarray(nrm.beta).reshape(1, -1), "float32")
 
-    def __call__(self, x, gamma=None, beta=None):
+    def __call__(self, x, gamma=None, beta=None, *, kcat: int = 0):
+        """``kcat`` (``KCAT_EMIT``, the consumer's :func:`operand_ktp`): return the consumer's split operand as a
+        :class:`KcatOperand` when the split-row kernel runs (else y as usual)."""
         import ttnn
 
         g = self.g if gamma is None else gamma
         b = self.b if beta is None else beta
         if self.mode == "fp32":
+            if self.fused:
+                from .ln_kernel import layer_norm_fp32_fused, supported
+
+                if self.split:
+                    from .ln_kernel import layer_norm_fp32_split, split_supported
+
+                    if split_supported(x):
+                        y = layer_norm_fp32_split(x, g, b, eps=C.LN_EPS, sfpu_bcast=self.sbc, kcat_ktp=kcat,
+                                                  memory_config=self.mem)
+                        return KcatOperand(y, kcat) if kcat else y
+                if supported(x):
+                    return layer_norm_fp32_fused(x, g, b, eps=C.LN_EPS, lean=self.fused == 2, sfpu_bcast=self.sbc,
+                                                 memory_config=self.mem)
             return layer_norm_fp32(x, g, b)
-        return ttnn.layer_norm(x, epsilon=C.LN_EPS, weight=g, bias=b, compute_kernel_config=self.cfg)
+        kw = {} if self.mem is None else {"memory_config": self.mem}
+        return ttnn.layer_norm(x, epsilon=C.LN_EPS, weight=g, bias=b, compute_kernel_config=self.cfg, **kw)
+
+    def can_transpose(self, x) -> bool:
+        """``LN_TR``: the fused one-core-per-row kernel runs for ``x`` (not the split-row form), so it can read the
+        residual / write the output transposed per entity."""
+        if not (self.tr and self.mode == "fp32" and self.fused and self.resid):
+            return False
+        from .ln_kernel import split_supported, supported
+
+        return supported(x) and not (self.split and split_supported(x))
+
+    def transposed(self, x, *, res=None):
+        """``LN_TR``: ``LN(x)`` (or, with ``res``, ``h = x + res`` and ``LN(h)``) with the output written as
+        ``[.., E, W, T]`` (the per-entity transpose); returns ``y^T`` or ``(h, y^T)``."""
+        from .ln_kernel import layer_norm_fp32_fused
+
+        return layer_norm_fp32_fused(x, self.g, self.b, eps=C.LN_EPS, residual=res, lean=self.fused == 2,
+                                     sfpu_bcast=self.sbc, out_t=True, memory_config=self.mem)
+
+    def residual_t(self, x, res_t_tensor):
+        """``LN_TR``: ``h = x + res^T`` (``res`` ``[.., E, W, T]``) and ``LN(h)`` -> ``(h, y)``."""
+        from .ln_kernel import layer_norm_fp32_fused
+
+        return layer_norm_fp32_fused(x, self.g, self.b, eps=C.LN_EPS, residual=res_t_tensor, lean=self.fused == 2,
+                                     sfpu_bcast=self.sbc, res_t=True, memory_config=self.mem)
+
+    def residual(self, x, res, rgate=None, gamma=None, beta=None, *, write_h: bool = True, kcat: int = 0):
+        """``h = x + res (* rgate)`` then ``(h, LN(h))``: one fused program (``LN_RESID``, ``tt/ln_kernel.py``)
+        when this LN runs the fused fp32 kernel, else the stock ``add`` (+ ``multiply``) and :meth:`__call__`.
+        ``h`` is None when ``write_h`` is False and the fused program runs."""
+        import ttnn
+
+        g = self.g if gamma is None else gamma
+        b = self.b if beta is None else beta
+        if self.mode == "fp32" and self.fused and self.resid:
+            from .ln_kernel import layer_norm_fp32_fused, supported
+
+            if supported(x) and supported(res) and tuple(x.shape) == tuple(res.shape):
+                if self.split:
+                    from .ln_kernel import layer_norm_fp32_split, split_supported
+
+                    if split_supported(x):
+                        h, y = layer_norm_fp32_split(x, g, b, eps=C.LN_EPS, residual=res, rgate=rgate,
+                                                     write_h=write_h, sfpu_bcast=self.sbc, kcat_ktp=kcat,
+                                                     memory_config=self.mem)
+                        return h, (KcatOperand(y, kcat) if kcat else y)
+                return layer_norm_fp32_fused(x, g, b, eps=C.LN_EPS, residual=res, rgate=rgate, write_h=write_h,
+                                             lean=self.fused == 2, sfpu_bcast=self.sbc, memory_config=self.mem)
+        h = ttnn.add(x, res if rgate is None else ttnn.multiply(res, rgate))
+        return h, self(h, gamma, beta)
 
 
 class Const:
     """A constant device tensor (upload once)."""
 
-    def __init__(self, build: Build, array: np.ndarray, dtype: str):
+    def __init__(self, build: Build, array: np.ndarray, dtype: str, memory_config=None):
         self.t = build.upload(array, dtype)
+        if memory_config is not None:
+            import ttnn
+
+            self.t = ttnn.to_memory_config(self.t, memory_config)
 
     def __call__(self):
         return self.t

code/tt_diffusion_planner/tt/ln_kernel.py

@@ -0,0 +1,285 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Fused fp32 LayerNorm as one ``ttnn.generic_op`` program (``LN_KERNEL``, OPT round 2 item 3).
+
+:func:`layer_norm_fp32_fused` computes exactly what :func:`tt.layers.layer_norm_fp32` computes with 7-9 stock
+programs (mean, subtract, square, mean, add eps, rsqrt, multiply, gamma, beta), in one program: the compute kernel
+(``kernels/ln32_compute.cpp``) issues the same SFPU LLK calls in the same order as the stock kernels, and every
+intermediate the stock graph writes to an fp32 DRAM tensor stays fp32 in L1 / DST, so the output is meant to be
+bit-identical (checked on the device: ``code/scripts/ln_kernel_check.py``).
+
+Layout: x is an fp32 TILE DRAM-interleaved tensor ``[..., R, W]`` (R a multiple of 32, W = 32 * Wt); the tile rows
+are split in contiguous blocks over ``min(rows, grid)`` cores (one core per tile row up to the grid). gamma / beta
+are ``[1, 1, 1, W]`` fp32 TILE rows (row 0 valid) or None.
+"""
+from __future__ import annotations
+
+import os
+import struct
+from typing import Any, Optional
+
+__all__ = ["layer_norm_fp32_fused", "supported"]
+
+_KDIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "kernels")
+TB = 4096          # fp32 tile bytes
+N_RT = 2           # per-core RT args of every kernel ([row0, n_rows] / [n_rows] + pad): a CT arg (probe P3)
+
+
+def _bits(v: float) -> int:
+    return int.from_bytes(struct.pack("<f", float(v)), "little")
+
+
+def supported(x: Any) -> bool:
+    """True when ``x`` is a device fp32 TILE interleaved tensor with tile-aligned rows and columns."""
+    import ttnn
+
+    try:
+        shp = list(x.padded_shape)
+        return (x.dtype == ttnn.float32 and x.layout == ttnn.TILE_LAYOUT and not x.is_sharded()
+                and shp[-1] % 32 == 0 and shp[-2] % 32 == 0 and int(x.shape[-1]) == shp[-1]
+                and hasattr(ttnn, "generic_op"))
+    except Exception:  # noqa: BLE001 - the host fake ttnn
+        return False
+
+
+def _cores(n: int, g):
+    """The first ``n`` cores of the grid in row-major order: a CoreRangeSet and the coordinate list."""
+    import ttnn
+
+    cs = [(i % g.x, i // g.x) for i in range(n)]
+    full, rem = divmod(n, g.x)
+    rs = []
+    if full:
+        rs.append(ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(g.x - 1, full - 1)))
+    if rem:
+        rs.append(ttnn.CoreRange(ttnn.CoreCoord(0, full), ttnn.CoreCoord(rem - 1, full)))
+    return ttnn.CoreRangeSet(set(rs)), cs
+
+
+def layer_norm_fp32_fused(x, gamma=None, beta=None, *, eps: float, residual=None, rgate=None, write_h: bool = True,
+                          lean: bool = True, sfpu_bcast: bool = False, res_t: bool = False, out_t: bool = False,
+                          memory_config=None):
+    """LayerNorm over the last dim of fp32 ``x`` -> fp32 tensor of x's shape (see the module docstring).
+
+    With ``residual`` (``LN_RESID``): ``h = x + residual (* rgate)`` first (the stock ``ttnn.add(x,
+    ttnn.multiply(residual, rgate))``, fp32 SFPU ops), then ``LN(h)``; returns ``(h, LN(h))`` (``h`` is None when
+    ``write_h`` is False). ``residual``: fp32 like ``x``; ``rgate``: a ``[1, 1, 1, W]`` fp32 row.
+    ``lean``: one unpacker / SFPU-binary init per phase instead of one per tile (same LLK math calls).
+    ``sfpu_bcast`` (``LN_SFPU_BCAST``): the SFPU row reduce writes the row statistics to every column
+    (``kernels/ln32_sfpu.h``) instead of the writer's RISC-V column fill (same values).
+    ``res_t`` / ``out_t`` (``LN_TR``, x ``[.., E, T, W]``): the residual comes as ``[.., E, W, T]`` (its per-entity
+    transpose) / the output is written as ``[.., E, W, T]``, the tiles transposed in the kernel with the stock
+    ``ttnn.transpose`` LLK (exact): the two transposes around the mixer's token-mixing MLP. ``memory_config``: of
+    the outputs (default DRAM; interleaved L1 for ``ENC_L1``)."""
+    import ttnn
+
+    dev = x.device()
+    shp = list(x.padded_shape)
+    W = shp[-1]
+    Wt = W // 32
+    rows = 1
+    for d in shp[:-1]:
+        rows *= d
+    rows //= 32
+    hr = int(residual is not None)
+    hrg = int(hr and rgate is not None)
+    wh = int(hr and write_h)
+    Tt = shp[-2] // 32
+    assert not (res_t and hrg), "LN_TR: no residual gate"
+    oshape = x.shape if not out_t else ttnn.Shape(list(x.shape)[:-2] + [shp[-1], shp[-2]])
+    omem = memory_config or ttnn.DRAM_MEMORY_CONFIG
+    out = ttnn.allocate_tensor_on_device(oshape, ttnn.float32, ttnn.TILE_LAYOUT, dev, omem)
+    h = (ttnn.allocate_tensor_on_device(x.shape, ttnn.float32, ttnn.TILE_LAYOUT, dev, omem)
+         if wh else None)
+    g = dev.compute_with_storage_grid_size()
+    n = min(rows, g.x * g.y)
+    crs, cs = _cores(n, g)
+    base, extra = divmod(rows, n)
+    rd, wr, cp = ttnn.RuntimeArgs(), ttnn.RuntimeArgs(), ttnn.RuntimeArgs()
+    r0 = 0
+    for i, (cx, cy) in enumerate(cs):
+        k = base + (1 if i < extra else 0)
+        rd[cx][cy] = [r0, k]
+        wr[cx][cy] = [r0, k]
+        cp[cx][cy] = [k, 0]
+        r0 += k
+    hg, hb = int(gamma is not None), int(beta is not None)
+
+    def acc(t):
+        return list(ttnn.TensorAccessorArgs(t).get_compile_time_args())
+
+    def cb(idx, pages):
+        return ttnn.CBDescriptor(total_size=pages * TB, core_ranges=crs, format_descriptors=[
+            ttnn.CBFormatDescriptor(buffer_index=idx, data_format=ttnn.float32, page_size=TB)])
+
+    cbs = [cb(0, 2 * Wt), cb(3, 1), cb(4, 2), cb(5, 2), cb(6, Wt), cb(16, 2 * Wt if Wt <= 4 else Wt)]
+    if hg:
+        cbs.append(cb(1, Wt))
+    if hb:
+        cbs.append(cb(2, Wt))
+    if hr:
+        cbs += [cb(7, Wt), cb(9, Wt)]
+    if hrg:
+        cbs.append(cb(8, Wt))
+    if wh:
+        cbs.append(cb(17, 2))
+    if out_t:
+        cbs.append(cb(18, 1))
+    um = [ttnn.UnpackToDestMode.Default] * 64
+    for i in (0, 1, 2, 3, 5, 6, 7, 8, 9, 18):
+        um[i] = ttnn.UnpackToDestMode.UnpackToDestFp32
+    ccfg = ttnn.ComputeConfigDescriptor(math_fidelity=ttnn.MathFidelity.HiFi4, fp32_dest_acc_en=True,
+                                        math_approx_mode=False)
+    ccfg.unpack_to_dest_mode = um
+    g_t = gamma if hg else x
+    b_t = beta if hb else x
+    r_t = residual if hr else x
+    rg_t = rgate if hrg else x
+    h_t = h if wh else out
+    reader_ct = [Wt, hg, hb, _bits(eps), N_RT, hr, hrg, int(bool(res_t)), Tt] + acc(x) + acc(g_t) + acc(b_t) + acc(
+        r_t) + acc(rg_t)
+    writer_ct = [Wt, N_RT, wh, int(bool(sfpu_bcast)), int(bool(out_t)), Tt] + acc(out) + acc(h_t)
+    compute_ct = [Wt, hg, hb, _bits(1.0 / W), N_RT, hr, hrg, wh, int(bool(lean)), int(bool(sfpu_bcast)),
+                  int(bool(res_t)), int(bool(out_t))]
+    SRC = ttnn.KernelDescriptor.SourceType.FILE_PATH
+
+    def kd(name, ct, rt, common, config):
+        return ttnn.KernelDescriptor(kernel_source=os.path.join(_KDIR, name), source_type=SRC, core_ranges=crs,
+                                     compile_time_args=ct, defines=[], runtime_args=rt, common_runtime_args=common,
+                                     config=config, compiler_include_paths=[_KDIR])
+
+    ks = [kd("ln32_reader.cpp", reader_ct, rd, [x.buffer_address(), g_t.buffer_address(), b_t.buffer_address(),
+                                                r_t.buffer_address(), rg_t.buffer_address()],
+             ttnn.ReaderConfigDescriptor()),
+          kd("ln32_writer.cpp", writer_ct, wr, [out.buffer_address(), h_t.buffer_address()],
+             ttnn.WriterConfigDescriptor()),
+          kd("ln32_compute.cpp", compute_ct, cp, [], ccfg)]
+    ins = [x] + ([gamma] if hg else []) + ([beta] if hb else []) + ([residual] if hr else []) + (
+        [rgate] if hrg else []) + ([h] if wh else [])
+    ttnn.generic_op(ins + [out], ttnn.ProgramDescriptor(kernels=ks, semaphores=[], cbs=cbs))
+    if hr:
+        return h, out
+    return out
+
+
+def reference_decomposition(x, gamma: Optional[Any] = None, beta: Optional[Any] = None, *, eps: float):
+    """The stock decomposition (``tt.layers.layer_norm_fp32``), for the device check."""
+    from .layers import layer_norm_fp32
+
+    return layer_norm_fp32(x, gamma, beta, eps=eps)
+
+
+def split_supported(x: Any) -> bool:
+    """The split-row form needs every tile row's Wt column tiles on distinct cores (rows * Wt <= grid cores)."""
+    if not supported(x):
+        return False
+    shp = list(x.padded_shape)
+    rows = 1
+    for d in shp[:-1]:
+        rows *= d
+    rows //= 32
+    g = x.device().compute_with_storage_grid_size()
+    wt = shp[-1] // 32
+    return 2 <= wt <= g.x * g.y and rows * wt <= g.x * g.y
+
+
+def layer_norm_fp32_split(x, gamma=None, beta=None, *, eps: float, residual=None, rgate=None, write_h: bool = True,
+                          sfpu_bcast: bool = False, kcat_ktp: int = 0, memory_config=None):
+    """:func:`layer_norm_fp32_fused` with each tile row spread over Wt cores (``LN_SPLIT``, ``kernels/ln32s_*.cpp``):
+    member j owns column tile j; the root (member 0) folds the gathered tiles in order and broadcasts the mean and
+    rstd, so the result is the same bit for bit. For few rows (the decoder's 11 tile rows: 88 cores instead of 11).
+    ``kcat_ktp`` (``KCAT_EMIT``): instead of y, write the split operand ``[y_hi | y_hi | y_lo | 1 | 0..]`` of the next
+    K-concatenated linear (``[..., R, 32 * kcat_ktp]``, the ``kcat_operand`` layout and LLK calls).
+    ``memory_config``: of the outputs (default DRAM; interleaved L1 for ``DEC_L1``)."""
+    import ttnn
+
+    dev = x.device()
+    shp = list(x.padded_shape)
+    W = shp[-1]
+    Wt = W // 32
+    omem = memory_config or ttnn.DRAM_MEMORY_CONFIG
+    rows = 1
+    for d in shp[:-1]:
+        rows *= d
+    rows //= 32
+    hg, hb = int(gamma is not None), int(beta is not None)
+    hr = int(residual is not None)
+    hrg = int(hr and rgate is not None)
+    wh = int(hr and write_h)
+    oshape = x.shape if not kcat_ktp else ttnn.Shape(list(x.shape)[:-1] + [32 * kcat_ktp])
+    out = ttnn.allocate_tensor_on_device(oshape, ttnn.float32, ttnn.TILE_LAYOUT, dev, omem)
+    h = (ttnn.allocate_tensor_on_device(x.shape, ttnn.float32, ttnn.TILE_LAYOUT, dev, omem)
+         if wh else None)
+    g = dev.compute_with_storage_grid_size()
+    G = min(rows, (g.x * g.y) // Wt)
+    n = G * Wt
+    crs, cs = _cores(n, g)
+    phys = [dev.worker_core_from_logical_core(ttnn.CoreCoord(cx, cy)) for cx, cy in cs]
+    rd, wr, cp = ttnn.RuntimeArgs(), ttnn.RuntimeArgs(), ttnn.RuntimeArgs()
+    n_rt = 6 + 2 * Wt
+    for gi in range(G):
+        n_rows = len(range(gi, rows, G))
+        root = phys[gi * Wt]
+        members = []
+        for m in range(Wt):
+            members += [phys[gi * Wt + m].x, phys[gi * Wt + m].y]
+        for j in range(Wt):
+            cx, cy = cs[gi * Wt + j]
+            args = [gi, n_rows, G, j, root.x, root.y] + members
+            rd[cx][cy] = args
+            wr[cx][cy] = args
+            cp[cx][cy] = [n_rows, int(j == 0)]
+
+    def acc(t):
+        return list(ttnn.TensorAccessorArgs(t).get_compile_time_args())
+
+    def cb(idx, pages):
+        return ttnn.CBDescriptor(total_size=pages * TB, core_ranges=crs, format_descriptors=[
+            ttnn.CBFormatDescriptor(buffer_index=idx, data_format=ttnn.float32, page_size=TB)])
+
+    cbs = [cb(0, 2), cb(3, 1), cb(4, 2), cb(5, 2), cb(6, 2), cb(10, Wt), cb(11, 2), cb(12, 1), cb(16, 2)]
+    if hg:
+        cbs.append(cb(1, 1))
+    if hb:
+        cbs.append(cb(2, 1))
+    if hr:
+        cbs += [cb(7, 2), cb(9, 2)]
+    if hrg:
+        cbs.append(cb(8, 1))
+    if wh:
+        cbs.append(cb(17, 2))
+    if kcat_ktp:
+        cbs += [cb(18, 1), cb(19, 2), cb(20, 1)] + ([cb(21, 1)] if kcat_ktp > 3 * Wt + 1 else [])
+    um = [ttnn.UnpackToDestMode.Default] * 64
+    for i in (0, 1, 2, 3, 5, 6, 7, 8, 9, 10, 18):
+        um[i] = ttnn.UnpackToDestMode.UnpackToDestFp32
+    ccfg = ttnn.ComputeConfigDescriptor(math_fidelity=ttnn.MathFidelity.HiFi4, fp32_dest_acc_en=True,
+                                        math_approx_mode=False)
+    ccfg.unpack_to_dest_mode = um
+    g_t = gamma if hg else x
+    b_t = beta if hb else x
+    r_t = residual if hr else x
+    rg_t = rgate if hrg else x
+    h_t = h if wh else out
+    reader_ct = [Wt, hg, hb, _bits(eps), n_rt, hr, hrg] + acc(x) + acc(g_t) + acc(b_t) + acc(r_t) + acc(rg_t)
+    writer_ct = [Wt, n_rt, wh, int(bool(sfpu_bcast)), int(bool(kcat_ktp)), int(kcat_ktp)] + acc(out) + acc(h_t)
+    compute_ct = [Wt, hg, hb, _bits(1.0 / W), 2, hr, hrg, wh, int(bool(sfpu_bcast)), int(bool(kcat_ktp))]
+    SRC = ttnn.KernelDescriptor.SourceType.FILE_PATH
+
+    def kd(name, ct, rt, common, config):
+        return ttnn.KernelDescriptor(kernel_source=os.path.join(_KDIR, name), source_type=SRC, core_ranges=crs,
+                                     compile_time_args=ct, defines=[], runtime_args=rt, common_runtime_args=common,
+                                     config=config, compiler_include_paths=[_KDIR])
+
+    ks = [kd("ln32s_reader.cpp", reader_ct, rd, [x.buffer_address(), g_t.buffer_address(), b_t.buffer_address(),
+                                                 r_t.buffer_address(), rg_t.buffer_address()],
+             ttnn.ReaderConfigDescriptor()),
+          kd("ln32s_writer.cpp", writer_ct, wr, [out.buffer_address(), h_t.buffer_address()],
+             ttnn.WriterConfigDescriptor()),
+          kd("ln32s_compute.cpp", compute_ct, cp, [], ccfg)]
+    sems = [ttnn.SemaphoreDescriptor(id=i, core_ranges=crs, initial_value=0) for i in range(2)]
+    ins = [x] + ([gamma] if hg else []) + ([beta] if hb else []) + ([residual] if hr else []) + (
+        [rgate] if hrg else []) + ([h] if wh else [])
+    ttnn.generic_op(ins + [out], ttnn.ProgramDescriptor(kernels=ks, semaphores=sems, cbs=cbs))
+    if hr:
+        return h, out
+    return out

code/tt_diffusion_planner/tt/model.py

@@ -33,7 +33,39 @@ from .decoder import STEP_KEYS, TtDecoder, TtTurnHead
 from .encoder import TtEncoder
 from .layers import Build, policy
 
-__all__ = ["TtDiffusionPlanner"]
+__all__ = ["TtDiffusionPlanner", "agent_buckets", "needed_rows", "bucket_rows"]
+
+
+def agent_buckets() -> tuple:
+    """``COMPACT``: the agent buckets of ``AGENT_BUCKETS`` (multiples of 32 below 352, sorted); () when off."""
+    k = T.KNOBS.read()
+    if not k.COMPACT:
+        return ()
+    out = sorted({int(v) for v in str(k.AGENT_BUCKETS).replace(" ", "").split(",") if v})
+    bad = [v for v in out if v % T.TILE or not 0 < v < T.AGENTS]
+    if bad:
+        raise ValueError(f"DIFFUSION_PLANNER_AGENT_BUCKETS: {bad} are not multiples of {T.TILE} below {T.AGENTS}")
+    return tuple(out)
+
+
+def needed_rows(prepared: Any) -> int:
+    """1 + the last decoder row a plan reads or attends to: the ego (row 0), the valid self-attention keys
+    (``agent_valid``) and the emitted neighbours (``neighbor_rows`` + 1). The rows past it are masked keys whose
+    outputs are never read, so dropping them leaves the needed rows' values unchanged."""
+    valid = np.flatnonzero(np.asarray(prepared.decoder.agent_valid, bool))
+    last = int(valid.max()) if valid.size else 0
+    tok = np.flatnonzero(np.asarray(prepared.features.valid["neighbor"], bool))   # valid neighbour tokens
+    if tok.size:
+        last = max(last, int(tok.max()) + 1)
+    emitted = np.asarray(prepared.neighbor_rows)
+    if emitted.size:
+        last = max(last, int(emitted.max()) + 1)
+    return last + 1
+
+
+def bucket_rows(prepared: Any, buckets: Sequence[int]) -> int:
+    n = needed_rows(prepared)
+    return next((int(b) for b in sorted(buckets) if b >= n), T.AGENTS)
 
 
 class TtDiffusionPlanner:
@@ -57,15 +89,24 @@ class TtDiffusionPlanner:
                            attn_fp32_acc=attn_fp32_acc, attn_matmul=attn_matmul)
         p = weights.params
         self.tables = P.step_tables(p, steps)
-        self.encoder = TtEncoder(self.build, p)
-        self.decoder = TtDecoder(self.build, p, self.tables)
+        self.buckets = agent_buckets() if hasattr(device, "compute_with_storage_grid_size") else ()
+        self.encoder = TtEncoder(self.build, p, nb_rows=self.buckets if self.build.compact_enc else ())
+        self.decoder = TtDecoder(self.build, p, self.tables, rows=(T.AGENTS,) + self.buckets)
         self.turn = TtTurnHead(self.build, p)
         self.runner = TraceRunner(device, num_command_queues=num_command_queues, name="diffusion-planner")
         warm = I.warmup_inputs()
         for name, shape in I.INPUT_SPECS.items():
             dtype = "bfloat16" if name in I.BF16_INPUTS else "float32"
             self.runner.add_input(name, init=warm[name], dtype=dtype, layout=ttnn.TILE_LAYOUT)
+        self.cs_pad = {}                        # INPUT_TRIM: zero columns that widen cs to STATE_COLS, per row count
+        if I.CS_COLS < T.STATE_COLS:
+            for r in (T.AGENTS,) + self.buckets:
+                self.cs_pad[r] = self.build.upload(np.zeros((r, T.STATE_COLS - I.CS_COLS), np.float32), "float32")
+        self.trim = I.CS_COLS < T.STATE_COLS
+        self._zero_inputs: set = set()          # INPUT_TRIM: inputs whose device buffer holds +0.0 (our last upload)
         self.runner.add_variant("plan", self._plan)
+        for r in self.buckets:                  # COMPACT: one trace per agent bucket (the decoder on r rows)
+            self.runner.add_variant(f"plan_r{r}", lambda ctx, r=r: self._plan(ctx, r))
         self.debug = bool(debug)
         if self.debug:
             self.runner.add_input("dbg_x", init=np.zeros((1, 1, T.AGENTS, T.STATE_COLS), np.float32),
@@ -80,11 +121,22 @@ class TtDiffusionPlanner:
         self.build_ms = (time.perf_counter() - t0) * 1e3
 
     # ---- traced functions --------------------------------------------------------------------------------------
-    def _plan(self, ctx):
-        enc = self.encoder.forward(ctx)
+    def _plan(self, ctx, rows: int = T.AGENTS):
+        """The whole plan; ``rows`` < 352 (``COMPACT``): the decoder on the first ``rows`` agents (the inputs'
+        leading rows, sliced in the trace)."""
+        import ttnn
+
+        enc = self.encoder.forward(ctx, nb=rows if rows != T.AGENTS else None)
         kv = self.decoder.cross_kv(enc)
-        self_mask = A.expand_key_bias(ctx["agent_key_row"], T.AGENTS)
-        final, ego = self.decoder.solve(ctx["y0"], ctx["cs"], kv, self_mask)
+        y0, cs, key = ctx["y0"], ctx["cs"], ctx["agent_key_row"]
+        if rows != T.AGENTS:
+            y0 = ttnn.slice(y0, [0, 0, 0, 0], [1, 1, rows, T.STATE_COLS])
+            cs = ttnn.slice(cs, [0, 0, 0, 0], [1, 1, rows, int(cs.shape[-1])])
+            key = ttnn.slice(key, [0, 0, 0, 0], [1, 1, 1, rows])
+        if int(cs.shape[-1]) < T.STATE_COLS:     # INPUT_TRIM: widen the one-tile cs (tile-aligned concat)
+            cs = ttnn.concat([cs, self.cs_pad[rows]], dim=-1)
+        self_mask = A.expand_key_bias(key, rows, memory_config=self.build.attn_mem())
+        final, ego = self.decoder.solve(y0, cs, kv, self_mask)
         logit = self.turn(final, enc)
         return pack_outputs({"final_x0": final, "logit": logit, "ego_steps": ego})
 
@@ -108,15 +160,54 @@ class TtDiffusionPlanner:
     def capture(self) -> None:
         self.runner.capture()
 
+    def filter_inputs(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
+        """``INPUT_TRIM``: drop the inputs that are all +0.0 while their device buffer already holds the zeros of this
+        model's previous upload (a value test, not a cache of the previous request: any non-zero array is always
+        uploaded). Every input passed on is recorded as uploaded."""
+        if not self.trim:
+            return inputs
+        out = {}
+        for name, a in inputs.items():
+            arr = a if isinstance(a, np.ndarray) else None
+            zero = arr is not None and arr.dtype == np.float32 and not arr.view(np.uint32).any()
+            if zero and name in self._zero_inputs:
+                continue
+            out[name] = a
+            if zero:
+                self._zero_inputs.add(name)
+            else:
+                self._zero_inputs.discard(name)
+        return out
+
     def _run(self, variant: str, inputs: Dict[str, Any], eager: bool):
+        inputs = self.filter_inputs(inputs)
         return self.runner.run_eager(variant, inputs=inputs) if eager else self.runner(variant, inputs=inputs)
 
-    def forward(self, prepared: Any, *, ego_steps: bool = True, eager: bool = False) -> Dict[str, Any]:
+    def rows_for(self, prepared: Any) -> int:
+        """The decoder rows a plan needs (``COMPACT``): the smallest agent bucket holding the ego, every valid
+        self-attention key and every emitted neighbour row; 352 without buckets or when none is large enough."""
+        return bucket_rows(prepared, self.buckets)
+
+    def variant_for(self, prepared: Any) -> str:
+        r = self.rows_for(prepared)
+        return "plan" if r == T.AGENTS else f"plan_r{r}"
+
+    def unpack(self, out: Dict[str, Any]) -> np.ndarray:
+        """The readback's ``final_x0`` (``[R, 324]``) -> ``[321, 324]``; the rows past R (not computed: masked keys,
+        never emitted) are zero."""
+        f = np.asarray(out["final_x0"], np.float32).reshape(-1, T.STATE_COLS)
+        if f.shape[0] < T.AGENTS:
+            f = np.concatenate([f, np.zeros((T.AGENTS - f.shape[0], T.STATE_COLS), np.float32)], 0)
+        return f[:C.MAX_NUM_AGENTS]
+
+    def forward(self, prepared: Any, *, ego_steps: bool = True, eager: bool = False,
+                variant: Optional[str] = None) -> Dict[str, Any]:
         """One plan: upload, replay, read (``eager=True``: the same graph without the trace, for bring-up and
         replay-vs-eager checks). ``final_x0`` ``[321, 81, 4]`` (normalised, prefix-constrained), ``logit`` ``[5]``,
-        ``denoising_steps``: the 11 iterates' ego rows as ``[1, 81, 4]`` arrays."""
-        out = self._run("plan", I.plan_inputs(prepared), eager)
-        final = out["final_x0"].reshape(T.AGENTS, T.STATE_COLS)[:C.MAX_NUM_AGENTS]
+        ``denoising_steps``: the 11 iterates' ego rows as ``[1, 81, 4]`` arrays. ``variant``: force a plan variant
+        (default :meth:`variant_for`)."""
+        out = self._run(variant or self.variant_for(prepared), I.plan_inputs(prepared), eager)
+        final = self.unpack(out)
         res = {"final_x0": final.reshape(C.MAX_NUM_AGENTS, C.OUTPUT_T + 1, C.POSE_DIM).astype(np.float32),
                "logit": out["logit"].reshape(-1)[:C.TURN_INDICATOR_OUTPUT_DIM].astype(np.float32)}
         if ego_steps:
@@ -180,7 +271,7 @@ class TtDiffusionPlanner:
                 "precision": self.build.policy.describe(),
                 "options": self.build.options(),
                 "uploaded_mb": round(self.build.uploaded_bytes / 2 ** 20, 2), "build_ms": round(self.build_ms, 1),
-                "tokens": T.TOKENS, "agents": T.AGENTS, "nfe": self.tables.nfe}
+                "tokens": T.TOKENS, "agents": T.AGENTS, "agent_buckets": list(self.buckets), "nfe": self.tables.nfe}
 
     def release(self) -> None:
         self.runner.release()

code/tt_diffusion_planner/tt/smask_kernel.py

@@ -0,0 +1,98 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Attention score scale + additive mask as one ``ttnn.generic_op`` (``ATTN_SMASK``, OPT round 2 item 4b).
+
+``scale_mask(s, scale, mask)``: fp32 TILE scores ``[1, H, Sq, Sk]`` and a ``[1, 1, Sq, Sk]`` mask (bf16 or fp32,
+broadcast over the heads) -> ``s * scale + mask`` fp32, the two stock ``binary_ng`` programs of
+``tt/attention.py`` (``ttnn.multiply(s, scale)``, ``ttnn.add(., mask)``) in one pass over the scores
+(``kernels/smask_*.cpp``): bit-identical, half the DRAM traffic and one program less per attention. A bf16 mask
+sends the stock fp32 + bf16 add down binary_ng's FPU path, which truncates the scores to TF32 (measured); the
+kernel reproduces that truncation (the mask holds 0 / -inf only, so the add itself is exact).
+The mask-plane tiles are split in contiguous ranges over the grid; each core streams the H score tiles of its mask
+tiles.
+"""
+from __future__ import annotations
+
+import os
+import struct
+from typing import Any
+
+__all__ = ["scale_mask", "supported"]
+
+_KDIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "kernels")
+TB = 4096
+N_RT = 2
+
+
+def supported(s: Any, mask: Any) -> bool:
+    import ttnn
+
+    try:
+        ss, ms = list(s.padded_shape), list(mask.padded_shape)
+        return (hasattr(ttnn, "generic_op") and s.dtype == ttnn.float32 and s.layout == ttnn.TILE_LAYOUT
+                and mask.layout == ttnn.TILE_LAYOUT and mask.dtype in (ttnn.bfloat16, ttnn.float32)
+                and not s.is_sharded() and not mask.is_sharded() and len(ss) == 4 and len(ms) == 4
+                and ss[0] == 1 and ss[1] % 2 == 0 and ms[0] == 1 and ms[1] == 1 and ss[2:] == ms[2:])
+    except Exception:  # noqa: BLE001 - the host fake ttnn
+        return False
+
+
+def scale_mask(s: Any, scale: float, mask: Any):
+    import ttnn
+
+    dev = s.device()
+    _, H, Sq, Sk = list(s.padded_shape)
+    plane = (Sq // 32) * (Sk // 32)
+    out = ttnn.allocate_tensor_on_device(s.shape, ttnn.float32, ttnn.TILE_LAYOUT, dev, ttnn.DRAM_MEMORY_CONFIG)
+    g = dev.compute_with_storage_grid_size()
+    n = min(plane, g.x * g.y)
+    cs = [(i % g.x, i // g.x) for i in range(n)]
+    full, rem = divmod(n, g.x)
+    rs = []
+    if full:
+        rs.append(ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(g.x - 1, full - 1)))
+    if rem:
+        rs.append(ttnn.CoreRange(ttnn.CoreCoord(0, full), ttnn.CoreCoord(rem - 1, full)))
+    crs = ttnn.CoreRangeSet(set(rs))
+    base, extra = divmod(plane, n)
+    rd, wr, cp = ttnn.RuntimeArgs(), ttnn.RuntimeArgs(), ttnn.RuntimeArgs()
+    m0 = 0
+    for i, (cx, cy) in enumerate(cs):
+        k = base + (1 if i < extra else 0)
+        rd[cx][cy] = [m0, k]
+        wr[cx][cy] = [m0, k]
+        cp[cx][cy] = [k, 0]
+        m0 += k
+    mbf = mask.dtype == ttnn.bfloat16
+    tbm = 2048 if mbf else 4096
+
+    def cb(idx, pages, dt, tb):
+        return ttnn.CBDescriptor(total_size=pages * tb, core_ranges=crs, format_descriptors=[
+            ttnn.CBFormatDescriptor(buffer_index=idx, data_format=dt, page_size=tb)])
+
+    cbs = [cb(0, 2 * H, ttnn.float32, TB), cb(1, 2, mask.dtype, tbm), cb(2, 1, ttnn.float32, TB),
+           cb(16, 4, ttnn.float32, TB)]
+    um = [ttnn.UnpackToDestMode.Default] * 64
+    um[0] = um[2] = ttnn.UnpackToDestMode.UnpackToDestFp32
+    if not mbf:
+        um[1] = ttnn.UnpackToDestMode.UnpackToDestFp32
+    ccfg = ttnn.ComputeConfigDescriptor(math_fidelity=ttnn.MathFidelity.HiFi4, fp32_dest_acc_en=True,
+                                        math_approx_mode=False)
+    ccfg.unpack_to_dest_mode = um
+    bits = int.from_bytes(struct.pack("<f", float(scale)), "little")
+    SRC = ttnn.KernelDescriptor.SourceType.FILE_PATH
+
+    def acc(t):
+        return list(ttnn.TensorAccessorArgs(t).get_compile_time_args())
+
+    def kd(name, ct, rt, common, config):
+        return ttnn.KernelDescriptor(kernel_source=os.path.join(_KDIR, name), source_type=SRC, core_ranges=crs,
+                                     compile_time_args=ct, defines=[], runtime_args=rt, common_runtime_args=common,
+                                     config=config)
+
+    ks = [kd("smask_reader.cpp", [H, plane, bits, N_RT, tbm] + acc(s) + acc(mask), rd,
+             [s.buffer_address(), mask.buffer_address()], ttnn.ReaderConfigDescriptor()),
+          kd("smask_writer.cpp", [H, plane, N_RT] + acc(out), wr, [out.buffer_address()],
+             ttnn.WriterConfigDescriptor()),
+          kd("smask_compute.cpp", [H, N_RT, int(mbf)], cp, [], ccfg)]
+    ttnn.generic_op([s, mask, out], ttnn.ProgramDescriptor(kernels=ks, semaphores=[], cbs=cbs))
+    return out