code/alpamayo_tt/EXPERT_NOTES.md

# Action expert on one Blackhole chip — implementation notes

Code: `alpamayo_tt/expert.py` (device graph, trace, Euler loop), `alpamayo_tt/action_proj.py` (host action_in_proj /
action_out_proj / rope tables), `alpamayo_tt/expert_ref.py` (torch oracle of one layer with HF semantics),
`tests/test_expert.py`.

## What runs where

| Piece | Where | Notes |
| --- | --- | --- |
| `action_in_proj` (Fourier 60 -> MLP 512 -> 1536 -> LayerNorm) | host, torch fp32 | 64x1536 output uploaded as bf16 (196 KB) per step |
| 64 decoder layers + final RMSNorm | device, one trace | inputs: `x_dev` (emb), `cos_dev`/`sin_dev`, `mask_dev`, the 64 VLM KV caches |
| `action_out_proj` (1536 -> 2) | host, torch fp32 | reads the final-normed hidden (196 KB) |
| Euler loop (10 steps) | host | `x <- x + dt * v`, `t = linspace(0, 1, 11)[:-1]` |

The reference model runs its `action_in_proj` in bf16 **with bf16-rounded Fourier frequencies** (`freqs` is an
`nn.Module` buffer, so `from_pretrained(dtype=bf16)` rounds 59.948 -> 60.0 etc.). Using exact frequencies gives a
systematic 0.07 max-abs error vs the golden `emb`; with the rounded frequencies (default) the error is 0.009 and
PCC 0.999995 (the remainder is the reference's bf16 MLP arithmetic).

## Layouts (batch 1, 64 expert tokens)

| Tensor | Shape / dtype | Notes |
| --- | --- | --- |
| residual stream `x` | `[1, 1, 64, 1536]` bf16 TILE DRAM | 2 tile rows |
| `wqkv` | `[1536, 4096]` bfp8 (or bf16) | HF `cat(q, k, v).T`; q heads 0..15, then k heads 0..7, then v heads 0..7 |
| `wo` | `[2048, 1536]` | `w1`/`w3`: `[1536, 6144]`, `w2`: `[6144, 1536]`, all HF `.T` |
| RMSNorm gammas | `[1, 1, 1, dim]` bf16 | `ttnn.rms_norm` requires gamma width == input width in this tt-metal |
| q / k / v heads | `[1, 16, 64, 128]`, `[1, 8, 64, 128]` | `nlp_create_qkv_heads(transpose_k_heads=False)` |
| cos / sin | `[1, 1, 64, 128]` bf16 | HF rotate_half tables for positions `i + rope_delta + S`, theta 5e6 |
| VLM KV cache (input) | `[1, 8, Smax, 128]` bf16 or bfp8 TILE DRAM per layer | HF head order, rows `[0, S)` valid |
| mask | `[1, 1, 128, Smax]` bfp4 (decode mode) / `[1, 1, 64 or 128, Smax]` bf16 (standard) | additive, 0 / -1e9 |

Per layer (19 ttnn ops): rms_norm -> linear(wqkv) -> nlp_create_qkv_heads -> rms_norm(q), rms_norm(k) (per-head,
eps 1e-6) -> rotary_embedding_hf(q), (k) -> fill_cache(K), fill_cache(V) -> attention -> nlp_concat_heads ->
linear(wo) -> add -> rms_norm -> linear(w1, silu), linear(w3) -> mul -> linear(w2) -> add.
`TTExpert.layer_fn` is the swap point for a fused kernel (signature `(layer_idx, x, K_cache, V_cache) -> x`).

## Attention over the VLM cache: masking approach

The expert's layer *i* attends non-causally to all `S` cached VLM positions of layer *i* plus its own 64 tokens.
Instead of slicing / concatenating (every copy of the ~19 MB cache costs ~40-75 us), the own roped K and V are
written **into the cache itself**, rows `[Smax-64, Smax)`, with `ttnn.fill_cache(cache, k, batch_idx=0,
update_idx=Smax-64)` (tile-aligned offset write, ~4 us). One attention op then reads the whole cache and an
additive mask hides rows `[S, Smax-64)` (garbage / unused). Attention is permutation invariant over keys, so
placing the own tokens at the tail is exact. Consequences:

* the traced graph does not depend on `S`; per trajectory only `cos/sin` (positions) and the mask are rewritten
  from host (`set_context`, ~1 ms for the bfp4 mask of 128 x 8192);
* **the last 64 rows of every VLM KV cache are expert scratch**: the text model must never store tokens there
  (`S <= Smax - 64`, asserted). With `max_seq_len = 8192` and prompts of ~4.6K + <= 256 CoC tokens this is far away;
* rows `[S, Smax-64)` may contain anything (tests fill them with 7.0).

Two attention kernels are implemented (`attn_mode`):

* **`decode` (default)** — `ttnn.transformer.scaled_dot_product_attention_decode`, the flash-decode kernel that
  splits the K/V sequence across cores and reduces. The cache `[1, 8, Smax, 128]` is *viewed* (zero copy) as
  `[8, 1, Smax, 128]` = 8 users with one kv head each; the queries `[1, 16, 64, 128]` are viewed as
  `[1, 8, 128, 128]` = per user 128 "heads" (the 2 GQA q heads x 64 tokens; tile order is identical, so both views
  are free). Mask `[1, 1, 128, Smax]` bfp4, broadcast over users. `k_chunk_size=128`, 16 cores per user.
  1024 "heads" in a single user would overflow L1 (tested), the 8-user view does not.
* **`standard`** — `ttnn.transformer.scaled_dot_product_attention(is_causal=False)` with GQA q `[1, 16, 64, 128]`
  (or the 2 q heads of a group folded onto rows, `fold_heads=True`). Every (head, q-chunk) work item re-reads the
  entire K/V, so it is either per-core latency bound (few chunks) or DRAM bound (many chunks): ~520 us at
  Smax = 8192 regardless of chunking.

Measured attention time per layer (bf16 cache, S = 4700, trace-timed):

| Kernel | Smax = 8192 | Smax = 5120 |
| --- | --- | --- |
| decode, mask bfp4, k_chunk 128 | 172 us | 149 us |
| decode, bfp8 cache | 143 us | 126 us |
| standard, GQA q, q_chunk 32, mask bfp4 | 516 us | 322 us |
| standard, folded heads, q_chunk 128 | 708 us | 460 us |

## Matmul program configs (M = 64 rows, bfp8 weights, trace-timed)

1D multicast configs (`MatmulMultiCoreReuseMultiCast1DProgramConfig`, in0 broadcast, N split over the 11x10 grid,
`per_core_M=2`) vs the ttnn default heuristics:

| Matmul | shape | default | tuned | tuned config |
| --- | --- | --- | --- | --- |
| wqkv | 1536 x 4096 | 38 us | 25 us | in0_block_w 8, per_core_N 2 |
| wo | 2048 x 1536 | 43 us | 19 us | in0_block_w 8, per_core_N 1 |
| w1 (silu fused), w3 | 1536 x 6144 | 45 us each | 35-38 us each | in0_block_w 8, per_core_N 2 |
| w2 | 6144 x 1536 | 113 us | 41 us | in0_block_w 48, per_core_N 1 |

Total weight traffic per layer 40.5 MB bfp8 -> ~158 us = ~255 GB/s average (DRAM peak ~512 GB/s).

## Trace safety (fixed a second-run corruption seen in the integrated pipeline)

A trace bakes in the addresses of the buffers its ops use; its intermediates are placed in memory that is free at
capture time. A device buffer allocated *after* some trace was captured can therefore alias that trace's
intermediates and is overwritten whenever the trace replays (verified directly on this tt-metal, see
`test_sample_repeat_with_foreign_trace`). In the pipeline the text decode trace is captured before the expert first
runs, so every expert input buffer (`x_dev`, `cos_dev`, `sin_dev`, `mask_dev` sized `cfg.max_seq_len`) is allocated in
`TTExpert.__init__`, and `sample()` re-uploads cos/sin/mask at its start (~1 ms) instead of trusting the cached
context across trajectories. The expert trace output is read back immediately after each replay. The expert trace
itself is keyed on the KV cache and mask buffer addresses and re-captured if they change. Rule for integrators:
allocate all persistent buffers (weights, caches, model inputs) before the first trace capture in the process.

## Accuracy (`tests/test_expert.py`, golden = ref_out/030c760c_5100000_ref.pt, S = 4592)

| Test | Result |
| --- | --- |
| torch layer oracle vs transformers `Qwen3VLTextModel` (random weights) | PCC 1.0000000, max abs 1.4e-6 |
| `action_in_proj` vs golden `emb`, steps 0..9 | max abs 9.0e-3, PCC 0.999995 (exact freqs: 6.9e-2 / 0.99953) |
| layer 0, device bfp8 / bf16 vs torch oracle (golden `exp.layer0.in`, golden cache rows + 13 synthetic) | PCC 0.99993 / 0.99994 |
| layer 31 | PCC 0.999997 / 0.999997 |
| layer 63 | PCC 0.999998 / 0.999998 |
| same, vs golden `exp.layer{0,31,63}.out` (13 of 4592 cache rows synthetic) | PCC 0.99993 / 1.00000 / 0.99994 |
| traced 64-layer step vs eager | bitwise identical |
| `sample()` repeated with a foreign trace replayed in between | bitwise identical |

bf16 weights are not measurably better than bfp8 here (bf16 activations / bf16 cache dominate), so bfp8 is the default
(2.55 GB for the 64 layers).

## Timing (chip 2, S = 4700, Smax = 8192, bf16 cache, bfp8 weights, trace-timed)

| | |
| --- | --- |
| traced 64-layer step (device, incl. 196 KB emb upload + hidden download) | **32.5 ms** (min 32.3) |
| full `step()` incl. host `action_in_proj` / `action_out_proj` | 33.5 ms |
| 10-step `sample()` | ~335 ms |
| one layer, back-to-back trace | 497 us |

Per-op device time of one layer (us):

| op | us | | op | us |
| --- | --- | --- | --- | --- |
| rms_norm | 17.3 | | concat_heads | 27.7 |
| wqkv (1536x4096 bfp8) | 25.6 | | wo (2048x1536) | 18.6 |
| nlp_create_qkv_heads | 52.8 | | add | 6.5 |
| q_norm / k_norm | 6.5 / 5.8 | | rms_norm2 | 17.4 |
| rope q / k | 11.0 / 9.1 | | w1+silu / w3 (1536x6144) | 40.2 / 37.3 |
| fill_cache K / V | 3.9 / 3.8 | | mul | 10.9 |
| attention (decode kernel) | 175.3 | | w2 (6144x1536) | 44.6 |
| | | | add2 | 6.7 |
| **sum** | **521** | | | |

Where the time goes: attention 34%, weight matmuls 32% (40.5 MB/layer at ~245 GB/s), head split + concat 15%, the
other 9 small ops 19%. The goal of < 30 ms/step as an op graph is missed by ~8%; what was tried and rejected:
batched per-head matmuls to skip `nlp_create_qkv_heads` (382 us), batched o_proj + sum (62 us vs 46 us), the
DRAM-sharded decode matmul (requires 32 rows), SDPA LoFi / exp-approx / other k-chunks (within 3%), 1024-"head"
decode SDPA (L1 overflow). Remaining op-graph options: a bfp8 KV cache (attention 143 us, -2 ms/step), a 5120-row
cache (-1.5 ms/step), fusing rms_norm+residual. The fused-layer kernel (megakernel stream) replaces the whole layer
via `TTExpert.layer_fn`.

## Known gaps

* The golden dump only holds the VLM cache for the prefill (4579 rows); the 13 generated CoC positions are
  missing, so layer tests compare against the torch oracle with synthetic rows for the generated span and only
  report (not assert) the PCC vs the golden layer output (0.99993-1.0 anyway). End-to-end parity (per-step `v`,
  final action) is covered by the integrated pipeline (teacher-forced trajectories 0.02-0.31 m from the reference).
* `S` must satisfy `S <= Smax - 64` (tail rows are scratch). If the integrator uses a smaller `max_seq_len`
  (5120 fits the trajectory-profile prompts) the attention gets ~15% faster.
* The trace bakes in the buffer addresses of the 64 KV caches and of the mask; it is re-captured automatically
  when they change (`_trace_key`). See "Trace safety" above.
* `nlp_create_qkv_heads` (53 us) and `nlp_concat_heads` (28 us) are slow for 64-row inputs (few cores busy); no
  cheaper ttnn formulation was found — a fused kernel is the way out.
* Only batch 1.