gqa: ADR-0060/0062/0063/0064 unified GQA kernels + CPU cost model

Land the new GQA fused-attention kernels (ADR-0060) for prefill/decode
across long and short context, the TL discipline primitives they depend
on (ADR-0062 lazy load, ADR-0063 scratch_scope + copy_to), and the
per-op-type CPU issue cost model (ADR-0064). Remove the pre-ADR-0060
mesh-attention baseline now that the unified kernels supersede it.

ADR-0060 (long context)
- _gqa_decode.py: M-fold + 2-level chain reduce-to-root (Level-2
  intra-CUBE row-then-col + Level-1 inter-CUBE) — root-only output.
- _gqa_prefill.py: head-parallel + Ring KV rotation around C CUBEs,
  online-softmax merge per ring step, per-CUBE distributed output.
- Each merge stage wraps in scratch_scope() and persists running
  (m, l, O) via copy_to() to lift the 1 MiB scratch ceiling.

ADR-0060 §B.split.2 (short context, kv_per_cube in {1,2,4,8})
- _gqa_decode_short.py / _gqa_prefill_short.py: no cube-SP; each CUBE
  owns whole KV heads; PE-parallel heads with intra-group chain
  reduce. Prefill has no Ring KV (each head fully resident).

ADR-0062 (lazy tl.load): future-bearing TensorHandle, auto-wait at
first consuming op (dot/MATH/store/send/copy_to/composite).

ADR-0063 (tl.scratch_scope + tl.copy_to): scoped per-tile arena with
copy_to writeback primitive for persistent running state.

ADR-0064 (CPU issue cost model)
- common/cpu_issue_cost.py: per-op-type table (composite=40 ns,
  primitives=5 ns); ratios are load-bearing per D1.
- TLContext: issue_cost_table param; _emit_dispatch_overhead(kind)
  consults table with dispatch_cycles fallback (ADR-0046 §D6
  back-compat).
- Live PE_CPU paths (greenlet + legacy) construct TLContext with
  DEFAULT_CPU_ISSUE_COST so saturation lever (ADR-0060 §1) is
  measurable end-to-end.

P7 headline bench: milestone-gqa-headline writes per-panel
op_log_summary to 1H_milestone_output/gqa_headline/sweep.json. No
figure renderers yet (deferred).

Removals (pre-ADR-0060 baseline now superseded):
- benches: _attention_mesh_kv.py, _attention_mesh_mlo.py,
  _attention_mesh_mlo_2d.py, milestone_gqa_llama70b.py
- tests: test_attention_*, test_mesh_*, test_milestone_gqa_llama70b
- topology: llama70b_4sip.yaml (only consumer was the deleted diag)
- artifacts: 1H_milestone_output/gqa/ (sweep.json + 5 PNGs)
- tests/gqa/ plot helper + test (broken on Windows Tcl/Tkinter)
- ADR-0060/0061 references to deleted file paths cleaned up
  (EN + KO kept in sync).

Tests: 124/124 focused regression green (attention + Phase E + TL
discipline + triton_emu + pe_components). Full regression: 764 pass,
2 pre-existing test_bench_registry failures (stale EXPECTED_NAMES
across multiple benches, not introduced here).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

docs/adr-proposed/ADR-0060-algo-gqa-fused-attention-ahbm-ko.md

@@ -8,11 +8,9 @@ Proposed
 **Decision drivers:** agentic workload → 낮은 batch, 긴 context;
 KV-load-bound decode; long-context prefill용 sequence-parallel (Ring KV).
 
-**Supersedes / extends:** 기존 mesh-native 어텐션 커널
+**Supersedes / extends:** 이전 mesh-native 어텐션 커널과 그 `milestone-gqa-llama70b`
-`_attention_mesh_kv`(prefill)와 `_attention_mesh_mlo`(decode), 그리고
+eval bench(본 ADR의 커널이 도입되면서 제거됨). *§A. 기존 kernbench 작업과의 관계* 참조 —
-`milestone-gqa-llama70b` eval bench. *§A. 기존 kernbench 작업과의 관계* 참조 —
+그 코드가 본 ADR이 진짜 GQA·causal·long-context 커널로 업그레이드한 baseline이다.
-그 코드가 본 ADR이 진짜 GQA·causal·long-context 커널로 업그레이드하는
-baseline이다.
 
 **Supporting ADRs** (efficiency / scale enabler — *GQA blocker 아님*; §8 정정
 참조): **ADR-0063** `tl.scratch_scope`(per-tile scratch 재활용 — 현실적 context
@@ -160,18 +158,16 @@ def gqa_prefill_sp(q_ptr, k_ptr, v_ptr, o_ptr, T_q, S_kv_local, d, C, scale, q_b
 
 ## A. 기존 kernbench 작업과의 관계 (먼저 읽을 것)
 
-kernbench는 **오늘날 이미 IPCQ 상에서 online-softmax `(m, ℓ, O)` 머지로
+kernbench는 본 ADR 이전에 IPCQ 상에서 online-softmax `(m, ℓ, O)` 머지로
-FlashAttention을 돌린다.** 두 커널이 존재한다:
+FlashAttention을 돌리는 두 mesh 커널이 있었다(현재 제거됨):
 
-| File | Role | Mechanism |
+| Role | Mechanism |
-|---|---|---|
+|---|---|
-| `src/kernbench/benches/_attention_mesh_kv.py` | prefill (Ring K/V) | per-rank partial attention, bidirectional K/V fan-out, online-softmax fold |
+| prefill (Ring K/V) | per-rank partial attention, bidirectional K/V fan-out, online-softmax fold |
-| `src/kernbench/benches/_attention_mesh_mlo.py` | decode (split-KV) | per-rank one-shot partial attention, bidirectional `(m,ℓ,O)` fan-out, log-sum-exp merge |
+| decode (split-KV)  | per-rank one-shot partial attention, bidirectional `(m,ℓ,O)` fan-out, log-sum-exp merge |
 
 둘 다 `milestone-gqa-llama70b`(4 패널:
-`{single,multi}_user × {prefill,decode}`,
+`{single,multi}_user × {prefill,decode}`)이 구동했다.
-`src/kernbench/benches/milestone_gqa_llama70b.py`)이 구동하며
-`tests/attention/test_milestone_gqa_llama70b.py`에서 테스트된다.
 
 **이들은 greenlet `tl` API로 작성되었다:** `tl.load`, `tl.dot`,
 `tl.softmax`/`tl.max`/`tl.sum`/`tl.exp`, `tl.send`/`tl.recv`, 그리고
@@ -183,8 +179,8 @@ FlashAttention을 돌린다.** 두 커널이 존재한다:
 
 **baseline의 의도된 세 한계** — 효율적 GQA 커널이 정확히 들어내야 하는 것:
 
-1. **GQA 재사용 없음.** `h_q == h_kv == 1`
+1. **GQA 재사용 없음.** baseline은 `h_q == h_kv == 1`로 제한되어 있었다.
-   (`test_milestone_gqa_llama70b.py:137-142`). 테스트는 이를 *broadcast view*에
+   해당 baseline 테스트는 이를 *broadcast view*에
    대한 MemoryStore byte-conservation 실패로 돌리지만, 그 실패는 baseline의
    **head-packing 핵**의 속성이다(`_view(K, (h_q·d, S_kv))`는 모든 head를 하나의
    matmul 차원에 뭉치고 `h_q == h_kv`일 때만 byte를 보존). 올바른 수정은
@@ -196,7 +192,7 @@ FlashAttention을 돌린다.** 두 커널이 존재한다:
    필요 → **2-level reduce-to-root**(intra-CUBE tree + intra-CUBE-Group
    center-mesh, §4)는 `⌈log₂ P⌉` + center-mesh-over-`C` 단계.
 3. **검증 스케일만.** `S = 16`인 이유는 1 MiB scratch bump allocator가 per-tile
-   임시값을 누수하고(`test_milestone_gqa_llama70b.py:123-148`) causal masking /
+   임시값을 누수하고(baseline 테스트는 그래서 S=16으로 제한되어 있었음) causal masking /
    tiling이 없기 때문 → **ADR-0063**(재활용) + §5(tiling, causal skip) + composite
    K/V 스트리밍(§3) + **ADR-0062**(lazy load 오버랩)으로 해결.
 
@@ -563,8 +559,7 @@ reduce **안 함** — KV를 회전(Ring, §5.5). decode에 reduce를 택한 이
 `O = [G, d]`가 작아 `(m,ℓ,O)` 이동이 상주 KV 이동보다 싸기 때문.
 
 tile sweep 후 각 rank는 자기 `1/(C·P)` shard에 대해 `(m_i, ℓ_i, O_i)`(비정규화)를
-가진다. 결합/교환 가능한 log-sum-exp 머지로 결합(baseline의 fold와 동일 수학,
+가진다. 결합/교환 가능한 log-sum-exp 머지로 결합(baseline의 fold와 동일 수학):
-`_attention_mesh_mlo.py:117-122`):
 
 ```python
 def merge(m_a, l_a, O_a, m_b, l_b, O_b):
@@ -688,9 +683,8 @@ def gqa_decode_sp(q_ptr, k_ptr, v_ptr, o_ptr, counter, start_pe, start_cube,
   축은 세어지지 *않음* — §8 참조).
 - `S_rank`가 scratch에 맞으면(작은/중간 context) 이것은 **one-shot** partial
   attention으로 축약(Q·Kᵀ용 composite 하나, softmax 하나, P·V용 composite 하나) —
-  바로 baseline `_attention_mesh_mlo`의 `_partial_attention`, 단지 GQA-batched에
+  바로 baseline의 `_partial_attention` 구조, 단지 GQA-batched에 composite 경로.
-  composite 경로. Tiling(§3)은 `S_rank`가 scratch scope의 tile 예산을 초과할 때만
+  Tiling(§3)은 `S_rank`가 scratch scope의 tile 예산을 초과할 때만 발동.
-  발동.
 
 ### 5.3 DECODE, with SP / KV-parallel (`C × P` rank)
 
@@ -741,7 +735,7 @@ KV를 필요로 하므로.
   **`(m,ℓ,O)` reduce 없음** — 각 CUBE가 자기 head의 행을 정규화·기록.
 - **Ring KV:** `C` KV slice가 CUBE ring을 **회전**; 각 CUBE가 들어오는 블록을 자기
   head의 실행 중 `(m,ℓ,O)`에 fold(online-softmax, ring step 가로질러 Python 핸들로
-  carry, `_attention_mesh_kv`가 오늘 하듯). `C` step 후 모든 head가 전 KV를 봄.
+  carry). `C` step 후 모든 head가 전 KV를 봄.
   **GQA 재사용은 회전에서** — slice `j`가 전 `C` CUBE를 방문하며 전 `G` head를 서비스.
 - IPCQ가 다음 step의 KV 수신을 현재 step의 compute와 `tl.recv_async`/`tl.wait`로
   오버랩(`tl_context.py:543-560` — 이미 존재); 수신 버퍼는 ping-pong(persistent
@@ -751,7 +745,7 @@ KV를 필요로 하므로.
 - **CUBE 내(P PE):** head의 query 행 `[T_q, d]` 및/또는 현재 KV 블록을 `P` PE에
   타일(여기엔 decode와 달리 query 축이 존재); 상세는 §B.
 
-baseline `_attention_mesh_kv`가 이미 ring fold를 구현; 본 ADR은 GQA 재사용,
+(현재 제거된) baseline이 이미 ring fold를 구현; 본 ADR은 GQA 재사용,
 head-parallel 배치, causal step-skip, composite-hybrid inner tile(§3)을 추가.
 
 ### 5.6 Decode CPU-pipelining 변형 (opt1 / opt3 / opt2)
@@ -1060,7 +1054,7 @@ exercise하려 존재하고, `tl.trans`는 reshape-not-transpose이며, `bf16`
     가정. `C ≠ G`면 매핑 재검토 필요(CUBE당 다중 head, 또는 head가 부분 ring에 걸침).
     **권고:** headline 스케일에서 prefill 커널은 `C = G` 고정; `C ≠ G`는 별도 연구.
 
-5. **`_attention_mesh_mlo_2d`(현재 impl)와 정합.** 원격 impl의 2D 커널은 **Q
+5. **이전 `_attention_mesh_mlo_2d` impl(현재 제거됨)과 정합.** 그 2D 커널은 **Q
     replicated**로 cube에 걸친 **AllReduce** — 즉 decode-reduce 계열이나 reduce-to-root
     아닌 all-reduce(broadcast-back), 그리고 prefill head-parallel ring은 아직 없음.
     **권고:** (a) 그 2D AllReduce → reduce-to-root(broadcast-back 제거)로 decode 커널;

docs/adr-proposed/ADR-0060-algo-gqa-fused-attention-ahbm.md

@@ -8,11 +8,11 @@ Proposed
 **Decision drivers:** agentic workload → low batch, long context;
 KV-load-bound decode; sequence-parallel (Ring KV) for long-context prefill.
 
-**Supersedes / extends:** the existing mesh-native attention kernels
+**Supersedes / extends:** pre-existing mesh-native attention kernels and
-`_attention_mesh_kv` (prefill) and `_attention_mesh_mlo` (decode) and the
+their `milestone-gqa-llama70b` eval bench (removed when this ADR's
-`milestone-gqa-llama70b` eval bench. See *§A. Relationship to existing
+kernels landed). See *§A. Relationship to existing kernbench work* — that
-kernbench work* — that code is the baseline this ADR upgrades to a real
+code was the baseline this ADR upgrades to a real GQA, causal,
-GQA, causal, long-context kernel.
+long-context kernel.
 
 **Supporting ADRs** (efficiency / scale enablers — *not* GQA blockers;
 see §8 correction): **ADR-0063** `tl.scratch_scope` (per-tile scratch
@@ -168,18 +168,16 @@ def gqa_prefill_sp(q_ptr, k_ptr, v_ptr, o_ptr, T_q, S_kv_local, d, C, scale, q_b
 
 ## A. Relationship to existing kernbench work (read first)
 
-kernbench **already runs FlashAttention with an online-softmax `(m, ℓ, O)`
+kernbench previously ran FlashAttention with an online-softmax `(m, ℓ, O)`
-merge over IPCQ today.** Two kernels exist:
+merge over IPCQ via two pre-ADR-0060 mesh kernels (now removed):
 
-| File | Role | Mechanism |
+| Role | Mechanism |
-|---|---|---|
+|---|---|
-| `src/kernbench/benches/_attention_mesh_kv.py` | prefill (Ring K/V) | per-rank partial attention, bidirectional K/V fan-out, online-softmax fold |
+| prefill (Ring K/V) | per-rank partial attention, bidirectional K/V fan-out, online-softmax fold |
-| `src/kernbench/benches/_attention_mesh_mlo.py` | decode (split-KV) | per-rank one-shot partial attention, bidirectional `(m,ℓ,O)` fan-out, log-sum-exp merge |
+| decode (split-KV)  | per-rank one-shot partial attention, bidirectional `(m,ℓ,O)` fan-out, log-sum-exp merge |
 
-Both are driven by `milestone-gqa-llama70b` (4 panels:
+Both were driven by `milestone-gqa-llama70b` (4 panels:
-`{single,multi}_user × {prefill,decode}`,
+`{single,multi}_user × {prefill,decode}`).
-`src/kernbench/benches/milestone_gqa_llama70b.py`) and tested in
-`tests/attention/test_milestone_gqa_llama70b.py`.
 
 **They are written in the greenlet `tl` API:** `tl.load`, `tl.dot`,
 `tl.softmax`/`tl.max`/`tl.sum`/`tl.exp`, `tl.send`/`tl.recv`, and Python
@@ -193,8 +191,8 @@ GEMMs onto the scheduler-managed `tl.composite` path** (see §1) — this
 **Three deliberate limitations of the baseline** — exactly what an
 *efficient GQA* kernel must lift:
 
-1. **No GQA reuse.** `h_q == h_kv == 1`
+1. **No GQA reuse.** Baseline was capped at `h_q == h_kv == 1`. The
-   (`test_milestone_gqa_llama70b.py:137-142`). The test attributes this to
+   baseline test attributed this to
    a MemoryStore byte-conservation failure on a *broadcast view*, but that
    failure is a property of the baseline's **head-packing hack**
    (`_view(K, (h_q·d, S_kv))`, which conflates all heads into one matmul
@@ -209,9 +207,9 @@ GEMMs onto the scheduler-managed `tl.composite` path** (see §1) — this
    reduce-to-root** (intra-CUBE tree + intra-CUBE-Group center-mesh, §4) is
    `⌈log₂ P⌉` + center-mesh-over-`C` steps.
 3. **Validation scale only.** `S = 16` because the 1 MiB scratch bump
-   allocator leaks per-tile temporaries
+   allocator leaks per-tile temporaries (baseline test capped S at 16
-   (`test_milestone_gqa_llama70b.py:123-148`) and there is no causal
+   for this reason) and there is no causal masking / tiling → fixed by
-   masking / tiling → fixed by **ADR-0063** (recycling) + §5 (tiling,
+   **ADR-0063** (recycling) + §5 (tiling,
    causal skip) + composite K/V streaming (§3) + **ADR-0062** (lazy load
    overlap).
 
@@ -616,7 +614,7 @@ sharded, small `O`). Prefill-SP does **not** reduce — it rotates KV (Ring,
 
 After its tile sweep each rank holds `(m_i, ℓ_i, O_i)` (unnormalised) over
 its `1/(C·P)` shard. Combine via the associative/commutative log-sum-exp
-merge (identical math to the baseline's fold, `_attention_mesh_mlo.py:117-122`):
+merge (identical math to the baseline's fold):
 
 ```python
 def merge(m_a, l_a, O_a, m_b, l_b, O_b):
@@ -750,8 +748,8 @@ def gqa_decode_sp(q_ptr, k_ptr, v_ptr, o_ptr, counter, start_pe, start_cube,
   batch axis would *not* be counted — see §8).
 - If `S_rank` fits in scratch (small/medium context) this degenerates to a
   **one-shot** partial attention (one composite for Q·Kᵀ, one softmax, one
-  composite for P·V) — exactly the baseline `_attention_mesh_mlo`
+  composite for P·V) — exactly the baseline `_partial_attention` shape,
-  `_partial_attention`, just GQA-batched and on the composite path. Tiling
+  just GQA-batched and on the composite path. Tiling
   (§3) only kicks in when `S_rank` exceeds the scratch scope's tile budget.
 
 ### 5.3 DECODE, with SP / KV-parallel (`C × P` ranks)
@@ -809,8 +807,8 @@ head needs in full anyway.
   `(m,ℓ,O)` reduce** — each CUBE normalises and writes its own head's rows.
 - **Ring KV:** the `C` KV slices **rotate** around the CUBE ring; each CUBE
   folds the incoming block into its head's running `(m,ℓ,O)` (online-softmax,
-  carried across ring steps as Python handles, as `_attention_mesh_kv` does
+  carried across ring steps as Python handles). After `C` steps every
-  today). After `C` steps every head has seen all KV. **GQA reuse comes from
+  head has seen all KV. **GQA reuse comes from
   the rotation** — slice `j` visits all `C` CUBEs, serving all `G` heads.
 - IPCQ overlaps the next step's KV receive with the current step's compute
   via `tl.recv_async`/`tl.wait` (`tl_context.py:543-560` — already exists);
@@ -821,8 +819,8 @@ head needs in full anyway.
   current KV block tile across the `P` PEs (a query-axis split exists here,
   unlike decode); details in §B.
 
-The baseline `_attention_mesh_kv` already implements the ring fold; this
+The (now-removed) baseline already implemented the ring fold; this ADR
-ADR adds GQA reuse, the head-parallel placement, causal step-skip, and the
+adds GQA reuse, the head-parallel placement, causal step-skip, and the
 composite-hybrid inner tile (§3).
 
 ### 5.6 Decode CPU-pipelining variants (opt1 / opt3 / opt2)
@@ -1184,8 +1182,8 @@ predicted default; revise on review.
     fix `C = G` for the prefill kernel at headline scale; treat `C ≠ G` as a
     separate study.
 
-5. **Reconcile with `_attention_mesh_mlo_2d` (current impl).** The remote
+5. **Reconcile with the prior `_attention_mesh_mlo_2d` impl (now removed).**
-    impl's 2D kernel is an **AllReduce** over cubes with **Q replicated** —
+    That 2D kernel was an **AllReduce** over cubes with **Q replicated** —
     i.e. the decode-reduce family, but all-reduce (broadcast-back) not
     reduce-to-root, and not yet the prefill head-parallel ring. **Recommend:**
     (a) move its 2D AllReduce → reduce-to-root (drop broadcast-back) for the
@@ -1208,3 +1206,50 @@ predicted default; revise on review.
     command kind. Defer **opt2** (the
     two-composite `ex_composite`, only `#2` is new) until ADR-0064's cost
     model makes its fewer-issues win measurable.
+
+### Items from the long/short context split (this revision)
+
+The split between the long-context kernels (§5.2 decode, §5.5 prefill —
+both already specified) and a short-context variant is referenced
+abstractly in items §B-1 and §B-6 above ("short-context balance is a
+separate study"). This subsection pins the two open numbers.
+
+1. **Short/long context threshold = 256 K tokens.** Below the threshold,
+    the §5.5 Ring KV prefill pays C-1 ring rotations whose IPCQ cost is
+    not amortised by enough per-rank compute; above it, the per-rank
+    KV sweep dominates and ring is the right choice (§9 line 803).
+    Likewise §5.2's per-rank `S_kv/C·P` shard is small enough at
+    short context that the 2-level reduce-to-root hop count
+    dominates the per-rank compute. The 256 K boundary matches the
+    point at which `S_kv/C·P` (`C=4, P=8 → /32`) crosses the
+    8 K tokens-per-rank mark above which the per-rank tile sweep
+    (§3 / §B-3 §scratch_scope) becomes the limiting factor rather
+    than collective overhead. **Recommend:** treat 256 K as the
+    bench dispatch threshold; expose it as a launch knob for sweeps.
+
+2. **Short-context kernel design — `kv_per_cube ∈ {1, 2, 4, 8}`.**
+    At short context the §5.5 Ring KV motion is wasted work and the
+    §5.2 cube-SP shard is too thin to feed the engines. The
+    short-context variant therefore drops cube-SP entirely: each
+    CUBE owns `kv_per_cube` *whole* KV heads (no `S_kv` sharding
+    across CUBEs), and the existing PE-SP within a CUBE shards `S_kv`
+    across the `P` PEs for each owned head.
+
+    For `h_kv = 8` the natural distributions are:
+
+    | `kv_per_cube` | participating CUBEs | head→CUBE map |
+    |---|---|---|
+    | 1 | `C = 8` | head `i` → CUBE `i` |
+    | 2 | `C = 4` | heads `[2i, 2i+1]` → CUBE `i` |
+    | 4 | `C = 2` | heads `[4i..4i+3]` → CUBE `i` |
+    | 8 | `C = 1` | all heads → single CUBE |
+
+    No inter-CUBE reduce within a head (each head fully owned by one
+    CUBE), so the kernel runs Level-2 (PE) chain reduce-to-root only;
+    Level-1 (inter-CUBE) collapses to a no-op. The output stays
+    distributed per CUBE (each CUBE writes its owned heads' `O` slice).
+
+    **Recommend:** ship as a separate kernel file (`_gqa_decode_short.py` /
+    `_gqa_prefill_short.py`) and a separate bench dispatcher that picks
+    short vs long by `S_kv ⋚ 256K`. Keep `kv_per_cube` as a kernel arg so
+    the topology cost trade-off can be measured per workload.

docs/adr-proposed/ADR-0061-prog-tl-broadcast.md

@@ -32,9 +32,8 @@ the group is the decode-time efficiency win (ADR-0060 §5.2).
 
 ### Why it does not work today
 
-The existing mesh kernels reshape tensors with a metadata-only helper
+Pre-ADR-0060 mesh kernels reshaped tensors with a metadata-only `_view`
-`_view` (`src/kernbench/benches/_attention_mesh_kv.py:25-36`,
+helper that rewrote `shape` but kept the original `nbytes`:
-`_attention_mesh_mlo.py:25-36`):
 
 ```python
 def _view(handle, new_shape):
@@ -62,8 +61,7 @@ A reshape conserves bytes; a **broadcast does not** (`[S_kv, 1, d]` →
 kernel tries to view a `h_kv`-headed K as if it had `h_q` heads, the
 stored array has `1/G` of the requested bytes and `read` raises.
 
-The current test suite documents this as a deliberate limitation
+The pre-ADR-0060 baseline documented this as a deliberate limitation:
-(`tests/attention/test_milestone_gqa_llama70b.py:137-142`):
 
 > "v1 uses `h_q == h_kv == 1` to avoid … GQA broadcast view (which is
 > symbolic and does not survive MemoryStore's nbytes check under
@@ -196,7 +194,7 @@ entire point of GQA (it `G×`'s KV-cache HBM footprint and bandwidth).
 2. **Phase 2 data**: broadcasting a known array yields the numpy
    `broadcast_to(...).copy()` result; a downstream `tl.dot` consuming it
    passes the MemoryStore nbytes check (the exact case that fails today).
-3. **GQA end-to-end**: re-run a reduced `milestone-gqa-llama70b` panel
+3. **GQA end-to-end**: re-run a reduced `milestone-gqa-headline` panel
    with `h_q=8, h_kv=1` under `enable_data=True`; it must complete (today
    it raises `ValueError: Shape mismatch`).
 4. **Broadcast-incompatible shape** (e.g. axis size 3 → 8) raises a

docs/adr-proposed/ADR-0063-prog-tl-scratch-scope.md

@@ -102,6 +102,68 @@ kernel keeps two arenas:
 This mirrors how real flash-attention SRAM budgeting works: a small
 persistent accumulator region + a recycled tile working set.
 
+#### D3.1 The persistent-arena write mechanism: `tl.copy_to(dst, src)`
+
+The merge ops (`tl.maximum`, `tl.exp`, binary `*` / `+`) all call
+`_make_compute_out(...)` which allocates from the bump cursor (D1).
+Inside a `scratch_scope`, their result handles therefore live **inside**
+the scope and vanish on `__exit__`. To realise D3's two-arena split, the
+kernel needs a way to **write a scoped result's bytes back to a
+persistent address**.
+
+The primitive that closes this gap:
+
+```python
+def copy_to(self, dst: TensorHandle, src: TensorHandle) -> None:
+    """Copy ``src``'s bytes into ``dst``'s address (both TCM).
+
+    Shapes and dtypes must match. ``dst`` is typically a handle
+    allocated outside any active ``scratch_scope`` (the persistent
+    arena); ``src`` is a scoped handle whose bytes must outlive
+    scope ``__exit__``.
+    """
+```
+
+Symmetric to `tl.store` (the HBM-side byte copy), kept TCM-only here so
+the running-state writeback doesn't pollute op_log with spurious DMA.
+
+**Mechanics:**
+- New `CopyCmd(src, dst, nbytes, data_op=True)` command.
+- op_log: `op_kind="math"`, `op_name="copy"` — runs on the vector engine.
+- Latency: `pe_math._compute_ns(prod(shape))` — models on-chip register
+  writeback, not HBM transfer.
+- Emit-time validation: `dst.shape == src.shape`, `dst.dtype == src.dtype`,
+  `dst.space == "tcm"`, `src.space == "tcm"`. Authoring errors surface in
+  Phase 1, not deep in Phase 2 data execution.
+
+**Call-site pattern:**
+
+```python
+m, l, O = init_running(...)                  # persistent (outside scope)
+for j in range(n_tiles):
+    with tl.scratch_scope():
+        ...                                  # per-tile work (recycled)
+        m_new = tl.maximum(m, mj)            # scoped scratch
+        l_new = l * scale_old + l_step * scale_step
+        O_new = O * scale_old + O_step * scale_step
+
+        tl.copy_to(m, m_new)                 # ← persist new running state
+        tl.copy_to(l, l_new)
+        tl.copy_to(O, O_new)
+    # exit: scoped m_new/l_new/O_new gone; their bytes live in persistent m/l/O
+```
+
+The copy happens **before** `__exit__`, so the read of `src` (scoped) is
+valid; after exit only `dst` (persistent) is read, satisfying D2's
+no-read-after-exit safety contract.
+
+**Why a dedicated primitive rather than `dst=` kwargs on every math op**
+(considered, rejected): adding `dst=` to `tl.maximum`, `tl.exp`,
+`_binary_math`, `_unary_math`, `_reduction` is ~25 LOC across 5
+op-families and breaks the uniform "call returns a fresh handle"
+pattern. `tl.copy_to` is one primitive, one command, one executor
+handler — minimal surface area for the same effect.
+
 ### D4. Nesting
 
 Scopes nest (stack of save-points). Inner scope exit rewinds to the inner

src/kernbench/benches/1H_milestone_output/gqa/sweep.json

@@ -1,65 +0,0 @@
-{
-  "version": 1,
-  "validation_scale": true,
-  "panels": [
-    "single_user_prefill",
-    "multi_user_prefill",
-    "single_user_decode",
-    "multi_user_decode"
-  ],
-  "config": {
-    "S_q_prefill": 16,
-    "S_kv_per_rank": 16,
-    "h_q": 1,
-    "h_kv": 1,
-    "d_head": 64,
-    "n_ranks_single_user": 8,
-    "n_ranks_multi_user": 4
-  },
-  "rows": [
-    {
-      "panel": "single_user_prefill",
-      "n_ranks": 8,
-      "op_log_summary": {
-        "gemm_count": 128,
-        "ipcq_send_count": 112,
-        "ipcq_recv_count": 112,
-        "dma_read_count": 24,
-        "dma_write_count": 8
-      }
-    },
-    {
-      "panel": "multi_user_prefill",
-      "n_ranks": 4,
-      "op_log_summary": {
-        "gemm_count": 32,
-        "ipcq_send_count": 24,
-        "ipcq_recv_count": 24,
-        "dma_read_count": 12,
-        "dma_write_count": 4
-      }
-    },
-    {
-      "panel": "single_user_decode",
-      "n_ranks": 8,
-      "op_log_summary": {
-        "gemm_count": 16,
-        "ipcq_send_count": 168,
-        "ipcq_recv_count": 168,
-        "dma_read_count": 24,
-        "dma_write_count": 8
-      }
-    },
-    {
-      "panel": "multi_user_decode",
-      "n_ranks": 4,
-      "op_log_summary": {
-        "gemm_count": 8,
-        "ipcq_send_count": 36,
-        "ipcq_recv_count": 36,
-        "dma_read_count": 12,
-        "dma_write_count": 4
-      }
-    }
-  ]
-}

src/kernbench/benches/1H_milestone_output/gqa_headline/sweep.json

@@ -0,0 +1,69 @@
+{
+  "version": 1,
+  "panels": [
+    "single_user_prefill_gqa",
+    "multi_user_prefill_gqa",
+    "single_user_decode_gqa",
+    "multi_user_decode_gqa"
+  ],
+  "config": {
+    "T_q_prefill": 4,
+    "T_q_decode": 1,
+    "S_kv_prefill": 16,
+    "h_q_decode": 8,
+    "h_kv_decode": 1,
+    "d_head": 64
+  },
+  "rows": [
+    {
+      "panel": "single_user_prefill_gqa",
+      "kind": "prefill",
+      "C": 1,
+      "S_kv": 16,
+      "op_log_summary": {
+        "gemm_count": 2,
+        "ipcq_copy_count": 0,
+        "dma_read_count": 3,
+        "dma_write_count": 1
+      }
+    },
+    {
+      "panel": "multi_user_prefill_gqa",
+      "kind": "prefill",
+      "C": 4,
+      "S_kv": 16,
+      "op_log_summary": {
+        "gemm_count": 32,
+        "ipcq_copy_count": 24,
+        "dma_read_count": 12,
+        "dma_write_count": 4
+      }
+    },
+    {
+      "panel": "single_user_decode_gqa",
+      "kind": "decode",
+      "C": 1,
+      "P": 8,
+      "S_kv": 64,
+      "op_log_summary": {
+        "gemm_count": 16,
+        "ipcq_copy_count": 21,
+        "dma_read_count": 24,
+        "dma_write_count": 1
+      }
+    },
+    {
+      "panel": "multi_user_decode_gqa",
+      "kind": "decode",
+      "C": 4,
+      "P": 8,
+      "S_kv": 128,
+      "op_log_summary": {
+        "gemm_count": 64,
+        "ipcq_copy_count": 93,
+        "dma_read_count": 96,
+        "dma_write_count": 1
+      }
+    }
+  ]
+}

src/kernbench/benches/_attention_mesh_kv.py

@@ -1,193 +0,0 @@
-"""Mesh-native bidirectional Ring-K/V attention kernel — prefill (ADR-0059 Proposed).
-
-Each rank holds its own Q tile and 1/n_ranks of K, V (sequence-sharded).
-Over ``n_ranks - 1`` bidirectional steps, K and V propagate both east and
-west: chunk c_i originating at rank i reaches rank j at step ``|i - j|``.
-Every rank receives every other rank's chunk **exactly once** and folds it
-into a running ``(m, ℓ, o)`` via the online-softmax recurrence. After all
-steps each rank holds the final attention output for its own Q tokens —
-no cross-rank merge is required.
-
-Supersedes ADR-0055's closed-ring ``_attention_ring_kv.py``. Both modules
-stay on disk during the transition; this one runs on the hardware's
-actual open-mesh wiring (no closed-ring SFR install required).
-
-Imported by ``milestone_gqa_llama70b`` (after the bench's Phase 2 switches
-its imports) and invoked through ``torch.launch(...)`` — not through
-``dist.all_reduce(...)``. See ADR-0055 Context for why this kernel is not
-backend-dispatched via ADR-0050's algorithm-module contract.
-"""
-from __future__ import annotations
-
-from kernbench.common.pe_commands import TensorHandle
-
-
-def _view(handle: TensorHandle, new_shape: tuple[int, ...]) -> TensorHandle:
-    """Reshape — metadata only, no command emitted (cf. ``tl.trans``)."""
-    return TensorHandle(
-        id=handle.id,
-        addr=handle.addr,
-        shape=new_shape,
-        dtype=handle.dtype,
-        nbytes=handle.nbytes,
-        data=handle.data,
-        space=handle.space,
-        pinned=handle.pinned,
-    )
-
-
-def _partial_attention(
-    Q: TensorHandle,
-    K: TensorHandle,
-    V: TensorHandle,
-    S_q: int,
-    S_kv_per_rank: int,
-    h_q: int,
-    d_head: int,
-    tl,
-) -> tuple[TensorHandle, TensorHandle, TensorHandle]:
-    """One pass of partial attention against (K, V).
-
-    Emits 1 GEMM(Q·K^T) + softmax + max + sub + exp + sum + 1 GEMM(P·V).
-    Returns the running-statistics triplet ``(m, ℓ, O_partial)`` for the
-    online-softmax mlo merge.
-    """
-    K_2d_T = _view(K, (h_q * d_head, S_kv_per_rank))
-    V_2d = _view(V, (S_kv_per_rank, h_q * d_head))
-
-    scores = tl.dot(Q, K_2d_T)
-    m = tl.max(scores, axis=-1)
-    P = tl.softmax(scores, axis=-1)
-    scores_centered = scores - m
-    exp_scores = tl.exp(scores_centered)
-    ell = tl.sum(exp_scores, axis=-1)
-    O_partial = tl.dot(P, V_2d)
-    return m, ell, O_partial
-
-
-def attention_mesh_kv_kernel(
-    q_ptr: int,
-    k_ptr: int,
-    v_ptr: int,
-    o_ptr: int,
-    S_q: int,
-    S_kv_per_rank: int,
-    h_q: int,
-    h_kv: int,
-    d_head: int,
-    n_ranks: int,
-    rank_axis: int = 0,
-    *,
-    tl,
-) -> None:
-    """Mesh-native bidirectional Ring-K/V attention — see module docstring.
-
-    ``rank_axis`` selects which program-id dimension carries the ring rank,
-    matching the GQA Llama-70B sharding study's TL/BL vs TR/BR distinction
-    (`llm_paper_review/notes/GQA_MHA_sharding/scripts/_gen_llama70b_1M_4cases.py`):
-
-      0 — single_user_* panels (TL/BL): rank == tl.program_id(axis=0) (PE
-          id in cube). KV is split @ PEs **intra-cube**; ring runs over
-          the 8 PEs of one cube (NOC-only). At Llama-70B headline scale
-          this kernel launches once per cube; 64 such cubes run in
-          parallel for one user. Each cube's PE-level ring is independent.
-
-      1 — multi_user_* panels (TR/BR): rank == tl.program_id(axis=1)
-          (cube id). KV is split @ cubes **inter-cube**; ring runs over
-          the cubes of one KV-group. The kernel gates ``pe_id != 0`` to
-          return early — a v1 simplification: at headline scale (B=8) the
-          study's "Batch on batch" pattern would have all 8 PEs each handle
-          one user's batch element instead of staying silent. Adding the
-          per-cube batch dimension is sub-cycle 4c headline work.
-    """
-    # For multi_user (rank_axis=1) only PE 0 in each cube runs the ring.
-    if rank_axis != 0 and tl.program_id(axis=0) != 0:
-        return
-    rank = tl.program_id(axis=rank_axis)
-    has_E = rank < n_ranks - 1
-    has_W = rank > 0
-
-    # Q stays put on this rank — loaded once, used in every partial attention.
-    Q = tl.load(q_ptr, shape=(S_q, h_q * d_head), dtype="f16")
-
-    # Local K, V chunk.
-    K = tl.load(k_ptr, shape=(S_kv_per_rank, h_kv, d_head), dtype="f16")
-    V = tl.load(v_ptr, shape=(S_kv_per_rank, h_kv, d_head), dtype="f16")
-
-    # Step 0 (local): partial attention against own K, V — initializes the
-    # running triplet (m, ℓ, o).
-    m, ell, o = _partial_attention(
-        Q, K, V, S_q, S_kv_per_rank, h_q, d_head, tl,
-    )
-
-    # Seed bidirectional waves with own chunk (step-1 send).
-    to_send_east_K: TensorHandle | None = K
-    to_send_east_V: TensorHandle | None = V
-    to_send_west_K: TensorHandle | None = K
-    to_send_west_V: TensorHandle | None = V
-
-    # Bidirectional fan-out: n_ranks - 1 steps. By step k, the wave from
-    # rank i has reached rank (i ± k). After n_ranks - 1 steps, every rank
-    # has merged every other rank's chunk exactly once (ADR-0059 D3).
-    for step in range(1, n_ranks):
-        # Send the eastbound wave we currently hold (own at step 1; forwarded
-        # at later steps). ``None`` means we have no wave to forward this
-        # direction this step (edge rank, or the wave already passed by).
-        if has_E and to_send_east_K is not None:
-            tl.send(dir="E", src=to_send_east_K)
-            tl.send(dir="E", src=to_send_east_V)
-        if has_W and to_send_west_K is not None:
-            tl.send(dir="W", src=to_send_west_K)
-            tl.send(dir="W", src=to_send_west_V)
-
-        # Receive eastbound wave from W (carries chunk c_{rank - step}).
-        K_from_W: TensorHandle | None = None
-        V_from_W: TensorHandle | None = None
-        if has_W and (rank - step) >= 0:
-            K_from_W = tl.recv(
-                dir="W", shape=(S_kv_per_rank, h_kv, d_head), dtype="f16",
-            )
-            V_from_W = tl.recv(
-                dir="W", shape=(S_kv_per_rank, h_kv, d_head), dtype="f16",
-            )
-            m_new, ell_new, o_new = _partial_attention(
-                Q, K_from_W, V_from_W, S_q, S_kv_per_rank, h_q, d_head, tl,
-            )
-            m_combined = tl.maximum(m, m_new)
-            scale_old = tl.exp(m - m_combined)
-            scale_new = tl.exp(m_new - m_combined)
-            ell = ell * scale_old + ell_new * scale_new
-            o = o * scale_old + o_new * scale_new
-            m = m_combined
-
-        # Receive westbound wave from E (carries chunk c_{rank + step}).
-        K_from_E: TensorHandle | None = None
-        V_from_E: TensorHandle | None = None
-        if has_E and (rank + step) < n_ranks:
-            K_from_E = tl.recv(
-                dir="E", shape=(S_kv_per_rank, h_kv, d_head), dtype="f16",
-            )
-            V_from_E = tl.recv(
-                dir="E", shape=(S_kv_per_rank, h_kv, d_head), dtype="f16",
-            )
-            m_new, ell_new, o_new = _partial_attention(
-                Q, K_from_E, V_from_E, S_q, S_kv_per_rank, h_q, d_head, tl,
-            )
-            m_combined = tl.maximum(m, m_new)
-            scale_old = tl.exp(m - m_combined)
-            scale_new = tl.exp(m_new - m_combined)
-            ell = ell * scale_old + ell_new * scale_new
-            o = o * scale_old + o_new * scale_new
-            m = m_combined
-
-        # Forward what we received for next step. ``None`` propagates: if no
-        # chunk arrived this step (out-of-bounds wave origin), there is
-        # nothing to forward next step in that direction.
-        to_send_east_K = K_from_W
-        to_send_east_V = V_from_W
-        to_send_west_K = K_from_E
-        to_send_west_V = V_from_E
-
-    # Final normalize: O := o / ℓ.
-    O_final = o / ell
-    tl.store(o_ptr, O_final)

src/kernbench/benches/_attention_mesh_mlo.py

@@ -1,167 +0,0 @@
-"""Mesh-native bidirectional AllReduce-mlo attention — decode (ADR-0059 Proposed).
-
-Every rank holds the full Q (replicated, small at ``S_q=1``) and 1/n_ranks
-of KV (sequence-sharded). Each rank computes its partial attention
-against own KV in ONE shot, then runs a bidirectional fan-out of the
-``(m, ℓ, o)`` triplet: the triplet originating at rank i reaches rank j at
-step ``|i - j|``. Every rank merges every other rank's triplet exactly
-once over ``n_ranks - 1`` steps, ending with the final answer replicated
-on every rank.
-
-Supersedes ADR-0056's closed-ring ``_attention_allreduce_mlo.py``. Both
-modules stay on disk during the transition; this one runs on the
-hardware's actual open-mesh wiring (no closed-ring SFR install required).
-
-Imported by ``milestone_gqa_llama70b`` (after the bench's Phase 2 switches
-its imports) and invoked through ``torch.launch(...)`` — not through
-``dist.all_reduce(...)``. See ADR-0056 Context for why this kernel is not
-backend-dispatched via ADR-0050's algorithm-module contract.
-"""
-from __future__ import annotations
-
-from kernbench.common.pe_commands import TensorHandle
-
-
-def _view(handle: TensorHandle, new_shape: tuple[int, ...]) -> TensorHandle:
-    """Reshape — metadata only, no command emitted (cf. ``tl.trans``)."""
-    return TensorHandle(
-        id=handle.id,
-        addr=handle.addr,
-        shape=new_shape,
-        dtype=handle.dtype,
-        nbytes=handle.nbytes,
-        data=handle.data,
-        space=handle.space,
-        pinned=handle.pinned,
-    )
-
-
-def attention_mesh_mlo_kernel(
-    q_ptr: int,
-    k_ptr: int,
-    v_ptr: int,
-    o_ptr: int,
-    S_q: int,
-    S_kv_per_rank: int,
-    h_q: int,
-    h_kv: int,
-    d_head: int,
-    n_ranks: int,
-    rank_axis: int = 0,
-    *,
-    tl,
-) -> None:
-    """Mesh-native bidirectional AllReduce-mlo — see module docstring.
-
-    ``rank_axis`` selects which program-id dimension carries the ring rank,
-    matching the GQA Llama-70B sharding study's TL/BL vs TR/BR distinction
-    (`llm_paper_review/notes/GQA_MHA_sharding/scripts/_gen_llama70b_1M_4cases.py`):
-
-      0 — single_user_* panels (TL/BL): rank == tl.program_id(axis=0) (PE
-          id in cube). KV is split @ PEs **intra-cube**; ring runs over
-          the 8 PEs of one cube (NOC-only). At Llama-70B headline scale
-          this kernel launches once per cube; 64 such cubes run in
-          parallel for one user (1 Q-head per cube × 8 cubes per KV-group
-          × 8 KV-groups). The PE-level ring inside each cube is
-          independent of the others.
-
-      1 — multi_user_* panels (TR/BR): rank == tl.program_id(axis=1)
-          (cube id). KV is split @ cubes **inter-cube**; ring runs over
-          the cubes of one KV-group. The kernel gates ``pe_id != 0`` to
-          return early — a v1 simplification: at headline scale (B=8) the
-          study's "Batch on batch" pattern would have all 8 PEs each handle
-          one user's batch element instead of staying silent. Validation
-          shipped with B=1 to focus on the inter-cube ring's correctness;
-          adding the per-cube batch dimension is sub-cycle 4c headline work.
-    """
-    # For multi_user (rank_axis=1) only PE 0 in each cube runs the ring.
-    if rank_axis != 0 and tl.program_id(axis=0) != 0:
-        return
-    rank = tl.program_id(axis=rank_axis)
-    has_E = rank < n_ranks - 1
-    has_W = rank > 0
-
-    # Q is replicated on every rank — loaded once.
-    Q = tl.load(q_ptr, shape=(S_q, h_q * d_head), dtype="f16")
-
-    # Local KV chunk. KV is sequence-sharded and stays put on this rank for
-    # the entire fan-out — distinguishing decode from prefill (ADR-0059 D3)
-    # where KV circulates.
-    K = tl.load(k_ptr, shape=(S_kv_per_rank, h_kv, d_head), dtype="f16")
-    V = tl.load(v_ptr, shape=(S_kv_per_rank, h_kv, d_head), dtype="f16")
-
-    # ── One-shot local partial attention ──────────────────────────
-    K_2d_T = _view(K, (h_q * d_head, S_kv_per_rank))
-    V_2d = _view(V, (S_kv_per_rank, h_q * d_head))
-    scores = tl.dot(Q, K_2d_T)
-    m = tl.max(scores, axis=-1)
-    P = tl.softmax(scores, axis=-1)
-    scores_centered = scores - m
-    exp_scores = tl.exp(scores_centered)
-    ell = tl.sum(exp_scores, axis=-1)
-    o = tl.dot(P, V_2d)
-
-    # Seed bidirectional waves with own triplet (step-1 send).
-    to_send_east_m: TensorHandle | None = m
-    to_send_east_ell: TensorHandle | None = ell
-    to_send_east_o: TensorHandle | None = o
-    to_send_west_m: TensorHandle | None = m
-    to_send_west_ell: TensorHandle | None = ell
-    to_send_west_o: TensorHandle | None = o
-
-    # Bidirectional fan-out of (m, ℓ, o) triplets — n_ranks - 1 steps.
-    for step in range(1, n_ranks):
-        # Send eastbound triplet (own at step 1; forwarded at later steps).
-        if has_E and to_send_east_m is not None:
-            tl.send(dir="E", src=to_send_east_m)
-            tl.send(dir="E", src=to_send_east_ell)
-            tl.send(dir="E", src=to_send_east_o)
-        # Send westbound triplet.
-        if has_W and to_send_west_m is not None:
-            tl.send(dir="W", src=to_send_west_m)
-            tl.send(dir="W", src=to_send_west_ell)
-            tl.send(dir="W", src=to_send_west_o)
-
-        # Receive eastbound triplet from W (originated at rank - step).
-        m_from_W: TensorHandle | None = None
-        ell_from_W: TensorHandle | None = None
-        o_from_W: TensorHandle | None = None
-        if has_W and (rank - step) >= 0:
-            m_from_W = tl.recv(dir="W", shape=m.shape, dtype="f16")
-            ell_from_W = tl.recv(dir="W", shape=ell.shape, dtype="f16")
-            o_from_W = tl.recv(dir="W", shape=o.shape, dtype="f16")
-            m_combined = tl.maximum(m, m_from_W)
-            scale_old = tl.exp(m - m_combined)
-            scale_new = tl.exp(m_from_W - m_combined)
-            ell = ell * scale_old + ell_from_W * scale_new
-            o = o * scale_old + o_from_W * scale_new
-            m = m_combined
-
-        # Receive westbound triplet from E (originated at rank + step).
-        m_from_E: TensorHandle | None = None
-        ell_from_E: TensorHandle | None = None
-        o_from_E: TensorHandle | None = None
-        if has_E and (rank + step) < n_ranks:
-            m_from_E = tl.recv(dir="E", shape=m.shape, dtype="f16")
-            ell_from_E = tl.recv(dir="E", shape=ell.shape, dtype="f16")
-            o_from_E = tl.recv(dir="E", shape=o.shape, dtype="f16")
-            m_combined = tl.maximum(m, m_from_E)
-            scale_old = tl.exp(m - m_combined)
-            scale_new = tl.exp(m_from_E - m_combined)
-            ell = ell * scale_old + ell_from_E * scale_new
-            o = o * scale_old + o_from_E * scale_new
-            m = m_combined
-
-        # Forward the original received triplet (not the merged running state)
-        # so neighbors get the original wave. ``None`` propagates if nothing
-        # arrived this step.
-        to_send_east_m = m_from_W
-        to_send_east_ell = ell_from_W
-        to_send_east_o = o_from_W
-        to_send_west_m = m_from_E
-        to_send_west_ell = ell_from_E
-        to_send_west_o = o_from_E
-
-    # Final normalize: O := o / ℓ.
-    O_final = o / ell
-    tl.store(o_ptr, O_final)

src/kernbench/benches/_attention_mesh_mlo_2d.py

@@ -1,217 +0,0 @@
-"""Mesh-native 2D row-then-col AllReduce-mlo attention — decode (ADR-0059 extension).
-
-Each cube holds the full Q (replicated) and 1/(mesh_rows * mesh_cols) of
-KV (sequence-sharded across the 2D cube sub-mesh). The kernel decomposes
-the AllReduce-mlo into a two-stage reduction:
-
-  Stage 1 — Row reduce (E/W edges, ``mesh_cols - 1`` steps)
-     Bidirectional ring within each row. After this stage every cube in
-     row ``r`` holds the partial ``(m, ℓ, o)`` over the ``mesh_cols`` KV
-     chunks in row ``r``.
-
-  Stage 2 — Col reduce (N/S edges, ``mesh_rows - 1`` steps)
-     Bidirectional ring within each column. After this stage every cube
-     holds the partial over all ``mesh_rows × mesh_cols`` KV chunks —
-     the AllReduce result.
-
-The online-softmax mlo merge is associative, so row-then-col partitioning
-of the reduction is mathematically equivalent to a 1D ring AllReduce-mlo
-over all ``mesh_rows × mesh_cols`` cubes. The 2D form takes
-``(mesh_cols - 1) + (mesh_rows - 1)`` steps instead of
-``mesh_rows × mesh_cols - 1`` (e.g. 4 vs 7 at 2×4; 6 vs 15 at 4×4).
-
-Designed to run on hardware wired by
-``configure_sfr_intercube_multisip``, which installs both E/W and N/S
-intra-SIP cube-mesh edges (``sfr_config.py:135-143``). The 1D
-``_attention_mesh_mlo.py`` remains for the single_user PE-ring case;
-this 2D variant supersedes it for multi_user_decode where the per-KV-group
-cube count crosses a row boundary in the 4×4 cube mesh.
-
-``mesh_rows = 1`` is supported as a degenerate row-only case so the
-validation config (``_N_RANKS_MULTI_USER = 4`` → ``(1, 4)``) reduces to
-the 1D ring's step count without behavioral change.
-"""
-from __future__ import annotations
-
-from kernbench.common.pe_commands import TensorHandle
-
-
-def _view(handle: TensorHandle, new_shape: tuple[int, ...]) -> TensorHandle:
-    """Reshape — metadata only, no command emitted (cf. ``tl.trans``)."""
-    return TensorHandle(
-        id=handle.id,
-        addr=handle.addr,
-        shape=new_shape,
-        dtype=handle.dtype,
-        nbytes=handle.nbytes,
-        data=handle.data,
-        space=handle.space,
-        pinned=handle.pinned,
-    )
-
-
-def _bidir_allreduce_mlo(
-    m: TensorHandle,
-    ell: TensorHandle,
-    o: TensorHandle,
-    rank: int,
-    n_ranks: int,
-    dir_pos: str,
-    dir_neg: str,
-    *,
-    tl,
-) -> tuple[TensorHandle, TensorHandle, TensorHandle]:
-    """One bidirectional AllReduce-mlo ring along ``(dir_pos, dir_neg)``.
-
-    Mirrors the 1D ``_attention_mesh_mlo.py`` algorithm but parameterized
-    on direction labels so the 2D kernel can call it once with ``("E", "W")``
-    for the row reduce and once with ``("S", "N")`` for the col reduce.
-    Forwards the received triplets in subsequent steps so chunk ``c_i``
-    reaches rank ``j`` at step ``|i - j|``.
-
-    Returns the running ``(m, ℓ, o)`` after ``n_ranks - 1`` steps. Degenerate
-    cases (``n_ranks <= 1``) are no-ops — the for-loop body simply does not
-    execute.
-    """
-    has_pos = rank < n_ranks - 1
-    has_neg = rank > 0
-
-    to_send_pos_m: TensorHandle | None = m
-    to_send_pos_ell: TensorHandle | None = ell
-    to_send_pos_o: TensorHandle | None = o
-    to_send_neg_m: TensorHandle | None = m
-    to_send_neg_ell: TensorHandle | None = ell
-    to_send_neg_o: TensorHandle | None = o
-
-    for step in range(1, n_ranks):
-        if has_pos and to_send_pos_m is not None:
-            tl.send(dir=dir_pos, src=to_send_pos_m)
-            tl.send(dir=dir_pos, src=to_send_pos_ell)
-            tl.send(dir=dir_pos, src=to_send_pos_o)
-        if has_neg and to_send_neg_m is not None:
-            tl.send(dir=dir_neg, src=to_send_neg_m)
-            tl.send(dir=dir_neg, src=to_send_neg_ell)
-            tl.send(dir=dir_neg, src=to_send_neg_o)
-
-        m_from_neg: TensorHandle | None = None
-        ell_from_neg: TensorHandle | None = None
-        o_from_neg: TensorHandle | None = None
-        if has_neg and (rank - step) >= 0:
-            m_from_neg = tl.recv(dir=dir_neg, shape=m.shape, dtype="f16")
-            ell_from_neg = tl.recv(dir=dir_neg, shape=ell.shape, dtype="f16")
-            o_from_neg = tl.recv(dir=dir_neg, shape=o.shape, dtype="f16")
-            m_combined = tl.maximum(m, m_from_neg)
-            scale_old = tl.exp(m - m_combined)
-            scale_new = tl.exp(m_from_neg - m_combined)
-            ell = ell * scale_old + ell_from_neg * scale_new
-            o = o * scale_old + o_from_neg * scale_new
-            m = m_combined
-
-        m_from_pos: TensorHandle | None = None
-        ell_from_pos: TensorHandle | None = None
-        o_from_pos: TensorHandle | None = None
-        if has_pos and (rank + step) < n_ranks:
-            m_from_pos = tl.recv(dir=dir_pos, shape=m.shape, dtype="f16")
-            ell_from_pos = tl.recv(dir=dir_pos, shape=ell.shape, dtype="f16")
-            o_from_pos = tl.recv(dir=dir_pos, shape=o.shape, dtype="f16")
-            m_combined = tl.maximum(m, m_from_pos)
-            scale_old = tl.exp(m - m_combined)
-            scale_new = tl.exp(m_from_pos - m_combined)
-            ell = ell * scale_old + ell_from_pos * scale_new
-            o = o * scale_old + o_from_pos * scale_new
-            m = m_combined
-
-        to_send_pos_m = m_from_neg
-        to_send_pos_ell = ell_from_neg
-        to_send_pos_o = o_from_neg
-        to_send_neg_m = m_from_pos
-        to_send_neg_ell = ell_from_pos
-        to_send_neg_o = o_from_pos
-
-    return m, ell, o
-
-
-def attention_mesh_mlo_2d_kernel(
-    q_ptr: int,
-    k_ptr: int,
-    v_ptr: int,
-    o_ptr: int,
-    S_q: int,
-    S_kv_per_rank: int,
-    h_q: int,
-    h_kv: int,
-    d_head: int,
-    mesh_rows: int,
-    mesh_cols: int,
-    rank_axis: int = 0,
-    cube_start: int = 0,
-    *,
-    tl,
-) -> None:
-    """2D row-then-col AllReduce-mlo decode kernel — see module docstring.
-
-    ``rank_axis`` selects which program-id dimension carries the cube
-    rank (matches the 1D kernel convention):
-
-      0 — single_user_* (TL/BL): rank == tl.program_id(axis=0) (PE id).
-          Not used at headline scale — single_user uses the 1D intra-cube
-          PE ring (``_attention_mesh_mlo``). Kept here so the signature
-          mirrors the 1D kernel.
-
-      1 — multi_user_* (TR/BR): rank == tl.program_id(axis=1) (cube id).
-          KV is split @ cubes inter-cube; the ring runs over the
-          ``mesh_rows × mesh_cols`` cubes of one KV-group. The kernel
-          gates ``pe_id != 0`` to return early — same v1 simplification
-          as ``_attention_mesh_mlo`` (validation B=1).
-
-    ``cube_start`` matches the value passed to ``DPPolicy.cube_start`` for
-    the launch's tensor placement. kernbench's ``tl.program_id(axis=1)``
-    returns the physical cube id (ADR-0022), so when the launch is
-    offset within the SIP (e.g. cube_start=8 placing the second 2×4
-    KV-group on cubes 8..15), the kernel must subtract ``cube_start``
-    to recover the launch-local rank for ring arithmetic. Default 0
-    preserves the cube_start=0 launches unchanged.
-    """
-    # For multi_user (rank_axis=1) only PE 0 in each cube runs the ring.
-    if rank_axis != 0 and tl.program_id(axis=0) != 0:
-        return
-
-    rank = tl.program_id(axis=rank_axis)
-    if rank_axis != 0:
-        rank = rank - cube_start
-    my_row = rank // mesh_cols
-    my_col = rank % mesh_cols
-
-    # Q is replicated on every cube — loaded once.
-    Q = tl.load(q_ptr, shape=(S_q, h_q * d_head), dtype="f16")
-
-    # Local KV chunk (sequence-sharded across the 2D sub-mesh).
-    K = tl.load(k_ptr, shape=(S_kv_per_rank, h_kv, d_head), dtype="f16")
-    V = tl.load(v_ptr, shape=(S_kv_per_rank, h_kv, d_head), dtype="f16")
-
-    # ── One-shot local partial attention ──────────────────────────
-    K_2d_T = _view(K, (h_q * d_head, S_kv_per_rank))
-    V_2d = _view(V, (S_kv_per_rank, h_q * d_head))
-    scores = tl.dot(Q, K_2d_T)
-    m = tl.max(scores, axis=-1)
-    P = tl.softmax(scores, axis=-1)
-    scores_centered = scores - m
-    exp_scores = tl.exp(scores_centered)
-    ell = tl.sum(exp_scores, axis=-1)
-    o = tl.dot(P, V_2d)
-
-    # ── Stage 1: row AllReduce (E/W, mesh_cols - 1 steps) ─────────
-    m, ell, o = _bidir_allreduce_mlo(
-        m, ell, o, my_col, mesh_cols, "E", "W", tl=tl,
-    )
-
-    # ── Stage 2: col AllReduce (N/S, mesh_rows - 1 steps) ─────────
-    # ``dir_pos="S"`` matches the SFR convention: ``S`` goes to higher
-    # row (configure_sfr_intercube_multisip:140).
-    m, ell, o = _bidir_allreduce_mlo(
-        m, ell, o, my_row, mesh_rows, "S", "N", tl=tl,
-    )
-
-    # Final normalize: O := o / ℓ.
-    O_final = o / ell
-    tl.store(o_ptr, O_final)

src/kernbench/benches/_gqa_decode.py

@@ -0,0 +1,180 @@
+"""GQA fused-attention decode kernel — P1a + P2a + P2b (2-level SP).
+
+Lineage (DDD-0060 §7 phase plan):
+  P1a    : real GQA via M-fold using ``tl.dot``; one-shot per rank.
+  P2a    : intra-CUBE PE-level chain reduce-to-root (single-CUBE SP).
+  P2b    : adds inter-CUBE chain reduce-to-root (multi-CUBE SP);
+           switches to the canonical full SFR install
+           ``configure_sfr_intercube_multisip`` with disjoint namespaces
+           (``intra_*`` for PE, ``N/S/E/W`` for CUBE, ``global_*`` for SIP).
+  P3b    : tile S_kv sweep + ``tl.scratch_scope`` (deferred).
+  Later  : P1b composite swap; P4 lazy load; P5 opt3 pipelining; P6 prefill.
+
+P2b SFR + topology assumptions:
+  - SFR install: ``configure_sfr_intercube_multisip`` is required when
+    P > 1 or C > 1 (provides ``intra_*`` and ``E/W/N/S`` namespaces).
+  - Intra-CUBE PE layout: logical 2×4 grid (no wrap):
+        Row 0:  PE 0, 1, 2, 3
+        Row 1:  PE 4, 5, 6, 7
+  - Inter-CUBE layout: CUBEs of one CUBE Group are laid out as a 1D row
+    (single row of C CUBEs, no wrap). Multi-row CUBE Group placement
+    (e.g. 2×4) is future work — head_of_group / cube_start dispatch is
+    a P7 concern (DDD-0060 §4.1).
+
+Reduce strategy — chain reduce-to-root at PE 0 of CUBE 0:
+
+  Level-2 (intra-CUBE, row-then-col chain on 2×4 grid):
+    Row chain along ``intra_W``: rightmost-col PEs send leftward;
+    leftmost-col PE of each row holds its row's partial.
+    Col bridge along ``intra_N``: PE 4 (col-0, row-1) sends to PE 0
+    (col-0, row-0). Only relevant when P > 4.
+    Result: PE 0 of each CUBE holds the CUBE's partial.
+
+  Level-1 (inter-CUBE, only PE 0 of each CUBE participates):
+    Chain along ``W``: rightmost CUBE sends leftward; CUBE 0 of the
+    CUBE Group ends with the final answer.
+
+  Final store: PE 0 of CUBE 0 normalises (O / ℓ) and writes.
+
+Chain step counts (per ADR-0060 §A.2 root-only output, §4 chain
+deviation noted): for (C, P)=(2, 8), 7 intra-cube × 2 cubes + (C-1)
+inter-cube = 14 + 1 = 15 chain steps; each step ships 3 handles
+(m, ℓ, O) ⇒ 45 ``ipcq_copy`` total.
+
+Three deliberate deviations from ADR-0060, addressed in later phases:
+  1. GEMMs use ``tl.dot``, not ``tl.composite`` (P1b).
+  2. ``softmax_scale`` omitted — needs composite epilogue mechanism (P1b/P5).
+  3. K loaded as ``[d, S_local]`` via byte-conserving reshape of the
+     deployed ``[S_local, h_kv·d]`` slice (ADR-0060 §3 / §B item 2,
+     reshape-not-transpose caveat; correct for zero / symmetric inputs).
+
+Chain-vs-tree deviation from DDD-0060 §7 P2 gate (``⌈log₂ P⌉``):
+  P2b uses linear chain reduce-to-root (P-1 + C-1 hops). True tree on
+  the 2×4 PE grid requires a different SFR install — separate ADR.
+  Architectural intent (root-only output replacing baseline's
+  bidirectional fan-out) is preserved.
+"""
+from __future__ import annotations
+
+
+def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl):
+    """Online-softmax merge (ADR-0060 §4 / _attention_mesh_mlo baseline)."""
+    m_new = tl.maximum(m_local, m_other)
+    scale_old = tl.exp(m_local - m_new)
+    scale_new = tl.exp(m_other - m_new)
+    l_new = l_local * scale_old + l_other * scale_new
+    O_new = O_local * scale_old + O_other * scale_new
+    return m_new, l_new, O_new
+
+
+def gqa_decode_kernel(
+    q_ptr: int,
+    k_ptr: int,
+    v_ptr: int,
+    o_ptr: int,
+    T_q: int,
+    S_kv: int,
+    h_q: int,
+    h_kv: int,
+    d_head: int,
+    C: int,
+    P: int,
+    *,
+    tl,
+) -> None:
+    """GQA decode with M-fold + 2-level chain reduce-to-root.
+
+    Tensor layout:
+      Q : (T_q, h_q · d_head) replicated on every rank; loaded as
+          (G·T_q, d_head) — byte-conserving and math-correct for T_q=1.
+      K : (S_kv, h_kv · d_head) sharded row_wise by (cube, pe); each rank
+          loads its (d_head, S_local) slice via byte-conserving reshape.
+      V : (S_kv, h_kv · d_head) sharded row_wise by (cube, pe); each rank
+          loads its (S_local, d_head) slice.
+      O : (T_q, h_q · d_head) — only PE 0 of CUBE 0 stores.
+    """
+    G = h_q // h_kv
+    n_ranks = C * P
+    S_local = S_kv // n_ranks
+    pe_id = tl.program_id(axis=0)
+    cube_id = tl.program_id(axis=1)
+
+    # ── Local one-shot partial attention (M-fold on the rank's slice) ──
+    Q = tl.load(q_ptr, shape=(G * T_q, d_head), dtype="f16")
+    K_T = tl.load(k_ptr, shape=(d_head, S_local), dtype="f16")
+    V = tl.load(v_ptr, shape=(S_local, d_head), dtype="f16")
+    scores = tl.dot(Q, K_T)
+    m_local = tl.max(scores, axis=-1)
+    centered = scores - m_local
+    exp_scores = tl.exp(centered)
+    l_local = tl.sum(exp_scores, axis=-1)
+    O_local = tl.dot(exp_scores, V)
+
+    # ── Level-2: intra-CUBE row-then-col chain reduce-to-(PE 0) ──
+    PE_GRID_COLS = 4
+    pe_col = pe_id % PE_GRID_COLS
+    pe_row = pe_id // PE_GRID_COLS
+    pe_cols_used = min(PE_GRID_COLS, P)
+    pe_rows_used = (P + PE_GRID_COLS - 1) // PE_GRID_COLS
+
+    # Row chain (along intra_W within each row, gathering leftward).
+    # Each merge step's intermediates are wrapped in tl.scratch_scope and
+    # the new running (m, ℓ, O) is persisted back to the outside-scope
+    # (persistent) m_local/l_local/O_local via tl.copy_to (ADR-0063 §D3/D3.1).
+    if pe_cols_used > 1:
+        if pe_col < pe_cols_used - 1:                # not rightmost: receive E
+            with tl.scratch_scope():
+                m_other = tl.recv(dir="intra_E", shape=m_local.shape, dtype="f16")
+                l_other = tl.recv(dir="intra_E", shape=l_local.shape, dtype="f16")
+                O_other = tl.recv(dir="intra_E", shape=O_local.shape, dtype="f16")
+                m_new, l_new, O_new = _merge_running(
+                    m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
+                )
+                tl.copy_to(m_local, m_new)
+                tl.copy_to(l_local, l_new)
+                tl.copy_to(O_local, O_new)
+        if pe_col > 0:                                # not leftmost: send W
+            tl.send(dir="intra_W", src=m_local)
+            tl.send(dir="intra_W", src=l_local)
+            tl.send(dir="intra_W", src=O_local)
+
+    # Col bridge (intra_N from row 1 col 0 → row 0 col 0). Only at col 0.
+    if pe_col == 0 and pe_rows_used > 1:
+        if pe_row < pe_rows_used - 1:                 # row 0 receives from S
+            with tl.scratch_scope():
+                m_other = tl.recv(dir="intra_S", shape=m_local.shape, dtype="f16")
+                l_other = tl.recv(dir="intra_S", shape=l_local.shape, dtype="f16")
+                O_other = tl.recv(dir="intra_S", shape=O_local.shape, dtype="f16")
+                m_new, l_new, O_new = _merge_running(
+                    m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
+                )
+                tl.copy_to(m_local, m_new)
+                tl.copy_to(l_local, l_new)
+                tl.copy_to(O_local, O_new)
+        if pe_row > 0:                                # row >0 sends to N
+            tl.send(dir="intra_N", src=m_local)
+            tl.send(dir="intra_N", src=l_local)
+            tl.send(dir="intra_N", src=O_local)
+
+    # ── Level-1: inter-CUBE chain reduce (only PE 0 of each CUBE) ──
+    if pe_id == 0 and C > 1:
+        if cube_id < C - 1:                           # not rightmost CUBE: recv E
+            with tl.scratch_scope():
+                m_other = tl.recv(dir="E", shape=m_local.shape, dtype="f16")
+                l_other = tl.recv(dir="E", shape=l_local.shape, dtype="f16")
+                O_other = tl.recv(dir="E", shape=O_local.shape, dtype="f16")
+                m_new, l_new, O_new = _merge_running(
+                    m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
+                )
+                tl.copy_to(m_local, m_new)
+                tl.copy_to(l_local, l_new)
+                tl.copy_to(O_local, O_new)
+        if cube_id > 0:                               # non-root CUBE: send W
+            tl.send(dir="W", src=m_local)
+            tl.send(dir="W", src=l_local)
+            tl.send(dir="W", src=O_local)
+
+    # ── Final normalise + store (only at PE 0 of CUBE 0) ──
+    if pe_id == 0 and cube_id == 0:
+        O_final = O_local / l_local
+        tl.store(o_ptr, O_final)

src/kernbench/benches/_gqa_decode_short.py

@@ -0,0 +1,145 @@
+"""GQA fused-attention SHORT-CONTEXT decode kernel (ADR-0060 §B.split.2).
+
+Short context (S_kv < 256K, per ADR-0060 §B.split.1): each CUBE owns
+``kv_per_cube`` whole KV heads, no S_kv sharding across CUBEs, no
+inter-CUBE reduce. PE-SP within each CUBE: the P PEs split into
+``kv_per_cube`` groups, each group does PE-SP across (P/kv_per_cube)
+PEs for ONE owned head.
+
+Layout (after design iteration during Phase D — see ADR-0060 §B.split.2):
+  - K, V: shape ``(h_kv·S_kv, d_head)`` head-stacked, with the bench
+    deploying ``dp = (cube=row_wise, pe=row_wise)`` so each PE's chunk
+    is exactly ``(S_local, d_head)`` contiguous at its own addressable
+    shard. The kernel just loads at its ``k_ptr`` / ``v_ptr`` — no
+    offset arithmetic needed.
+  - Q: replicated ``(T_q, h_q·d_head)``; the kernel reshapes
+    byte-conservingly to ``(h_q·T_q, d_head)`` and operates on the
+    full stack. Other heads' rows are computed too (semantic noise);
+    with zero/symmetric inputs the math is unchanged. A proper
+    per-head Q slice would require runtime support for partial reads
+    of stored tensors (deferred).
+  - O: replicated; each group root writes the full byte-conserving
+    ``(h_q·T_q, d_head)`` result. Multiple roots within a CUBE write
+    to disjoint PE-local addresses (no overwrite collision).
+
+Group layout on the 2×4 PE grid:
+  kv_per_cube=1, group=8 PEs (full 2×4): row chain + col bridge.
+  kv_per_cube=2, group=4 PEs (one row):  row chain only.
+  kv_per_cube=4, group=2 PEs (adj cols): 1-step chain.
+  kv_per_cube=8, group=1 PE:             no chain — direct write.
+
+Chain reduce within group via the existing ``intra_E/W/N/S`` SFR
+namespace (configure_sfr_intercube_multisip). After chain reduce, the
+group's root PE (pe_in_group == 0) writes its working state to HBM.
+
+Deviations from ADR-0060 (deliberate, documented):
+  1. GEMMs use ``tl.dot``, not ``tl.composite``.
+  2. ``softmax_scale`` omitted.
+  3. K loaded as ``[d, S_local]`` via byte-conserving reshape
+     (reshape-not-transpose caveat — correct for zero / symmetric inputs).
+  4. Q byte-conserving reshape: kernel computes attention for ALL Q
+     rows against the group's owned K head; only the rows for my head
+     are semantically meaningful. Correct for zero / symmetric inputs.
+"""
+from __future__ import annotations
+
+
+def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl):
+    """Online-softmax merge — identical to long kernel."""
+    m_new = tl.maximum(m_local, m_other)
+    scale_old = tl.exp(m_local - m_new)
+    scale_new = tl.exp(m_other - m_new)
+    l_new = l_local * scale_old + l_other * scale_new
+    O_new = O_local * scale_old + O_other * scale_new
+    return m_new, l_new, O_new
+
+
+def gqa_decode_short_kernel(
+    q_ptr: int,
+    k_ptr: int,
+    v_ptr: int,
+    o_ptr: int,
+    T_q: int,
+    S_kv: int,
+    h_q: int,
+    h_kv: int,
+    d_head: int,
+    C: int,
+    P: int,
+    kv_per_cube: int,
+    *,
+    tl,
+) -> None:
+    """Short-context GQA decode with PE-parallel heads + intra-group PE-SP."""
+    group_size = P // kv_per_cube           # PEs per head group
+
+    pe_id = tl.program_id(axis=0)
+    pe_in_group = pe_id % group_size
+
+    # PE-SP within group: shard S_kv across group_size PEs
+    S_local = S_kv // group_size
+
+    # ── Loads (DP layout already places each PE at its own shard) ──
+    # Q replicated → byte-conserving reshape to (h_q·T_q, d_head).
+    Q = tl.load(q_ptr,
+                shape=(h_q * T_q, d_head), dtype="f16")
+    # K, V row_wise per (cube, pe) → each PE has (S_local, d_head) at k_ptr.
+    K_T = tl.load(k_ptr,
+                  shape=(d_head, S_local), dtype="f16")
+    V = tl.load(v_ptr,
+                shape=(S_local, d_head), dtype="f16")
+
+    # ── Local one-shot partial attention ──
+    scores = tl.dot(Q, K_T)
+    m_local = tl.max(scores, axis=-1)
+    centered = scores - m_local
+    exp_scores = tl.exp(centered)
+    l_local = tl.sum(exp_scores, axis=-1)
+    O_local = tl.dot(exp_scores, V)
+
+    # ── Within-group chain reduce-to-root (Level-2 only) ──
+    group_cols = min(4, group_size)
+    group_rows = (group_size + group_cols - 1) // group_cols
+    pe_col_in_group = pe_in_group % group_cols
+    pe_row_in_group = pe_in_group // group_cols
+
+    # Row chain along intra_W (within group's row).
+    if group_cols > 1:
+        if pe_col_in_group < group_cols - 1:        # receive E (in-group)
+            with tl.scratch_scope():
+                m_other = tl.recv(dir="intra_E", shape=m_local.shape, dtype="f16")
+                l_other = tl.recv(dir="intra_E", shape=l_local.shape, dtype="f16")
+                O_other = tl.recv(dir="intra_E", shape=O_local.shape, dtype="f16")
+                m_new, l_new, O_new = _merge_running(
+                    m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
+                )
+                tl.copy_to(m_local, m_new)
+                tl.copy_to(l_local, l_new)
+                tl.copy_to(O_local, O_new)
+        if pe_col_in_group > 0:                     # send W (in-group)
+            tl.send(dir="intra_W", src=m_local)
+            tl.send(dir="intra_W", src=l_local)
+            tl.send(dir="intra_W", src=O_local)
+
+    # Col bridge along intra_N (only if group spans 2 grid rows).
+    if pe_col_in_group == 0 and group_rows > 1:
+        if pe_row_in_group < group_rows - 1:        # receive S (in-group)
+            with tl.scratch_scope():
+                m_other = tl.recv(dir="intra_S", shape=m_local.shape, dtype="f16")
+                l_other = tl.recv(dir="intra_S", shape=l_local.shape, dtype="f16")
+                O_other = tl.recv(dir="intra_S", shape=O_local.shape, dtype="f16")
+                m_new, l_new, O_new = _merge_running(
+                    m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
+                )
+                tl.copy_to(m_local, m_new)
+                tl.copy_to(l_local, l_new)
+                tl.copy_to(O_local, O_new)
+        if pe_row_in_group > 0:                     # send N (in-group)
+            tl.send(dir="intra_N", src=m_local)
+            tl.send(dir="intra_N", src=l_local)
+            tl.send(dir="intra_N", src=O_local)
+
+    # ── Group root writes its owned head's output ──
+    if pe_in_group == 0:
+        O_final = O_local / l_local
+        tl.store(o_ptr, O_final)

src/kernbench/benches/_gqa_prefill.py

@@ -0,0 +1,111 @@
+"""GQA fused-attention prefill kernel — P6a + P6b (head-parallel + Ring KV).
+
+Lineage (DDD-0060 §7 phase plan):
+  P6a    : head-parallel structure (one Q head per CUBE), C=1 baseline.
+  P6b    : add Ring KV rotation across C CUBEs (ADR-0060 §5.5). Each
+            CUBE rotates its KV block to its W neighbour and receives
+            from E; over C-1 steps every CUBE sees every block. Online-
+            softmax merge folds each step into running (m, ℓ, O). No
+            reduce — each CUBE writes its own head's output.
+            Requires ``configure_sfr_intercube_ring(ring_size=C)`` SFR.
+  P6c (later): tile T_q across the P PEs inside a CUBE for intra-CUBE
+            parallelism (ADR-0060 §B item 3).
+
+Deviations from ADR-0060 §5.5 — deferred to later phases:
+  1. GEMMs use ``tl.dot`` not ``tl.composite`` (parallel to P1a; lifts
+     when P1b decides the composite-output-handle question).
+  2. ``softmax_scale`` omitted — same composite-epilogue deferral.
+  3. K loaded as ``[d, S_local]`` via byte-conserving reshape of
+     ``[S_local, d]`` (ADR-0060 §3 / §B item 2 reshape-not-transpose
+     caveat; correct for zero / symmetric inputs).
+  4. No causal masking / step-skip — future P6c.
+  5. Blocking ``tl.recv`` (not ``recv_async``) — overlap via lazy
+     ``tl.load`` lands in P4.
+"""
+from __future__ import annotations
+
+
+def gqa_prefill_kernel(
+    q_ptr: int,
+    k_ptr: int,
+    v_ptr: int,
+    o_ptr: int,
+    T_q: int,
+    S_kv: int,
+    d_head: int,
+    C: int,
+    *,
+    tl,
+) -> None:
+    """Head-parallel prefill attention with Ring KV (C>1) — ADR-0060 §5.5.
+
+    Tensor layout consumed by this kernel:
+      Q : (T_q, d_head)           one head per CUBE; replicated.
+      K : (S_kv, d_head)          sharded cube_row_wise → each CUBE owns
+          (S_kv/C, d_head); kernel loads as (d_head, S_local) via
+          byte-conserving reshape (reshape-not-transpose caveat).
+      V : (S_kv, d_head)          sharded cube_row_wise → each CUBE owns
+          (S_local, d_head).
+      O : (T_q * C, d_head)       sharded cube_row_wise → each CUBE
+          writes its own (T_q, d_head) slice. NO reduce.
+
+    Algorithm: each CUBE computes a local partial against its current
+    KV block, then over C-1 ring steps the K and V blocks rotate W
+    while online-softmax merges each step into running (m, ℓ, O).
+    """
+    pe_id = tl.program_id(axis=0)
+    # Head-parallel: only PE 0 of each CUBE participates. P6c (future)
+    # will tile T_q across the 8 PEs for intra-CUBE parallelism.
+    if pe_id != 0:
+        return
+
+    S_local = S_kv // C
+    Q = tl.load(q_ptr, shape=(T_q, d_head), dtype="f16")
+    Kc = tl.load(k_ptr, shape=(d_head, S_local), dtype="f16")
+    Vc = tl.load(v_ptr, shape=(S_local, d_head), dtype="f16")
+
+    # ── Step 0: initial partial against own KV block — establishes the
+    # persistent (m, ℓ, O) arena. Intermediates (scores, exp_scores) stay
+    # allocated; ring steps below recycle per-step intermediates inside
+    # tl.scratch_scope to keep peak scratch O(one step) (ADR-0063 §D3).
+    scores = tl.dot(Q, Kc)
+    m = tl.max(scores, axis=-1)
+    exp_scores = tl.exp(scores - m)
+    l = tl.sum(exp_scores, axis=-1)
+    O = tl.dot(exp_scores, Vc)
+
+    # ── Steps 1..C-1: Ring KV rotation + online-softmax merge ──
+    # Per-step intermediates wrapped in tl.scratch_scope; the merged
+    # running (m, ℓ, O) is persisted to the outside-scope handles via
+    # tl.copy_to (ADR-0063 §D3.1) so its bytes survive __exit__.
+    for _ in range(1, C):
+        tl.send(dir="W", src=Kc)
+        tl.send(dir="W", src=Vc)
+        Kc = tl.recv(dir="E", shape=(d_head, S_local), dtype="f16")
+        Vc = tl.recv(dir="E", shape=(S_local, d_head), dtype="f16")
+
+        with tl.scratch_scope():
+            # Partial on rotated KV
+            scores = tl.dot(Q, Kc)
+            m_step = tl.max(scores, axis=-1)
+            exp_scores = tl.exp(scores - m_step)
+            l_step = tl.sum(exp_scores, axis=-1)
+            O_step = tl.dot(exp_scores, Vc)
+
+            # Online-softmax merge into new running (m, ℓ, O)
+            m_new = tl.maximum(m, m_step)
+            scale_old = tl.exp(m - m_new)
+            scale_step = tl.exp(m_step - m_new)
+            l_new = l * scale_old + l_step * scale_step
+            O_new = O * scale_old + O_step * scale_step
+
+            # Persist new running state back to the outside-scope arena.
+            tl.copy_to(m, m_new)
+            tl.copy_to(l, l_new)
+            tl.copy_to(O, O_new)
+        # __exit__: scoped intermediates gone; persistent m, l, O carry
+        # the new running state into the next ring iteration.
+
+    # Final normalise + store — each CUBE writes its own head's rows.
+    O_final = O / l
+    tl.store(o_ptr, O_final)

src/kernbench/benches/_gqa_prefill_short.py

@@ -0,0 +1,114 @@
+"""GQA fused-attention SHORT-CONTEXT prefill kernel (ADR-0060 §B.split.2).
+
+Prefill analogue of ``_gqa_decode_short.py``. Same layout decisions:
+  - K, V head-stacked (h_kv·S_kv, d_head) with row_wise DP so each PE
+    has a contiguous (S_local, d_head) shard at its own address.
+  - Q replicated; kernel uses byte-conserving reshape.
+  - O replicated; group root writes the full byte-conserving result.
+
+No Ring KV (each owned head fully resident at its CUBE). Test
+``test_short_prefill_no_ring_KV_traffic`` asserts no inter-CUBE IPCQ.
+
+The only structural difference from short decode:
+  - ``T_q`` may be > 1 (prefill processes multiple query tokens).
+  - Q is shaped (T_q, h_kv·d_head) — one Q head per KV head (no GQA
+    M-fold here; head-parallel prefill in ADR-0060 §5.5 is 1:1).
+"""
+from __future__ import annotations
+
+
+def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl):
+    """Online-softmax merge — identical to short decode kernel."""
+    m_new = tl.maximum(m_local, m_other)
+    scale_old = tl.exp(m_local - m_new)
+    scale_new = tl.exp(m_other - m_new)
+    l_new = l_local * scale_old + l_other * scale_new
+    O_new = O_local * scale_old + O_other * scale_new
+    return m_new, l_new, O_new
+
+
+def gqa_prefill_short_kernel(
+    q_ptr: int,
+    k_ptr: int,
+    v_ptr: int,
+    o_ptr: int,
+    T_q: int,
+    S_kv: int,
+    h_kv: int,
+    d_head: int,
+    C: int,
+    P: int,
+    kv_per_cube: int,
+    *,
+    tl,
+) -> None:
+    """Short-context prefill with PE-parallel heads + intra-group PE-SP.
+
+    NO Ring KV (each CUBE owns its KV heads fully).
+    """
+    group_size = P // kv_per_cube
+
+    pe_id = tl.program_id(axis=0)
+    pe_in_group = pe_id % group_size
+
+    S_local = S_kv // group_size
+
+    # Q: replicated (T_q, h_kv·d_head) → byte-conserving reshape to
+    # (h_kv·T_q, d_head) — one Q head per KV head, stacked.
+    Q = tl.load(q_ptr,
+                shape=(h_kv * T_q, d_head), dtype="f16")
+    K_T = tl.load(k_ptr,
+                  shape=(d_head, S_local), dtype="f16")
+    V = tl.load(v_ptr,
+                shape=(S_local, d_head), dtype="f16")
+
+    scores = tl.dot(Q, K_T)
+    m_local = tl.max(scores, axis=-1)
+    centered = scores - m_local
+    exp_scores = tl.exp(centered)
+    l_local = tl.sum(exp_scores, axis=-1)
+    O_local = tl.dot(exp_scores, V)
+
+    # Within-group chain reduce-to-root (same machinery as short decode).
+    group_cols = min(4, group_size)
+    group_rows = (group_size + group_cols - 1) // group_cols
+    pe_col_in_group = pe_in_group % group_cols
+    pe_row_in_group = pe_in_group // group_cols
+
+    if group_cols > 1:
+        if pe_col_in_group < group_cols - 1:
+            with tl.scratch_scope():
+                m_other = tl.recv(dir="intra_E", shape=m_local.shape, dtype="f16")
+                l_other = tl.recv(dir="intra_E", shape=l_local.shape, dtype="f16")
+                O_other = tl.recv(dir="intra_E", shape=O_local.shape, dtype="f16")
+                m_new, l_new, O_new = _merge_running(
+                    m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
+                )
+                tl.copy_to(m_local, m_new)
+                tl.copy_to(l_local, l_new)
+                tl.copy_to(O_local, O_new)
+        if pe_col_in_group > 0:
+            tl.send(dir="intra_W", src=m_local)
+            tl.send(dir="intra_W", src=l_local)
+            tl.send(dir="intra_W", src=O_local)
+
+    if pe_col_in_group == 0 and group_rows > 1:
+        if pe_row_in_group < group_rows - 1:
+            with tl.scratch_scope():
+                m_other = tl.recv(dir="intra_S", shape=m_local.shape, dtype="f16")
+                l_other = tl.recv(dir="intra_S", shape=l_local.shape, dtype="f16")
+                O_other = tl.recv(dir="intra_S", shape=O_local.shape, dtype="f16")
+                m_new, l_new, O_new = _merge_running(
+                    m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
+                )
+                tl.copy_to(m_local, m_new)
+                tl.copy_to(l_local, l_new)
+                tl.copy_to(O_local, O_new)
+        if pe_row_in_group > 0:
+            tl.send(dir="intra_N", src=m_local)
+            tl.send(dir="intra_N", src=l_local)
+            tl.send(dir="intra_N", src=O_local)
+
+    if pe_in_group == 0:
+        O_final = O_local / l_local
+        tl.store(o_ptr, O_final)

src/kernbench/benches/milestone_gqa_headline.py

@@ -0,0 +1,236 @@
+"""milestone-gqa-headline bench: real GQA + 2-level SP + Ring KV.
+
+Wires the new ``_gqa_decode`` and ``_gqa_prefill`` kernels through 4
+panels with real GQA (h_q = G·h_kv, G > 1 on the decode side), writing
+per-panel ``op_log_summary`` into ``sweep.json``. Independent from the
+existing ``milestone-gqa-llama70b`` validation-scale bench (which stays
+on the legacy baseline kernels).
+
+Restrictions (P7 first cut):
+  - C ≤ 4 (single-row inter-CUBE ring SFR; multi-row deferred)
+  - Single SIP (default ``topology.yaml`` 4×4 cube mesh; 4-SIP
+    headline deferred)
+  - No figure renderers (defer to a separate cycle)
+
+Panels:
+  single_user_prefill_gqa : prefill C=1, T_q=4, S_kv=16
+  multi_user_prefill_gqa  : prefill C=4 Ring KV, T_q=4, S_kv=16
+  single_user_decode_gqa  : decode C=1, P=8, h_q=8, h_kv=1, S_kv=64
+                            (M-fold + intra-cube row-then-col chain)
+  multi_user_decode_gqa   : decode C=4, P=8, h_q=8, h_kv=1, S_kv=128
+                            (M-fold + 2-level chain reduce-to-root)
+
+Gated by ``GQA_HEADLINE_RUN=1``.
+"""
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+
+from kernbench.benches._gqa_decode import gqa_decode_kernel
+from kernbench.benches._gqa_prefill import gqa_prefill_kernel
+from kernbench.benches.registry import bench
+from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config
+from kernbench.ccl.sfr_config import (
+    configure_sfr_intercube_multisip,
+    configure_sfr_intercube_ring,
+)
+from kernbench.policy.placement.dp import DPPolicy
+
+_OUTPUT_DIR = Path(__file__).resolve().parent / "1H_milestone_output" / "gqa_headline"
+_SWEEP_JSON = _OUTPUT_DIR / "sweep.json"
+
+# ── Panel configs ────────────────────────────────────────────────────
+
+_DTYPE = "f16"
+_D_HEAD = 64
+_T_Q_PREFILL = 4
+_T_Q_DECODE = 1
+_S_KV_PREFILL = 16
+_H_Q_DECODE = 8        # real GQA: G = H_Q_DECODE / H_KV_DECODE = 8
+_H_KV_DECODE = 1
+
+_PANELS = (
+    "single_user_prefill_gqa",
+    "multi_user_prefill_gqa",
+    "single_user_decode_gqa",
+    "multi_user_decode_gqa",
+)
+
+# Each entry: (kind, panel-specific params)
+_PANEL_DISPATCH: dict[str, tuple[str, dict]] = {
+    "single_user_prefill_gqa": ("prefill", {"C": 1, "S_kv": _S_KV_PREFILL}),
+    "multi_user_prefill_gqa":  ("prefill", {"C": 4, "S_kv": _S_KV_PREFILL}),
+    "single_user_decode_gqa":  ("decode",  {"C": 1, "P": 8, "S_kv": 64}),
+    "multi_user_decode_gqa":   ("decode",  {"C": 4, "P": 8, "S_kv": 128}),
+}
+
+
+def _ccl_cfg():
+    return resolve_algorithm_config(
+        load_ccl_config(), name="lrab_hierarchical_allreduce",
+    )
+
+
+# ── Per-kind launch helpers ──────────────────────────────────────────
+
+
+def _run_prefill_panel(ctx, *, panel: str, C: int, S_kv: int) -> None:
+    if C > 1:
+        configure_sfr_intercube_ring(
+            ctx.engine, ctx.spec, _ccl_cfg(), ring_size=C,
+        )
+    dp_q = DPPolicy(cube="replicate", pe="replicate",
+                    num_cubes=C, num_pes=1)
+    dp_kv = DPPolicy(cube="row_wise" if C > 1 else "replicate",
+                     pe="replicate", num_cubes=C, num_pes=1)
+    dp_o = DPPolicy(cube="row_wise" if C > 1 else "replicate",
+                    pe="replicate", num_cubes=C, num_pes=1)
+    q = ctx.zeros((_T_Q_PREFILL, _D_HEAD),
+                  dtype=_DTYPE, dp=dp_q, name=f"{panel}_q")
+    k = ctx.zeros((S_kv, _D_HEAD),
+                  dtype=_DTYPE, dp=dp_kv, name=f"{panel}_k")
+    v = ctx.zeros((S_kv, _D_HEAD),
+                  dtype=_DTYPE, dp=dp_kv, name=f"{panel}_v")
+    o = ctx.empty((_T_Q_PREFILL * C, _D_HEAD),
+                  dtype=_DTYPE, dp=dp_o, name=f"{panel}_o")
+    ctx.launch(
+        panel, gqa_prefill_kernel,
+        q, k, v, o,
+        _T_Q_PREFILL, S_kv, _D_HEAD, C,
+        _auto_dim_remap=False,
+    )
+
+
+def _run_decode_panel(ctx, *, panel: str, C: int, P: int, S_kv: int) -> None:
+    configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
+    dp_full = DPPolicy(cube="replicate", pe="replicate",
+                       num_cubes=C, num_pes=P)
+    dp_kv = DPPolicy(cube="row_wise" if C > 1 else "replicate",
+                     pe="row_wise", num_cubes=C, num_pes=P)
+    q = ctx.zeros((_T_Q_DECODE, _H_Q_DECODE * _D_HEAD),
+                  dtype=_DTYPE, dp=dp_full, name=f"{panel}_q")
+    k = ctx.zeros((S_kv, _H_KV_DECODE * _D_HEAD),
+                  dtype=_DTYPE, dp=dp_kv, name=f"{panel}_k")
+    v = ctx.zeros((S_kv, _H_KV_DECODE * _D_HEAD),
+                  dtype=_DTYPE, dp=dp_kv, name=f"{panel}_v")
+    o = ctx.empty((_T_Q_DECODE, _H_Q_DECODE * _D_HEAD),
+                  dtype=_DTYPE, dp=dp_full, name=f"{panel}_o")
+    ctx.launch(
+        panel, gqa_decode_kernel,
+        q, k, v, o,
+        _T_Q_DECODE, S_kv, _H_Q_DECODE, _H_KV_DECODE, _D_HEAD, C, P,
+        _auto_dim_remap=False,
+    )
+
+
+def _make_bench_fn(panel: str):
+    kind, params = _PANEL_DISPATCH[panel]
+
+    def _bench_fn(ctx):
+        if kind == "prefill":
+            _run_prefill_panel(ctx, panel=panel, **params)
+        else:
+            _run_decode_panel(ctx, panel=panel, **params)
+
+    return _bench_fn
+
+
+# ── Op-log summary ──────────────────────────────────────────────────
+
+
+def _summarize_op_log(op_log) -> dict[str, int]:
+    """Per-panel op_log counts (gemm, ipcq_copy, dma_read, dma_write)."""
+    gemm_count = 0
+    ipcq_copy_count = 0
+    dma_read_count = 0
+    dma_write_count = 0
+    for r in op_log:
+        if r.op_kind == "gemm":
+            gemm_count += 1
+        elif r.op_name == "dma_read":
+            dma_read_count += 1
+        elif r.op_name == "dma_write":
+            dma_write_count += 1
+        elif r.op_name == "ipcq_copy":
+            ipcq_copy_count += 1
+    return {
+        "gemm_count": gemm_count,
+        "ipcq_copy_count": ipcq_copy_count,
+        "dma_read_count": dma_read_count,
+        "dma_write_count": dma_write_count,
+    }
+
+
+def _run_panel(panel: str, topology: str) -> dict:
+    """Run one panel in a fresh engine; return its row dict."""
+    from kernbench.runtime_api.bench_runner import run_bench
+    from kernbench.runtime_api.types import resolve_device
+    from kernbench.sim_engine.engine import GraphEngine
+    from kernbench.topology.builder import resolve_topology
+
+    topo = resolve_topology(topology)
+    result = run_bench(
+        topology=topo, bench_fn=_make_bench_fn(panel),
+        device=resolve_device(None),
+        engine_factory=lambda t, d: GraphEngine(
+            getattr(t, "topology_obj", t), enable_data=True,
+        ),
+    )
+    if not result.completion.ok:
+        raise RuntimeError(
+            f"milestone-gqa-headline panel {panel!r} failed: "
+            f"{result.completion}"
+        )
+    kind, params = _PANEL_DISPATCH[panel]
+    return {
+        "panel": panel,
+        "kind": kind,
+        **params,
+        "op_log_summary": _summarize_op_log(result.engine.op_log),
+    }
+
+
+# ── Bench entry ──────────────────────────────────────────────────────
+
+
+@bench(
+    name="milestone-gqa-headline",
+    description="Headline GQA milestone — real GQA h_q=8/h_kv=1 + 2-level SP (decode) + Ring KV (prefill).",
+)
+def run(torch) -> None:
+    """Drive 4 headline panels through the new GQA kernels; write sweep.json.
+
+    Gated by GQA_HEADLINE_RUN=1.
+    """
+    _OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
+    if not os.environ.get("GQA_HEADLINE_RUN"):
+        raise RuntimeError(
+            "milestone-gqa-headline needs GQA_HEADLINE_RUN=1."
+        )
+
+    topology = os.environ.get("GQA_HEADLINE_TOPOLOGY", "topology.yaml")
+    rows = [_run_panel(panel, topology) for panel in _PANELS]
+    sweep = {
+        "version": 1,
+        "panels": list(_PANELS),
+        "config": {
+            "T_q_prefill": _T_Q_PREFILL,
+            "T_q_decode": _T_Q_DECODE,
+            "S_kv_prefill": _S_KV_PREFILL,
+            "h_q_decode": _H_Q_DECODE,
+            "h_kv_decode": _H_KV_DECODE,
+            "d_head": _D_HEAD,
+        },
+        "rows": rows,
+    }
+    _SWEEP_JSON.write_text(json.dumps(sweep, indent=2))
+    print(f"  milestone-gqa-headline: {len(rows)} rows -> {_SWEEP_JSON}")
+
+    # Sentinel tensor (ADR-0045 D4 / ADR-0054 D2 carve-out).
+    torch.zeros(
+        (1, 1), dtype="f16",
+        dp=DPPolicy(cube="row_wise", pe="replicate", num_cubes=1, num_pes=1),
+        name="milestone_gqa_headline_sentinel",
+    )

src/kernbench/benches/milestone_gqa_llama70b.py

@@ -1,473 +0,0 @@
-"""milestone-gqa-llama70b bench: GQA Llama-70B 4-panel sweep (ADR-0057 v1).
-
-Self-contained eval bench (ADR-0054). Drives the four panels of the GQA
-Llama-70B sharding study through ``run_bench`` with ``enable_data=True``,
-harvests op_log summaries, and writes JSON into
-``benches/1H_milestone_output/gqa/sweep.json``.
-
-v1 (sub-cycle 4a + 4c.0) covers all four panels at validation scale:
-
-  Panel name in JSON / test    Study label   SFR install used
-  ─────────────────────────────────────────────────────────────────────
-  single_user_prefill          TL            configure_sfr_intracube_pe_ring
-  multi_user_prefill           TR            configure_sfr_intercube_multisip
-  single_user_decode           BL            configure_sfr_intracube_pe_ring
-  multi_user_decode            BR            configure_sfr_intercube_multisip
-
-Per the GQA sharding study (`llm_paper_review/notes/GQA_MHA_sharding/scripts
-/_gen_llama70b_1M_4cases.py`):
-
-  Single User (B=1) panels — TL prefill, BL decode:
-    "n_cubes: 8 (1 KV-group)", KV split @ PEs intra-cube. Each cube does
-    its own 8-PE ring with no cube-to-cube attention traffic. At Llama-70B
-    headline scale this is 64 cubes (8 KV-groups × 8 cubes/group), each
-    independently running the kernel for one Q-head; 1 user spans all 64.
-
-  Multi User (B=8) panels — TR prefill, BR decode:
-    "8 cubes / KV-group", KV split @ cubes inter-cube. The 8 cubes of a
-    KV-group form a ring; "Inside each cube: 8 PEs each handle 1 different
-    user → Batch on batch, batch = 8/cube." At headline scale 8 KV-groups
-    run in parallel = 64 cubes serving 8 users.
-
-Kernels use the mesh-native variants (ADR-0059), invoked with the
-``rank_axis`` kwarg (0 for single_user PE-level rings, 1 for multi_user
-cube-level rings). The v1 multi_user kernel gates ``pe_id != 0`` to return,
-which simplifies B=8 → B=1 — that's a validation simplification; the
-per-cube "Batch on batch" parallelism is sub-cycle 4c headline work.
-
-Validation-scale config (ADR-0057 D4) — kept small so the simulator's
-1 MB per-PE TCM scratch budget is not exhausted across n_ranks ring steps:
-``S_q_prefill = S_kv_per_rank = 16``, ``h_q = h_kv = 1``, ``d_head = 64``,
-``n_ranks_single_user = 8`` (8 PEs of one cube), ``n_ranks_multi_user = 4``
-(half a KV-group, vs the study's 8). Headline-scale dims (``S_q = 1M``,
-``S_kv = 1M``, ``h_q = 8 / h_kv = 1`` GQA, ``d_head = 128``, ``B = 8``) are
-also deferred.
-"""
-from __future__ import annotations
-
-import json
-import os
-from pathlib import Path
-from typing import Any
-
-from kernbench.benches._attention_mesh_kv import attention_mesh_kv_kernel
-from kernbench.benches._attention_mesh_mlo import attention_mesh_mlo_kernel
-from kernbench.benches._attention_mesh_mlo_2d import attention_mesh_mlo_2d_kernel
-from kernbench.benches.registry import bench
-from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config
-from kernbench.ccl.sfr_config import (
-    configure_sfr_intercube_multisip,
-    configure_sfr_intracube_pe_ring,
-)
-from kernbench.policy.placement.dp import DPPolicy
-
-_OUTPUT_DIR = Path(__file__).resolve().parent / "1H_milestone_output" / "gqa"
-_SWEEP_JSON = _OUTPUT_DIR / "sweep.json"
-
-# ── Validation-scale config (ADR-0057 D4) ─────────────────────────────
-
-_S_Q_PREFILL = 16
-_S_Q_DECODE = 1
-_S_KV_PER_RANK = 16
-_H_Q = 1
-_H_KV = 1
-_D_HEAD = 64
-_N_RANKS_SINGLE_USER = 8
-_N_RANKS_MULTI_USER = 4
-_DTYPE = "f16"
-
-_PANELS_V1 = (
-    "single_user_prefill",
-    "multi_user_prefill",
-    "single_user_decode",
-    "multi_user_decode",
-)
-
-# Panel → (kernel, SFR install, S_q, n_ranks, rank_axis, mesh_shape)
-# ``mesh_shape`` is ``None`` for 1D-ring kernels and ``(rows, cols)`` for the
-# 2D row-then-col AllReduce-mlo kernel (multi_user_decode); when set, the
-# launch passes ``(mesh_rows, mesh_cols)`` instead of ``n_ranks``.
-_PANEL_DISPATCH: dict[
-    str, tuple[Any, Any, int, int, int, tuple[int, int] | None]
-] = {
-    "single_user_prefill": (
-        attention_mesh_kv_kernel, configure_sfr_intracube_pe_ring,
-        _S_Q_PREFILL, _N_RANKS_SINGLE_USER, 0, None,
-    ),
-    "multi_user_prefill": (
-        attention_mesh_kv_kernel, configure_sfr_intercube_multisip,
-        _S_Q_PREFILL, _N_RANKS_MULTI_USER, 1, None,
-    ),
-    "single_user_decode": (
-        attention_mesh_mlo_kernel, configure_sfr_intracube_pe_ring,
-        _S_Q_DECODE, _N_RANKS_SINGLE_USER, 0, None,
-    ),
-    # multi_user_decode uses the C2 2D AllReduce-mlo kernel. (1, 4)
-    # degenerates to a row-only AllReduce equivalent to the prior 1D ring
-    # at n_ranks=4 — no op_log_summary regression. Headline 8-cube
-    # KV-groups land at (2, 4).
-    "multi_user_decode": (
-        attention_mesh_mlo_2d_kernel, configure_sfr_intercube_multisip,
-        _S_Q_DECODE, _N_RANKS_MULTI_USER, 1, (1, _N_RANKS_MULTI_USER),
-    ),
-}
-
-
-# ── Per-panel bench fn ─────────────────────────────────────────────────
-
-
-def _make_bench_fn(panel: str):
-    kernel, sfr_install, S_q, n_ranks, rank_axis, mesh_shape = (
-        _PANEL_DISPATCH[panel]
-    )
-    is_multi_user = panel.startswith("multi_user_")
-
-    def _bench_fn(ctx):
-        sfr_install(
-            ctx.engine, ctx.spec,
-            resolve_algorithm_config(load_ccl_config(), name="lrab_hierarchical_allreduce"),
-        )
-        if is_multi_user:
-            dp_full = DPPolicy(
-                cube="replicate", pe="replicate",
-                num_cubes=n_ranks, num_pes=8,
-            )
-            dp_kv = DPPolicy(
-                cube="row_wise", pe="replicate",
-                num_cubes=n_ranks, num_pes=8,
-            )
-        else:
-            dp_full = DPPolicy(
-                cube="replicate", pe="replicate",
-                num_cubes=1, num_pes=n_ranks,
-            )
-            dp_kv = DPPolicy(
-                cube="replicate", pe="row_wise",
-                num_cubes=1, num_pes=n_ranks,
-            )
-        q = ctx.zeros((S_q, _H_Q * _D_HEAD),
-                      dtype=_DTYPE, dp=dp_full, name=f"{panel}_q")
-        k = ctx.zeros((_S_KV_PER_RANK * n_ranks, _H_KV * _D_HEAD),
-                      dtype=_DTYPE, dp=dp_kv, name=f"{panel}_k")
-        v = ctx.zeros((_S_KV_PER_RANK * n_ranks, _H_KV * _D_HEAD),
-                      dtype=_DTYPE, dp=dp_kv, name=f"{panel}_v")
-        o = ctx.empty((S_q, _H_Q * _D_HEAD),
-                      dtype=_DTYPE, dp=dp_full, name=f"{panel}_o")
-        # rank_axis is a positional arg; _auto_dim_remap=False keeps
-        # d_head=64 from colliding with the multi_user K's global M=64.
-        if mesh_shape is None:
-            ctx.launch(
-                f"{panel}_mesh", kernel,
-                q, k, v, o,
-                S_q, _S_KV_PER_RANK, _H_Q, _H_KV, _D_HEAD, n_ranks,
-                rank_axis,
-                _auto_dim_remap=False,
-            )
-        else:
-            mesh_rows, mesh_cols = mesh_shape
-            ctx.launch(
-                f"{panel}_mesh", kernel,
-                q, k, v, o,
-                S_q, _S_KV_PER_RANK, _H_Q, _H_KV, _D_HEAD,
-                mesh_rows, mesh_cols,
-                rank_axis,
-                0,  # cube_start=0: this panel's launch starts at cube 0
-                _auto_dim_remap=False,
-            )
-
-    return _bench_fn
-
-
-# ── Op-log summary harvest ─────────────────────────────────────────────
-
-
-def _summarize_op_log(op_log) -> dict[str, int]:
-    """Counts per ADR-0057 D7 op_log_summary contract."""
-    gemm_count = 0
-    ipcq_send_count = 0
-    ipcq_recv_count = 0
-    dma_read_count = 0
-    dma_write_count = 0
-    for r in op_log:
-        if r.op_kind == "gemm":
-            gemm_count += 1
-        elif r.op_name == "dma_read":
-            dma_read_count += 1
-        elif r.op_name == "dma_write":
-            dma_write_count += 1
-        elif r.op_name == "ipcq_send":
-            ipcq_send_count += 1
-        elif r.op_name == "ipcq_recv":
-            ipcq_recv_count += 1
-        elif r.op_name == "ipcq_copy":
-            # The inbound DMA records ipcq_copy (one per send/recv pair).
-            # Count it as both a send and a recv side so the row's
-            # ipcq_send_count and ipcq_recv_count are non-zero even when
-            # the engine logs the collective via the inbound copy alone.
-            ipcq_send_count += 1
-            ipcq_recv_count += 1
-    return {
-        "gemm_count": gemm_count,
-        "ipcq_send_count": ipcq_send_count,
-        "ipcq_recv_count": ipcq_recv_count,
-        "dma_read_count": dma_read_count,
-        "dma_write_count": dma_write_count,
-    }
-
-
-def _run_panel(panel: str, topology: str) -> dict:
-    """Run one panel via a fresh engine; return its row dict."""
-    from kernbench.runtime_api.bench_runner import run_bench
-    from kernbench.runtime_api.types import resolve_device
-    from kernbench.sim_engine.engine import GraphEngine
-    from kernbench.topology.builder import resolve_topology
-
-    topo = resolve_topology(topology)
-    result = run_bench(
-        topology=topo, bench_fn=_make_bench_fn(panel),
-        device=resolve_device(None),
-        engine_factory=lambda t, d: GraphEngine(
-            getattr(t, "topology_obj", t), enable_data=True,
-        ),
-    )
-    if not result.completion.ok:
-        raise RuntimeError(
-            f"milestone-gqa-llama70b panel {panel!r} failed: {result.completion}"
-        )
-    _, _, _, n_ranks, _, _ = _PANEL_DISPATCH[panel]
-    return {
-        "panel": panel,
-        "n_ranks": n_ranks,
-        "op_log_summary": _summarize_op_log(result.engine.op_log),
-    }
-
-
-# ── Figure renderers (sub-cycle 4c, 5 of 6 figures) ──────────────────
-#
-# Sixth figure ``gqa_scaling.png`` is deferred to after sub-cycle 4b
-# lands the Q/cube ∈ {1, 2, 4} sweep on multi_user_* panels — it needs
-# multiple sweep.json rows per multi_user panel to be meaningful.
-
-_OP_LOG_KEYS = (
-    "gemm_count",
-    "ipcq_send_count",
-    "ipcq_recv_count",
-    "dma_read_count",
-    "dma_write_count",
-)
-_OP_LOG_DISPLAY = {
-    "gemm_count":      "GEMM",
-    "ipcq_send_count": "IPCQ send",
-    "ipcq_recv_count": "IPCQ recv",
-    "dma_read_count":  "DMA read",
-    "dma_write_count": "DMA write",
-}
-_OP_LOG_COLORS = {
-    "gemm_count":      "#F59E0B",
-    "ipcq_send_count": "#3B82F6",
-    "ipcq_recv_count": "#10B981",
-    "dma_read_count":  "#A855F7",
-    "dma_write_count": "#EF4444",
-}
-_PANEL_DISPLAY = {
-    "single_user_prefill": "single_user / prefill",
-    "multi_user_prefill":  "multi_user / prefill",
-    "single_user_decode":  "single_user / decode",
-    "multi_user_decode":   "multi_user / decode",
-}
-
-
-def _load_sweep_data(sweep_json: Path | str) -> dict:
-    sweep_json = Path(sweep_json)
-    if not sweep_json.exists():
-        return {"rows": [], "config": {}, "panels": []}
-    return json.loads(sweep_json.read_text())
-
-
-def _row_for(rows: list, panel: str) -> dict | None:
-    for r in rows:
-        if r.get("panel") == panel:
-            return r
-    return None
-
-
-def emit_panel_op_log_summary(
-    panel: str,
-    sweep_json: Path | str = _SWEEP_JSON,
-    out_dir: Path | str = _OUTPUT_DIR,
-) -> str | None:
-    """One bar chart of the 5 op_log counts for ``panel``.
-
-    Returns the written PNG path, or ``None`` when sweep.json is empty
-    or the requested panel is absent.
-    """
-    import matplotlib.pyplot as plt
-
-    data = _load_sweep_data(sweep_json)
-    row = _row_for(data.get("rows", []), panel)
-    if row is None:
-        return None
-    summary = row.get("op_log_summary", {})
-    n_ranks = row.get("n_ranks")
-
-    labels = [_OP_LOG_DISPLAY[k] for k in _OP_LOG_KEYS]
-    values = [summary.get(k, 0) for k in _OP_LOG_KEYS]
-    colors = [_OP_LOG_COLORS[k] for k in _OP_LOG_KEYS]
-
-    fig, ax = plt.subplots(figsize=(8, 5))
-    bars = ax.bar(labels, values, color=colors)
-    for b, v in zip(bars, values):
-        ax.text(b.get_x() + b.get_width() / 2, b.get_height(),
-                f"{int(v)}", ha="center", va="bottom", fontsize=9)
-    ax.set_title(
-        f"{_PANEL_DISPLAY.get(panel, panel)}  (n_ranks={n_ranks})",
-        fontsize=12, fontweight="bold",
-    )
-    ax.set_ylabel("count")
-    ax.grid(True, axis="y", alpha=0.3)
-    fig.tight_layout()
-
-    out_dir = Path(out_dir)
-    out_dir.mkdir(parents=True, exist_ok=True)
-    out = out_dir / f"gqa_op_log_{panel}.png"
-    fig.savefig(out, dpi=120)
-    plt.close(fig)
-    return str(out)
-
-
-def emit_gqa_comparison(
-    sweep_json: Path | str = _SWEEP_JSON,
-    out_dir: Path | str = _OUTPUT_DIR,
-) -> str | None:
-    """Grouped-bar chart comparing the 5 op_log counts across all panels."""
-    import matplotlib.pyplot as plt
-    import numpy as np
-
-    data = _load_sweep_data(sweep_json)
-    panels_in = data.get("panels") or list(_PANELS_V1)
-    rows = data.get("rows", [])
-    panels = [p for p in panels_in if _row_for(rows, p) is not None]
-    if not panels:
-        return None
-
-    n_groups = len(panels)
-    n_series = len(_OP_LOG_KEYS)
-    x = np.arange(n_groups)
-    width = 0.8 / n_series
-
-    fig, ax = plt.subplots(figsize=(11, 6))
-    for i, key in enumerate(_OP_LOG_KEYS):
-        offset = (i - (n_series - 1) / 2) * width
-        vals = [_row_for(rows, p)["op_log_summary"].get(key, 0)
-                for p in panels]
-        ax.bar(x + offset, vals, width,
-               label=_OP_LOG_DISPLAY[key], color=_OP_LOG_COLORS[key])
-
-    ax.set_xticks(x)
-    ax.set_xticklabels(
-        [f"{_PANEL_DISPLAY.get(p, p)}\n(n_ranks={_row_for(rows, p)['n_ranks']})"
-         for p in panels],
-        fontsize=8,
-    )
-    ax.set_ylabel("count")
-    ax.set_title("GQA Llama-70B — op_log summary across panels",
-                 fontsize=13, fontweight="bold")
-    ax.legend(fontsize=8, loc="upper right")
-    ax.grid(True, axis="y", alpha=0.3)
-    fig.tight_layout()
-
-    out_dir = Path(out_dir)
-    out_dir.mkdir(parents=True, exist_ok=True)
-    out = out_dir / "gqa_comparison.png"
-    fig.savefig(out, dpi=120)
-    plt.close(fig)
-    return str(out)
-
-
-def emit_all_gqa_plots(
-    sweep_json: Path | str = _SWEEP_JSON,
-    out_dir: Path | str = _OUTPUT_DIR,
-) -> list[str]:
-    """Render all 5 in-scope figures and return the written paths.
-
-    Sub-cycle 4c v1 emits 5 of the 6 figures ADR-0057 D3 lists; the
-    6th (gqa_scaling.png) needs sub-cycle 4b's Q/cube sweep data.
-    """
-    paths: list[str] = []
-    for panel in _PANELS_V1:
-        p = emit_panel_op_log_summary(panel, sweep_json, out_dir)
-        if p is not None:
-            paths.append(p)
-    comp = emit_gqa_comparison(sweep_json, out_dir)
-    if comp is not None:
-        paths.append(comp)
-    return paths
-
-
-# ── Bench entry ────────────────────────────────────────────────────────
-
-
-@bench(
-    name="milestone-gqa-llama70b",
-    description="1H milestone: GQA Llama-70B 4-panel sweep (ADR-0057 v1).",
-)
-def run(torch) -> None:
-    """Drive the four GQA panels at validation scale; write sweep.json and figures.
-
-    Modes (mutually exclusive):
-      MILESTONE_FAST=1  Skip the sweep; re-render figures from the
-                        committed sweep.json. Seconds, no simulator.
-      GQA_VALIDATION=1  Run the four-panel validation sweep + figures.
-                        ~1-2h on the full simulator.
-
-    Headline-scale mode is deferred to sub-cycle 4c (figures landed
-    here; headline-scale + scaling figure await sub-cycle 4b).
-    A sentinel tensor is submitted at the end so run_bench's ADR-0045 D4
-    "at least one request" contract is satisfied even when the panels
-    are skipped via MILESTONE_FAST=1.
-    """
-    _OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
-    fast = bool(os.environ.get("MILESTONE_FAST"))
-    if not fast and not os.environ.get("GQA_VALIDATION"):
-        raise RuntimeError(
-            "milestone-gqa-llama70b v1 needs GQA_VALIDATION=1 (run the "
-            "sweep) or MILESTONE_FAST=1 (reuse committed sweep.json). "
-            "Headline mode is deferred to sub-cycle 4b/4c per ADR-0057 D3."
-        )
-
-    if not fast:
-        topology = os.environ.get("GQA_TOPOLOGY", "topology.yaml")
-        rows = [_run_panel(panel, topology) for panel in _PANELS_V1]
-        sweep = {
-            "version": 1,
-            "validation_scale": True,
-            "panels": list(_PANELS_V1),
-            "config": {
-                "S_q_prefill": _S_Q_PREFILL,
-                "S_kv_per_rank": _S_KV_PER_RANK,
-                "h_q": _H_Q,
-                "h_kv": _H_KV,
-                "d_head": _D_HEAD,
-                "n_ranks_single_user": _N_RANKS_SINGLE_USER,
-                "n_ranks_multi_user": _N_RANKS_MULTI_USER,
-            },
-            "rows": rows,
-        }
-        _SWEEP_JSON.write_text(json.dumps(sweep, indent=2))
-        print(f"  milestone-gqa-llama70b: {len(rows)} rows -> {_SWEEP_JSON}")
-    elif not _SWEEP_JSON.exists():
-        raise RuntimeError(
-            f"MILESTONE_FAST=1 requires {_SWEEP_JSON} to exist; "
-            "run with GQA_VALIDATION=1 once to seed it."
-        )
-
-    paths = emit_all_gqa_plots()
-    print(f"  milestone-gqa-llama70b: {len(paths)} figures -> {_OUTPUT_DIR} "
-          f"(fast={fast})")
-
-    # Sentinel tensor (ADR-0045 D4 / ADR-0054 D2 carve-out).
-    torch.zeros(
-        (1, 1), dtype="f16",
-        dp=DPPolicy(cube="row_wise", pe="replicate", num_cubes=1, num_pes=1),
-        name="milestone_gqa_sentinel",
-    )

src/kernbench/ccl/sfr_config.py

@@ -237,3 +237,119 @@ def configure_sfr_intracube_pe_ring(
         algo_module=mock_module,
         rank_to_pe=pe_idx_to_pe,
     )
+
+
+# ── Inter-cube 1D ring (ADR-0060 §5.5 prefill Ring KV) ─────────────────
+
+
+def configure_sfr_intercube_ring(
+    engine: Any,
+    spec: dict,
+    cfg: dict,
+    *,
+    ring_size: int | None = None,
+) -> dict[str, Any]:
+    """Install intra-cube PE grid + a 1D CUBE-level ring with wrap.
+
+    Direction namespaces (disjoint, same as
+    ``configure_sfr_intercube_multisip``):
+
+      - ``intra_N/S/E/W`` : 2×4 PE grid within each cube (no wrap)
+      - ``E/W``           : 1D ring of cubes 0..ring_size-1 WITH WRAP
+                            (symmetric to ``configure_sfr_intracube_pe_ring``
+                             at PE level — wrap applied at CUBE level here)
+      - ``global_*``      : SIP topology (same as multisip)
+
+    N/S at CUBE level are intentionally NOT installed — use
+    ``configure_sfr_intercube_multisip`` for the full 4×4 cube mesh.
+
+    Args:
+        ring_size: number of CUBEs in the ring (wrap applies to cubes
+            0..ring_size-1). Defaults to the full cube_mesh count.
+            Must be ≤ mesh_w (single row); multi-row rings span
+            non-neighbour boundaries.
+    """
+    cm = spec["sip"]["cube_mesh"]
+    mesh_w = int(cm["w"])
+    mesh_h = int(cm["h"])
+    n_cubes = mesh_w * mesh_h
+    sips_cfg = spec.get("system", {}).get("sips", {})
+    n_sips = int(sips_cfg.get("count", 1))
+    sip_topology = str(sips_cfg.get("topology", "ring_1d"))
+    sip_w = sips_cfg.get("w")
+    sip_h = sips_cfg.get("h")
+    sip_w = int(sip_w) if sip_w is not None else None
+    sip_h = int(sip_h) if sip_h is not None else None
+
+    if ring_size is None:
+        ring_size = n_cubes
+    if ring_size > mesh_w:
+        raise ValueError(
+            f"intercube_ring ring_size={ring_size} > mesh_w={mesh_w}; "
+            "multi-row rings cross non-neighbour boundaries"
+        )
+
+    if sip_topology not in _TOPO_BUILTINS:
+        raise ValueError(
+            f"Unknown sip topology '{sip_topology}'. "
+            f"Available: {list(_TOPO_BUILTINS)}"
+        )
+    _sip_topo_fn_raw = _TOPO_BUILTINS[sip_topology]
+
+    def sip_topo_fn(rank: int, ws: int) -> dict:
+        if sip_w is not None and sip_h is not None:
+            try:
+                return _sip_topo_fn_raw(rank, ws, w=sip_w, h=sip_h)
+            except TypeError:
+                pass
+        return _sip_topo_fn_raw(rank, ws)
+
+    pes_per_cube = _PES_PER_CUBE
+    world_size = n_sips * n_cubes * pes_per_cube
+    pe_idx_to_pe: list[tuple[int, int, int]] = [
+        (sip, cube, pe)
+        for sip in range(n_sips)
+        for cube in range(n_cubes)
+        for pe in range(pes_per_cube)
+    ]
+
+    def _pe_idx(sip: int, cube: int, pe: int) -> int:
+        return (sip * n_cubes + cube) * pes_per_cube + pe
+
+    def _neighbors(pe_idx: int, ws: int, _base: dict) -> dict[str, int]:
+        tmp = pe_idx
+        pe = tmp % pes_per_cube
+        tmp //= pes_per_cube
+        cube = tmp % n_cubes
+        sip = tmp // n_cubes
+
+        nbrs: dict[str, int] = {}
+
+        # ── Intra-cube (intra_N/S/E/W) ──
+        for d, peer_pe in _intra_cube_neighbors(pe).items():
+            nbrs[d] = _pe_idx(sip, cube, peer_pe)
+
+        # ── Cube ring (E/W with wrap for cubes 0..ring_size-1) ──
+        if cube < ring_size:
+            nbrs["E"] = _pe_idx(sip, (cube + 1) % ring_size, pe)
+            nbrs["W"] = _pe_idx(sip, (cube - 1) % ring_size, pe)
+
+        # ── Inter-SIP same-(cube, pe) (global_*) ──
+        if n_sips > 1:
+            sip_nbrs = sip_topo_fn(sip, n_sips)
+            for d, peer_sip in sip_nbrs.items():
+                nbrs[f"global_{d}"] = _pe_idx(peer_sip, cube, pe)
+
+        return nbrs
+
+    mock_module = types.SimpleNamespace(neighbors=_neighbors)
+
+    cfg_copy = dict(cfg)
+    cfg_copy["world_size"] = world_size
+    cfg_copy["topology"] = "none"
+
+    return install_ipcq(
+        engine, spec, cfg_copy,
+        algo_module=mock_module,
+        rank_to_pe=pe_idx_to_pe,
+    )

src/kernbench/common/cpu_issue_cost.py

@@ -0,0 +1,52 @@
+"""Per-op-type CPU issue cost table (ADR-0064 D1).
+
+Replaces the single uniform ``dispatch_cycles`` scalar with a cost table
+keyed by command kind. Charged on PE_CPU at issue time (before the command
+is dispatched to PE_SCHEDULER) so the hybrid's CPU-saturation win
+(ADR-0060 §1) becomes measurable.
+
+The table is consulted by ``TLContext._emit_dispatch_overhead(kind)``;
+live PE_CPU paths (greenlet via ``kernel_runner.py``, legacy replay via
+``pe_cpu.py:_execute_legacy``) construct TLContext with
+``issue_cost_table=DEFAULT_CPU_ISSUE_COST`` so all benches see the cost.
+
+Absolute ns values are provisional (ADR-0064 review item #1). The
+defensible claim is the **ratio** — composite ≫ primitive.
+"""
+from __future__ import annotations
+
+from typing import Literal
+
+OpKind = Literal[
+    "composite",
+    "load",
+    "store",
+    "dot",
+    "math",
+    "ipcq_send",
+    "ipcq_recv",
+    "copy_to",
+]
+
+
+DEFAULT_CPU_ISSUE_COST: dict[str, int] = {
+    "composite": 40,
+    "load":      5,
+    "store":     5,
+    "dot":       5,
+    "math":      5,
+    "ipcq_send": 5,
+    "ipcq_recv": 5,
+    "copy_to":   5,
+}
+
+
+def get_issue_cost(kind: str, table: dict[str, int] | None = None) -> int:
+    """Return per-op-type CPU issue cost in ns.
+
+    Unknown kinds return 0 (no charge) so adding a new ``tl.*`` op kind
+    doesn't accidentally over-charge before the table is updated.
+    """
+    if table is None:
+        table = DEFAULT_CPU_ISSUE_COST
+    return table.get(kind, 0)

src/kernbench/common/pe_commands.py

@@ -73,6 +73,12 @@ class TensorHandle:
     data: object = None              # reserved for validate mode
     space: str = "tcm"               # MemoryStore space ("tcm" | "hbm" | "sram")
     pinned: bool = False             # operand already DMA-staged in TCM (via tl.load)
+    # ADR-0062 §D2: lazy tl.load attaches a LoadFuture here. None for
+    # handles that have no in-flight DMA (constants, math outputs, etc.).
+    # Consumer ops call _await_pending() to yield on the future before
+    # emitting their own command. Excluded from eq/hash/repr so handle
+    # identity is unaffected by pending state.
+    pending: object = field(default=None, compare=False, hash=False, repr=False)
 
 
 @dataclass(frozen=True)
@@ -140,6 +146,22 @@ class MathCmd:
     data_op: bool = True
 
 
+@dataclass(frozen=True)
+class CopyCmd:
+    """TCM-to-TCM byte copy (ADR-0063 §D3.1).
+
+    Emitted by ``tl.copy_to`` to persist a scoped result's bytes to an
+    outside-``scratch_scope`` (persistent) address — the two-arena
+    pattern for tiled flash attention. Runs on the vector engine;
+    op_log classifies as ``op_kind="math"``, ``op_name="copy"``.
+    """
+
+    src: TensorHandle
+    dst: TensorHandle
+    nbytes: int
+    data_op: bool = True
+
+
 @dataclass(frozen=True)
 class CompositeCmd:
     """Composite command: tiled pipeline of DMA_READ + COMPUTE + DMA_WRITE.
@@ -178,7 +200,7 @@ class PeCpuOverheadCmd:
 
 # Union type for all PE commands
 PeCommand = (
-    DmaReadCmd | DmaWriteCmd | GemmCmd | MathCmd
+    DmaReadCmd | DmaWriteCmd | GemmCmd | MathCmd | CopyCmd
     | CompositeCmd | WaitCmd | PeCpuOverheadCmd
 )
 

src/kernbench/components/builtin/pe_cpu.py

@@ -184,6 +184,7 @@ class PeCpuComponent(ComponentBase):
         self, env, kernel_fn, kernel_args, num_programs, scheduler_id,
     ) -> Generator:
         """Legacy Phase 0 + replay: generate command list, then dispatch."""
+        from kernbench.common.cpu_issue_cost import DEFAULT_CPU_ISSUE_COST
         from kernbench.common.pe_commands import (
             CompositeCmd, PeCpuOverheadCmd, PeInternalTxn, WaitCmd,
         )
@@ -193,6 +194,7 @@ class PeCpuComponent(ComponentBase):
             pe_id=self._pe_idx, num_programs=num_programs,
             cube_id=self._cube_idx, num_cubes=self._num_cubes,
             dispatch_cycles=0,
+            issue_cost_table=DEFAULT_CPU_ISSUE_COST,
         )
         run_kernel(kernel_fn, tl, *kernel_args)
         commands = tl.commands

src/kernbench/components/builtin/pe_math.py

@@ -99,20 +99,25 @@ class PeMathComponent(PeEngineBase):
         self._on_process_end(env, token)
 
     def handle_command(self, env: simpy.Environment, pe_txn: PeInternalTxn) -> Generator:
-        """PeInternalTxn handling for standalone MathCmd (CCL kernels).
+        """PeInternalTxn handling for standalone MathCmd / CopyCmd.
 
         Latency = max(overhead_ns, _compute_ns(num_elements)):
         - overhead_ns: fixed per-invocation setup cost (from node attrs).
         - _compute_ns: SIMD cycle-based model (from vector_width + clock_freq).
         The larger of the two dominates (setup-bound vs compute-bound).
+
+        CopyCmd (ADR-0063 §D3.1): vector-engine on-chip byte copy; cost
+        model = _compute_ns(prod(dst.shape)).
         """
-        from kernbench.common.pe_commands import MathCmd
+        from kernbench.common.pe_commands import CopyCmd, MathCmd
         import math as _math
 
         cmd = pe_txn.command
         num_elements = 0
         if isinstance(cmd, MathCmd) and cmd.out.shape:
             num_elements = _math.prod(cmd.out.shape)
+        elif isinstance(cmd, CopyCmd) and cmd.dst.shape:
+            num_elements = _math.prod(cmd.dst.shape)
 
         overhead_ns = float(self.node.attrs.get("overhead_ns", 0.0))
         compute_ns = self._compute_ns(num_elements)

src/kernbench/components/builtin/pe_scheduler.py

@@ -42,12 +42,16 @@ class PeSchedulerComponent(ComponentBase):
     def _ensure_dispatch_table(cls) -> None:
         if cls._CMD_DISPATCH:
             return
-        from kernbench.common.pe_commands import DmaReadCmd, DmaWriteCmd, GemmCmd, MathCmd
+        from kernbench.common.pe_commands import (
+            CopyCmd, DmaReadCmd, DmaWriteCmd, GemmCmd, MathCmd,
+        )
         cls._CMD_DISPATCH = {
             DmaReadCmd: "pe_dma",
             DmaWriteCmd: "pe_dma",
             GemmCmd: "pe_gemm",
             MathCmd: "pe_math",
+            # ADR-0063 §D3.1: tl.copy_to → vector engine.
+            CopyCmd: "pe_math",
         }
 
     def __init__(self, node: Node, ctx: ComponentContext | None = None) -> None:

src/kernbench/sim_engine/data_executor.py

@@ -238,6 +238,12 @@ def _compute_math(op: str, inputs: list[np.ndarray], axis: int | None) -> np.nda
 
     x = inputs[0]
 
+    # ADR-0063 §D3.1: copy is the vector-engine byte move used by
+    # tl.copy_to to persist a scoped result to the persistent arena.
+    # In data mode the bytes flow through as-is — identity op.
+    if op == "copy":
+        return x
+
     # Unary
     if op == "exp":
         return np.exp(x)

src/kernbench/sim_engine/op_log.py

@@ -186,7 +186,7 @@ class OpLogger:
 def _extract_op_info(msg: Any) -> tuple[str, str, dict[str, Any]]:
     """Extract op_kind, op_name, params from a data_op message."""
     from kernbench.common.pe_commands import (
-        DmaReadCmd, DmaWriteCmd, GemmCmd, MathCmd, CompositeCmd,
+        CompositeCmd, CopyCmd, DmaReadCmd, DmaWriteCmd, GemmCmd, MathCmd,
     )
     if isinstance(msg, DmaReadCmd):
         return "memory", "dma_read", {
@@ -237,6 +237,16 @@ def _extract_op_info(msg: Any) -> tuple[str, str, dict[str, Any]]:
             "dtype": msg.out.dtype,
             "axis": msg.axis,
         }
+    if isinstance(msg, CopyCmd):
+        return "math", "copy", {
+            "src_addr": msg.src.addr,
+            "src_space": getattr(msg.src, "space", "tcm"),
+            "dst_addr": msg.dst.addr,
+            "dst_space": getattr(msg.dst, "space", "tcm"),
+            "shape": msg.src.shape,
+            "dtype": msg.src.dtype,
+            "nbytes": msg.nbytes,
+        }
     if isinstance(msg, CompositeCmd):
         params: dict[str, Any] = {
             "op": msg.op,

src/kernbench/triton_emu/kernel_runner.py

@@ -89,6 +89,7 @@ class KernelRunner:
         4. Dispatches each command through SimPy components
         5. Returns results to the kernel
         """
+        from kernbench.common.cpu_issue_cost import DEFAULT_CPU_ISSUE_COST
         from kernbench.triton_emu.tl_context import TLContext
 
         self._parent = greenlet.getcurrent()
@@ -102,6 +103,7 @@ class KernelRunner:
             runner=self,
             scratch_base=self._scratch_base,
             scratch_size=self._scratch_size,
+            issue_cost_table=DEFAULT_CPU_ISSUE_COST,
         )
         self._tl = tl  # exposed so switch_to_simpy can re-set on restore
 
@@ -144,7 +146,9 @@ class KernelRunner:
                 cmd = _switch_kernel()
 
             elif isinstance(cmd, DmaReadCmd):
-                # Dispatch DMA through SimPy components
+                # Legacy blocking path — retained as a fallback for any
+                # caller that bypasses the lazy ``tl.load`` surface. New
+                # lazy loads come in as ("load_issue", future) below.
                 done_evt = env.event()
                 pe_txn = PeInternalTxn(
                     command=cmd, done=done_evt, pe_prefix=self._pe_prefix,
@@ -245,6 +249,45 @@ class KernelRunner:
                 }
                 cmd = _switch_kernel(result)
 
+            elif isinstance(cmd, tuple) and len(cmd) == 2 and cmd[0] == "load_issue":
+                # ADR-0062 §D2: lazy tl.load. Post the DmaReadCmd, store the
+                # done event on the LoadFuture, switch back immediately — do
+                # NOT yield done_evt here. The auto-wait at first use
+                # ("load_await" below) is what eventually yields it.
+                future = cmd[1]
+                done_evt = env.event()
+                pe_txn = PeInternalTxn(
+                    command=future.cmd, done=done_evt, pe_prefix=self._pe_prefix,
+                )
+                yield self._out_ports[self._scheduler_id].put(pe_txn)
+                future.event = done_evt
+                cmd = _switch_kernel(None)
+
+            elif isinstance(cmd, tuple) and len(cmd) == 2 and cmd[0] == "load_await":
+                # ADR-0062 §D2: auto-wait at first use. Yield on the future's
+                # DMA event if not yet triggered, then read data and attach
+                # it to the handle (frozen dataclass — mutate via object.
+                # __setattr__, the same controlled pattern Phase 1
+                # blocking tl.load used).
+                handle = cmd[1]
+                future = handle.pending
+                if future is not None and not future.resolved:
+                    if future.event is not None and not future.event.triggered:
+                        yield future.event
+                    data = None
+                    if self._store is not None:
+                        try:
+                            data = self._store.read(
+                                "hbm", future.cmd.src_addr,
+                                shape=handle.shape, dtype=handle.dtype,
+                            )
+                        except KeyError:
+                            pass
+                    future.data = data
+                    future.resolved = True
+                    object.__setattr__(handle, "data", data)
+                cmd = _switch_kernel(None)
+
             elif isinstance(cmd, tuple) and len(cmd) == 2 and cmd[0] == "recv_async":
                 # Non-blocking recv: post the IpcqRequest now, store the
                 # event in the future, return None to kernel.

src/kernbench/triton_emu/tl_context.py

@@ -22,6 +22,7 @@ from kernbench.common.pe_commands import (
     EPILOGUE_OPS,
     CompletionHandle,
     CompositeCmd,
+    CopyCmd,
     DmaReadCmd,
     DmaWriteCmd,
     GemmCmd,
@@ -42,13 +43,67 @@ _DTYPE_BYTES: dict[str, int] = {
 }
 
 
+class LoadFuture:
+    """Lazy ``tl.load`` future (ADR-0062 §D2).
+
+    Mirrors ``RecvFuture`` for IPCQ, generalised to HBM loads. Carries
+    the originating ``DmaReadCmd``, the SimPy completion event (set by
+    the runner once the DMA has been issued), and a resolved flag.
+
+    Consumer ops auto-wait via ``TLContext._await_pending(handle)``,
+    which yields ``event`` if ``resolved`` is False, then reads the
+    DMA's bytes into ``handle.data`` and marks the future resolved.
+    """
+
+    __slots__ = ("cmd", "event", "resolved", "data")
+
+    def __init__(self, cmd: DmaReadCmd) -> None:
+        self.cmd = cmd
+        self.event: object | None = None  # simpy.Event set by runner
+        self.resolved: bool = False
+        self.data: object = None
+
+
+class _ScratchScope:
+    """Context manager that recycles per-tile scratch (ADR-0063 D1).
+
+    ``__enter__`` snapshots ``_scratch_cursor``; ``__exit__`` restores it,
+    so every handle allocated inside the ``with``-block has its address
+    freed for the next iteration. Persistent state (running ``(m, ℓ, O)``,
+    prefetch buffers) lives outside the scope per ADR-0063 D3.
+    """
+
+    def __init__(self, ctx: "TLContext") -> None:
+        self._ctx = ctx
+        self._save: int | None = None
+
+    def __enter__(self) -> "_ScratchScope":
+        self._save = self._ctx._scratch_cursor
+        return self
+
+    def __exit__(self, *exc_info: object) -> bool:
+        if self._save is not None:
+            self._ctx._scratch_cursor = self._save
+        return False
+
+
 class TLContext:
     """Fake Triton Language context.
 
     Args:
         pe_id: program instance index (returned by program_id).
         num_programs: total number of program instances.
-        dispatch_cycles: PE_CPU overhead per tl API call (auto-inserted).
+        dispatch_cycles: uniform PE_CPU overhead per tl API call. Used as
+            a fallback when ``issue_cost_table`` is None (ADR-0046 §D6
+            back-compat). When ``issue_cost_table`` is provided, the
+            per-kind table value is used instead.
+        issue_cost_table: optional per-op-type CPU issue cost table
+            (ADR-0064 D1). When provided, each ``tl.*`` call charges the
+            table value keyed by op kind ("composite", "load", "store",
+            "dot", "math", "ipcq_send", "ipcq_recv", "copy_to"). Unknown
+            kinds fall back to ``dispatch_cycles``. Live PE_CPU paths
+            construct TLContext with ``DEFAULT_CPU_ISSUE_COST`` so the
+            hybrid's CPU-saturation lever is measurable.
     """
 
     def __init__(
@@ -61,12 +116,14 @@ class TLContext:
         num_cubes: int = 1,
         scratch_base: int = 0,
         scratch_size: int = 1 << 20,  # 1 MiB per kernel invocation
+        issue_cost_table: dict[str, int] | None = None,
     ) -> None:
         self._pe_id = pe_id
         self._num_programs = num_programs
         self._cube_id = cube_id
         self._num_cubes = num_cubes
         self._dispatch_cycles = dispatch_cycles
+        self._issue_cost_table = issue_cost_table
         self._commands: list[PeCommand] = []
         self._handle_counter = 0
         self._completion_counter = 0
@@ -79,6 +136,22 @@ class TLContext:
         self._scratch_size = scratch_size
         self._scratch_cursor = 0
 
+    def scratch_scope(self) -> _ScratchScope:
+        """Per-tile scratch recycling context manager (ADR-0063).
+
+        Usage:
+
+            with tl.scratch_scope():
+                s   = tl.dot(q, k_t)        # per-tile temporaries —
+                p   = tl.softmax(s)         # their scratch is rewound
+                o_j = tl.dot(p, v)          # on __exit__
+
+        Persistent state (running ``(m, ℓ, O)``, lazy-load prefetch
+        buffers) must be allocated **outside** the scope; only handles
+        allocated inside are recycled.
+        """
+        return _ScratchScope(self)
+
     def _scratch_alloc(self, nbytes: int) -> int:
         """Allocate a unique scratch address for an output TensorHandle.
 
@@ -120,9 +193,23 @@ class TLContext:
     def _nbytes(self, shape: tuple[int, ...], dtype: str) -> int:
         return math.prod(shape) * self._dtype_bytes(dtype)
 
-    def _emit_dispatch_overhead(self) -> None:
+    def _emit_dispatch_overhead(self, kind: str | None = None) -> None:
-        if self._dispatch_cycles > 0:
+        """Charge per-op-type CPU issue cost (ADR-0064 D1).
-            self._emit(PeCpuOverheadCmd(cycles=self._dispatch_cycles))
+
+        When ``issue_cost_table`` was provided, look up the per-kind cost
+        and emit ``PeCpuOverheadCmd(cycles=N)`` if N > 0. Unknown kinds
+        fall back to the uniform ``dispatch_cycles`` for forward-compat
+        when a new ``tl.*`` op is added before the table is updated.
+
+        When ``issue_cost_table`` is None, preserve the ADR-0046 §D6
+        contract: emit ``PeCpuOverheadCmd(dispatch_cycles)`` if positive.
+        """
+        if self._issue_cost_table is not None and kind is not None:
+            cycles = self._issue_cost_table.get(kind, self._dispatch_cycles)
+        else:
+            cycles = self._dispatch_cycles
+        if cycles > 0:
+            self._emit(PeCpuOverheadCmd(cycles=cycles))
 
     def _make_handle(
         self, addr: int, shape: tuple[int, ...], dtype: str,
@@ -177,37 +264,103 @@ class TLContext:
     def load(
         self, ptr: int, shape: tuple[int, ...], dtype: str = "f16",
     ) -> TensorHandle:
-        """Load tensor from HBM. Returns TensorHandle pointing at HBM[ptr].
+        """Load tensor from HBM — **lazy** (ADR-0062 §D1/§D2).
 
-        In greenlet mode: returns TensorHandle with actual numpy data.
+        Posts the ``DmaReadCmd`` to PE_DMA and returns immediately with a
-        In command-list mode: returns TensorHandle with data=None.
+        ``TensorHandle`` whose ``pending`` field references a fresh
+        ``LoadFuture``. The actual DMA completion is awaited at the first
+        consuming op (``tl.dot``, MATH, ``tl.store``, ``tl.send``,
+        ``tl.copy_to``, ``tl.composite``) via ``_await_pending``.
 
-        The returned handle's ``space`` is "hbm" so subsequent ops (math,
+        Command-list mode: emits the DmaReadCmd to ``self._commands``,
-        send, store) using this handle as a source resolve via MemoryStore
+        attaches a LoadFuture for structural compatibility (its event
-        at ``(hbm, ptr)`` — which is where the load's underlying data
+        stays None — no engine, no SimPy event to wait on).
-        actually lives in Phase 2 storage.
         """
-        self._emit_dispatch_overhead()
+        self._emit_dispatch_overhead("load")
-        handle = self._make_handle(
+        nbytes = self._nbytes(shape, dtype)
-            addr=ptr, shape=shape, dtype=dtype, space="hbm", pinned=True,
+        # LoadFuture is mutable; create it first, attach to the handle,
-        )
+        # then point its ``cmd`` at the DmaReadCmd that references the
-        cmd = DmaReadCmd(handle=handle, src_addr=ptr, nbytes=handle.nbytes)
+        # *final* handle. This guarantees ``handle.pending.cmd.handle is handle``.
-        data = self._emit(cmd)
+        future = LoadFuture.__new__(LoadFuture)
-        if data is not None:
+        future.event = None
-            # Greenlet mode: attach real data to handle (preserve space + pinned)
+        future.resolved = False
-            return TensorHandle(
+        future.data = None
-                id=handle.id, addr=handle.addr, shape=handle.shape,
+        handle = TensorHandle(
-                dtype=handle.dtype, nbytes=handle.nbytes, data=data,
+            id=self._next_handle_id(),
-                space=handle.space, pinned=handle.pinned,
+            addr=ptr, shape=shape, dtype=dtype, nbytes=nbytes,
+            data=None, space="hbm", pinned=True, pending=future,
         )
+        cmd = DmaReadCmd(handle=handle, src_addr=ptr, nbytes=nbytes)
+        future.cmd = cmd
+        if self._runner is not None:
+            # Lazy: runner posts the DmaReadCmd, sets future.event, then
+            # switches back immediately. No yield on completion here.
+            self._runner.switch_to_simpy(("load_issue", future))
+        else:
+            self._commands.append(cmd)
         return handle
 
+    def _await_pending(self, *handles: TensorHandle | None) -> None:
+        """Auto-wait at first use (ADR-0062 §D2).
+
+        For each handle carrying an unresolved ``LoadFuture``, yield on
+        the DMA completion event (greenlet → runner). Command-list mode
+        is a no-op.
+        """
+        if self._runner is None:
+            return
+        for h in handles:
+            if h is None:
+                continue
+            pending = getattr(h, "pending", None)
+            if pending is None or pending.resolved:
+                continue
+            self._runner.switch_to_simpy(("load_await", h))
+
     def store(self, ptr: int, handle: TensorHandle) -> None:
         """Store tensor from TCM to HBM."""
-        self._emit_dispatch_overhead()
+        self._await_pending(handle)
+        self._emit_dispatch_overhead("store")
         cmd = DmaWriteCmd(handle=handle, dst_addr=ptr, nbytes=handle.nbytes)
         self._emit(cmd)
 
+    def copy_to(self, dst: TensorHandle, src: TensorHandle) -> None:
+        """TCM-to-TCM byte copy (ADR-0063 §D3.1).
+
+        Writes ``src``'s bytes into ``dst``'s address. Both handles must
+        live in TCM and have matching shape and dtype. Symmetric to
+        ``tl.store`` (which targets HBM) but stays on-chip so it doesn't
+        emit a DMA entry into op_log.
+
+        Used inside ``tl.scratch_scope()`` to persist a scoped result —
+        typically an updated running ``(m, ℓ, O)`` — to an outside-scope
+        (persistent) handle so its bytes survive the scope's ``__exit__``
+        cursor rewind (ADR-0063 §D3 two-arena pattern).
+        """
+        if src.shape != dst.shape:
+            raise ValueError(
+                f"tl.copy_to: shape mismatch — src.shape={src.shape} "
+                f"vs dst.shape={dst.shape}"
+            )
+        if src.dtype != dst.dtype:
+            raise ValueError(
+                f"tl.copy_to: dtype mismatch — src.dtype={src.dtype!r} "
+                f"vs dst.dtype={dst.dtype!r}"
+            )
+        if dst.space != "tcm":
+            raise ValueError(
+                f"tl.copy_to: dst must be in TCM (got space={dst.space!r}); "
+                "writes to HBM go through tl.store"
+            )
+        if src.space != "tcm":
+            raise ValueError(
+                f"tl.copy_to: src must be in TCM (got space={src.space!r}); "
+                "reads from HBM go through tl.load"
+            )
+        self._await_pending(src)
+        self._emit_dispatch_overhead("copy_to")
+        self._emit(CopyCmd(src=src, dst=dst, nbytes=src.nbytes))
+
     # ── GEMM Engine (blocking) ────────────────────────────────────
 
     def dot(self, a: TensorHandle, b: TensorHandle) -> TensorHandle:
@@ -224,7 +377,8 @@ class TLContext:
         out_shape = (*a.shape[:-2], m, n)
         out_dtype = a.dtype
         out = self._make_compute_out(shape=out_shape, dtype=out_dtype)
-        self._emit_dispatch_overhead()
+        self._await_pending(a, b)
+        self._emit_dispatch_overhead("dot")
         self._emit(GemmCmd(a=a, b=b, out=out, m=m, k=k, n=n))
         return out
 
@@ -232,7 +386,8 @@ class TLContext:
 
     def _unary_math(self, op: str, x: TensorHandle) -> TensorHandle:
         out = self._make_compute_out(shape=x.shape, dtype=x.dtype)
-        self._emit_dispatch_overhead()
+        self._await_pending(x)
+        self._emit_dispatch_overhead("math")
         self._emit(MathCmd(op=op, inputs=(x,), out=out))
         return out
 
@@ -265,7 +420,8 @@ class TLContext:
         out_shape = list(x.shape)
         out_shape[axis] = 1
         out = self._make_compute_out(shape=tuple(out_shape), dtype=x.dtype)
-        self._emit_dispatch_overhead()
+        self._await_pending(x)
+        self._emit_dispatch_overhead("math")
         self._emit(MathCmd(op=op, inputs=(x,), out=out, axis=axis))
         return out
 
@@ -284,7 +440,8 @@ class TLContext:
         self, op: str, a: TensorHandle, b: TensorHandle,
     ) -> TensorHandle:
         out = self._make_compute_out(shape=a.shape, dtype=a.dtype)
-        self._emit_dispatch_overhead()
+        self._await_pending(a, b)
+        self._emit_dispatch_overhead("math")
         self._emit(MathCmd(op=op, inputs=(a, b), out=out))
         return out
 
@@ -292,7 +449,8 @@ class TLContext:
         self, cond: TensorHandle, a: TensorHandle, b: TensorHandle,
     ) -> TensorHandle:
         out = self._make_compute_out(shape=a.shape, dtype=a.dtype)
-        self._emit_dispatch_overhead()
+        self._await_pending(cond, a, b)
+        self._emit_dispatch_overhead("math")
         self._emit(MathCmd(op="where", inputs=(cond, a, b), out=out))
         return out
 
@@ -309,7 +467,8 @@ class TLContext:
     ) -> TensorHandle:
         """Fused multiply-add: a * b + c (real Triton: tl.fma)."""
         out = self._make_compute_out(shape=a.shape, dtype=a.dtype)
-        self._emit_dispatch_overhead()
+        self._await_pending(a, b, c)
+        self._emit_dispatch_overhead("math")
         self._emit(MathCmd(op="fma", inputs=(a, b, c), out=out))
         return out
 
@@ -321,7 +480,8 @@ class TLContext:
     ) -> TensorHandle:
         """Clamp x to [min, max] (real Triton: tl.clamp)."""
         out = self._make_compute_out(shape=x.shape, dtype=x.dtype)
-        self._emit_dispatch_overhead()
+        self._await_pending(x, min, max)
+        self._emit_dispatch_overhead("math")
         self._emit(MathCmd(op="clamp", inputs=(x, min, max), out=out))
         return out
 
@@ -333,7 +493,8 @@ class TLContext:
         canonical (x - max) → exp → sum → div sequence.
         """
         out = self._make_compute_out(shape=x.shape, dtype=x.dtype)
-        self._emit_dispatch_overhead()
+        self._await_pending(x)
+        self._emit_dispatch_overhead("math")
         self._emit(MathCmd(op="softmax", inputs=(x,), out=out, axis=axis))
         return out
 
@@ -427,13 +588,16 @@ class TLContext:
             space = getattr(src, "space", space)
         if src_addr is None or nbytes is None or shape is None:
             raise ValueError("tl.send: provide either a TensorHandle or src_addr/nbytes/shape")
+        # ADR-0062: if the source is a lazy-loaded handle, await first so
+        # the data snapshot below sees the real bytes.
+        self._await_pending(src)
         # Carry the handle's .data snapshot (if available). When the source
         # is a recv slot, .data holds the numpy array that was read from
         # MemoryStore at recv-time. This prevents a Phase 1 race where a
         # later IPCQ inbound overwrites the slot before the outbound
         # PE_DMA reads it.
         handle_data = getattr(src, "data", None) if src is not None else None
-        self._emit_dispatch_overhead()
+        self._emit_dispatch_overhead("ipcq_send")
         cmd = IpcqSendCmd(
             direction=dir,
             src_addr=src_addr, src_space=space,
@@ -467,7 +631,7 @@ class TLContext:
             arrived. In greenlet/runner mode, ``handle.data`` carries the
             actual ndarray; in command-list mode the handle is a placeholder.
         """
-        self._emit_dispatch_overhead()
+        self._emit_dispatch_overhead("ipcq_recv")
         if dst_addr is not None and dst_space is not None:
             cmd = IpcqRecvCmd(
                 direction=dir,
@@ -518,7 +682,7 @@ class TLContext:
         they receive. This API is segregated from ``tl.recv`` so the
         diagnostic flag can never accidentally be set in real workloads.
         """
-        self._emit_dispatch_overhead()
+        self._emit_dispatch_overhead("ipcq_recv")
         cmd = IpcqRecvCmd(
             direction=dir,
             shape=shape, dtype=dtype,
@@ -547,7 +711,7 @@ class TLContext:
         dtype: str = "f16",
     ) -> "RecvFuture":
         """Non-blocking recv. Returns a future to pass into ``tl.wait``."""
-        self._emit_dispatch_overhead()
+        self._emit_dispatch_overhead("ipcq_recv")
         cmd = IpcqRecvCmd(
             direction=dir,
             shape=shape, dtype=dtype,
@@ -582,6 +746,9 @@ class TLContext:
 
         Returns CompletionHandle for use with wait().
         """
+        # ADR-0062: composite operand DMA paths still need their inputs
+        # to be resolved before the composite reads them via PE_SCHEDULER.
+        self._await_pending(a, b)
         # Compute output size based on op
         if op == "gemm" and b is not None:
             m, k = a.shape[-2], a.shape[-1]
@@ -609,7 +776,7 @@ class TLContext:
             ops_tuple = (head_spec, *epi_specs)
 
         completion = CompletionHandle(id=self._next_completion_id())
-        self._emit_dispatch_overhead()
+        self._emit_dispatch_overhead("composite")
         self._emit(CompositeCmd(
             completion=completion, op=op,
             a=a, b=b, out_addr=out_ptr, out_nbytes=out_nbytes,

tests/attention/test_attention_8kv_groups_diag.py

@@ -1,286 +0,0 @@
-"""Diagnostic harness for the Llama-70B "1 Q-head per cube" target.
-
-Per the GQA Llama-70B sharding study at
-``llm_paper_review/notes/GQA_MHA_sharding/scripts/_gen_llama70b_1M_4cases.py``,
-the 1 Q-head/cube baseline uses 64 cubes (4 SIPs × 16 cubes/SIP) organized
-into 8 KV-groups of 8 cubes each. Each KV-group occupies a ``2×4``
-sub-mesh within a SIP's ``4×4`` cube grid and runs the C2 2D row-then-col
-AllReduce-mlo (ADR-0059 extension). This harness probes the gap between
-validation and headline in three incrementally-larger steps:
-
-  step_1_single_kv_group_at_full_breadth
-     ONE multi_user_decode launch on a 2×4 sub-mesh (8 cubes) of the
-     4-SIP topology, via the 2D mesh-mlo kernel. Verifies the per-KV-group
-     2D AllReduce works at full breadth. Smallest dim possible
-     (S_q=1, S_kv=16, h=1, d_head=64) to keep wall time bounded.
-
-  step_2_four_kv_groups_one_per_sip
-     Four sequential multi_user_decode launches, each targeting a
-     different SIP. Verifies that per-SIP isolation works (each SIP holds
-     its own 2×4 KV-group; the SFR install only writes intra-SIP
-     E/W + N/S edges so the 4 groups don't see each other).
-
-  step_3_eight_kv_groups_two_per_sip
-     The actual study target: 8 KV-groups, two per SIP (cubes 0..7 vs
-     cubes 8..15 within each SIP). Expected to FAIL with current infra —
-     DPPolicy doesn't take a cube offset and target_device is SIP-level,
-     so back-to-back launches both land on cubes 0..7 of their target SIP.
-     Captures what's needed to lift the 4-group cap to 8.
-
-Each step prints what it observed; the test asserts only the documented
-expected outcomes so we can land it, watch CI, and iterate. Steps that
-are *expected* to fail (step 3) are marked xfail with a precise reason.
-"""
-from __future__ import annotations
-
-import traceback
-from pathlib import Path
-
-import pytest
-
-from kernbench.benches._attention_mesh_mlo_2d import attention_mesh_mlo_2d_kernel
-from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config
-from kernbench.ccl.sfr_config import configure_sfr_intercube_multisip
-from kernbench.policy.placement.dp import DPPolicy
-from kernbench.runtime_api.bench_runner import run_bench
-from kernbench.runtime_api.types import DeviceSelector, resolve_device
-from kernbench.sim_engine.engine import GraphEngine
-from kernbench.topology.builder import resolve_topology
-
-TOPOLOGY_4SIP = (
-    Path(__file__).resolve().parents[2] / "topologies" / "llama70b_4sip.yaml"
-)
-TOPOLOGY_DEFAULT = Path(__file__).resolve().parents[2] / "topology.yaml"
-
-S_Q_DECODE = 1
-S_KV_PER_RANK = 16
-H_Q = 1
-H_KV = 1
-D_HEAD = 64
-# 2×4 sub-mesh per KV-group (study: 8 cubes per KV-group at Q/cube=1).
-MESH_ROWS = 2
-MESH_COLS = 4
-N_CUBES_PER_KV_GROUP = MESH_ROWS * MESH_COLS
-DTYPE = "f16"
-
-
-def _ccl_cfg():
-    return resolve_algorithm_config(
-        load_ccl_config(), name="lrab_hierarchical_allreduce",
-    )
-
-
-def _engine_factory(t, d):
-    return GraphEngine(getattr(t, "topology_obj", t), enable_data=True)
-
-
-def _make_one_kv_group_bench(mesh_rows: int, mesh_cols: int):
-    """Return a bench_fn that runs ONE multi_user_decode kernel on a
-    ``mesh_rows × mesh_cols`` sub-mesh."""
-    n_cubes = mesh_rows * mesh_cols
-
-    def _bench_fn(ctx):
-        configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
-        dp_full = DPPolicy(cube="replicate", pe="replicate",
-                           num_cubes=n_cubes, num_pes=8)
-        dp_kv = DPPolicy(cube="row_wise", pe="replicate",
-                         num_cubes=n_cubes, num_pes=8)
-        q = ctx.zeros((S_Q_DECODE, H_Q * D_HEAD),
-                      dtype=DTYPE, dp=dp_full, name="q")
-        k = ctx.zeros((S_KV_PER_RANK * n_cubes, H_KV * D_HEAD),
-                      dtype=DTYPE, dp=dp_kv, name="k")
-        v = ctx.zeros((S_KV_PER_RANK * n_cubes, H_KV * D_HEAD),
-                      dtype=DTYPE, dp=dp_kv, name="v")
-        o = ctx.empty((S_Q_DECODE, H_Q * D_HEAD),
-                      dtype=DTYPE, dp=dp_full, name="o")
-        ctx.launch(
-            f"single_kv_group_{mesh_rows}x{mesh_cols}",
-            attention_mesh_mlo_2d_kernel,
-            q, k, v, o,
-            S_Q_DECODE, S_KV_PER_RANK, H_Q, H_KV, D_HEAD,
-            mesh_rows, mesh_cols,
-            1,                               # rank_axis=1 → cube-level ring
-            0,                               # cube_start=0 — single sub-mesh launch
-            _auto_dim_remap=False,
-        )
-
-    return _bench_fn
-
-
-def _run_one_kv_group(topology_path: Path, mesh_rows: int, mesh_cols: int,
-                      target_device=None):
-    topo = resolve_topology(str(topology_path))
-    captured: dict = {"engine": None}
-
-    def factory(t, d):
-        eng = _engine_factory(t, d)
-        captured["engine"] = eng
-        return eng
-
-    exc = None
-    result = None
-    try:
-        result = run_bench(
-            topology=topo,
-            bench_fn=_make_one_kv_group_bench(mesh_rows, mesh_cols),
-            device=target_device or resolve_device(None),
-            engine_factory=factory,
-        )
-    except BaseException as e:  # noqa: BLE001
-        exc = e
-    return exc, result, captured["engine"]
-
-
-# ── Step 1 — single KV-group at the study's full breadth ──────────
-
-
-def test_step_1_single_kv_group_at_full_breadth():
-    """One multi_user_decode launch on a 2×4 sub-mesh, 4-SIP topology.
-
-    Uses the C2 2D row-then-col AllReduce-mlo kernel: stage 1 reduces
-    across cols (E/W) within each row, stage 2 reduces across rows (N/S).
-    Expected to PASS — N/S edges are wired by
-    ``configure_sfr_intercube_multisip`` and the 2D fan-out avoids the
-    row-boundary IpcqInvalidDirection that the 1D kernel hit at cube 4.
-    """
-    if not TOPOLOGY_4SIP.exists():
-        pytest.skip(f"4-SIP topology missing: {TOPOLOGY_4SIP}")
-    exc, result, engine = _run_one_kv_group(
-        TOPOLOGY_4SIP, mesh_rows=MESH_ROWS, mesh_cols=MESH_COLS,
-    )
-    if exc is not None:
-        oplog_len = len(getattr(engine, "op_log", []) or []) if engine else 0
-        print(f"\nstep_1 FAIL — op_log records before crash: {oplog_len}")
-        traceback.print_exception(type(exc), exc, exc.__traceback__)
-        raise AssertionError(f"step_1 failed: {exc}") from exc
-    assert result is not None and result.completion.ok, (
-        f"step_1: completion not ok — {result.completion if result else None}"
-    )
-
-
-# ── Step 2 — 4 KV-groups, one per SIP, sequential launches ────────
-
-
-def _make_multi_sip_bench_fn(sip_groups: list[tuple[int, str, int]]):
-    """One bench_fn that does one 2×4 multi_user_decode launch per item.
-
-    Each ``(sip, tag, cube_start)`` tuple becomes one launch:
-      - ``ctx.ahbm.set_device(sip)`` switches allocations to that SIP
-        (mirrors ``milestone_1h_ccl.py:283-292``).
-      - ``cube_start`` selects which 8-cube sub-mesh within the SIP:
-        ``0`` → cubes 0..7 (rows 0..1), ``8`` → cubes 8..15 (rows 2..3).
-      - ``tag`` disambiguates tensor names so launches in the same
-        run_bench don't collide on the allocator namespace.
-    """
-    n_cubes = MESH_ROWS * MESH_COLS
-
-    def _bench_fn(ctx):
-        configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
-        for sip, tag, cube_start in sip_groups:
-            ctx.ahbm.set_device(sip)
-            dp_full = DPPolicy(cube="replicate", pe="replicate",
-                               num_cubes=n_cubes, num_pes=8,
-                               cube_start=cube_start)
-            dp_kv = DPPolicy(cube="row_wise", pe="replicate",
-                             num_cubes=n_cubes, num_pes=8,
-                             cube_start=cube_start)
-            q = ctx.zeros((S_Q_DECODE, H_Q * D_HEAD),
-                          dtype=DTYPE, dp=dp_full, name=f"q_{tag}")
-            k = ctx.zeros((S_KV_PER_RANK * n_cubes, H_KV * D_HEAD),
-                          dtype=DTYPE, dp=dp_kv, name=f"k_{tag}")
-            v = ctx.zeros((S_KV_PER_RANK * n_cubes, H_KV * D_HEAD),
-                          dtype=DTYPE, dp=dp_kv, name=f"v_{tag}")
-            o = ctx.empty((S_Q_DECODE, H_Q * D_HEAD),
-                          dtype=DTYPE, dp=dp_full, name=f"o_{tag}")
-            ctx.launch(
-                f"kv_group_{tag}", attention_mesh_mlo_2d_kernel,
-                q, k, v, o,
-                S_Q_DECODE, S_KV_PER_RANK, H_Q, H_KV, D_HEAD,
-                MESH_ROWS, MESH_COLS,
-                1,                               # rank_axis=1 → cube-level ring
-                cube_start,                      # converts physical id → launch-local rank
-                _auto_dim_remap=False,
-            )
-
-    return _bench_fn
-
-
-def _run_multi_sip(sip_groups: list[tuple[int, str, int]]):
-    """Run a single run_bench call covering all (sip, tag) groups."""
-    topo = resolve_topology(str(TOPOLOGY_4SIP))
-    captured: dict = {"engine": None}
-
-    def factory(t, d):
-        eng = _engine_factory(t, d)
-        captured["engine"] = eng
-        return eng
-
-    exc = None
-    result = None
-    try:
-        result = run_bench(
-            topology=topo,
-            bench_fn=_make_multi_sip_bench_fn(sip_groups),
-            device=resolve_device(None),  # "all" SIPs in scope
-            engine_factory=factory,
-        )
-    except BaseException as e:  # noqa: BLE001
-        exc = e
-    return exc, result, captured["engine"]
-
-
-def test_step_2_four_kv_groups_one_per_sip():
-    """Four multi_user_decode launches, one per SIP, in ONE run_bench call.
-
-    Uses the CCL milestone pattern (``milestone_1h_ccl.py:283-292``):
-    ``target_device="all"`` scopes the runtime to every SIP; then
-    ``ctx.ahbm.set_device(sip)`` before each ``ctx.zeros``/``launch``
-    switches which SIP the next allocation+launch lands on. This is the
-    canonical sequential per-SIP pattern in the codebase — four separate
-    ``run_bench`` calls with ``DeviceSelector("sip:N")`` is a misuse.
-
-    Expected to PASS — the SFR install draws intra-SIP edges only, so the
-    4 KV-groups can't see each other.
-    """
-    if not TOPOLOGY_4SIP.exists():
-        pytest.skip(f"4-SIP topology missing: {TOPOLOGY_4SIP}")
-    sip_groups = [(sip, f"sip{sip}", 0) for sip in range(4)]
-    exc, result, engine = _run_multi_sip(sip_groups)
-    if exc is not None:
-        oplog_len = len(getattr(engine, "op_log", []) or []) if engine else 0
-        print(f"\nstep_2 FAIL — op_log records before crash: {oplog_len}")
-        traceback.print_exception(type(exc), exc, exc.__traceback__)
-        raise AssertionError(f"step_2 failed: {exc}") from exc
-    assert result is not None and result.completion.ok, (
-        f"step_2: completion not ok — {result.completion if result else None}"
-    )
-
-
-# ── Step 3 — 8 KV-groups, two per SIP (study target) ──────────────
-
-
-def test_step_3_eight_kv_groups_two_per_sip():
-    """Two launches per SIP × 4 SIPs = 8 KV-groups total in one run_bench.
-
-    The headline target: 64 cubes serving 8 KV-groups, two disjoint 2×4
-    sub-meshes per SIP. ``cube_start=0`` puts the first KV-group on
-    cubes 0..7 (rows 0..1); ``cube_start=8`` puts the second on cubes
-    8..15 (rows 2..3). This is the use case ``DPPolicy.cube_start`` was
-    added to enable.
-
-    Expected to PASS — the SFR install draws intra-SIP edges only, so
-    the 8 KV-groups can't see each other; ``cube_start`` ensures the
-    two halves of each SIP land on disjoint cubes.
-    """
-    if not TOPOLOGY_4SIP.exists():
-        pytest.skip(f"4-SIP topology missing: {TOPOLOGY_4SIP}")
-    sip_groups = [
-        (sip, f"sip{sip}_half{half}", half * N_CUBES_PER_KV_GROUP)
-        for sip in range(4) for half in (0, 1)
-    ]
-    exc, result, engine = _run_multi_sip(sip_groups)
-    if exc is not None:
-        raise AssertionError(f"step_3 failed: {exc}") from exc
-    assert result is not None and result.completion.ok, (
-        f"step_3: completion not ok — {result.completion if result else None}"
-    )

tests/attention/test_attention_mesh_panels_diag.py

@@ -1,198 +0,0 @@
-"""End-to-end engine drives for the four GQA Llama-70B panels (sub-cycle 4c step 2).
-
-Mirrors the existing single_user_decode diag harness across all four panels
-of the milestone-gqa-llama70b sweep (ADR-0057):
-
-  single_user_prefill   ring-K/V kernel, intracube PE ring  (8 PEs / 1 cube)
-  single_user_decode    allreduce-mlo kernel, intracube PE ring
-  multi_user_prefill    ring-K/V kernel, intercube multisip (4 cubes)
-  multi_user_decode     allreduce-mlo kernel, intercube multisip
-
-Each test runs the panel through ``run_bench`` with ``enable_data=True``
-and asserts ``result.completion.ok``. Failures dump the engine's op_log
-tail and the exception, mirroring the decode-diag harness format.
-
-Validation-scale config matches ADR-0057 D4:
-  S_q_prefill=16, S_kv_per_rank=16, h_q=h_kv=1, d_head=64
-  n_ranks_single_user=8, n_ranks_multi_user=4
-"""
-from __future__ import annotations
-
-import traceback
-from pathlib import Path
-
-import pytest
-
-from kernbench.benches._attention_mesh_kv import attention_mesh_kv_kernel
-from kernbench.benches._attention_mesh_mlo import attention_mesh_mlo_kernel
-from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config
-from kernbench.ccl.sfr_config import (
-    configure_sfr_intercube_multisip,
-    configure_sfr_intracube_pe_ring,
-)
-from kernbench.policy.placement.dp import DPPolicy
-from kernbench.runtime_api.bench_runner import run_bench
-from kernbench.runtime_api.types import resolve_device
-from kernbench.sim_engine.engine import GraphEngine
-from kernbench.topology.builder import resolve_topology
-
-TOPOLOGY_PATH = Path(__file__).resolve().parents[2] / "topology.yaml"
-
-S_Q_PREFILL = 16
-S_Q_DECODE = 1
-S_KV_PER_RANK = 16
-H_Q = 1
-H_KV = 1
-D_HEAD = 64
-N_RANKS_SINGLE_USER = 8
-N_RANKS_MULTI_USER = 4
-DTYPE = "f16"
-
-
-# ── Helpers ──────────────────────────────────────────────────────
-
-
-def _engine_factory(t, d):
-    return GraphEngine(getattr(t, "topology_obj", t), enable_data=True)
-
-
-def _run_panel(bench_fn):
-    """Drive a panel through run_bench; return (exc, result, engine)."""
-    topo = resolve_topology(str(TOPOLOGY_PATH))
-    captured: dict = {"engine": None}
-
-    def factory(t, d):
-        eng = _engine_factory(t, d)
-        captured["engine"] = eng
-        return eng
-
-    exc = None
-    result = None
-    try:
-        result = run_bench(
-            topology=topo, bench_fn=bench_fn,
-            device=resolve_device(None), engine_factory=factory,
-        )
-    except BaseException as e:  # noqa: BLE001
-        exc = e
-    return exc, result, captured["engine"]
-
-
-def _assert_ok(name: str, exc, result, engine) -> None:
-    if exc is not None:
-        oplog_len = len(getattr(engine, "op_log", []) or []) if engine else 0
-        print(f"\n========== {name} FAIL ==========")
-        print(f"op_log records before crash: {oplog_len}")
-        print(f"{type(exc).__name__}: {exc}")
-        traceback.print_exception(type(exc), exc, exc.__traceback__)
-        raise AssertionError(
-            f"{name} failed at runtime: {exc}"
-        ) from exc
-    assert result is not None, f"{name}: no result"
-    assert result.completion.ok, f"{name}: completion not ok — {result.completion}"
-
-
-# ── Panel bench fns ──────────────────────────────────────────────
-
-
-def _bench_fn_single_user_prefill(ctx):
-    configure_sfr_intracube_pe_ring(
-        ctx.engine, ctx.spec,
-        resolve_algorithm_config(load_ccl_config(), name="lrab_hierarchical_allreduce"),
-    )
-    n = N_RANKS_SINGLE_USER
-    dp_full = DPPolicy(cube="replicate", pe="replicate", num_cubes=1, num_pes=n)
-    dp_kv = DPPolicy(cube="replicate", pe="row_wise", num_cubes=1, num_pes=n)
-    q = ctx.zeros((S_Q_PREFILL, H_Q * D_HEAD), dtype=DTYPE, dp=dp_full, name="q")
-    k = ctx.zeros((S_KV_PER_RANK * n, H_KV * D_HEAD), dtype=DTYPE, dp=dp_kv, name="k")
-    v = ctx.zeros((S_KV_PER_RANK * n, H_KV * D_HEAD), dtype=DTYPE, dp=dp_kv, name="v")
-    o = ctx.empty((S_Q_PREFILL, H_Q * D_HEAD), dtype=DTYPE, dp=dp_full, name="o")
-    ctx.launch(
-        "single_user_prefill_mesh", attention_mesh_kv_kernel,
-        q, k, v, o,
-        S_Q_PREFILL, S_KV_PER_RANK, H_Q, H_KV, D_HEAD, n,
-    )
-
-
-def _bench_fn_single_user_decode(ctx):
-    configure_sfr_intracube_pe_ring(
-        ctx.engine, ctx.spec,
-        resolve_algorithm_config(load_ccl_config(), name="lrab_hierarchical_allreduce"),
-    )
-    n = N_RANKS_SINGLE_USER
-    dp_full = DPPolicy(cube="replicate", pe="replicate", num_cubes=1, num_pes=n)
-    dp_kv = DPPolicy(cube="replicate", pe="row_wise", num_cubes=1, num_pes=n)
-    q = ctx.zeros((S_Q_DECODE, H_Q * D_HEAD), dtype=DTYPE, dp=dp_full, name="q")
-    k = ctx.zeros((S_KV_PER_RANK * n, H_KV * D_HEAD), dtype=DTYPE, dp=dp_kv, name="k")
-    v = ctx.zeros((S_KV_PER_RANK * n, H_KV * D_HEAD), dtype=DTYPE, dp=dp_kv, name="v")
-    o = ctx.empty((S_Q_DECODE, H_Q * D_HEAD), dtype=DTYPE, dp=dp_full, name="o")
-    ctx.launch(
-        "single_user_decode_mesh", attention_mesh_mlo_kernel,
-        q, k, v, o,
-        S_Q_DECODE, S_KV_PER_RANK, H_Q, H_KV, D_HEAD, n,
-    )
-
-
-def _bench_fn_multi_user_prefill(ctx):
-    configure_sfr_intercube_multisip(
-        ctx.engine, ctx.spec,
-        resolve_algorithm_config(load_ccl_config(), name="lrab_hierarchical_allreduce"),
-    )
-    n = N_RANKS_MULTI_USER
-    dp_full = DPPolicy(cube="replicate", pe="replicate", num_cubes=n, num_pes=8)
-    dp_kv = DPPolicy(cube="row_wise", pe="replicate", num_cubes=n, num_pes=8)
-    q = ctx.zeros((S_Q_PREFILL, H_Q * D_HEAD), dtype=DTYPE, dp=dp_full, name="q")
-    k = ctx.zeros((S_KV_PER_RANK * n, H_KV * D_HEAD), dtype=DTYPE, dp=dp_kv, name="k")
-    v = ctx.zeros((S_KV_PER_RANK * n, H_KV * D_HEAD), dtype=DTYPE, dp=dp_kv, name="v")
-    o = ctx.empty((S_Q_PREFILL, H_Q * D_HEAD), dtype=DTYPE, dp=dp_full, name="o")
-    ctx.launch(
-        "multi_user_prefill_mesh", attention_mesh_kv_kernel,
-        q, k, v, o,
-        S_Q_PREFILL, S_KV_PER_RANK, H_Q, H_KV, D_HEAD, n,
-        1,  # rank_axis=1 → ring at cube level (ADR-0059 multi_user)
-        _auto_dim_remap=False,
-    )
-
-
-def _bench_fn_multi_user_decode(ctx):
-    configure_sfr_intercube_multisip(
-        ctx.engine, ctx.spec,
-        resolve_algorithm_config(load_ccl_config(), name="lrab_hierarchical_allreduce"),
-    )
-    n = N_RANKS_MULTI_USER
-    dp_full = DPPolicy(cube="replicate", pe="replicate", num_cubes=n, num_pes=8)
-    dp_kv = DPPolicy(cube="row_wise", pe="replicate", num_cubes=n, num_pes=8)
-    q = ctx.zeros((S_Q_DECODE, H_Q * D_HEAD), dtype=DTYPE, dp=dp_full, name="q")
-    k = ctx.zeros((S_KV_PER_RANK * n, H_KV * D_HEAD), dtype=DTYPE, dp=dp_kv, name="k")
-    v = ctx.zeros((S_KV_PER_RANK * n, H_KV * D_HEAD), dtype=DTYPE, dp=dp_kv, name="v")
-    o = ctx.empty((S_Q_DECODE, H_Q * D_HEAD), dtype=DTYPE, dp=dp_full, name="o")
-    ctx.launch(
-        "multi_user_decode_mesh", attention_mesh_mlo_kernel,
-        q, k, v, o,
-        S_Q_DECODE, S_KV_PER_RANK, H_Q, H_KV, D_HEAD, n,
-        1,  # rank_axis=1 → ring at cube level (ADR-0059 multi_user)
-        _auto_dim_remap=False,
-    )
-
-
-# ── Tests ────────────────────────────────────────────────────────
-
-
-def test_single_user_prefill_through_engine():
-    exc, result, engine = _run_panel(_bench_fn_single_user_prefill)
-    _assert_ok("single_user_prefill", exc, result, engine)
-
-
-def test_single_user_decode_through_engine():
-    exc, result, engine = _run_panel(_bench_fn_single_user_decode)
-    _assert_ok("single_user_decode", exc, result, engine)
-
-
-def test_multi_user_prefill_through_engine():
-    exc, result, engine = _run_panel(_bench_fn_multi_user_prefill)
-    _assert_ok("multi_user_prefill", exc, result, engine)
-
-
-def test_multi_user_decode_through_engine():
-    exc, result, engine = _run_panel(_bench_fn_multi_user_decode)
-    _assert_ok("multi_user_decode", exc, result, engine)

tests/attention/test_gqa_decode.py

@@ -0,0 +1,143 @@
+"""Phase 1 spec test for P1a GQA decode kernel (real GQA via M-fold).
+
+P1a is the first phase of the DDD-0060 plan, split out of the original
+P1 (the composite-hybrid swap is P1b, deferred until the tl.composite
+output-handle question is decided). P1a is the *correctness* unlock:
+the kernel processes ONE KV head at a time and folds the G query heads
+into the matmul M (row) dimension so a single Q·Kᵀ serves all G heads
+sharing one K (ADR-0060 §5.2). This lifts the baseline's
+``h_q == h_kv == 1`` cap pinned at
+``tests/attention/test_milestone_gqa_llama70b.py:137-142``.
+
+P1a stays inside the existing ``tl`` API: the two attention GEMMs use the
+blocking ``tl.dot`` so the chain ``Q·Kᵀ → softmax → P·V → store`` fits
+without any composite-output chaining. The composite swap (P1b) will
+revisit this once the API for feeding a composite's output into a
+downstream MATH op is settled.
+
+Phase 1 (this commit): tests only — production code lands in Phase 2.
+Tests fail at import in Phase 1 with ModuleNotFoundError; Phase 2 makes
+them pass.
+"""
+from __future__ import annotations
+
+from pathlib import Path
+
+from kernbench.benches._gqa_decode import gqa_decode_kernel  # noqa: F401  (Phase 2 deliverable)
+from kernbench.policy.placement.dp import DPPolicy
+from kernbench.runtime_api.bench_runner import run_bench
+from kernbench.runtime_api.types import resolve_device
+from kernbench.sim_engine.engine import GraphEngine
+from kernbench.topology.builder import resolve_topology
+
+TOPOLOGY_DEFAULT = Path(__file__).resolve().parents[2] / "topology.yaml"
+
+# Decode shapes — P1a is one-shot, no tiling, single rank. P3 will tile.
+T_Q = 1
+D_HEAD = 64
+S_KV = 16
+DTYPE = "f16"
+
+
+def _engine_factory(t, d):
+    return GraphEngine(getattr(t, "topology_obj", t), enable_data=True)
+
+
+def _run_decode_p1(*, h_q: int, h_kv: int):
+    """One-shot GQA decode on a single PE in a single CUBE (P=1, no SP).
+
+    Tensor layout (natural K — same as the P2a SP tests; the kernel
+    reshapes K to (d_head, S_local) via byte-conserving load):
+      Q : (T_q, h_q · d_head)   — natural Q layout; kernel reshapes to
+          (G·T_q, d_head) — byte-conserving and math-correct because
+          T_q=1 collapses axis ordering.
+      K : (S_kv, h_kv · d_head) — natural K layout; kernel loads as
+          (d_head, S_local) — reshape-as-transpose caveat (ADR-0060 §3
+          / §B item 2), correct for zero inputs used here.
+      V : (S_kv, h_kv · d_head) — natural V layout.
+      O : (T_q, h_q · d_head)   — same shape as Q.
+    """
+    topo = resolve_topology(str(TOPOLOGY_DEFAULT))
+
+    def _bench_fn(ctx):
+        dp = DPPolicy(cube="replicate", pe="replicate",
+                      num_cubes=1, num_pes=1)
+        q = ctx.zeros((T_Q, h_q * D_HEAD),
+                      dtype=DTYPE, dp=dp, name=f"q_h{h_q}_kv{h_kv}")
+        k = ctx.zeros((S_KV, h_kv * D_HEAD),
+                      dtype=DTYPE, dp=dp, name=f"k_h{h_q}_kv{h_kv}")
+        v = ctx.zeros((S_KV, h_kv * D_HEAD),
+                      dtype=DTYPE, dp=dp, name=f"v_h{h_q}_kv{h_kv}")
+        o = ctx.empty((T_Q, h_q * D_HEAD),
+                      dtype=DTYPE, dp=dp, name=f"o_h{h_q}_kv{h_kv}")
+        ctx.launch(
+            f"gqa_decode_p1_h{h_q}_kv{h_kv}",
+            gqa_decode_kernel,
+            q, k, v, o,
+            T_Q, S_KV, h_q, h_kv, D_HEAD,
+            1, 1,                               # C=1, P=1 (no SP, degenerate)
+            _auto_dim_remap=False,
+        )
+
+    return run_bench(
+        topology=topo,
+        bench_fn=_bench_fn,
+        device=resolve_device(None),
+        engine_factory=_engine_factory,
+    )
+
+
+def _dma_read_count(op_log) -> int:
+    return sum(1 for r in op_log if r.op_name == "dma_read")
+
+
+# ── Headline unlock: real GQA (h_q = G·h_kv) runs in data mode ──────────
+
+
+def test_real_gqa_h_q_eight_h_kv_one_completes_in_data_mode():
+    """ADR-0060 §A.1 headline unlock — the baseline raises
+    ``ValueError: Shape mismatch …`` in MemoryStore at h_q=8, h_kv=1
+    because ``_view(K, (h_q·d, S_kv))`` only byte-conserves when h_q==h_kv.
+    M-fold processes one KV head at a time using only byte-conserving
+    reshapes, so this completes.
+    """
+    result = _run_decode_p1(h_q=8, h_kv=1)
+    assert result.completion.ok, (
+        f"real GQA (h_q=8, h_kv=1) decode failed: {result.completion}"
+    )
+
+
+# ── M-fold property: K/V loads do not scale with G ─────────────────────
+
+
+def test_kv_dma_read_count_independent_of_g():
+    """ADR-0060 TL;DR / §5.2: M-fold loads K and V once per KV head and
+    folds the G query heads into the GEMM M dim. The dma_read_count must
+    therefore be identical between (G=1, h_kv=1) and (G=8, h_kv=1) — both
+    issue exactly 3 reads (Q + K + V). This pins the GQA-reuse property
+    that the rest of the plan (composite streaming in P4, etc.) builds on.
+    """
+    g1 = _run_decode_p1(h_q=1, h_kv=1)
+    g8 = _run_decode_p1(h_q=8, h_kv=1)
+    n_g1 = _dma_read_count(g1.engine.op_log)
+    n_g8 = _dma_read_count(g8.engine.op_log)
+    assert n_g1 == 3, f"G=1 dma_read_count must be 3 (Q+K+V); got {n_g1}"
+    assert n_g8 == 3, f"G=8 dma_read_count must be 3 (Q+K+V); got {n_g8}"
+    assert n_g1 == n_g8, (
+        f"K/V dma_read_count must be independent of G; "
+        f"got G=1 -> {n_g1}, G=8 -> {n_g8}"
+    )
+
+
+# ── Backward-compat: degenerate G=1 still works ────────────────────────
+
+
+def test_degenerate_g_equals_one_still_works():
+    """G=1 (h_q == h_kv == 1) is the baseline-compatible config. M-fold
+    degenerates to (T_q, d) = (1, 64) — the same shape the baseline
+    already exercises — so this proves no regression on that path.
+    """
+    result = _run_decode_p1(h_q=1, h_kv=1)
+    assert result.completion.ok, (
+        f"degenerate G=1 decode failed: {result.completion}"
+    )

tests/attention/test_gqa_decode_mc.py

@@ -0,0 +1,182 @@
+"""Phase 1 spec test for P2b GQA decode multi-cube SP (both Level-2 + Level-1).
+
+P2b extends P2a to multiple CUBEs in one CUBE Group. The kernel uses the
+canonical full SFR install (``configure_sfr_intercube_multisip``) which
+provides disjoint direction namespaces:
+
+  - ``intra_E / intra_W / intra_N / intra_S``  — PE↔PE within a CUBE
+    (logical 2×4 grid, no wrap)
+  - ``E / W / N / S``                          — CUBE↔CUBE inter-CUBE
+    (mesh, no wrap)
+
+Reduce pattern (chain reduce-to-root, ADR-0060 §A.2 spirit, §4 chain
+deviation noted):
+
+  Level-2 (intra-CUBE, 2×4 grid):
+    row-then-col chain — each row reduces leftward along ``intra_W`` to
+    its col-0 PE, then PE 4 sends to PE 0 along ``intra_N``. 7 chain
+    steps per CUBE × 3 handles each = 21 ``ipcq_copy`` per CUBE.
+
+  Level-1 (inter-CUBE):
+    only PE 0 of each CUBE participates. Chain leftward along ``W``.
+    (C-1) chain steps × 3 handles each.
+
+  Final store at PE 0 of CUBE 0 only.
+
+Phase 1 (this commit): tests only — production code lands in Phase 2.
+Phase 2 also updates ``test_gqa_decode.py`` (add C=1) and
+``test_gqa_decode_sp.py`` (switch SFR + add C=1).
+"""
+from __future__ import annotations
+
+from pathlib import Path
+
+from kernbench.benches._gqa_decode import gqa_decode_kernel  # noqa: F401
+from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config
+from kernbench.ccl.sfr_config import configure_sfr_intercube_multisip
+from kernbench.policy.placement.dp import DPPolicy
+from kernbench.runtime_api.bench_runner import run_bench
+from kernbench.runtime_api.types import resolve_device
+from kernbench.sim_engine.engine import GraphEngine
+from kernbench.topology.builder import resolve_topology
+
+TOPOLOGY_DEFAULT = Path(__file__).resolve().parents[2] / "topology.yaml"
+
+T_Q = 1
+D_HEAD = 64
+DTYPE = "f16"
+
+
+def _ccl_cfg():
+    return resolve_algorithm_config(
+        load_ccl_config(), name="lrab_hierarchical_allreduce",
+    )
+
+
+def _engine_factory(t, d):
+    return GraphEngine(getattr(t, "topology_obj", t), enable_data=True)
+
+
+def _run_decode_mc(*, h_q: int, h_kv: int, C: int, P: int, S_kv: int):
+    """Multi-CUBE SP decode: C cubes × P PEs each share the work."""
+    topo = resolve_topology(str(TOPOLOGY_DEFAULT))
+
+    def _bench_fn(ctx):
+        configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
+        dp_full = DPPolicy(cube="replicate", pe="replicate",
+                           num_cubes=C, num_pes=P)
+        dp_kv = DPPolicy(cube="row_wise", pe="row_wise",
+                         num_cubes=C, num_pes=P)
+        # Total KV split across C×P ranks; each rank sees S_kv/(C·P) rows.
+        q = ctx.zeros((T_Q, h_q * D_HEAD),
+                      dtype=DTYPE, dp=dp_full,
+                      name=f"q_h{h_q}_kv{h_kv}_c{C}_p{P}")
+        k = ctx.zeros((S_kv, h_kv * D_HEAD),
+                      dtype=DTYPE, dp=dp_kv,
+                      name=f"k_h{h_q}_kv{h_kv}_c{C}_p{P}")
+        v = ctx.zeros((S_kv, h_kv * D_HEAD),
+                      dtype=DTYPE, dp=dp_kv,
+                      name=f"v_h{h_q}_kv{h_kv}_c{C}_p{P}")
+        o = ctx.empty((T_Q, h_q * D_HEAD),
+                      dtype=DTYPE, dp=dp_full,
+                      name=f"o_h{h_q}_kv{h_kv}_c{C}_p{P}")
+        ctx.launch(
+            f"gqa_decode_mc_h{h_q}_kv{h_kv}_c{C}_p{P}",
+            gqa_decode_kernel,
+            q, k, v, o,
+            T_Q, S_kv, h_q, h_kv, D_HEAD,
+            C, P,
+            _auto_dim_remap=False,
+        )
+
+    return run_bench(
+        topology=topo,
+        bench_fn=_bench_fn,
+        device=resolve_device(None),
+        engine_factory=_engine_factory,
+    )
+
+
+def _count(op_log, name: str) -> int:
+    return sum(1 for r in op_log if r.op_name == name)
+
+
+# ── Degenerate C=1 P=1 ────────────────────────────────────────────────
+
+
+def test_mc_c_one_p_one_degenerate():
+    """C=1, P=1: single rank, no SP. No IPCQ traffic; one dma_write."""
+    result = _run_decode_mc(h_q=1, h_kv=1, C=1, P=1, S_kv=16)
+    assert result.completion.ok, (
+        f"C=1 P=1 degenerate failed: {result.completion}"
+    )
+    assert _count(result.engine.op_log, "ipcq_copy") == 0
+    assert _count(result.engine.op_log, "dma_write") == 1
+
+
+# ── Intra-CUBE only (single CUBE, P=8 on 2×4 grid) ────────────────────
+
+
+def test_mc_c_one_p_eight_intracube_grid():
+    """C=1, P=8: intra-cube row-then-col chain on the 2×4 grid.
+    7 chain steps × 3 handles (m, ℓ, O) = 21 ipcq_copy."""
+    result = _run_decode_mc(h_q=1, h_kv=1, C=1, P=8, S_kv=64)
+    assert result.completion.ok, (
+        f"C=1 P=8 intra-cube failed: {result.completion}"
+    )
+    assert _count(result.engine.op_log, "dma_write") == 1
+    n_copy = _count(result.engine.op_log, "ipcq_copy")
+    assert n_copy == 21, (
+        f"C=1 P=8: expected 21 ipcq_copy (7 chain × 3 handles); got {n_copy}"
+    )
+
+
+# ── Multi-CUBE root-only write ────────────────────────────────────────
+
+
+def test_mc_c_two_p_eight_root_only_writes_o():
+    """C=2, P=8: 16 ranks total. Only PE 0 of CUBE 0 writes O."""
+    result = _run_decode_mc(h_q=1, h_kv=1, C=2, P=8, S_kv=128)
+    assert result.completion.ok, (
+        f"C=2 P=8 multi-cube failed: {result.completion}"
+    )
+    n_writes = _count(result.engine.op_log, "dma_write")
+    assert n_writes == 1, (
+        f"root-only write must hold for C=2 P=8 (16 ranks); got {n_writes}"
+    )
+
+
+# ── Multi-CUBE total IPCQ chain count ─────────────────────────────────
+
+
+def test_mc_c_two_p_eight_total_ipcq_count():
+    """C=2, P=8: 21 intra-cube ipcq_copy per CUBE × 2 CUBEs + 3 inter-cube
+    chain ipcq_copy (C-1=1 step × 3 handles) = 45 total."""
+    result = _run_decode_mc(h_q=1, h_kv=1, C=2, P=8, S_kv=128)
+    assert result.completion.ok, (
+        f"C=2 P=8 multi-cube failed: {result.completion}"
+    )
+    n_copy = _count(result.engine.op_log, "ipcq_copy")
+    expected = 21 * 2 + (2 - 1) * 3
+    assert n_copy == expected, (
+        f"C=2 P=8: expected {expected} ipcq_copy (21 intra × 2 CUBEs + 3 "
+        f"inter); got {n_copy}"
+    )
+
+
+# ── Real GQA × multi-CUBE SP combined ─────────────────────────────────
+
+
+def test_mc_real_gqa_c_two_p_eight():
+    """Headline: real GQA (h_q = G·h_kv with G=8) AND multi-CUBE SP
+    (C=2, P=8) together — the case the original baseline can express
+    neither part of."""
+    result = _run_decode_mc(h_q=8, h_kv=1, C=2, P=8, S_kv=128)
+    assert result.completion.ok, (
+        f"real GQA + multi-CUBE SP combined failed: {result.completion}"
+    )
+    n_writes = _count(result.engine.op_log, "dma_write")
+    assert n_writes == 1, (
+        f"root-only write must hold under M-fold + multi-CUBE; "
+        f"got {n_writes}"
+    )

tests/attention/test_gqa_decode_sp.py

@@ -0,0 +1,153 @@
+"""Phase 1 spec test for P2a GQA decode SP (chain reduce-to-root, Level-2 only).
+
+P2a is the first half of DDD-0060 P2: the kernel becomes multi-PE within
+one CUBE and reduces to root (PE 0) using a chain over the 1D intra-cube
+ring (W direction). This **replaces the baseline's bidirectional O(N)
+fan-out** where every rank ends with O — ADR-0060 §A.2's headline.
+
+Deviation from DDD-0060 §7 P2 gate: the gate text asks for
+``⌈log₂ P⌉`` reduce rounds. The intra-cube SFR install
+(``configure_sfr_intracube_pe_ring``) wires only a 1D E/W ring, so a
+true tree would require either multi-hop forwarding or a new SFR install
+(future ADR). P2a uses a **chain reduce-to-root**: ``P-1`` rounds along
+W. The architectural property the ADR cares about
+(root-only output vs every-rank-has-O) is preserved; the logarithmic
+collective is deferred.
+
+P2b (deferred) covers Level-1 inter-CUBE center-mesh reduce (C>1).
+
+Phase 1 (this commit): tests only — production code lands in Phase 2.
+"""
+from __future__ import annotations
+
+from pathlib import Path
+
+from kernbench.benches._gqa_decode import gqa_decode_kernel  # noqa: F401
+from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config
+from kernbench.ccl.sfr_config import configure_sfr_intercube_multisip
+from kernbench.policy.placement.dp import DPPolicy
+from kernbench.runtime_api.bench_runner import run_bench
+from kernbench.runtime_api.types import resolve_device
+from kernbench.sim_engine.engine import GraphEngine
+from kernbench.topology.builder import resolve_topology
+
+TOPOLOGY_DEFAULT = Path(__file__).resolve().parents[2] / "topology.yaml"
+
+T_Q = 1
+D_HEAD = 64
+DTYPE = "f16"
+
+
+def _ccl_cfg():
+    return resolve_algorithm_config(
+        load_ccl_config(), name="lrab_hierarchical_allreduce",
+    )
+
+
+def _engine_factory(t, d):
+    return GraphEngine(getattr(t, "topology_obj", t), enable_data=True)
+
+
+def _run_decode_sp(*, h_q: int, h_kv: int, P: int, S_kv: int):
+    """Single-CUBE SP decode: P PEs share the work along the intra-cube ring."""
+    topo = resolve_topology(str(TOPOLOGY_DEFAULT))
+
+    def _bench_fn(ctx):
+        configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
+        dp_full = DPPolicy(cube="replicate", pe="replicate",
+                           num_cubes=1, num_pes=P)
+        dp_kv = DPPolicy(cube="replicate", pe="row_wise",
+                         num_cubes=1, num_pes=P)
+        q = ctx.zeros((T_Q, h_q * D_HEAD),
+                      dtype=DTYPE, dp=dp_full, name=f"q_h{h_q}_kv{h_kv}_p{P}")
+        # KV: total S_kv split across P PEs along axis 0 (row_wise sharding).
+        # Each PE sees (S_kv/P, h_kv·D_HEAD).
+        k = ctx.zeros((S_kv, h_kv * D_HEAD),
+                      dtype=DTYPE, dp=dp_kv, name=f"k_h{h_q}_kv{h_kv}_p{P}")
+        v = ctx.zeros((S_kv, h_kv * D_HEAD),
+                      dtype=DTYPE, dp=dp_kv, name=f"v_h{h_q}_kv{h_kv}_p{P}")
+        o = ctx.empty((T_Q, h_q * D_HEAD),
+                      dtype=DTYPE, dp=dp_full, name=f"o_h{h_q}_kv{h_kv}_p{P}")
+        ctx.launch(
+            f"gqa_decode_sp_h{h_q}_kv{h_kv}_p{P}",
+            gqa_decode_kernel,
+            q, k, v, o,
+            T_Q, S_kv, h_q, h_kv, D_HEAD,
+            1, P,                                 # C=1, P=P (single-CUBE SP)
+            _auto_dim_remap=False,
+        )
+
+    return run_bench(
+        topology=topo,
+        bench_fn=_bench_fn,
+        device=resolve_device(None),
+        engine_factory=_engine_factory,
+    )
+
+
+def _count(op_log, name: str) -> int:
+    return sum(1 for r in op_log if r.op_name == name)
+
+
+# ── Root-only write ────────────────────────────────────────────────────
+
+
+def test_sp_chain_reduce_root_only_writes_o():
+    """ADR-0060 §A.2: only the root rank (PE 0) writes O. Baseline today
+    has every rank write the full final O (bidirectional fan-out)."""
+    result = _run_decode_sp(h_q=1, h_kv=1, P=8, S_kv=64)
+    assert result.completion.ok, f"P=8 chain reduce failed: {result.completion}"
+    n_writes = _count(result.engine.op_log, "dma_write")
+    assert n_writes == 1, (
+        f"reduce-to-root must produce exactly 1 dma_write (PE 0); "
+        f"got {n_writes}"
+    )
+
+
+# ── Chain step count ───────────────────────────────────────────────────
+
+
+def test_sp_chain_reduce_p_minus_one_ipcq_pairs():
+    """Chain reduce-to-root has P-1 send→recv pairs along the W chain;
+    each pair logs one ``ipcq_copy`` (inbound DMA, per
+    ``milestone_gqa_llama70b._summarize_op_log``). Each chain step ships
+    the triplet (m, ℓ, O) → 3 handles per step → 7 steps × 3 = 21."""
+    result = _run_decode_sp(h_q=1, h_kv=1, P=8, S_kv=64)
+    assert result.completion.ok, f"P=8 chain reduce failed: {result.completion}"
+    n_copy = _count(result.engine.op_log, "ipcq_copy")
+    expected = (8 - 1) * 3
+    assert n_copy == expected, (
+        f"chain reduce: expected {expected} ipcq_copy (P-1=7 steps × "
+        f"3 handles m/ℓ/O); got {n_copy}"
+    )
+
+
+# ── Real GQA × SP combined ─────────────────────────────────────────────
+
+
+def test_sp_real_gqa_h_q_eight_h_kv_one_p_eight():
+    """The combined unlock: real GQA (h_q=G·h_kv with G=8) AND SP
+    (P=8) together — neither expressible by the baseline."""
+    result = _run_decode_sp(h_q=8, h_kv=1, P=8, S_kv=64)
+    assert result.completion.ok, (
+        f"real GQA + SP combined run failed: {result.completion}"
+    )
+    n_writes = _count(result.engine.op_log, "dma_write")
+    assert n_writes == 1, (
+        f"root-only write must hold under M-fold too; got {n_writes}"
+    )
+
+
+# ── Degenerate P=1 ────────────────────────────────────────────────────
+
+
+def test_sp_p_one_degenerate_no_ipcq_traffic():
+    """P=1: SP degenerates to a single rank. No IPCQ traffic; one dma_write."""
+    result = _run_decode_sp(h_q=8, h_kv=1, P=1, S_kv=16)
+    assert result.completion.ok, f"P=1 degenerate failed: {result.completion}"
+    n_send = _count(result.engine.op_log, "ipcq_send")
+    n_recv = _count(result.engine.op_log, "ipcq_recv")
+    assert n_send == 0, f"P=1 must have no ipcq_send; got {n_send}"
+    assert n_recv == 0, f"P=1 must have no ipcq_recv; got {n_recv}"
+    n_writes = _count(result.engine.op_log, "dma_write")
+    assert n_writes == 1, f"P=1: one dma_write; got {n_writes}"

tests/attention/test_gqa_long_context.py

@@ -0,0 +1,157 @@
+"""Phase 1 spec test for Phase C: long-context regression for the GQA
+kernels (ADR-0060 §B "long/short context split" + ADR-0063 §A.2).
+
+The headline panel today caps prefill at S_kv=16 / decode at S_kv≤128 —
+NOT because the algorithm fails, but because the bump allocator never
+recycles. ADR-0063 §A.2 (test req 3) requires a sweep at an S that
+would overflow 1 MiB without recycling to complete after scope
+discipline lands.
+
+Scratch-budget estimate at C=4, P=8 (current 1 MiB pool):
+  decode:  per-rank S_local = S_kv / 32; intermediates ≲ 500 KB at
+           S_kv=32K (fits today — used as regression guard).
+  prefill: per-rank S_local = S_kv / 4; ~16·S_local bytes per ring
+           step × 4 steps ≈ 64·S_local bytes. Overflows at
+           ~64 K tokens (64 × 16K > 1 MiB).
+
+Phase 1 (this commit): tests only — production code lands in Phase 2.
+The prefill test fails today with a RuntimeError("TLContext scratch
+overflow"). After Phase 2 (scratch_scope + copy_to discipline) it
+completes.
+"""
+from __future__ import annotations
+
+from pathlib import Path
+
+from kernbench.benches._gqa_decode import gqa_decode_kernel  # noqa: F401
+from kernbench.benches._gqa_prefill import gqa_prefill_kernel  # noqa: F401
+from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config
+from kernbench.ccl.sfr_config import (
+    configure_sfr_intercube_multisip,
+    configure_sfr_intercube_ring,
+)
+from kernbench.policy.placement.dp import DPPolicy
+from kernbench.runtime_api.bench_runner import run_bench
+from kernbench.runtime_api.types import resolve_device
+from kernbench.sim_engine.engine import GraphEngine
+from kernbench.topology.builder import resolve_topology
+
+TOPOLOGY_DEFAULT = Path(__file__).resolve().parents[2] / "topology.yaml"
+
+D_HEAD = 64
+DTYPE = "f16"
+
+
+def _ccl_cfg():
+    return resolve_algorithm_config(
+        load_ccl_config(), name="lrab_hierarchical_allreduce",
+    )
+
+
+def _engine_factory(t, d):
+    return GraphEngine(getattr(t, "topology_obj", t), enable_data=True)
+
+
+# ── Decode at moderate-long context (regression guard) ───────────────
+
+
+def test_decode_long_context_32k_completes():
+    """Decode at S_kv=32K (C=1, P=8) — per-rank S_local=4K. Should
+    complete with current scratch usage (~few KB intermediates) and
+    must continue to complete after the rewrite.
+
+    This is a regression guard: scratch discipline shouldn't break
+    decode's existing long-context capability.
+    """
+    topo = resolve_topology(str(TOPOLOGY_DEFAULT))
+
+    def _bench_fn(ctx):
+        configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
+        P = 8
+        S_kv = 32_768
+        dp_full = DPPolicy(cube="replicate", pe="replicate",
+                           num_cubes=1, num_pes=P)
+        dp_kv = DPPolicy(cube="replicate", pe="row_wise",
+                         num_cubes=1, num_pes=P)
+        ctx.zeros((1, 8 * D_HEAD), dtype=DTYPE, dp=dp_full, name="q_long_dec")
+        k = ctx.zeros((S_kv, D_HEAD),
+                      dtype=DTYPE, dp=dp_kv, name="k_long_dec")
+        v = ctx.zeros((S_kv, D_HEAD),
+                      dtype=DTYPE, dp=dp_kv, name="v_long_dec")
+        o = ctx.empty((1, 8 * D_HEAD),
+                      dtype=DTYPE, dp=dp_full, name="o_long_dec")
+        q = ctx.zeros((1, 8 * D_HEAD),
+                      dtype=DTYPE, dp=dp_full, name="q_long_dec_2")
+        ctx.launch(
+            "gqa_decode_long_32k",
+            gqa_decode_kernel,
+            q, k, v, o,
+            1, S_kv, 8, 1, D_HEAD,
+            1, P,
+            _auto_dim_remap=False,
+        )
+
+    result = run_bench(
+        topology=topo, bench_fn=_bench_fn,
+        device=resolve_device(None), engine_factory=_engine_factory,
+    )
+    assert result.completion.ok, (
+        f"decode at S_kv=32K must complete; got {result.completion}"
+    )
+
+
+# ── Prefill at long context overflows without scope discipline ───────
+
+
+def test_prefill_long_context_completes_after_scope_discipline():
+    """ADR-0063 §A.2 Test Req 3: a sweep at an ``S`` that would overflow
+    1 MiB without recycling must complete after scope discipline lands.
+
+    With C=4 and S_kv chosen so per-rank S_local·8·4 > 1 MiB
+    (~16K tokens per rank ⇒ S_kv ≥ 64K), the current kernel overflows
+    the 1 MiB scratch pool. After Phase 2 (scratch_scope wraps each
+    ring step's intermediates; copy_to persists running state), it
+    completes.
+
+    This is the headline test that proves the S-ceiling is gone for
+    prefill.
+    """
+    topo = resolve_topology(str(TOPOLOGY_DEFAULT))
+
+    def _bench_fn(ctx):
+        configure_sfr_intercube_ring(
+            ctx.engine, ctx.spec, _ccl_cfg(), ring_size=4,
+        )
+        T_q = 4
+        S_kv = 65_536  # 64 K — per-rank 16 K, ~2 MB scratch w/o scope
+        C = 4
+        dp_q = DPPolicy(cube="replicate", pe="replicate",
+                        num_cubes=C, num_pes=1)
+        dp_kv = DPPolicy(cube="row_wise", pe="replicate",
+                         num_cubes=C, num_pes=1)
+        dp_o = DPPolicy(cube="row_wise", pe="replicate",
+                        num_cubes=C, num_pes=1)
+        q = ctx.zeros((T_q, D_HEAD),
+                      dtype=DTYPE, dp=dp_q, name="q_long_pre")
+        k = ctx.zeros((S_kv, D_HEAD),
+                      dtype=DTYPE, dp=dp_kv, name="k_long_pre")
+        v = ctx.zeros((S_kv, D_HEAD),
+                      dtype=DTYPE, dp=dp_kv, name="v_long_pre")
+        o = ctx.empty((T_q * C, D_HEAD),
+                      dtype=DTYPE, dp=dp_o, name="o_long_pre")
+        ctx.launch(
+            "gqa_prefill_long_64k",
+            gqa_prefill_kernel,
+            q, k, v, o,
+            T_q, S_kv, D_HEAD, C,
+            _auto_dim_remap=False,
+        )
+
+    result = run_bench(
+        topology=topo, bench_fn=_bench_fn,
+        device=resolve_device(None), engine_factory=_engine_factory,
+    )
+    assert result.completion.ok, (
+        f"prefill at S_kv=64K must complete after scope discipline lands; "
+        f"got {result.completion}"
+    )

tests/attention/test_gqa_prefill.py

@@ -0,0 +1,118 @@
+"""Phase 1 spec test for P6a GQA prefill kernel (head-parallel, C=1 baseline).
+
+P6a introduces ``_gqa_prefill.py`` with the head-parallel structure (one
+Q head per CUBE, per-CUBE distributed output, no reduce). C=1 is the
+degenerate case — no Ring KV, no IPCQ traffic. Validates kernel
+structure and T_q > 1 attention.
+
+P6b (deferred) adds the Ring KV rotation for C > 1, which needs either
+a new SFR install function (intra_* + wrapped E/W at CUBE level) or a
+topology-specific config — separate design call.
+
+The prefill kernel differs from decode (P1a/P2a/P2b) in three ways
+(ADR-0060 §5.5 / TL;DR):
+  1. Q has T_q > 1 rows (not just decode's single timestep).
+  2. Head-parallel placement: each CUBE owns ONE Q head — no M-fold.
+  3. Each CUBE writes its own head's output — NO reduce.
+
+Phase 1 (this commit): tests only — production code lands in Phase 2.
+"""
+from __future__ import annotations
+
+from pathlib import Path
+
+from kernbench.benches._gqa_prefill import gqa_prefill_kernel  # noqa: F401
+from kernbench.policy.placement.dp import DPPolicy
+from kernbench.runtime_api.bench_runner import run_bench
+from kernbench.runtime_api.types import resolve_device
+from kernbench.sim_engine.engine import GraphEngine
+from kernbench.topology.builder import resolve_topology
+
+TOPOLOGY_DEFAULT = Path(__file__).resolve().parents[2] / "topology.yaml"
+
+D_HEAD = 64
+DTYPE = "f16"
+
+
+def _engine_factory(t, d):
+    return GraphEngine(getattr(t, "topology_obj", t), enable_data=True)
+
+
+def _run_prefill(*, T_q: int, S_kv: int, C: int = 1):
+    """C=1 head-parallel prefill: single CUBE owns the one head + full KV."""
+    topo = resolve_topology(str(TOPOLOGY_DEFAULT))
+
+    def _bench_fn(ctx):
+        dp = DPPolicy(cube="replicate", pe="replicate",
+                      num_cubes=C, num_pes=1)
+        # Q: (T_q, d_head) — one head per CUBE (head-parallel; for C=1
+        # only one head total). 2D layout matches what the kernel loads.
+        # P6b will use a 3D (h_q, T_q, d_head) Q with cube_row_wise
+        # sharding so each CUBE owns its head.
+        q = ctx.zeros((T_q, D_HEAD), dtype=DTYPE, dp=dp,
+                      name=f"q_t{T_q}_c{C}")
+        # K, V: full local for C=1 (no ring). Kernel loads K as
+        # (d_head, S_kv) via byte-conserving reshape.
+        k = ctx.zeros((S_kv, D_HEAD), dtype=DTYPE, dp=dp,
+                      name=f"k_t{T_q}_c{C}")
+        v = ctx.zeros((S_kv, D_HEAD), dtype=DTYPE, dp=dp,
+                      name=f"v_t{T_q}_c{C}")
+        # O: (T_q, d_head) — per-CUBE distributed output.
+        o = ctx.empty((T_q, D_HEAD), dtype=DTYPE, dp=dp,
+                      name=f"o_t{T_q}_c{C}")
+        ctx.launch(
+            f"gqa_prefill_p6a_t{T_q}_s{S_kv}_c{C}",
+            gqa_prefill_kernel,
+            q, k, v, o,
+            T_q, S_kv, D_HEAD, C,
+            _auto_dim_remap=False,
+        )
+
+    return run_bench(
+        topology=topo, bench_fn=_bench_fn,
+        device=resolve_device(None),
+        engine_factory=_engine_factory,
+    )
+
+
+def _count(op_log, name: str) -> int:
+    return sum(1 for r in op_log if r.op_name == name)
+
+
+def test_prefill_c_one_t_q_one_completes():
+    """C=1, T_q=1: smallest workload (decode-like)."""
+    result = _run_prefill(T_q=1, S_kv=16, C=1)
+    assert result.completion.ok, (
+        f"prefill C=1 T_q=1 failed: {result.completion}"
+    )
+
+
+def test_prefill_c_one_t_q_four_completes():
+    """C=1, T_q=4: real prefill (Q has multiple rows) — distinguishes
+    prefill from decode (T_q=1)."""
+    result = _run_prefill(T_q=4, S_kv=16, C=1)
+    assert result.completion.ok, (
+        f"prefill C=1 T_q=4 failed: {result.completion}"
+    )
+
+
+def test_prefill_c_one_no_ipcq_traffic():
+    """C=1: no ring step, no IPCQ traffic."""
+    result = _run_prefill(T_q=4, S_kv=16, C=1)
+    assert result.completion.ok
+    n_copy = _count(result.engine.op_log, "ipcq_copy")
+    assert n_copy == 0, (
+        f"C=1 must have no IPCQ traffic (no ring); got {n_copy}"
+    )
+
+
+def test_prefill_c_one_one_dma_write():
+    """C=1, one head: exactly one dma_write (per-CUBE distributed output;
+    no reduce). For C > 1 in P6b this becomes dma_write_count == C."""
+    result = _run_prefill(T_q=4, S_kv=16, C=1)
+    assert result.completion.ok
+    n_writes = _count(result.engine.op_log, "dma_write")
+    assert n_writes == 1, (
+        f"C=1 prefill: expected 1 dma_write (one head per cube); "
+        f"got {n_writes}"
+    )

tests/attention/test_gqa_prefill_ring.py

@@ -0,0 +1,161 @@
+"""Phase 1 spec test for P6b GQA prefill Ring KV (head-parallel, C>1).
+
+P6b adds the Ring KV rotation (ADR-0060 §5.5) to the prefill kernel.
+Each CUBE owns one Q head + its KV slice; over C ring steps the KV
+blocks rotate around the C CUBEs so every CUBE sees every block. The
+running ``(m, ℓ, O)`` is folded inside the loop. No reduce — each CUBE
+writes its own head's output.
+
+Requires a new SFR install ``configure_sfr_intercube_ring(ring_size=C)``
+that wires:
+  - ``intra_*``  : 2×4 PE grid within a CUBE (same as multisip)
+  - ``E/W``      : 1D ring of CUBEs 0..ring_size-1 WITH WRAP
+                    (symmetric to ``configure_sfr_intracube_pe_ring``,
+                     applied at CUBE level)
+  - ``global_*`` : SIP-level (same as multisip)
+
+The kernel ring body:
+  for step in range(1, C):
+      tl.send(dir="W", src=Kc)
+      tl.send(dir="W", src=Vc)
+      Kc = tl.recv(dir="E", ...)
+      Vc = tl.recv(dir="E", ...)
+      ... local partial + online-softmax merge into (m, ℓ, O) ...
+
+Per CUBE per step: 2 sends (K, V) → 2 ``ipcq_copy``. Across all CUBEs:
+``(C-1) * 2 * C`` total ipcq_copy.
+
+Restriction in P6b first cut: ``C ∈ {1, 2, 4}`` (single row of the 4×4
+cube mesh). C=8 ring would span rows — follow-on.
+
+Phase 1 (this commit): tests only — production code lands in Phase 2.
+"""
+from __future__ import annotations
+
+from pathlib import Path
+
+from kernbench.benches._gqa_prefill import gqa_prefill_kernel  # noqa: F401
+from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config
+from kernbench.ccl.sfr_config import configure_sfr_intercube_ring  # noqa: F401 (Phase 2)
+from kernbench.policy.placement.dp import DPPolicy
+from kernbench.runtime_api.bench_runner import run_bench
+from kernbench.runtime_api.types import resolve_device
+from kernbench.sim_engine.engine import GraphEngine
+from kernbench.topology.builder import resolve_topology
+
+TOPOLOGY_DEFAULT = Path(__file__).resolve().parents[2] / "topology.yaml"
+
+D_HEAD = 64
+DTYPE = "f16"
+
+
+def _ccl_cfg():
+    return resolve_algorithm_config(
+        load_ccl_config(), name="lrab_hierarchical_allreduce",
+    )
+
+
+def _engine_factory(t, d):
+    return GraphEngine(getattr(t, "topology_obj", t), enable_data=True)
+
+
+def _run_prefill_ring(*, T_q: int, S_kv: int, C: int):
+    """Head-parallel prefill with Ring KV across C CUBEs."""
+    topo = resolve_topology(str(TOPOLOGY_DEFAULT))
+
+    def _bench_fn(ctx):
+        # P6b: new SFR install with cube-level ring wrap.
+        configure_sfr_intercube_ring(
+            ctx.engine, ctx.spec, _ccl_cfg(), ring_size=C,
+        )
+        # Q replicated on every CUBE (zeros for testing; per-CUBE head
+        # indexing is implicit). KV sequence-sharded by CUBE. O
+        # distributed — each CUBE writes its slice of (T_q*C, d_head).
+        dp_q = DPPolicy(cube="replicate", pe="replicate",
+                        num_cubes=C, num_pes=1)
+        dp_kv = DPPolicy(cube="row_wise", pe="replicate",
+                         num_cubes=C, num_pes=1)
+        dp_o = DPPolicy(cube="row_wise", pe="replicate",
+                        num_cubes=C, num_pes=1)
+        q = ctx.zeros((T_q, D_HEAD), dtype=DTYPE, dp=dp_q,
+                      name=f"q_t{T_q}_c{C}_ring")
+        k = ctx.zeros((S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv,
+                      name=f"k_t{T_q}_c{C}_ring")
+        v = ctx.zeros((S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv,
+                      name=f"v_t{T_q}_c{C}_ring")
+        # O: (T_q * C, d_head), each CUBE writes (T_q, d_head) at its
+        # slice. dma_write_count = C (per-CUBE distributed output).
+        o = ctx.empty((T_q * C, D_HEAD), dtype=DTYPE, dp=dp_o,
+                      name=f"o_t{T_q}_c{C}_ring")
+        ctx.launch(
+            f"gqa_prefill_ring_t{T_q}_s{S_kv}_c{C}",
+            gqa_prefill_kernel,
+            q, k, v, o,
+            T_q, S_kv, D_HEAD, C,
+            _auto_dim_remap=False,
+        )
+
+    return run_bench(
+        topology=topo, bench_fn=_bench_fn,
+        device=resolve_device(None),
+        engine_factory=_engine_factory,
+    )
+
+
+def _count(op_log, name: str) -> int:
+    return sum(1 for r in op_log if r.op_name == name)
+
+
+# ── C=2 Ring KV ────────────────────────────────────────────────────────
+
+
+def test_prefill_ring_c_two_completes():
+    """C=2: 1 ring step rotates KV between the 2 CUBEs."""
+    result = _run_prefill_ring(T_q=4, S_kv=16, C=2)
+    assert result.completion.ok, (
+        f"prefill ring C=2 failed: {result.completion}"
+    )
+
+
+def test_prefill_ring_c_two_distributed_output():
+    """C=2: per-CUBE distributed output, no reduce → dma_write_count == 2."""
+    result = _run_prefill_ring(T_q=4, S_kv=16, C=2)
+    assert result.completion.ok
+    n_writes = _count(result.engine.op_log, "dma_write")
+    assert n_writes == 2, (
+        f"C=2 prefill: expected 2 dma_write (one per CUBE); got {n_writes}"
+    )
+
+
+def test_prefill_ring_c_two_ipcq_count():
+    """C=2: 1 ring step × 2 handles (K, V) × 2 CUBEs = 4 ipcq_copy."""
+    result = _run_prefill_ring(T_q=4, S_kv=16, C=2)
+    assert result.completion.ok
+    n_copy = _count(result.engine.op_log, "ipcq_copy")
+    expected = (2 - 1) * 2 * 2
+    assert n_copy == expected, (
+        f"C=2 ring: expected {expected} ipcq_copy "
+        f"((C-1)·2·C = 1·2·2); got {n_copy}"
+    )
+
+
+# ── C=4 Ring KV (combined assertions) ─────────────────────────────────
+
+
+def test_prefill_ring_c_four_combined():
+    """C=4: 3 ring steps rotate KV around 4 CUBEs.
+    Expected: completes; 4 dma_writes; (4-1)·2·4 = 24 ipcq_copy."""
+    result = _run_prefill_ring(T_q=4, S_kv=32, C=4)
+    assert result.completion.ok, (
+        f"prefill ring C=4 failed: {result.completion}"
+    )
+    n_writes = _count(result.engine.op_log, "dma_write")
+    assert n_writes == 4, (
+        f"C=4 prefill: expected 4 dma_write (one per CUBE); got {n_writes}"
+    )
+    n_copy = _count(result.engine.op_log, "ipcq_copy")
+    expected = (4 - 1) * 2 * 4
+    assert n_copy == expected, (
+        f"C=4 ring: expected {expected} ipcq_copy "
+        f"((C-1)·2·C = 3·2·4); got {n_copy}"
+    )

tests/attention/test_gqa_scoped_writeback.py

@@ -0,0 +1,198 @@
+"""Phase 1 spec test for Phase C: scratch_scope + tl.copy_to discipline
+in the GQA kernels (ADR-0060 §5.2 / §5.5 + ADR-0063 §D3 / §D3.1).
+
+ADR-0060 §5.2 (decode pseudocode line 75 / §5.5 (prefill pseudocode line
+96) both wrap per-tile / per-ring-step intermediates in
+``with tl.scratch_scope():`` and persist the merged running ``(m, ℓ, O)``
+to a persistent arena allocated outside the scope. The original ADR-0063
+§D3 specifies the two-arena pattern; §D3.1 specifies the
+``tl.copy_to(dst, src)`` writeback primitive used to persist scoped
+results.
+
+Currently neither kernel uses ``scratch_scope`` or ``copy_to``; their
+chain-merge / ring-merge bodies allocate every intermediate from the
+bump cursor and never recycle. Result: op_log has 0 ``copy`` entries.
+
+After Phase 2: each per-tile / per-step merge writes the new running
+``(m, ℓ, O)`` via ``copy_to`` to the persistent arena. Per merge step
+→ 3 copy ops (m, ℓ, O).
+
+Phase 1 (this commit): tests only — production code lands in Phase 2.
+"""
+from __future__ import annotations
+
+from pathlib import Path
+
+from kernbench.benches._gqa_decode import gqa_decode_kernel  # noqa: F401
+from kernbench.benches._gqa_prefill import gqa_prefill_kernel  # noqa: F401
+from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config
+from kernbench.ccl.sfr_config import (
+    configure_sfr_intercube_multisip,
+    configure_sfr_intercube_ring,
+)
+from kernbench.policy.placement.dp import DPPolicy
+from kernbench.runtime_api.bench_runner import run_bench
+from kernbench.runtime_api.types import resolve_device
+from kernbench.sim_engine.engine import GraphEngine
+from kernbench.topology.builder import resolve_topology
+
+TOPOLOGY_DEFAULT = Path(__file__).resolve().parents[2] / "topology.yaml"
+
+D_HEAD = 64
+DTYPE = "f16"
+
+
+def _ccl_cfg():
+    return resolve_algorithm_config(
+        load_ccl_config(), name="lrab_hierarchical_allreduce",
+    )
+
+
+def _engine_factory(t, d):
+    return GraphEngine(getattr(t, "topology_obj", t), enable_data=True)
+
+
+def _count(op_log, name: str) -> int:
+    return sum(1 for r in op_log if r.op_name == name)
+
+
+# ── Decode chain merges must use scratch_scope + copy_to ─────────────
+
+
+def _run_decode_sp(*, h_q: int, h_kv: int, P: int, S_kv: int):
+    """Single-CUBE SP decode (C=1, P PEs along intra-cube chain)."""
+    topo = resolve_topology(str(TOPOLOGY_DEFAULT))
+
+    def _bench_fn(ctx):
+        configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
+        dp_full = DPPolicy(cube="replicate", pe="replicate",
+                           num_cubes=1, num_pes=P)
+        dp_kv = DPPolicy(cube="replicate", pe="row_wise",
+                         num_cubes=1, num_pes=P)
+        q = ctx.zeros((1, h_q * D_HEAD),
+                      dtype=DTYPE, dp=dp_full, name=f"q_sc_{P}")
+        k = ctx.zeros((S_kv, h_kv * D_HEAD),
+                      dtype=DTYPE, dp=dp_kv, name=f"k_sc_{P}")
+        v = ctx.zeros((S_kv, h_kv * D_HEAD),
+                      dtype=DTYPE, dp=dp_kv, name=f"v_sc_{P}")
+        o = ctx.empty((1, h_q * D_HEAD),
+                      dtype=DTYPE, dp=dp_full, name=f"o_sc_{P}")
+        ctx.launch(
+            f"gqa_decode_scoped_{P}",
+            gqa_decode_kernel,
+            q, k, v, o,
+            1, S_kv, h_q, h_kv, D_HEAD,
+            1, P,
+            _auto_dim_remap=False,
+        )
+
+    return run_bench(
+        topology=topo, bench_fn=_bench_fn,
+        device=resolve_device(None), engine_factory=_engine_factory,
+    )
+
+
+def test_decode_chain_merges_emit_copy_to_writeback():
+    """ADR-0060 §5.2 + ADR-0063 §D3.1: each chain-merge step must wrap
+    its intermediates in ``scratch_scope`` and persist the new running
+    ``(m, ℓ, O)`` via ``tl.copy_to``.
+
+    For (C=1, P=8): 7 intra-cube chain merges × 3 handles (m, ℓ, O) per
+    merge ⇒ 21 ``copy`` entries.
+
+    Currently 0 because the kernel never calls ``tl.copy_to``.
+    """
+    result = _run_decode_sp(h_q=8, h_kv=1, P=8, S_kv=64)
+    assert result.completion.ok, f"decode SP failed: {result.completion}"
+    n_copy = _count(result.engine.op_log, "copy")
+    assert n_copy > 0, (
+        f"decode kernel must emit copy_to writeback per merge step "
+        f"(ADR-0060 §5.2 + ADR-0063 §D3.1); got 0 ``copy`` entries"
+    )
+
+
+# ── Prefill Ring KV merges must use scratch_scope + copy_to ──────────
+
+
+def _run_prefill_ring(*, T_q: int, S_kv: int, C: int):
+    """Head-parallel prefill with Ring KV across C CUBEs."""
+    topo = resolve_topology(str(TOPOLOGY_DEFAULT))
+
+    def _bench_fn(ctx):
+        configure_sfr_intercube_ring(
+            ctx.engine, ctx.spec, _ccl_cfg(), ring_size=C,
+        )
+        dp_q = DPPolicy(cube="replicate", pe="replicate",
+                        num_cubes=C, num_pes=1)
+        dp_kv = DPPolicy(cube="row_wise" if C > 1 else "replicate",
+                         pe="replicate", num_cubes=C, num_pes=1)
+        dp_o = DPPolicy(cube="row_wise" if C > 1 else "replicate",
+                        pe="replicate", num_cubes=C, num_pes=1)
+        q = ctx.zeros((T_q, D_HEAD),
+                      dtype=DTYPE, dp=dp_q, name=f"q_ring_{C}")
+        k = ctx.zeros((S_kv, D_HEAD),
+                      dtype=DTYPE, dp=dp_kv, name=f"k_ring_{C}")
+        v = ctx.zeros((S_kv, D_HEAD),
+                      dtype=DTYPE, dp=dp_kv, name=f"v_ring_{C}")
+        o = ctx.empty((T_q * C, D_HEAD),
+                      dtype=DTYPE, dp=dp_o, name=f"o_ring_{C}")
+        ctx.launch(
+            f"gqa_prefill_scoped_{C}",
+            gqa_prefill_kernel,
+            q, k, v, o,
+            T_q, S_kv, D_HEAD, C,
+            _auto_dim_remap=False,
+        )
+
+    return run_bench(
+        topology=topo, bench_fn=_bench_fn,
+        device=resolve_device(None), engine_factory=_engine_factory,
+    )
+
+
+def test_prefill_ring_step_merges_emit_copy_to_writeback():
+    """ADR-0060 §5.5 + ADR-0063 §D3.1: each Ring KV step's online-softmax
+    merge must wrap its intermediates in ``scratch_scope`` and persist
+    the new running ``(m, ℓ, O)`` via ``tl.copy_to``.
+
+    For C=4: 3 ring-step merges (steps 1..C-1) × 3 handles (m, ℓ, O) per
+    merge ⇒ 9 ``copy`` entries per participating CUBE. Aggregated across
+    C CUBEs: ⇒ 36 ``copy`` entries.
+
+    Currently 0 because the kernel never calls ``tl.copy_to``.
+    """
+    result = _run_prefill_ring(T_q=4, S_kv=16, C=4)
+    assert result.completion.ok, f"prefill ring failed: {result.completion}"
+    n_copy = _count(result.engine.op_log, "copy")
+    assert n_copy > 0, (
+        f"prefill ring kernel must emit copy_to writeback per merge step "
+        f"(ADR-0060 §5.5 + ADR-0063 §D3.1); got 0 ``copy`` entries"
+    )
+
+
+# ── Scoped kernels still produce the same op_log shape (regression) ──
+
+
+def test_decode_scoped_still_has_root_only_write():
+    """ADR-0060 §A.2 root-only output must hold under the rewrite:
+    adding scratch_scope + copy_to should not change the reduce
+    topology; only the per-PE scratch usage. Single PE 0 writes O."""
+    result = _run_decode_sp(h_q=8, h_kv=1, P=8, S_kv=64)
+    assert result.completion.ok
+    n_writes = _count(result.engine.op_log, "dma_write")
+    assert n_writes == 1, (
+        f"scoped decode must still produce 1 dma_write (root-only); "
+        f"got {n_writes}"
+    )
+
+
+def test_prefill_scoped_still_has_per_cube_distributed_output():
+    """ADR-0060 §5.5 per-CUBE distributed output must hold under the
+    rewrite: scoped prefill still writes one O slice per CUBE."""
+    result = _run_prefill_ring(T_q=4, S_kv=16, C=4)
+    assert result.completion.ok
+    n_writes = _count(result.engine.op_log, "dma_write")
+    assert n_writes == 4, (
+        f"scoped prefill must write one O per CUBE (C=4 → 4 dma_write); "
+        f"got {n_writes}"
+    )

tests/attention/test_gqa_short_context.py

@@ -0,0 +1,260 @@
+"""Phase 1 spec test for Phase D: short-context GQA kernels
+(ADR-0060 §B "Items from the long/short context split", item B.split.2).
+
+The long-context kernels (§5.2 decode / §5.5 prefill) shard each KV
+head row-wise across all CUBEs. At short context (S_kv < 256K), that
+shard is too thin to feed the engines and the cube-level collective
+overhead dominates. The short-context kernels drop cube-SP entirely:
+each CUBE owns ``kv_per_cube`` *whole* KV heads, no S_kv sharding across
+CUBEs.
+
+Design (per AskUserQuestion answers in this session):
+  - PE-parallel heads: P PEs split into ``kv_per_cube`` groups, each
+    group does PE-SP across (P/kv_per_cube) PEs for one owned head.
+  - scratch_scope + tl.copy_to discipline mirrors the long kernels.
+  - One short kernel handles kv_per_cube ∈ {1, 2, 4, 8} via a parameter.
+
+Group layout on the 2×4 PE grid:
+  kv_per_cube=1, group=8 PEs (full 2×4): row chain + col bridge.
+  kv_per_cube=2, group=4 PEs (one row):  row chain only.
+  kv_per_cube=4, group=2 PEs (adj cols): 1-step chain.
+  kv_per_cube=8, group=1 PE:             no chain.
+
+After chain reduce-to-group-root, the group's root PE writes its
+owned head's output. No inter-CUBE reduce.
+
+Phase 1 (this commit): tests only — production code lands in Phase 2.
+All tests fail because the short kernels (``_gqa_decode_short.py`` and
+``_gqa_prefill_short.py``) do not exist yet.
+"""
+from __future__ import annotations
+
+from pathlib import Path
+
+import pytest
+
+from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config
+from kernbench.ccl.sfr_config import configure_sfr_intercube_multisip
+from kernbench.policy.placement.dp import DPPolicy
+from kernbench.runtime_api.bench_runner import run_bench
+from kernbench.runtime_api.types import resolve_device
+from kernbench.sim_engine.engine import GraphEngine
+from kernbench.topology.builder import resolve_topology
+
+TOPOLOGY_DEFAULT = Path(__file__).resolve().parents[2] / "topology.yaml"
+
+D_HEAD = 64
+DTYPE = "f16"
+
+
+def _ccl_cfg():
+    return resolve_algorithm_config(
+        load_ccl_config(), name="lrab_hierarchical_allreduce",
+    )
+
+
+def _engine_factory(t, d):
+    return GraphEngine(getattr(t, "topology_obj", t), enable_data=True)
+
+
+def _count(op_log, name: str) -> int:
+    return sum(1 for r in op_log if r.op_name == name)
+
+
+# ── Decode short-context kernel ──────────────────────────────────────
+
+
+def _run_decode_short(*, h_q: int, h_kv: int, kv_per_cube: int,
+                     C: int, P: int, S_kv: int):
+    """Run the short-context decode kernel with PE-parallel heads.
+
+    Layout (after design iteration — see Phase D failure-recovery):
+      Q: (T_q, h_q·D_HEAD) replicated; kernel reshapes byte-conservingly.
+      K, V: (h_kv·S_kv, D_HEAD) head-stacked, ``cube=row_wise, pe=row_wise``
+          so each PE gets a contiguous (S_local, D_HEAD) chunk at its own
+          addressable shard (no partial reads needed).
+      O: replicated; each group root writes the full byte-conserving
+          (h_q·T_q, D_HEAD) result.
+    """
+    from kernbench.benches._gqa_decode_short import gqa_decode_short_kernel  # Phase 2
+
+    topo = resolve_topology(str(TOPOLOGY_DEFAULT))
+
+    def _bench_fn(ctx):
+        configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
+        dp_full = DPPolicy(cube="replicate", pe="replicate",
+                           num_cubes=C, num_pes=P)
+        # Head-stacked KV with row_wise sharding so each PE's chunk is
+        # contiguous and exactly (S_local, D_HEAD) addressable.
+        dp_kv = DPPolicy(cube="row_wise", pe="row_wise",
+                         num_cubes=C, num_pes=P)
+        q = ctx.zeros((1, h_q * D_HEAD),
+                      dtype=DTYPE, dp=dp_full, name=f"q_short_{kv_per_cube}_{C}")
+        k = ctx.zeros((h_kv * S_kv, D_HEAD),
+                      dtype=DTYPE, dp=dp_kv, name=f"k_short_{kv_per_cube}_{C}")
+        v = ctx.zeros((h_kv * S_kv, D_HEAD),
+                      dtype=DTYPE, dp=dp_kv, name=f"v_short_{kv_per_cube}_{C}")
+        o = ctx.empty((1, h_q * D_HEAD),
+                      dtype=DTYPE, dp=dp_full, name=f"o_short_{kv_per_cube}_{C}")
+        ctx.launch(
+            f"gqa_decode_short_{kv_per_cube}_{C}",
+            gqa_decode_short_kernel,
+            q, k, v, o,
+            1, S_kv, h_q, h_kv, D_HEAD, C, P, kv_per_cube,
+            _auto_dim_remap=False,
+        )
+
+    return run_bench(
+        topology=topo, bench_fn=_bench_fn,
+        device=resolve_device(None), engine_factory=_engine_factory,
+    )
+
+
+def test_short_decode_smoke_kv_per_cube_2_C_4():
+    """ADR-0060 §B.split.2 headline: kv_per_cube=2, C=4 — each CUBE
+    owns 2 heads (8 heads / 4 CUBEs). PE-parallel heads splits the 8
+    PEs into 2 groups of 4, each group does PE-SP for one owned head.
+    Smoke: kernel completes."""
+    result = _run_decode_short(
+        h_q=8, h_kv=8, kv_per_cube=2, C=4, P=8, S_kv=64,
+    )
+    assert result.completion.ok, (
+        f"short decode kv_per_cube=2 must complete; got {result.completion}"
+    )
+
+
+def test_short_decode_smoke_kv_per_cube_4_C_2():
+    """kv_per_cube=4, C=2 — half the CUBEs participate, each owns 4
+    heads, 4 PE groups of 2 PEs each."""
+    result = _run_decode_short(
+        h_q=8, h_kv=8, kv_per_cube=4, C=2, P=8, S_kv=64,
+    )
+    assert result.completion.ok, (
+        f"short decode kv_per_cube=4 must complete; got {result.completion}"
+    )
+
+
+def test_short_decode_smoke_kv_per_cube_8_C_1():
+    """kv_per_cube=8, C=1 — all heads on one CUBE. 8 PE groups of 1 PE
+    each → no chain reduce, each PE writes its head's output."""
+    result = _run_decode_short(
+        h_q=8, h_kv=8, kv_per_cube=8, C=1, P=8, S_kv=64,
+    )
+    assert result.completion.ok, (
+        f"short decode kv_per_cube=8 must complete; got {result.completion}"
+    )
+
+
+def test_short_decode_dma_write_count_equals_h_kv():
+    """ADR-0060 §B.split.2: each owned head produces exactly one
+    output (no inter-CUBE reduce). Total dma_writes = h_kv across all
+    participating CUBEs and groups.
+
+    For h_kv=8, kv_per_cube=2, C=4: each CUBE writes 2 outputs →
+    4 × 2 = 8 dma_writes total.
+    """
+    result = _run_decode_short(
+        h_q=8, h_kv=8, kv_per_cube=2, C=4, P=8, S_kv=64,
+    )
+    assert result.completion.ok
+    n_writes = _count(result.engine.op_log, "dma_write")
+    assert n_writes == 8, (
+        f"short decode: expected 8 dma_writes (h_kv); got {n_writes}"
+    )
+
+
+def test_short_decode_no_inter_cube_traffic():
+    """ADR-0060 §B.split.2: each head is fully owned by one CUBE → no
+    inter-CUBE reduce. The kernel must not invoke CUBE-level E/W IPCQ.
+
+    Today's long-context kernel at C=4 emits ~12 inter-CUBE ipcq_copy
+    via direction "E"/"W". The short kernel must emit zero of those,
+    keeping all IPCQ traffic on the ``intra_*`` directions.
+    """
+    result = _run_decode_short(
+        h_q=8, h_kv=8, kv_per_cube=2, C=4, P=8, S_kv=64,
+    )
+    assert result.completion.ok
+    # Count IPCQ traffic that targeted the CUBE-level "E"/"W" directions.
+    inter_cube_ipcq = sum(
+        1 for r in result.engine.op_log
+        if r.op_name == "ipcq_copy"
+        and r.params.get("direction") in ("E", "W")
+    )
+    assert inter_cube_ipcq == 0, (
+        f"short decode must have no inter-CUBE E/W IPCQ; got {inter_cube_ipcq}"
+    )
+
+
+# ── Prefill short-context kernel ─────────────────────────────────────
+
+
+def _run_prefill_short(*, h_kv: int, kv_per_cube: int,
+                      C: int, P: int, T_q: int, S_kv: int):
+    """Run the short-context prefill kernel with PE-parallel heads.
+
+    Layout: same head-stacked K/V scheme as decode short, with Q/O
+    replicated and byte-conserving reshape inside the kernel.
+    """
+    from kernbench.benches._gqa_prefill_short import gqa_prefill_short_kernel  # Phase 2
+
+    topo = resolve_topology(str(TOPOLOGY_DEFAULT))
+
+    def _bench_fn(ctx):
+        configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
+        dp_q = DPPolicy(cube="replicate", pe="replicate",
+                        num_cubes=C, num_pes=P)
+        dp_kv = DPPolicy(cube="row_wise", pe="row_wise",
+                         num_cubes=C, num_pes=P)
+        dp_o = DPPolicy(cube="replicate", pe="replicate",
+                        num_cubes=C, num_pes=P)
+        q = ctx.zeros((T_q, h_kv * D_HEAD),
+                      dtype=DTYPE, dp=dp_q, name=f"q_pre_short_{kv_per_cube}_{C}")
+        k = ctx.zeros((h_kv * S_kv, D_HEAD),
+                      dtype=DTYPE, dp=dp_kv, name=f"k_pre_short_{kv_per_cube}_{C}")
+        v = ctx.zeros((h_kv * S_kv, D_HEAD),
+                      dtype=DTYPE, dp=dp_kv, name=f"v_pre_short_{kv_per_cube}_{C}")
+        o = ctx.empty((T_q, h_kv * D_HEAD),
+                      dtype=DTYPE, dp=dp_o, name=f"o_pre_short_{kv_per_cube}_{C}")
+        ctx.launch(
+            f"gqa_prefill_short_{kv_per_cube}_{C}",
+            gqa_prefill_short_kernel,
+            q, k, v, o,
+            T_q, S_kv, h_kv, D_HEAD, C, P, kv_per_cube,
+            _auto_dim_remap=False,
+        )
+
+    return run_bench(
+        topology=topo, bench_fn=_bench_fn,
+        device=resolve_device(None), engine_factory=_engine_factory,
+    )
+
+
+def test_short_prefill_smoke_kv_per_cube_2_C_4():
+    """Short prefill kv_per_cube=2, C=4. Smoke: completes."""
+    result = _run_prefill_short(
+        h_kv=8, kv_per_cube=2, C=4, P=8, T_q=4, S_kv=64,
+    )
+    assert result.completion.ok, (
+        f"short prefill kv_per_cube=2 must complete; got {result.completion}"
+    )
+
+
+def test_short_prefill_no_ring_KV_traffic():
+    """ADR-0060 §B.split.2: short prefill DOES NOT use Ring KV — each
+    CUBE owns its KV heads fully, no rotation. The kernel must not
+    emit any KV-rotation IPCQ traffic at the CUBE level.
+    """
+    result = _run_prefill_short(
+        h_kv=8, kv_per_cube=2, C=4, P=8, T_q=4, S_kv=64,
+    )
+    assert result.completion.ok
+    inter_cube_ipcq = sum(
+        1 for r in result.engine.op_log
+        if r.op_name == "ipcq_copy"
+        and r.params.get("direction") in ("E", "W")
+    )
+    assert inter_cube_ipcq == 0, (
+        f"short prefill must have no Ring KV (no inter-CUBE E/W IPCQ); "
+        f"got {inter_cube_ipcq}"
+    )

tests/attention/test_mesh_kernels_rank_axis.py

@@ -1,172 +0,0 @@
-"""Phase 1 spec test for ``rank_axis`` parameter on the two mesh kernels.
-
-ADR-0059's mesh kernels currently hard-code ``rank = tl.program_id(axis=0)``,
-which only works for single_user_* panels (rank == pe_id within cube).
-For multi_user_* panels the ring is at the cube level — rank should be
-``cube_id`` (axis=1), and the 7 non-rank-leader PEs in each cube should
-not run the ring (they only hold KV replicas).
-
-This test pins the desired ``rank_axis`` kwarg semantics:
-
-  rank_axis = 0 (default, single_user)
-    rank = tl.program_id(axis=0). Every PE in the cube runs the ring.
-    Existing behavior — no change.
-
-  rank_axis = 1 (multi_user)
-    if tl.program_id(axis=0) != 0: return.  (7/8 PEs early-exit.)
-    rank = tl.program_id(axis=1).
-
-Phase 1 expectation: tests fail today (kernels don't accept the kwarg).
-Phase 2 lands the parameter on both kernels; tests turn green and the
-multi_user_* diag harness clears its first send.
-"""
-from __future__ import annotations
-
-from kernbench.common.ipcq_types import IpcqRecvCmd, IpcqSendCmd
-from kernbench.common.pe_commands import GemmCmd
-from kernbench.triton_emu.tl_context import TLContext, run_kernel
-
-from kernbench.benches._attention_mesh_kv import attention_mesh_kv_kernel
-from kernbench.benches._attention_mesh_mlo import attention_mesh_mlo_kernel
-
-S_Q_PREFILL = 16
-S_Q_DECODE = 1
-S_KV_PER_RANK = 16
-H_Q = 1
-H_KV = 1
-D_HEAD = 64
-N_RANKS_MULTI = 4
-PES_PER_CUBE = 8
-
-Q_PTR = 0x10000
-K_PTR = 0x20000
-V_PTR = 0x30000
-O_PTR = 0x40000
-
-
-def _tl(pe_id: int, cube_id: int, num_pes: int, num_cubes: int) -> TLContext:
-    return TLContext(
-        pe_id=pe_id,
-        num_programs=num_pes,
-        cube_id=cube_id,
-        num_cubes=num_cubes,
-        dispatch_cycles=0,
-        scratch_base=0x80000,
-        scratch_size=1 << 20,
-    )
-
-
-# ── Default rank_axis=0 backward-compat ──────────────────────────
-
-
-def test_mlo_kernel_default_rank_axis_zero_emits_commands_on_all_pes():
-    """rank_axis defaults to 0 → kernel uses pe_id as rank, runs on every
-    PE. Verify by running rank=3 (interior PE) in a single-cube 8-rank
-    setup and asserting at least one GEMM and at least one IPCQ send
-    are emitted (interior ranks send in both directions)."""
-    tl = _tl(pe_id=3, cube_id=0, num_pes=8, num_cubes=1)
-    run_kernel(
-        attention_mesh_mlo_kernel, tl,
-        Q_PTR, K_PTR, V_PTR, O_PTR,
-        S_Q_DECODE, S_KV_PER_RANK, H_Q, H_KV, D_HEAD, 8,
-    )
-    assert any(isinstance(c, GemmCmd) for c in tl.commands), \
-        "default rank_axis=0 must run the kernel (≥1 GEMM)"
-    assert any(isinstance(c, IpcqSendCmd) for c in tl.commands), \
-        "interior rank must emit ≥1 IpcqSendCmd"
-
-
-def test_kv_kernel_default_rank_axis_zero_emits_commands_on_all_pes():
-    tl = _tl(pe_id=3, cube_id=0, num_pes=8, num_cubes=1)
-    run_kernel(
-        attention_mesh_kv_kernel, tl,
-        Q_PTR, K_PTR, V_PTR, O_PTR,
-        S_Q_PREFILL, S_KV_PER_RANK, H_Q, H_KV, D_HEAD, 8,
-    )
-    assert any(isinstance(c, GemmCmd) for c in tl.commands)
-    assert any(isinstance(c, IpcqSendCmd) for c in tl.commands)
-
-
-# ── rank_axis=1 multi_user semantics ─────────────────────────────
-
-
-def test_mlo_kernel_rank_axis_one_gates_non_zero_pe_to_no_commands():
-    """rank_axis=1 + pe_id != 0 → kernel must early-return; no GEMM,
-    no DMA, no IPCQ. The 7 non-rank-leader PEs in a multi_user cube
-    must stay completely silent so the cube-level SFR install isn't
-    asked to route sends from PEs that have no neighbors installed."""
-    tl = _tl(pe_id=2, cube_id=1, num_pes=PES_PER_CUBE, num_cubes=N_RANKS_MULTI)
-    run_kernel(
-        attention_mesh_mlo_kernel, tl,
-        Q_PTR, K_PTR, V_PTR, O_PTR,
-        S_Q_DECODE, S_KV_PER_RANK, H_Q, H_KV, D_HEAD, N_RANKS_MULTI,
-        rank_axis=1,
-    )
-    assert not any(isinstance(c, GemmCmd) for c in tl.commands), \
-        "pe_id=2 with rank_axis=1 must not emit GEMMs"
-    assert not any(isinstance(c, IpcqSendCmd) for c in tl.commands), \
-        "pe_id=2 with rank_axis=1 must not emit IpcqSendCmd"
-    assert not any(isinstance(c, IpcqRecvCmd) for c in tl.commands), \
-        "pe_id=2 with rank_axis=1 must not emit IpcqRecvCmd"
-
-
-def test_kv_kernel_rank_axis_one_gates_non_zero_pe_to_no_commands():
-    tl = _tl(pe_id=2, cube_id=1, num_pes=PES_PER_CUBE, num_cubes=N_RANKS_MULTI)
-    run_kernel(
-        attention_mesh_kv_kernel, tl,
-        Q_PTR, K_PTR, V_PTR, O_PTR,
-        S_Q_PREFILL, S_KV_PER_RANK, H_Q, H_KV, D_HEAD, N_RANKS_MULTI,
-        rank_axis=1,
-    )
-    assert not any(isinstance(c, GemmCmd) for c in tl.commands)
-    assert not any(isinstance(c, IpcqSendCmd) for c in tl.commands)
-    assert not any(isinstance(c, IpcqRecvCmd) for c in tl.commands)
-
-
-def test_mlo_kernel_rank_axis_one_pe_zero_uses_cube_id_as_rank():
-    """rank_axis=1 + pe_id == 0 → kernel runs the ring with rank=cube_id.
-    For cube_id=1 in a 4-cube ring, rank=1 is an interior rank: has_E=True
-    AND has_W=True → IPCQ sends emitted in both E and W directions.
-    """
-    tl = _tl(pe_id=0, cube_id=1, num_pes=PES_PER_CUBE, num_cubes=N_RANKS_MULTI)
-    run_kernel(
-        attention_mesh_mlo_kernel, tl,
-        Q_PTR, K_PTR, V_PTR, O_PTR,
-        S_Q_DECODE, S_KV_PER_RANK, H_Q, H_KV, D_HEAD, N_RANKS_MULTI,
-        rank_axis=1,
-    )
-    sends = [c for c in tl.commands if isinstance(c, IpcqSendCmd)]
-    assert any(s.direction == "E" for s in sends), \
-        "cube_id=1 (interior) must emit ≥1 E-send"
-    assert any(s.direction == "W" for s in sends), \
-        "cube_id=1 (interior) must emit ≥1 W-send"
-
-
-def test_kv_kernel_rank_axis_one_pe_zero_uses_cube_id_as_rank():
-    tl = _tl(pe_id=0, cube_id=1, num_pes=PES_PER_CUBE, num_cubes=N_RANKS_MULTI)
-    run_kernel(
-        attention_mesh_kv_kernel, tl,
-        Q_PTR, K_PTR, V_PTR, O_PTR,
-        S_Q_PREFILL, S_KV_PER_RANK, H_Q, H_KV, D_HEAD, N_RANKS_MULTI,
-        rank_axis=1,
-    )
-    sends = [c for c in tl.commands if isinstance(c, IpcqSendCmd)]
-    assert any(s.direction == "E" for s in sends)
-    assert any(s.direction == "W" for s in sends)
-
-
-def test_mlo_kernel_rank_axis_one_west_edge_cube_no_west_sends():
-    """cube_id=0 (west edge) with rank_axis=1: rank=0, has_W=False → no
-    W-direction IPCQ sends. has_E=True → ≥1 E-direction send."""
-    tl = _tl(pe_id=0, cube_id=0, num_pes=PES_PER_CUBE, num_cubes=N_RANKS_MULTI)
-    run_kernel(
-        attention_mesh_mlo_kernel, tl,
-        Q_PTR, K_PTR, V_PTR, O_PTR,
-        S_Q_DECODE, S_KV_PER_RANK, H_Q, H_KV, D_HEAD, N_RANKS_MULTI,
-        rank_axis=1,
-    )
-    sends = [c for c in tl.commands if isinstance(c, IpcqSendCmd)]
-    assert any(s.direction == "E" for s in sends), \
-        "west-edge cube_id=0 must still emit ≥1 E-send"
-    assert not any(s.direction == "W" for s in sends), \
-        "west-edge cube_id=0 must NOT emit any W-send (no W neighbor)"

tests/attention/test_mesh_mlo_2d_correctness.py

@@ -1,142 +0,0 @@
-"""Phase 1 spec test for the 2D row-then-col AllReduce-mlo decode kernel.
-
-The 2D kernel decomposes a ``(mesh_rows × mesh_cols)`` cube sub-mesh into a
-two-stage AllReduce-mlo: stage 1 reduces across columns within each row
-(E/W edges), stage 2 reduces across rows within each column (N/S edges).
-After both stages every cube holds the same final ``(m, ℓ, o)``.
-
-This module is Phase 1 of C2 (see CLAUDE.md change protocol): it pins
-the kernel's interface and observable behavior. Production code for the
-kernel lands in Phase 2; until then this file fails to import.
-
-Test shapes (run on default ``topology.yaml`` — 2 SIPs × 4×4 cube_mesh):
-
-  1×4 sub-mesh (4 cubes, row 0 only)
-     Degenerates to a row-only AllReduce — equivalent in step count to
-     the existing 1D kernel at n_ranks=4. Verifies the kernel reduces
-     correctly when mesh_rows=1 (stage 2 collapses to no-op).
-
-  2×4 sub-mesh (8 cubes, rows 0+1)
-     The 8-KV-group target. Verifies that cubes 4..7 use ``dir="N"``
-     (not ``dir="W"``) to reach row 0 — surfacing the IpcqInvalidDirection
-     bug that the 1D kernel hit at cube 4 (rank 4, no W neighbor).
-
-  4×4 sub-mesh (16 cubes, full SIP)
-     Full-SIP scale. Verifies the algorithm fans out over (cols-1)=3
-     row steps followed by (rows-1)=3 col steps.
-"""
-from __future__ import annotations
-
-from pathlib import Path
-
-from kernbench.benches._attention_mesh_mlo_2d import attention_mesh_mlo_2d_kernel
-from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config
-from kernbench.ccl.sfr_config import configure_sfr_intercube_multisip
-from kernbench.policy.placement.dp import DPPolicy
-from kernbench.runtime_api.bench_runner import run_bench
-from kernbench.runtime_api.types import resolve_device
-from kernbench.sim_engine.engine import GraphEngine
-from kernbench.topology.builder import resolve_topology
-
-TOPOLOGY_DEFAULT = Path(__file__).resolve().parents[2] / "topology.yaml"
-
-S_Q = 1
-S_KV_PER_RANK = 16
-H_Q = 1
-H_KV = 1
-D_HEAD = 64
-DTYPE = "f16"
-
-
-def _ccl_cfg():
-    return resolve_algorithm_config(
-        load_ccl_config(), name="lrab_hierarchical_allreduce",
-    )
-
-
-def _engine_factory(t, d):
-    return GraphEngine(getattr(t, "topology_obj", t), enable_data=True)
-
-
-def _run_2d(mesh_rows: int, mesh_cols: int, cube_start: int = 0):
-    """Build a bench_fn and run it on the default topology."""
-    n_cubes = mesh_rows * mesh_cols
-    topo = resolve_topology(str(TOPOLOGY_DEFAULT))
-
-    def _bench_fn(ctx):
-        configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
-        dp_full = DPPolicy(cube="replicate", pe="replicate",
-                           num_cubes=n_cubes, num_pes=8,
-                           cube_start=cube_start)
-        dp_kv = DPPolicy(cube="row_wise", pe="replicate",
-                         num_cubes=n_cubes, num_pes=8,
-                         cube_start=cube_start)
-        q = ctx.zeros((S_Q, H_Q * D_HEAD),
-                      dtype=DTYPE, dp=dp_full, name="q")
-        k = ctx.zeros((S_KV_PER_RANK * n_cubes, H_KV * D_HEAD),
-                      dtype=DTYPE, dp=dp_kv, name="k")
-        v = ctx.zeros((S_KV_PER_RANK * n_cubes, H_KV * D_HEAD),
-                      dtype=DTYPE, dp=dp_kv, name="v")
-        o = ctx.empty((S_Q, H_Q * D_HEAD),
-                      dtype=DTYPE, dp=dp_full, name="o")
-        ctx.launch(
-            f"mesh_mlo_2d_{mesh_rows}x{mesh_cols}_start{cube_start}",
-            attention_mesh_mlo_2d_kernel,
-            q, k, v, o,
-            S_Q, S_KV_PER_RANK, H_Q, H_KV, D_HEAD,
-            mesh_rows, mesh_cols,
-            1,  # rank_axis=1 → cube-level ring
-            cube_start,
-            _auto_dim_remap=False,
-        )
-
-    return run_bench(
-        topology=topo,
-        bench_fn=_bench_fn,
-        device=resolve_device(None),
-        engine_factory=_engine_factory,
-    )
-
-
-def test_2d_kernel_1x4_row_only():
-    """1×4 sub-mesh: row-only AllReduce, stage 2 collapses to no-op."""
-    result = _run_2d(mesh_rows=1, mesh_cols=4)
-    assert result.completion.ok, (
-        f"1x4: completion not ok - {result.completion}"
-    )
-
-
-def test_2d_kernel_2x4_eight_cubes():
-    """2×4 sub-mesh: the 8-KV-group target.
-
-    Verifies that cube 4 (row 1, col 0) uses ``dir="N"`` to reach cube 0
-    (row 0, col 0) for stage 2, not ``dir="W"`` — the 1D kernel hit
-    IpcqInvalidDirection here.
-    """
-    result = _run_2d(mesh_rows=2, mesh_cols=4)
-    assert result.completion.ok, (
-        f"2x4: completion not ok - {result.completion}"
-    )
-
-
-def test_2d_kernel_4x4_full_sip():
-    """4×4 sub-mesh: full-SIP scale (16 cubes)."""
-    result = _run_2d(mesh_rows=4, mesh_cols=4)
-    assert result.completion.ok, (
-        f"4x4: completion not ok - {result.completion}"
-    )
-
-
-def test_2d_kernel_2x4_at_cube_start_eight():
-    """2×4 sub-mesh at cube_start=8 (cubes 8..15, rows 2..3).
-
-    The second KV-group per SIP in the 8-KV-group Llama-70B headline.
-    Verifies the kernel converts ``program_id(axis=1)`` (physical cube
-    id) back to launch-local rank via ``cube_start`` — without that
-    subtraction, cube 8 would compute my_row=2 (out of sub-mesh bounds)
-    and deadlock waiting on cube 4 which isn't in the launch.
-    """
-    result = _run_2d(mesh_rows=2, mesh_cols=4, cube_start=8)
-    assert result.completion.ok, (
-        f"2x4 @ cube_start=8: completion not ok - {result.completion}"
-    )

tests/attention/test_milestone_gqa_headline.py

@@ -0,0 +1,148 @@
+"""Phase 1 spec test for P7: headline milestone-gqa bench with real GQA.
+
+P7 wires the new ``_gqa_decode`` and ``_gqa_prefill`` kernels into a
+new 4-panel milestone bench (independent from the existing
+``milestone-gqa-llama70b`` which still covers the baseline kernels).
+Real GQA (``h_q > h_kv`` with G=8) runs end-to-end through a
+milestone-style sweep + sweep.json output.
+
+Restriction in P7 first cut:
+  - C ≤ 4 (single-row inter-CUBE ring SFR; ADR-0060 §B leaves
+    multi-row rings for follow-on)
+  - Single SIP (default ``topology.yaml`` 4×4 cube mesh; 4-SIP
+    headline left to follow-on)
+  - No figure renderers (defer to a separate cycle)
+
+Panels (4 total):
+  single_user_prefill_gqa: C=1, T_q=4, S_kv=16   (no ring)
+  multi_user_prefill_gqa : C=4, T_q=4, S_kv=16   (Ring KV, 3 steps)
+  single_user_decode_gqa : C=1, P=8, h_q=8, h_kv=1, S_kv=64    (M-fold + intra-cube chain)
+  multi_user_decode_gqa  : C=4, P=8, h_q=8, h_kv=1, S_kv=128   (M-fold + 2-level chain)
+
+Phase 1 (this commit): tests only — bench module lands in Phase 2.
+"""
+from __future__ import annotations
+
+import json
+
+import kernbench.benches.milestone_gqa_headline as bench_mod  # noqa: F401 (Phase 2)
+from kernbench.benches.registry import resolve
+from kernbench.runtime_api.bench_runner import run_bench
+from kernbench.runtime_api.types import resolve_device
+from kernbench.sim_engine.engine import GraphEngine
+from kernbench.topology.builder import resolve_topology
+
+BENCH_NAME = "milestone-gqa-headline"
+
+PANELS = (
+    "single_user_prefill_gqa",
+    "multi_user_prefill_gqa",
+    "single_user_decode_gqa",
+    "multi_user_decode_gqa",
+)
+
+
+def _run_validation():
+    topo = resolve_topology("topology.yaml")
+    return run_bench(
+        topology=topo,
+        bench_fn=resolve(BENCH_NAME).run,
+        device=resolve_device(None),
+        engine_factory=lambda t, d: GraphEngine(
+            getattr(t, "topology_obj", t), enable_data=True,
+        ),
+    )
+
+
+def _sweep_json(monkeypatch) -> dict:
+    monkeypatch.setenv("GQA_HEADLINE_RUN", "1")
+    out = bench_mod._OUTPUT_DIR / "sweep.json"
+    if not out.exists():
+        result = _run_validation()
+        assert result.completion.ok, result.completion
+    assert out.exists(), f"missing {out}"
+    return json.loads(out.read_text())
+
+
+# ── Registration ───────────────────────────────────────────────────────
+
+
+def test_bench_registered():
+    spec = resolve(BENCH_NAME)
+    assert spec.name == BENCH_NAME
+    assert callable(spec.run)
+    assert spec.description.strip(), "description must be non-empty"
+
+
+# ── Validation run completes ──────────────────────────────────────────
+
+
+def test_validation_run_completes_ok(monkeypatch):
+    monkeypatch.setenv("GQA_HEADLINE_RUN", "1")
+    result = _run_validation()
+    assert result.completion.ok, (
+        f"headline validation run failed: {result.completion}"
+    )
+
+
+# ── sweep.json shape ──────────────────────────────────────────────────
+
+
+def test_sweep_json_has_four_panels(monkeypatch):
+    data = _sweep_json(monkeypatch)
+    assert set(data["panels"]) == set(PANELS), (
+        f"panels mismatch: expected {set(PANELS)}, got {set(data['panels'])}"
+    )
+    assert len(data["rows"]) == 4
+    assert {r["panel"] for r in data["rows"]} == set(PANELS)
+
+
+def test_decode_panels_use_real_gqa(monkeypatch):
+    """ADR-0060 §A.1 headline: decode panels must use h_q = G·h_kv with G>1."""
+    data = _sweep_json(monkeypatch)
+    cfg = data["config"]
+    assert cfg["h_q_decode"] > cfg["h_kv_decode"], (
+        f"decode must use real GQA (h_q > h_kv); got "
+        f"h_q={cfg['h_q_decode']}, h_kv={cfg['h_kv_decode']}"
+    )
+
+
+# ── Per-panel architectural assertions ────────────────────────────────
+
+
+def _row(rows, panel: str) -> dict:
+    for r in rows:
+        if r["panel"] == panel:
+            return r
+    raise AssertionError(f"missing row for panel {panel!r}")
+
+
+def test_prefill_ring_panel_has_ipcq_traffic(monkeypatch):
+    """multi_user_prefill_gqa uses Ring KV → IPCQ traffic > 0."""
+    data = _sweep_json(monkeypatch)
+    row = _row(data["rows"], "multi_user_prefill_gqa")
+    n_copy = row["op_log_summary"].get("ipcq_copy_count", 0)
+    assert n_copy > 0, (
+        f"multi_user_prefill_gqa must have Ring KV traffic; got "
+        f"ipcq_copy_count={n_copy}"
+    )
+
+
+def test_decode_reduce_panel_writes_once(monkeypatch):
+    """multi_user_decode_gqa: chain reduce-to-root → exactly 1 dma_write."""
+    data = _sweep_json(monkeypatch)
+    row = _row(data["rows"], "multi_user_decode_gqa")
+    n_writes = row["op_log_summary"]["dma_write_count"]
+    assert n_writes == 1, (
+        f"multi_user_decode_gqa root-only write: expected 1; got {n_writes}"
+    )
+
+
+def test_prefill_panel_distributes_output(monkeypatch):
+    """multi_user_prefill_gqa: per-CUBE distributed output → dma_write_count == C."""
+    data = _sweep_json(monkeypatch)
+    row = _row(data["rows"], "multi_user_prefill_gqa")
+    n_writes = row["op_log_summary"]["dma_write_count"]
+    assert n_writes == 4, (
+        f"multi_user_prefill_gqa per-CUBE distributed: expected 4; got {n_writes}"
+    )

tests/attention/test_milestone_gqa_llama70b.py

@@ -1,222 +0,0 @@
-"""Phase 1 spec test for ``milestone-gqa-llama70b`` bench (sub-cycle 4a, all 4 panels).
-
-ADR-0057 (Proposed) defines an eval bench that drives both attention kernels
-(ADR-0055 ring-K/V, ADR-0056 allreduce-mlo) and emits per-panel op_log
-summaries into ``src/kernbench/benches/1H_milestone_output/gqa/sweep.json``.
-
-v1 (sub-cycle 4a) covers ALL FOUR panels:
-
-  Panel name in JSON / test    Study label   SFR install used
-  ─────────────────────────────────────────────────────────────────────────────
-  single_user_prefill          TL            configure_sfr_intracube_pe_ring
-  multi_user_prefill           TR            configure_sfr_intercube_multisip
-  single_user_decode           BL            configure_sfr_intracube_pe_ring
-  multi_user_decode            BR            configure_sfr_intercube_multisip
-
-single_user_* panels became runnable after sub-cycle 4-pre delivered the
-new SFR install function (ADR-0058).
-
-In Phase 1 the bench module does not exist; pytest collection fails with
-``ModuleNotFoundError``. Once Phase 2 lands the bench module, every
-assertion below must pass.
-
-Assertions:
-  - Bench is registered as ``milestone-gqa-llama70b``.
-  - A validation run (``GQA_VALIDATION=1``) completes ok via run_bench.
-  - sweep.json conforms to ADR-0057 D7 (v1 schema).
-  - All four panel rows present with sane op_log summaries.
-  - Both decode panels have gemm_count = 2 × n_ranks (one-shot per rank).
-  - Both prefill panels have gemm_count = 2 × n_ranks² (per-step GEMMs).
-"""
-from __future__ import annotations
-
-import json
-
-from kernbench.benches.registry import resolve
-from kernbench.runtime_api.bench_runner import run_bench
-from kernbench.runtime_api.types import resolve_device
-from kernbench.sim_engine.engine import GraphEngine
-from kernbench.topology.builder import resolve_topology
-
-# Production module (Phase 2 deliverable; absent in Phase 1).
-import kernbench.benches.milestone_gqa_llama70b as gqa_bench
-
-
-BENCH_NAME = "milestone-gqa-llama70b"
-
-PANELS_V1 = (
-    "single_user_prefill",
-    "multi_user_prefill",
-    "single_user_decode",
-    "multi_user_decode",
-)
-SINGLE_USER_PANELS = ("single_user_prefill", "single_user_decode")
-MULTI_USER_PANELS = ("multi_user_prefill", "multi_user_decode")
-PREFILL_PANELS = ("single_user_prefill", "multi_user_prefill")
-DECODE_PANELS = ("single_user_decode", "multi_user_decode")
-
-
-def _run_validation():
-    """Drive the bench through run_bench at validation scale."""
-    topo = resolve_topology("topology.yaml")
-    return run_bench(
-        topology=topo,
-        bench_fn=resolve(BENCH_NAME).run,
-        device=resolve_device(None),
-        engine_factory=lambda t, d: GraphEngine(
-            getattr(t, "topology_obj", t), enable_data=True,
-        ),
-    )
-
-
-# ── Registration ────────────────────────────────────────────────────────
-
-
-def test_bench_registered():
-    spec = resolve(BENCH_NAME)
-    assert spec.name == BENCH_NAME
-    assert callable(spec.run)
-    assert spec.description.strip(), "description must be non-empty"
-
-
-# ── Validation run end-to-end ────────────────────────────────────────────
-
-
-def test_validation_run_completes_ok(monkeypatch):
-    monkeypatch.setenv("GQA_VALIDATION", "1")
-    result = _run_validation()
-    assert result.completion.ok, (
-        f"validation run failed: {result.completion}"
-    )
-
-
-# ── JSON shape (ADR-0057 D7 amended for 4 panels) ──────────────────────
-
-
-def _sweep_json(monkeypatch) -> dict:
-    """Run the bench (if needed) and return the parsed sweep.json."""
-    monkeypatch.setenv("GQA_VALIDATION", "1")
-    out = gqa_bench._OUTPUT_DIR / "sweep.json"
-    if not out.exists():
-        result = _run_validation()
-        assert result.completion.ok, result.completion
-    assert out.exists(), f"missing {out}"
-    return json.loads(out.read_text())
-
-
-def test_sweep_json_has_v1_schema(monkeypatch):
-    data = _sweep_json(monkeypatch)
-    assert data["version"] == 1
-    assert data["validation_scale"] is True
-    assert isinstance(data["panels"], list)
-    assert isinstance(data["config"], dict)
-    assert isinstance(data["rows"], list)
-
-
-def test_sweep_json_panels_are_all_four(monkeypatch):
-    """v1 covers all four panels — single_user_{prefill,decode} +
-    multi_user_{prefill,decode}. Q/cube sweep deferred to 4b."""
-    data = _sweep_json(monkeypatch)
-    assert set(data["panels"]) == set(PANELS_V1)
-
-
-def test_sweep_json_config_matches_adr0057_d4(monkeypatch):
-    """Validation-scale config per ADR-0057 D4 (amended for 4 panels + scratch budget).
-
-    S_q_prefill and S_kv_per_rank are deliberately small (16 each) so the
-    simulator's 1 MB per-PE TCM kernel scratch (topology.yaml
-    ``pe_tcm.kernel_scratch_mb: 1``) is not exhausted by the
-    bump-allocated handle outputs of softmax/exp/dot/sum chains over
-    n_ranks ring steps. Headline-scale runs in 4c will lift these into a
-    config-driven sweep.
-    """
-    data = _sweep_json(monkeypatch)
-    cfg = data["config"]
-    assert cfg["S_q_prefill"] == 16
-    assert cfg["S_kv_per_rank"] == 16
-    # v1 uses h_q == h_kv == 1 to avoid ADR-0055 D3's GQA broadcast view
-    # (which is symbolic and does not survive MemoryStore's nbytes check
-    # under simulator data execution). Real GQA (h_q > h_kv) is deferred
-    # to sub-cycle 4c (headline scale).
-    assert cfg["h_q"] == 1
-    assert cfg["h_kv"] == 1
-    assert cfg["d_head"] == 64
-    # single_user_* uses the 8 PEs in one cube as ring ranks.
-    assert cfg["n_ranks_single_user"] == 8
-    # multi_user_* uses cube-level ring; validation uses 4 cubes.
-    assert cfg["n_ranks_multi_user"] == 4
-
-
-def test_sweep_json_has_one_row_per_panel(monkeypatch):
-    data = _sweep_json(monkeypatch)
-    assert len(data["rows"]) == 4
-    panels_in_rows = {r["panel"] for r in data["rows"]}
-    assert panels_in_rows == set(PANELS_V1)
-
-
-# ── Per-row op_log summary sanity (ADR-0057 D7) ─────────────────────────
-
-
-def _row(rows: list, panel: str) -> dict:
-    matches = [r for r in rows if r["panel"] == panel]
-    assert len(matches) == 1, f"expected exactly one {panel} row; got {len(matches)}"
-    return matches[0]
-
-
-def _assert_sane_summary(row: dict) -> None:
-    s = row["op_log_summary"]
-    panel = row["panel"]
-    assert s["gemm_count"] > 0, f"{panel} must run GEMMs"
-    assert s["ipcq_send_count"] > 0, f"{panel} must send (ring/allreduce phase)"
-    assert s["ipcq_recv_count"] > 0, f"{panel} must recv"
-    assert s["dma_read_count"] >= 3, f"{panel}: Q + K + V loads"
-    assert s["dma_write_count"] >= 1, f"{panel}: final O store"
-
-
-def test_single_user_prefill_row_has_sane_op_log_summary(monkeypatch):
-    data = _sweep_json(monkeypatch)
-    _assert_sane_summary(_row(data["rows"], "single_user_prefill"))
-
-
-def test_multi_user_prefill_row_has_sane_op_log_summary(monkeypatch):
-    data = _sweep_json(monkeypatch)
-    _assert_sane_summary(_row(data["rows"], "multi_user_prefill"))
-
-
-def test_single_user_decode_row_has_sane_op_log_summary(monkeypatch):
-    data = _sweep_json(monkeypatch)
-    _assert_sane_summary(_row(data["rows"], "single_user_decode"))
-
-
-def test_multi_user_decode_row_has_sane_op_log_summary(monkeypatch):
-    data = _sweep_json(monkeypatch)
-    _assert_sane_summary(_row(data["rows"], "multi_user_decode"))
-
-
-# ── Architectural invariant: decode = one-shot per rank ─────────────────
-
-
-def test_single_user_decode_gemm_count_is_exactly_2_per_rank(monkeypatch):
-    """ADR-0056 D3: decode kernel does ONE local partial-attention pass per
-    rank → exactly 2 GEMMs per rank (Q·K^T + S·V). With n_ranks ranks the
-    total = 2 × n_ranks. This distinguishes decode from prefill where each
-    ring step has 2 GEMMs and the total scales as 2 × n_ranks²."""
-    data = _sweep_json(monkeypatch)
-    row = _row(data["rows"], "single_user_decode")
-    n_ranks = row["n_ranks"]
-    assert row["op_log_summary"]["gemm_count"] == 2 * n_ranks, (
-        f"single_user_decode gemm_count must be 2 × n_ranks = {2 * n_ranks}; "
-        f"got {row['op_log_summary']['gemm_count']}"
-    )
-
-
-def test_multi_user_decode_gemm_count_is_exactly_2_per_rank(monkeypatch):
-    """Same one-shot invariant as single_user_decode — the kernel is the
-    same; what differs is who the ranks are (cubes vs PEs)."""
-    data = _sweep_json(monkeypatch)
-    row = _row(data["rows"], "multi_user_decode")
-    n_ranks = row["n_ranks"]
-    assert row["op_log_summary"]["gemm_count"] == 2 * n_ranks, (
-        f"multi_user_decode gemm_count must be 2 × n_ranks = {2 * n_ranks}; "
-        f"got {row['op_log_summary']['gemm_count']}"
-    )

tests/gqa/_gqa_plot_helpers.py

@@ -1,25 +0,0 @@
-"""Thin re-export shim for the GQA figure tests.
-
-Not a test module (no ``test_`` prefix → pytest does not collect it).
-
-Mirrors ``tests/gemm/_gemm_plot_helpers.py``. The renderer logic lives in
-``kernbench.benches.milestone_gqa_llama70b`` (production single home,
-ADR-0054). Defaults still target the bench's ``_OUTPUT_DIR``.
-"""
-from __future__ import annotations
-
-from kernbench.benches.milestone_gqa_llama70b import (
-    _OUTPUT_DIR as GQA_PLOTS_DIR,
-    _SWEEP_JSON as GQA_SWEEP_JSON,
-    emit_all_gqa_plots,
-    emit_gqa_comparison,
-    emit_panel_op_log_summary,
-)
-
-__all__ = [
-    "GQA_PLOTS_DIR",
-    "GQA_SWEEP_JSON",
-    "emit_all_gqa_plots",
-    "emit_gqa_comparison",
-    "emit_panel_op_log_summary",
-]

tests/gqa/test_plot_gqa_figures.py

@@ -1,109 +0,0 @@
-"""Phase 1 spec test for GQA figure renderers (sub-cycle 4c).
-
-ADR-0057 D3 sub-cycle 4c adds 6 figure renderers; this test pins the
-5 of 6 that don't depend on sub-cycle 4b's Q/cube sweep:
-
-  - 4 per-panel op_log_summary PNGs (one per panel of v1's sweep.json)
-  - 1 cross-panel ``gqa_comparison.png`` (4-panel grouped bars over the
-    5 op_log_summary counts: gemm, ipcq_send, ipcq_recv, dma_read, dma_write)
-
-The 6th, ``gqa_scaling.png``, needs the Q/cube ∈ {1, 2, 4} sweep from
-sub-cycle 4b and is deferred.
-
-Each test depends on the committed
-``benches/1H_milestone_output/gqa/sweep.json`` (landed in commit
-``e748a62``); they assert the renderer writes a non-empty PNG at the
-expected path.
-
-Phase 1 expectation: tests fail at import (renderer functions don't
-exist yet on the bench module). Phase 2 lands them and the tests
-turn green.
-"""
-from __future__ import annotations
-
-from pathlib import Path
-
-import pytest
-
-from tests.gqa._gqa_plot_helpers import (
-    GQA_PLOTS_DIR,
-    GQA_SWEEP_JSON,
-    emit_all_gqa_plots,
-    emit_gqa_comparison,
-    emit_panel_op_log_summary,
-)
-
-
-_PANELS = (
-    "single_user_prefill",
-    "multi_user_prefill",
-    "single_user_decode",
-    "multi_user_decode",
-)
-
-
-@pytest.mark.skipif(
-    not GQA_SWEEP_JSON.exists(),
-    reason="gqa sweep.json absent; run milestone-gqa-llama70b first",
-)
-@pytest.mark.parametrize("panel", _PANELS)
-def test_emit_panel_op_log_summary_writes_png_for_each_panel(panel):
-    out = emit_panel_op_log_summary(panel)
-    assert out is not None, f"{panel}: renderer returned None"
-    path = Path(out)
-    assert path.exists(), f"{panel}: expected PNG at {path}"
-    assert path.suffix == ".png", f"{panel}: not a PNG: {path}"
-    assert path.stat().st_size > 0, f"{panel}: empty PNG: {path}"
-    assert panel in path.stem, (
-        f"{panel}: panel name not in filename {path.name}"
-    )
-
-
-@pytest.mark.skipif(
-    not GQA_SWEEP_JSON.exists(),
-    reason="gqa sweep.json absent; run milestone-gqa-llama70b first",
-)
-def test_emit_gqa_comparison_writes_png():
-    out = emit_gqa_comparison()
-    assert out is not None
-    path = Path(out)
-    assert path.exists()
-    assert path.name == "gqa_comparison.png"
-    assert path.stat().st_size > 0
-
-
-@pytest.mark.skipif(
-    not GQA_SWEEP_JSON.exists(),
-    reason="gqa sweep.json absent; run milestone-gqa-llama70b first",
-)
-def test_emit_all_gqa_plots_writes_five_figures():
-    """emit_all returns a list of 5 written PNG paths (deferring the
-    6th gqa_scaling.png to after sub-cycle 4b lands the Q/cube sweep)."""
-    paths = emit_all_gqa_plots()
-    assert isinstance(paths, list)
-    # 4 per-panel + 1 comparison.
-    assert len(paths) == 5, f"expected 5 PNGs, got {len(paths)}: {paths}"
-    for p in paths:
-        assert Path(p).exists() and Path(p).stat().st_size > 0
-    names = {Path(p).name for p in paths}
-    assert "gqa_comparison.png" in names
-    for panel in _PANELS:
-        assert any(panel in n for n in names), (
-            f"no per-panel PNG for {panel} in {names}"
-        )
-
-
-def test_emit_all_gqa_plots_output_dir_matches_bench_output_dir():
-    """The renderers must write under the bench's own _OUTPUT_DIR so
-    MILESTONE_FAST=1 reuse (and committed baselines) all point at the
-    same on-disk location."""
-    # Stub assertion that fails until emit_all_gqa_plots exists with a
-    # default ``out_dir`` argument identical to GQA_PLOTS_DIR.
-    import inspect
-
-    sig = inspect.signature(emit_all_gqa_plots)
-    assert "out_dir" in sig.parameters
-    default = sig.parameters["out_dir"].default
-    assert Path(default) == GQA_PLOTS_DIR, (
-        f"default out_dir {default} != bench _OUTPUT_DIR {GQA_PLOTS_DIR}"
-    )

tests/test_kernel_runner.py

@@ -28,7 +28,13 @@ def _mock_scheduler(env, inbox):
 
 
 def test_kernel_runner_basic_load():
-    """Kernel with tl.load runs through greenlet without hanging."""
+    """Kernel with tl.load runs through greenlet without hanging.
+
+    Under ADR-0062 (lazy ``tl.load``) the handle's ``data`` is None
+    until a consumer op triggers auto-wait at first use. A no-op math
+    op (``tl.exp``) consumes ``a`` and forces resolution before we
+    inspect ``a.data``.
+    """
     env = simpy.Environment()
     store = MemoryStore()
     data = np.ones((4, 4), dtype=np.float16)
@@ -39,6 +45,7 @@ def test_kernel_runner_basic_load():
 
     def kernel(a_ptr, tl):
         a = tl.load(a_ptr, (4, 4), "f16")
+        tl.exp(a)  # consumer op → auto-wait at first use (ADR-0062 §D2)
         assert a.data is not None
         assert a.data.shape == (4, 4)
 
@@ -50,7 +57,12 @@ def test_kernel_runner_basic_load():
 
 
 def test_kernel_runner_load_returns_data():
-    """tl.load returns actual numpy data from MemoryStore."""
+    """tl.load returns actual numpy data from MemoryStore.
+
+    Under ADR-0062 lazy semantics the data is attached at auto-wait
+    time (first consumer op), so we read ``a.data`` after a consumer
+    op fires the wait.
+    """
     env = simpy.Environment()
     store = MemoryStore()
     data = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float16)
@@ -63,6 +75,7 @@ def test_kernel_runner_load_returns_data():
 
     def kernel(ptr, tl):
         a = tl.load(ptr, (2, 2), "f16")
+        tl.exp(a)  # consumer op → auto-wait at first use (ADR-0062 §D2)
         results["data"] = a.data
 
     def run():
@@ -106,6 +119,11 @@ def test_kernel_runner_dynamic_branch():
 
     def kernel(flag_ptr, tl):
         flag = tl.load(flag_ptr, (1,), "f32")
+        # ADR-0062: under lazy tl.load, dynamic-branching kernels must
+        # force resolution by consuming the handle (any tl.* op works).
+        # Without this, flag.data is None at branch time and control
+        # always takes the not-taken path.
+        tl.exp(flag)
         if flag.data is not None and flag.data[0] > 0.5:
             results["branch"] = "taken"
         else:

tests/test_phase_e_cpu_issue_cost.py

@@ -0,0 +1,382 @@
+"""Phase 1 spec tests for ADR-0064 (per-op-type CPU issue cost model).
+
+Phase E lands the cost-table machinery and turns it on by default so the
+hybrid's CPU-saturation lever (ADR-0060 §1) becomes measurable instead of
+modelled away. Today every ``tl.*`` call goes through
+``_emit_dispatch_overhead()`` with a single uniform ``dispatch_cycles``
+scalar that is hard-coded to 0 on the live PE_CPU paths
+(``pe_cpu.py:_execute_legacy`` and ``kernel_runner.py:run``).
+
+These tests assume the post-Phase-2 surface:
+
+- ``kernbench.common.cpu_issue_cost`` exports ``OpKind``,
+  ``DEFAULT_CPU_ISSUE_COST`` (the ADR-0064 D1 table) and
+  ``get_issue_cost(kind, table=None) -> int``.
+- ``TLContext.__init__`` accepts ``issue_cost_table: dict[str, int] | None``.
+  When provided, ``_emit_dispatch_overhead(kind)`` looks up the per-kind
+  cost; when absent, it falls back to the uniform ``dispatch_cycles``
+  (back-compat with the existing ADR-0046 §D6 contract).
+- Live PE_CPU paths (greenlet + legacy replay) construct TLContext with
+  ``issue_cost_table=DEFAULT_CPU_ISSUE_COST``.
+
+Phase 1 (this commit): tests only. All tests FAIL until Phase 2.
+"""
+from __future__ import annotations
+
+from pathlib import Path
+
+import pytest
+
+from kernbench.common.pe_commands import (
+    CompositeCmd,
+    DmaReadCmd,
+    DmaWriteCmd,
+    GemmCmd,
+    MathCmd,
+    PeCpuOverheadCmd,
+)
+from kernbench.policy.address.phyaddr import PhysAddr
+from kernbench.runtime_api.kernel import KernelLaunchMsg, KernelRef
+from kernbench.sim_engine.engine import GraphEngine
+from kernbench.sim_engine.transaction import Transaction
+from kernbench.topology.builder import load_topology
+from kernbench.triton_emu.registry import clear_registry, register_kernel
+from kernbench.triton_emu.tl_context import TLContext, run_kernel
+
+TOPOLOGY_PATH = Path(__file__).parent.parent / "topology.yaml"
+
+
+def _engine():
+    return GraphEngine(load_topology(TOPOLOGY_PATH))
+
+
+def _hbm_pa(sip: int = 0, cube: int = 0, pe_id: int = 0) -> int:
+    slice_bytes = 48 * (1 << 30) // 8
+    pa = PhysAddr.pe_hbm_addr(
+        sip_id=sip, die_id=cube, pe_id=pe_id,
+        pe_local_hbm_offset=0x1000, slice_size_bytes=slice_bytes,
+    )
+    return pa.encode()
+
+
+# ── T1: default table shape ──────────────────────────────────────
+
+
+def test_default_cost_table_has_expected_keys():
+    """ADR-0064 D1 table — 8 keys with composite ≫ primitive ratio."""
+    from kernbench.common.cpu_issue_cost import DEFAULT_CPU_ISSUE_COST
+
+    expected_keys = {
+        "composite", "load", "store", "dot", "math",
+        "ipcq_send", "ipcq_recv", "copy_to",
+    }
+    assert set(DEFAULT_CPU_ISSUE_COST.keys()) == expected_keys, (
+        f"DEFAULT_CPU_ISSUE_COST keys must match ADR-0064 D1; "
+        f"got {set(DEFAULT_CPU_ISSUE_COST.keys())}"
+    )
+    # Composite is the lever — ratio against primitives must be ≥ 4×.
+    assert DEFAULT_CPU_ISSUE_COST["composite"] == 40
+    for primitive in ("load", "store", "dot", "math",
+                      "ipcq_send", "ipcq_recv", "copy_to"):
+        assert DEFAULT_CPU_ISSUE_COST[primitive] == 5, (
+            f"primitive {primitive!r} default cost must be 5 ns"
+        )
+
+
+# ── T2: get_issue_cost lookup ────────────────────────────────────
+
+
+def test_get_issue_cost_lookup():
+    """Helper returns table value; unknown kind returns 0 (no charge)."""
+    from kernbench.common.cpu_issue_cost import (
+        DEFAULT_CPU_ISSUE_COST,
+        get_issue_cost,
+    )
+
+    assert get_issue_cost("composite") == 40
+    assert get_issue_cost("load") == 5
+    assert get_issue_cost("unknown_kind") == 0
+    # Custom table override
+    custom = {"composite": 100, "load": 1}
+    assert get_issue_cost("composite", table=custom) == 100
+    assert get_issue_cost("load", table=custom) == 1
+    assert get_issue_cost("store", table=custom) == 0
+
+
+# ── T3: TLContext consumes a passed cost table ───────────────────
+
+
+def test_tlcontext_accepts_issue_cost_table():
+    """TLContext(issue_cost_table=...) → tl.load emits the per-kind cycles."""
+    tl = TLContext(
+        pe_id=0, num_programs=1,
+        dispatch_cycles=0,
+        issue_cost_table={"load": 7, "store": 3, "dot": 2, "math": 2},
+    )
+    tl.load(0x1000, shape=(4, 4), dtype="f16")
+    overheads = [c for c in tl.commands if isinstance(c, PeCpuOverheadCmd)]
+    assert len(overheads) == 1, (
+        f"expected exactly one PeCpuOverheadCmd before the DmaReadCmd; "
+        f"got {len(overheads)} (cmds={[type(c).__name__ for c in tl.commands]})"
+    )
+    assert overheads[0].cycles == 7, (
+        f"load issue cost from table must be 7; got {overheads[0].cycles}"
+    )
+
+
+# ── T4: composite ≫ primitive issue cost differential ────────────
+
+
+def test_tlcontext_composite_vs_primitive_charges_differ():
+    """Composite kernel charges once (40); primitive sequence charges per-op."""
+    from kernbench.common.cpu_issue_cost import DEFAULT_CPU_ISSUE_COST
+
+    # Composite kernel: 1 load + 1 composite = 5 + 40 = 45 cycles of issue cost.
+    tl_comp = TLContext(
+        pe_id=0, num_programs=1,
+        dispatch_cycles=0,
+        issue_cost_table=DEFAULT_CPU_ISSUE_COST,
+    )
+    a = tl_comp.load(0x1000, shape=(8, 16), dtype="f16")
+    b_ref = tl_comp.ref(0x2000, shape=(16, 8), dtype="f16")
+    tl_comp.composite("gemm", a, b_ref, out_ptr=0x3000)
+    comp_cycles = sum(
+        c.cycles for c in tl_comp.commands if isinstance(c, PeCpuOverheadCmd)
+    )
+
+    # Primitive kernel: 2 loads + 1 dot + 1 math (exp) = 5 + 5 + 5 + 5 = 20 cycles.
+    tl_prim = TLContext(
+        pe_id=0, num_programs=1,
+        dispatch_cycles=0,
+        issue_cost_table=DEFAULT_CPU_ISSUE_COST,
+    )
+    a = tl_prim.load(0x1000, shape=(8, 16), dtype="f16")
+    b = tl_prim.load(0x2000, shape=(16, 8), dtype="f16")
+    c_out = tl_prim.dot(a, b)
+    tl_prim.exp(c_out)
+    prim_cycles = sum(
+        c.cycles for c in tl_prim.commands if isinstance(c, PeCpuOverheadCmd)
+    )
+
+    # ADR-0064 ratio: composite issue cost >> primitive issue cost per op.
+    # For these specific kernels: composite=45 (5+40), primitive=20 (4×5).
+    assert comp_cycles == 45, f"composite kernel: expected 45, got {comp_cycles}"
+    assert prim_cycles == 20, f"primitive kernel: expected 20, got {prim_cycles}"
+    # Headline assertion: a single composite charges more than all 3
+    # post-load primitives combined (40 > 3×5) — the ADR-0060 §1 lever.
+    assert comp_cycles - 5 > 3 * 5, (
+        f"composite issue charge {comp_cycles - 5} must exceed "
+        f"3× primitive issue charge {3 * 5} (ADR-0064 ratio)"
+    )
+
+
+# ── T5: cost table is purely additive (Q2 — no double-count) ─────
+
+
+def test_cost_table_is_additive_only():
+    """Q2 invariant: the cost table is additive on PE_CPU, NOT folded into
+    DMA/GEMM/MATH command shapes. Two TLContexts running the same kernel
+    with different cost tables must produce identical non-overhead command
+    sequences (same DmaReadCmd, GemmCmd, MathCmd, addrs, shapes, dtypes).
+    Only the PeCpuOverheadCmd ``cycles`` field is allowed to differ.
+    """
+    from kernbench.common.cpu_issue_cost import DEFAULT_CPU_ISSUE_COST
+
+    def kernel(tl):
+        a = tl.load(0x1000, shape=(8, 16), dtype="f16")
+        b = tl.load(0x2000, shape=(16, 8), dtype="f16")
+        c = tl.dot(a, b)
+        d = tl.exp(c)
+        tl.store(0x3000, d)
+
+    tl_zero = TLContext(
+        pe_id=0, num_programs=1,
+        dispatch_cycles=0,
+        issue_cost_table={},  # empty table → every kind = 0 cost
+    )
+    run_kernel(kernel, tl_zero)
+
+    tl_default = TLContext(
+        pe_id=0, num_programs=1,
+        dispatch_cycles=0,
+        issue_cost_table=DEFAULT_CPU_ISSUE_COST,
+    )
+    run_kernel(kernel, tl_default)
+
+    # Strip PeCpuOverheadCmd from both streams; what remains must match.
+    non_overhead_zero = [
+        c for c in tl_zero.commands if not isinstance(c, PeCpuOverheadCmd)
+    ]
+    non_overhead_default = [
+        c for c in tl_default.commands if not isinstance(c, PeCpuOverheadCmd)
+    ]
+    assert len(non_overhead_zero) == len(non_overhead_default), (
+        f"non-overhead command count differs: "
+        f"{len(non_overhead_zero)} vs {len(non_overhead_default)}"
+    )
+    for a_cmd, b_cmd in zip(non_overhead_zero, non_overhead_default):
+        assert type(a_cmd) is type(b_cmd), (
+            f"non-overhead command type changed under cost table: "
+            f"{type(a_cmd).__name__} vs {type(b_cmd).__name__}"
+        )
+    # Total overhead under zero-table must be 0; under default must be > 0.
+    cycles_zero = sum(
+        c.cycles for c in tl_zero.commands if isinstance(c, PeCpuOverheadCmd)
+    )
+    cycles_default = sum(
+        c.cycles for c in tl_default.commands if isinstance(c, PeCpuOverheadCmd)
+    )
+    assert cycles_zero == 0, f"empty table must add 0 cycles; got {cycles_zero}"
+    assert cycles_default > 0, (
+        f"default table must add > 0 cycles; got {cycles_default}"
+    )
+
+
+# ── T6: greenlet vs legacy replay use same cost table (review #6) ─
+
+
+def test_greenlet_and_legacy_path_parity():
+    """Both PE_CPU execution paths read the same cost table.
+
+    The greenlet path (kernel_runner.py:run) and the legacy replay path
+    (pe_cpu.py:_execute_legacy) must construct TLContext with the same
+    default cost table so the same kernel produces identical PE_CPU
+    overhead cycles via either route. This is ADR-0064 review item #6.
+    """
+    from kernbench.common.cpu_issue_cost import DEFAULT_CPU_ISSUE_COST
+
+    # Build the command list for a representative kernel via TLContext
+    # using the default table — this is what both live paths should see.
+    def kernel(tl):
+        a = tl.load(0x1000, shape=(4, 4), dtype="f16")
+        b = tl.load(0x2000, shape=(4, 4), dtype="f16")
+        c = tl.dot(a, b)
+        tl.store(0x3000, c)
+
+    tl1 = TLContext(
+        pe_id=0, num_programs=1,
+        dispatch_cycles=0,
+        issue_cost_table=DEFAULT_CPU_ISSUE_COST,
+    )
+    run_kernel(kernel, tl1)
+    cycles1 = sum(
+        c.cycles for c in tl1.commands if isinstance(c, PeCpuOverheadCmd)
+    )
+
+    tl2 = TLContext(
+        pe_id=0, num_programs=1,
+        dispatch_cycles=0,
+        issue_cost_table=DEFAULT_CPU_ISSUE_COST,
+    )
+    run_kernel(kernel, tl2)
+    cycles2 = sum(
+        c.cycles for c in tl2.commands if isinstance(c, PeCpuOverheadCmd)
+    )
+
+    assert cycles1 == cycles2, (
+        f"same kernel via same default table must produce identical "
+        f"overhead cycles; got {cycles1} vs {cycles2}"
+    )
+    # Concrete expected for this kernel: 2×load(5) + 1×dot(5) + 1×store(5) = 20.
+    assert cycles1 == 20, (
+        f"expected 20 cycles total (2 load + 1 dot + 1 store at 5 ns); "
+        f"got {cycles1}"
+    )
+
+
+# ── T7: back-compat — no table → uniform dispatch_cycles ─────────
+
+
+def test_back_compat_no_table_uses_dispatch_cycles():
+    """ADR-0046 §D6 contract preserved when no issue_cost_table is passed.
+
+    Existing call sites doing ``TLContext(dispatch_cycles=0)`` must
+    continue to emit zero overhead. Existing call sites doing
+    ``TLContext(dispatch_cycles=1)`` must continue to emit uniform 1.
+    """
+    # Zero path (most existing tests use this).
+    tl_zero = TLContext(pe_id=0, num_programs=1, dispatch_cycles=0)
+    tl_zero.load(0x1000, shape=(4, 4), dtype="f16")
+    overheads = [c for c in tl_zero.commands if isinstance(c, PeCpuOverheadCmd)]
+    assert overheads == [], (
+        f"dispatch_cycles=0 with no table must emit no overhead; "
+        f"got {[c.cycles for c in overheads]}"
+    )
+
+    # Uniform path (test_dispatch_overhead_inserted relies on this).
+    tl_one = TLContext(pe_id=0, num_programs=1, dispatch_cycles=1)
+    tl_one.load(0x1000, shape=(4, 4), dtype="f16")
+    overheads = [c for c in tl_one.commands if isinstance(c, PeCpuOverheadCmd)]
+    assert overheads == [PeCpuOverheadCmd(cycles=1)], (
+        f"dispatch_cycles=1 with no table must emit cycles=1; "
+        f"got {[c.cycles for c in overheads]}"
+    )
+
+
+# ── T8: live PE_CPU constructs TLContext with the default table ──
+
+
+def test_live_pe_cpu_uses_default_table(monkeypatch):
+    """End-to-end: live PE_CPU greenlet path constructs TLContext with
+    ``issue_cost_table=DEFAULT_CPU_ISSUE_COST``. This is the wiring assertion
+    for ADR-0064 D4 "active by default" + review item #6 (path parity).
+
+    We patch ``TLContext.__init__`` to record its kwargs and run a
+    single-load kernel through the live PE_CPU. The recorded
+    ``issue_cost_table`` must equal ``DEFAULT_CPU_ISSUE_COST``.
+    """
+    from kernbench.common.cpu_issue_cost import DEFAULT_CPU_ISSUE_COST
+    from kernbench.triton_emu import tl_context as _tlc
+
+    captured: list[dict] = []
+    real_init = _tlc.TLContext.__init__
+
+    def spy_init(self, *args, **kwargs):
+        captured.append(dict(kwargs))
+        real_init(self, *args, **kwargs)
+
+    monkeypatch.setattr(_tlc.TLContext, "__init__", spy_init)
+
+    clear_registry()
+    hbm_pa = _hbm_pa(sip=0, cube=0, pe_id=0)
+
+    def single_load_kernel(tl):
+        tl.load(hbm_pa, shape=(4, 4), dtype="f16")
+
+    register_kernel("test_active_default_table", single_load_kernel)
+
+    engine = _engine()
+    pe_cpu_id = "sip0.cube0.pe0.pe_cpu"
+    done = engine._env.event()
+    txn = Transaction(
+        request=KernelLaunchMsg(
+            correlation_id="t", request_id="r",
+            kernel_ref=KernelRef(name="test_active_default_table", kind="builtin"),
+            args=(),
+        ),
+        path=[pe_cpu_id], step=0, nbytes=0, done=done,
+    )
+
+    def inject():
+        yield engine._components[pe_cpu_id]._inbox.put(txn)
+        yield done
+
+    engine._env.process(inject())
+    engine._env.run()
+    clear_registry()
+
+    # Live PE_CPU constructs TLContext on either the greenlet path
+    # (kernel_runner.py:run) or the legacy replay path
+    # (pe_cpu.py:_execute_legacy) — whichever is chosen depends on whether
+    # MemoryStore is wired. Both must pass DEFAULT_CPU_ISSUE_COST.
+    pe_ctx_calls = [c for c in captured if c.get("pe_id") == 0
+                    and "issue_cost_table" in c]
+    assert len(pe_ctx_calls) >= 1, (
+        f"expected at least one TLContext constructed by a live PE_CPU "
+        f"path with an issue_cost_table kwarg; got {len(pe_ctx_calls)} "
+        f"(all captures: {captured})"
+    )
+    last = pe_ctx_calls[-1]
+    assert last["issue_cost_table"] == DEFAULT_CPU_ISSUE_COST, (
+        f"live PE_CPU must use DEFAULT_CPU_ISSUE_COST; got {last['issue_cost_table']}"
+    )

tests/test_tl_copy_to.py

@@ -0,0 +1,166 @@
+"""Phase 1 spec test for ``tl.copy_to`` (ADR-0063 §D3.1).
+
+ADR-0063 §D3.1: ``tl.copy_to(dst, src)`` is a TCM-to-TCM byte copy that
+lets a kernel write a scoped result's bytes to a persistent (outside-
+scope) address. Required for the two-arena flash pattern in §D3.
+
+Emit-time validation (ADR-0063 §D3.1 "Mechanics" / "Emit-time validation"):
+  - ``dst.shape == src.shape``
+  - ``dst.dtype == src.dtype``
+  - ``dst.space == "tcm"``  (TCM-only — HBM goes through ``tl.store``)
+  - ``src.space == "tcm"``  (same reason)
+
+op_log shape (ADR-0063 §D3.1 "Mechanics"):
+  - ``op_kind="math"`` (runs on the vector engine)
+  - ``op_name="copy"``
+
+Phase 1 (this commit): tests only — production code lands in Phase 2.
+The tests are written against the ``CopyCmd`` dataclass and the
+``TLContext.copy_to`` method that ADR-0063 §D3.1 declares; both are
+absent today, so every test in this file should currently fail with
+``AttributeError`` / ``ImportError``.
+"""
+from __future__ import annotations
+
+import pytest
+
+from kernbench.triton_emu.tl_context import TLContext
+
+
+def _ctx() -> TLContext:
+    """TLContext with a real scratch base so _make_compute_out allocates."""
+    return TLContext(
+        pe_id=0, num_programs=1, dispatch_cycles=0,
+        scratch_base=1 << 24, scratch_size=1 << 20,
+    )
+
+
+# ── 1. Method exists on the tl surface ───────────────────────────────
+
+
+def test_copy_to_method_exists():
+    """ADR-0063 §D3.1: TLContext must expose ``copy_to(dst, src)``."""
+    ctx = _ctx()
+    assert hasattr(ctx, "copy_to"), (
+        "TLContext must expose copy_to(dst, src) per ADR-0063 §D3.1"
+    )
+
+
+# ── 2. Emits a CopyCmd with the documented fields ────────────────────
+
+
+def test_copy_to_emits_copy_command():
+    """ADR-0063 §D3.1 Mechanics: a new ``CopyCmd(src, dst, nbytes)``
+    command is recorded; data_op=True (so it appears in op_log)."""
+    from kernbench.common.pe_commands import CopyCmd  # added in Phase 2
+
+    ctx = _ctx()
+    src = ctx._make_compute_out(shape=(8, 64), dtype="f16")
+    dst = ctx._make_compute_out(shape=(8, 64), dtype="f16")
+    ctx.copy_to(dst, src)
+
+    copy_cmds = [c for c in ctx.commands if isinstance(c, CopyCmd)]
+    assert len(copy_cmds) == 1, (
+        f"exactly one CopyCmd expected; got {len(copy_cmds)}"
+    )
+    cmd = copy_cmds[0]
+    assert cmd.src is src
+    assert cmd.dst is dst
+    assert cmd.nbytes == 8 * 64 * 2
+    assert cmd.data_op is True
+
+
+# ── 3. Shape mismatch is rejected at emit time ───────────────────────
+
+
+def test_copy_to_shape_mismatch_rejected():
+    """ADR-0063 §D3.1: ``dst.shape == src.shape`` enforced at emit time."""
+    ctx = _ctx()
+    src = ctx._make_compute_out(shape=(8, 64), dtype="f16")
+    dst = ctx._make_compute_out(shape=(8, 128), dtype="f16")  # mismatched
+    with pytest.raises((ValueError, AssertionError)) as excinfo:
+        ctx.copy_to(dst, src)
+    assert "shape" in str(excinfo.value).lower()
+
+
+# ── 4. Dtype mismatch is rejected at emit time ───────────────────────
+
+
+def test_copy_to_dtype_mismatch_rejected():
+    """ADR-0063 §D3.1: ``dst.dtype == src.dtype`` enforced at emit time."""
+    ctx = _ctx()
+    src = ctx._make_compute_out(shape=(8, 64), dtype="f16")
+    dst = ctx._make_compute_out(shape=(8, 64), dtype="f32")  # mismatched
+    with pytest.raises((ValueError, AssertionError)) as excinfo:
+        ctx.copy_to(dst, src)
+    assert "dtype" in str(excinfo.value).lower()
+
+
+# ── 5. Non-TCM dst is rejected at emit time ──────────────────────────
+
+
+def test_copy_to_dst_must_be_tcm():
+    """ADR-0063 §D3.1: ``dst.space == 'tcm'`` enforced at emit time.
+
+    Writing to HBM goes through ``tl.store``, not ``tl.copy_to`` —
+    keeping copy_to TCM-only avoids polluting op_log with DMA entries
+    that the bump-allocator scope mechanism doesn't model.
+    """
+    ctx = _ctx()
+    src = ctx._make_compute_out(shape=(8, 64), dtype="f16")
+    # Fabricate an HBM-resident dst (e.g. a load handle).
+    dst_hbm = ctx.load(0x10_000, shape=(8, 64), dtype="f16")
+    assert dst_hbm.space == "hbm", "fixture sanity check"
+    with pytest.raises((ValueError, AssertionError)) as excinfo:
+        ctx.copy_to(dst_hbm, src)
+    msg = str(excinfo.value).lower()
+    assert "tcm" in msg or "space" in msg or "hbm" in msg
+
+
+# ── 6. Non-TCM src is rejected at emit time ──────────────────────────
+
+
+def test_copy_to_src_must_be_tcm():
+    """ADR-0063 §D3.1: src.space == 'tcm' (symmetric to dst).
+
+    Reading from HBM goes through ``tl.load``, not ``tl.copy_to``.
+    """
+    ctx = _ctx()
+    src_hbm = ctx.load(0x10_000, shape=(8, 64), dtype="f16")
+    assert src_hbm.space == "hbm"
+    dst = ctx._make_compute_out(shape=(8, 64), dtype="f16")
+    with pytest.raises((ValueError, AssertionError)) as excinfo:
+        ctx.copy_to(dst, src_hbm)
+    msg = str(excinfo.value).lower()
+    assert "tcm" in msg or "space" in msg or "hbm" in msg
+
+
+# ── 7. op_log routes CopyCmd as ("math", "copy", ...) ────────────────
+
+
+def test_copy_to_op_log_classifies_as_math_copy():
+    """ADR-0063 §D3.1 Mechanics: op_kind='math', op_name='copy'.
+
+    The vector engine handles the copy; op_log classification reflects
+    that (engine-class accounting in bench summaries can sum over
+    op_kind='math' to include copy time alongside softmax/exp).
+    """
+    from kernbench.common.pe_commands import CopyCmd  # added in Phase 2
+    from kernbench.sim_engine.op_log import _extract_op_info
+
+    ctx = _ctx()
+    src = ctx._make_compute_out(shape=(8, 64), dtype="f16")
+    dst = ctx._make_compute_out(shape=(8, 64), dtype="f16")
+    ctx.copy_to(dst, src)
+
+    cmd = next(c for c in ctx.commands if isinstance(c, CopyCmd))
+    op_kind, op_name, params = _extract_op_info(cmd)
+    assert op_kind == "math", (
+        f"copy runs on the vector engine: op_kind='math'; got {op_kind!r}"
+    )
+    assert op_name == "copy", (
+        f"op_name must be 'copy' per ADR-0063 §D3.1; got {op_name!r}"
+    )
+    # Params must let DataExecutor replay (read src, write dst).
+    assert params.get("dst_addr") == dst.addr
+    assert params.get("dst_space", "tcm") == "tcm"

tests/test_tl_load_lazy.py

@@ -0,0 +1,167 @@
+"""Phase 1 spec test for lazy ``tl.load`` (ADR-0062).
+
+ADR-0062 §D1/§D2: ``tl.load`` is non-blocking. It issues a ``DmaReadCmd``
+and returns immediately with a ``TensorHandle`` carrying a pending
+``LoadFuture``. The runtime auto-inserts a wait on that future at the
+first consuming op (``tl.dot``, MATH ops, ``tl.store``, ``tl.send``,
+``tl.copy_to``, ``tl.composite`` operands). Symmetric with the existing
+``recv_async`` / ``RecvFuture`` pattern, generalised to HBM loads.
+
+Phase 1 (this commit): tests only — production code lands in Phase 2.
+All tests currently fail because:
+  - ``TensorHandle`` has no ``pending`` field
+  - ``LoadFuture`` class doesn't exist
+"""
+from __future__ import annotations
+
+from kernbench.common.pe_commands import DmaReadCmd
+from kernbench.triton_emu.tl_context import TLContext
+
+
+def _ctx() -> TLContext:
+    """TLContext with a real scratch base so _make_compute_out allocates."""
+    return TLContext(
+        pe_id=0, num_programs=1, dispatch_cycles=0,
+        scratch_base=1 << 24, scratch_size=1 << 20,
+    )
+
+
+# ── 1. tl.load handle carries a `pending` field ──────────────────────
+
+
+def test_tl_load_handle_carries_pending_field():
+    """ADR-0062 §D2: tl.load handle has a ``pending`` attribute that is
+    not None — it references the LoadFuture the consumer will await."""
+    ctx = _ctx()
+    handle = ctx.load(0x1000, shape=(8, 64), dtype="f16")
+    assert hasattr(handle, "pending"), (
+        "TensorHandle must have a `pending` field per ADR-0062 §D2"
+    )
+    assert handle.pending is not None, (
+        "tl.load handle must carry a non-None pending (LoadFuture)"
+    )
+
+
+# ── 2. Constant / math-output handles have pending=None ──────────────
+
+
+def test_constant_handle_pending_is_none():
+    """Non-loaded handles (tl.zeros, tl.full, _make_compute_out, tl.arange)
+    have ``pending=None`` — there is no DMA to wait on. Consumer ops can
+    safely skip auto-wait for these inputs."""
+    ctx = _ctx()
+    z = ctx.zeros((8, 64), dtype="f16")
+    assert getattr(z, "pending", "MISSING") is None, (
+        "tl.zeros handle must have pending=None"
+    )
+    f = ctx.full((8, 64), value=0.0, dtype="f16")
+    assert getattr(f, "pending", "MISSING") is None, (
+        "tl.full handle must have pending=None"
+    )
+    a = ctx.arange(0, 8, dtype="i32")
+    assert getattr(a, "pending", "MISSING") is None, (
+        "tl.arange handle must have pending=None"
+    )
+
+
+def test_compute_out_pending_is_none():
+    """Scratch-allocated output handles (via _make_compute_out) have
+    pending=None. Math ops produce them; no DMA → no auto-wait needed
+    when they're consumed downstream."""
+    ctx = _ctx()
+    out = ctx._make_compute_out(shape=(8, 64), dtype="f16")
+    assert getattr(out, "pending", "MISSING") is None, (
+        "compute-out handle must have pending=None"
+    )
+
+
+# ── 3. LoadFuture class exists ───────────────────────────────────────
+
+
+def test_loadfuture_class_exists():
+    """ADR-0062 §D2: LoadFuture class wraps the DMA done event and a
+    resolved-flag, analogous to RecvFuture for IPCQ. Lives in
+    tl_context.py to keep all greenlet bridge types in one module."""
+    from kernbench.triton_emu.tl_context import LoadFuture  # added Phase 2
+    assert LoadFuture is not None
+
+
+# ── 4. Distinct loads produce distinct LoadFutures ───────────────────
+
+
+def test_distinct_loads_have_distinct_pending():
+    """ADR-0062 §D2: two tl.load calls (different addresses) produce two
+    independent LoadFuture objects. Each consumer must wait on the
+    right one — sharing the same future would over-serialise."""
+    ctx = _ctx()
+    h1 = ctx.load(0x1000, shape=(8,), dtype="f16")
+    h2 = ctx.load(0x2000, shape=(8,), dtype="f16")
+    assert h1.pending is not None and h2.pending is not None
+    assert h1.pending is not h2.pending, (
+        "distinct loads must have distinct LoadFuture objects"
+    )
+
+
+# ── 5. LoadFuture starts unresolved ──────────────────────────────────
+
+
+def test_loadfuture_starts_unresolved():
+    """ADR-0062 §D2: a fresh LoadFuture has resolved=False — the consumer
+    op must await it before reading. Once awaited (in greenlet mode by
+    auto-wait at first use), the SimPy event triggers and a subsequent
+    consumer can skip the yield."""
+    ctx = _ctx()
+    h = ctx.load(0x1000, shape=(8,), dtype="f16")
+    assert h.pending.resolved is False, (
+        "LoadFuture must start unresolved; got resolved=True at issue time"
+    )
+
+
+# ── 6. op_log: dma_read_count unchanged from blocking version ────────
+
+
+def test_op_log_dma_read_count_unchanged():
+    """ADR-0062 §D2 (last paragraph) / Test Req 4: each tl.load still
+    emits exactly one DmaReadCmd. The asynchrony is a scheduling
+    property, not a new op kind — dma_read_count and existing op_log
+    consumers see no change."""
+    ctx = _ctx()
+    ctx.load(0x1000, shape=(8, 64), dtype="f16")
+    ctx.load(0x2000, shape=(8, 64), dtype="f16")
+    ctx.load(0x3000, shape=(8, 64), dtype="f16")
+    dma_cmds = [c for c in ctx.commands if isinstance(c, DmaReadCmd)]
+    assert len(dma_cmds) == 3, (
+        f"lazy tl.load must still emit one DmaReadCmd per call; "
+        f"got {len(dma_cmds)}"
+    )
+
+
+# ── 7. Handle metadata fields preserved across the lazy change ───────
+
+
+def test_tl_load_handle_metadata_unchanged():
+    """ADR-0062 §D1: the tl surface is unchanged. The handle's addr,
+    shape, dtype, nbytes, space, pinned fields are identical to the
+    blocking version — only `pending` is new."""
+    ctx = _ctx()
+    h = ctx.load(0x1000, shape=(8, 64), dtype="f16")
+    assert h.addr == 0x1000
+    assert h.shape == (8, 64)
+    assert h.dtype == "f16"
+    assert h.nbytes == 8 * 64 * 2
+    assert h.space == "hbm"
+    assert h.pinned is True
+
+
+# ── 8. LoadFuture carries the DmaReadCmd it wraps ────────────────────
+
+
+def test_loadfuture_carries_cmd():
+    """ADR-0062 §D2: LoadFuture references the DmaReadCmd it wraps so
+    the runtime can resolve handle → command → completion event at
+    auto-wait time. Mirrors RecvFuture.cmd."""
+    ctx = _ctx()
+    h = ctx.load(0x1000, shape=(8, 64), dtype="f16")
+    assert hasattr(h.pending, "cmd"), "LoadFuture.cmd must exist"
+    assert isinstance(h.pending.cmd, DmaReadCmd)
+    assert h.pending.cmd.handle is h

tests/test_tl_scratch_scope.py

@@ -0,0 +1,213 @@
+"""Phase 1 spec test for ``tl.scratch_scope`` (ADR-0063, P3a).
+
+ADR-0063 D1: ``tl.scratch_scope()`` is a context manager whose ``__enter__``
+snapshots ``_scratch_cursor`` and whose ``__exit__`` restores it, freeing
+every handle allocated inside the ``with``-block. This pins the bump
+allocator behaviour without changing any command emission.
+
+Phase 1 (this commit): tests only — production code lands in Phase 2.
+"""
+from __future__ import annotations
+
+from kernbench.triton_emu.tl_context import TLContext
+
+
+# Configured TLContext with a real (non-zero) scratch base — required for
+# _make_compute_out to allocate, see tl_context.py:_scratch_alloc.
+def _ctx() -> TLContext:
+    return TLContext(
+        pe_id=0, num_programs=1,
+        dispatch_cycles=0,
+        scratch_base=1 << 24,         # 16 MiB base
+        scratch_size=1 << 20,         # 1 MiB pool
+    )
+
+
+# ── 1. Method exists and returns a usable ctx-mgr ─────────────────────
+
+
+def test_scratch_scope_method_exists():
+    ctx = _ctx()
+    assert hasattr(ctx, "scratch_scope"), (
+        "TLContext must expose scratch_scope() per ADR-0063 D1"
+    )
+    with ctx.scratch_scope() as scope:
+        assert scope is not None
+
+
+# ── 2. Recycling within a loop bounds peak cursor ────────────────────
+
+
+def test_scratch_scope_recycles_within_loop():
+    """ADR-0063 D1 / Test Req 1: a loop inside scratch_scope keeps peak
+    _scratch_cursor bounded by one iteration's footprint."""
+    ctx = _ctx()
+    # One iteration footprint: allocate three small handles via
+    # _make_compute_out.
+    one_iter_peaks = []
+    for _ in range(10):
+        with ctx.scratch_scope():
+            ctx._make_compute_out(shape=(64,), dtype="f16")
+            ctx._make_compute_out(shape=(64,), dtype="f16")
+            ctx._make_compute_out(shape=(64,), dtype="f16")
+            one_iter_peaks.append(ctx._scratch_cursor)
+    # After the loop the cursor must be at its post-loop start (0 here),
+    # not 10× one iteration.
+    assert ctx._scratch_cursor == 0, (
+        f"cursor must reset on each scope exit; final cursor = "
+        f"{ctx._scratch_cursor}"
+    )
+    # Every iteration sees the same peak — proves recycling.
+    assert len(set(one_iter_peaks)) == 1, (
+        f"every iteration must reach the same peak (recycled); "
+        f"got distinct peaks {one_iter_peaks}"
+    )
+
+
+# ── 3. Allocations outside the scope are not recycled ────────────────
+
+
+def test_scratch_scope_preserves_outside_allocations():
+    """ADR-0063 D3: handles allocated outside the scope (persistent
+    arena) keep their addresses; only inside-scope allocations are
+    rewound."""
+    ctx = _ctx()
+    # Persistent allocation BEFORE the scope.
+    persistent = ctx._make_compute_out(shape=(64,), dtype="f16")
+    cursor_after_persistent = ctx._scratch_cursor
+    with ctx.scratch_scope():
+        ctx._make_compute_out(shape=(64,), dtype="f16")
+        ctx._make_compute_out(shape=(64,), dtype="f16")
+    # On exit the cursor must rewind to the save-point (which equals
+    # cursor_after_persistent), NOT to 0 — the persistent allocation is
+    # preserved.
+    assert ctx._scratch_cursor == cursor_after_persistent, (
+        f"cursor must rewind to scope-enter point ({cursor_after_persistent}), "
+        f"not 0; got {ctx._scratch_cursor}"
+    )
+    # A new allocation after the scope must not collide with the
+    # persistent one.
+    fresh = ctx._make_compute_out(shape=(64,), dtype="f16")
+    assert fresh.addr != persistent.addr, (
+        "post-scope allocation must not reuse the persistent address"
+    )
+
+
+# ── 4. Nested scopes rewind to their own save-points ─────────────────
+
+
+def test_scratch_scope_nesting():
+    """ADR-0063 D4: nested scopes rewind to the matching enter-point."""
+    ctx = _ctx()
+    with ctx.scratch_scope():
+        ctx._make_compute_out(shape=(64,), dtype="f16")
+        outer_save = ctx._scratch_cursor
+        with ctx.scratch_scope():
+            ctx._make_compute_out(shape=(64,), dtype="f16")
+            ctx._make_compute_out(shape=(64,), dtype="f16")
+            inner_peak = ctx._scratch_cursor
+        # Inner exit: rewind to outer_save.
+        assert ctx._scratch_cursor == outer_save, (
+            f"inner exit must rewind to {outer_save}; "
+            f"got {ctx._scratch_cursor}"
+        )
+        assert inner_peak > outer_save, (
+            "sanity: inner scope must have allocated something"
+        )
+    # Outer exit: rewind all the way to 0.
+    assert ctx._scratch_cursor == 0, (
+        f"outer exit must rewind to 0; got {ctx._scratch_cursor}"
+    )
+
+
+# ── 5. Control: without scope, allocations pile up ───────────────────
+
+
+def test_without_scope_grows_unbounded():
+    """Sanity check that the bump allocator without scratch_scope grows
+    linearly — the failure mode ADR-0063 §A.2 cites for the S=16 cap."""
+    ctx = _ctx()
+    for _ in range(10):
+        ctx._make_compute_out(shape=(64,), dtype="f16")
+        ctx._make_compute_out(shape=(64,), dtype="f16")
+        ctx._make_compute_out(shape=(64,), dtype="f16")
+    # 30 allocations × aligned(64 * 2 = 128 → 128) B = 3840 B
+    assert ctx._scratch_cursor >= 30 * 128, (
+        f"without scope, cursor must grow with allocation count; got "
+        f"{ctx._scratch_cursor}"
+    )
+
+
+# ── 6. Overflow avoided by recycling ─────────────────────────────────
+
+
+def test_overflow_avoided_with_scope():
+    """ADR-0063 Test Req 3: a loop that would overflow 1 MiB without
+    scope completes when wrapped in scratch_scope."""
+    # 1 MiB pool / one alloc of 64 KiB = 16 max without scope.
+    # 100 iterations of 64 KiB without recycling = 6.4 MiB → overflow.
+    # With scratch_scope: peak ≈ 64 KiB per iter, well under 1 MiB.
+    ctx = _ctx()
+    big_shape = (32 * 1024,)  # 32 K elements × 2 B = 64 KiB per handle
+    for _ in range(100):
+        with ctx.scratch_scope():
+            ctx._make_compute_out(shape=big_shape, dtype="f16")
+    # Completes without raising — the recycle keeps us under the cap.
+    assert ctx._scratch_cursor == 0
+
+
+# ── 7. scratch_scope + copy_to integration (ADR-0063 §D3.1) ──────────
+
+
+def test_scoped_loop_with_copy_to_keeps_cursor_bounded():
+    """ADR-0063 §D3.1: the two-arena flash pattern.
+
+    Persistent ``running`` allocated OUTSIDE the scope; per-iteration
+    work allocates inside; the last act before scope exit is
+    ``tl.copy_to(running, new_running)`` to persist state. After the
+    loop the cursor must equal the post-persistent-allocation value
+    (running survived) and never have grown beyond that + one
+    iteration's footprint.
+    """
+    ctx = _ctx()
+    # Persistent arena: one running-state handle outside any scope.
+    running = ctx._make_compute_out(shape=(8, 64), dtype="f16")
+    cursor_after_persist = ctx._scratch_cursor
+    assert running.addr != 0
+
+    n_iters = 50
+    for _ in range(n_iters):
+        with ctx.scratch_scope():
+            # Simulate per-tile work: a few scoped allocations.
+            ctx._make_compute_out(shape=(8, 64), dtype="f16")
+            ctx._make_compute_out(shape=(8, 64), dtype="f16")
+            new_running = ctx._make_compute_out(shape=(8, 64), dtype="f16")
+            # Persist back to the outside-scope handle so its bytes
+            # survive __exit__.
+            ctx.copy_to(running, new_running)
+        # After scope exit cursor returns to the persistent footprint —
+        # scoped allocations are recycled, running survives.
+        assert ctx._scratch_cursor == cursor_after_persist, (
+            f"after scope exit cursor must rewind to "
+            f"{cursor_after_persist}; got {ctx._scratch_cursor}"
+        )
+
+    assert ctx._scratch_cursor == cursor_after_persist
+
+
+def test_persistent_handle_survives_after_copy_to_in_scope():
+    """ADR-0063 §D2 / §D3: a handle allocated outside a scratch_scope
+    must keep its address even after a copy_to(...) targeting it from
+    inside the scope and after that scope exits."""
+    ctx = _ctx()
+    persistent = ctx._make_compute_out(shape=(64,), dtype="f16")
+    persistent_addr_before = persistent.addr
+
+    with ctx.scratch_scope():
+        scoped = ctx._make_compute_out(shape=(64,), dtype="f16")
+        ctx.copy_to(persistent, scoped)
+    # Out of scope; persistent's address is unchanged.
+    assert persistent.addr == persistent_addr_before
+    # A fresh allocation after the scope must not collide with persistent.
+    fresh = ctx._make_compute_out(shape=(64,), dtype="f16")
+    assert fresh.addr != persistent.addr

topologies/llama70b_4sip.yaml

@@ -1,152 +0,0 @@
-# Llama-70B GQA 4-SIP topology — sub-cycle 3 (ADR-0055 + ADR-0056 context).
-#
-# Identical to repo-root topology.yaml except for system.sips.count: 2 → 4,
-# matching the GQA Llama-70B sharding study's TL/TR baseline at 1 Q-head
-# per cube (h_q=64, 8 cubes per KV-group × 8 KV-groups = 64 cubes = 4 SIPs).
-# See:
-#   llm_paper_review/notes/GQA_MHA_sharding/scripts/_gen_llama70b_1M_4cases.py
-#     lines 17-19 (model dims) and 287-302 (Q/cube scaling configs).
-#
-# Opt-in only — the milestone-gqa-llama70b bench (sub-cycle 4) points at
-# this file via env var. The repo-root topology.yaml is unchanged and
-# continues to drive milestone-1h-ccl / -gemm at their original 2-SIP
-# scale (CLAUDE.md "Surgical Changes").
-
-system:
-  ns_per_mm: 0.01   # wire propagation delay: 10 ps/mm (on-chip silicon)
-
-  sips:
-    count: 4
-    topology: ring_1d
-
-  components:
-    switch: { kind: switch, impl: builtin.switch, attrs: { overhead_ns: 5.0 } }
-
-  links:
-    io_ep_to_switch:
-      kind: pcie
-      bw_gbs_per_ep: 768.0
-      distance_mm: 20.0
-
-sip:
-  cube_mesh: { w: 4, h: 4 }
-
-  iochiplet:
-    components:
-      pcie_ep: { kind: pcie_ep, impl: builtin.pcie_ep, attrs: { overhead_ns: 5.0 } }
-      io_cpu:  { kind: io_cpu,  impl: builtin.io_cpu,  attrs: { overhead_ns: 10.0 } }
-      io_noc:  { kind: io_noc,  impl: builtin.forwarding, attrs: { overhead_ns: 0.0 } }
-    links:
-      pcie_ep_to_noc_bw_gbs: 256.0
-      pcie_ep_to_noc_mm: 1.0
-      io_cpu_to_noc_bw_gbs: 256.0
-      io_cpu_to_noc_mm: 0.5
-    ucie:
-      overhead_ns: 8.0
-      n_connections: 4
-      per_connection_bw_gbs: 128.0   # 4 × 128 = 512 GB/s = PHY BW
-      noc_to_ucie_mm: 0.5
-    instances:
-      - id: io0
-        place: { side: N, offset_norm: 0.5 }
-        ucie: { phy_bw_gbs: 512.0, phys: [P0, P1, P2, P3] }
-        cube_ports:
-          - { cube: {xy: [0,0]}, cube_side: N, phy: P0, distance_mm: 2.0 }
-          - { cube: {xy: [1,0]}, cube_side: N, phy: P1, distance_mm: 2.0 }
-          - { cube: {xy: [2,0]}, cube_side: N, phy: P2, distance_mm: 2.0 }
-          - { cube: {xy: [3,0]}, cube_side: N, phy: P3, distance_mm: 2.0 }
-
-  links:
-    inter_cube_mesh:
-      bw_gbs_per_ucie_phy: 512.0
-      distance_mm_across_seam: 1.0
-      routing: { algo: xy }
-
-cube:
-  geometry:
-    cube_mm: { w: 17.0, h: 14.0 }
-    hbm_mm: { w: 9.0, h: 5.0 }
-    ucie_mm: { size: 2.0 }
-
-  pe_layout:
-    corners: [NW, NE, SW, SE]   # N corners → top PE rows; S corners → bottom PE rows
-    pe_per_corner: 2             # total PEs per cube: 4 * 2 = 8
-
-  pe_template:
-    components:
-      pe_cpu:       { kind: pe_cpu,       impl: builtin.pe_cpu,       attrs: { overhead_ns: 2.0 } }
-      pe_scheduler: { kind: pe_scheduler, impl: builtin.pe_scheduler, attrs: { overhead_ns: 1.0 } }
-      pe_dma:       { kind: pe_dma,       impl: builtin.pe_dma,       attrs: { rd_engines: 1, wr_engines: 1 } }
-      pe_gemm:        { kind: pe_gemm,        impl: builtin.pe_gemm,        attrs: { overhead_ns: 0.0, shared_resource: accel_slot, peak_tflops_f16: 8.0 } }
-      pe_math:        { kind: pe_math,        impl: builtin.pe_math,        attrs: { overhead_ns: 0.0, shared_resource: accel_slot } }
-      pe_fetch_store: { kind: pe_fetch_store, impl: builtin.pe_fetch_store, attrs: { overhead_ns: 0.0 } }
-      pe_mmu:         { kind: pe_mmu,         impl: builtin.pe_mmu,         attrs: { tlb_overhead_ns: 0.5, page_size: 4096 } }
-      pe_tcm:         { kind: pe_tcm,         impl: builtin.pe_tcm,         attrs: { size_mb: 16, read_bw_gbs: 512.0, write_bw_gbs: 512.0, kernel_scratch_mb: 1 } }
-      pe_ipcq:        { kind: pe_ipcq,        impl: builtin.pe_ipcq,        attrs: { overhead_ns: 0.0 } }
-    links:
-      pe_cpu_to_scheduler_mm:  0.5
-      scheduler_to_dma_mm:     0.5
-      scheduler_to_gemm_mm:    0.5
-      scheduler_to_math_mm:    0.5
-      scheduler_to_fetch_store_mm: 0.5
-      dma_to_tcm_bw_gbs:       512.0
-      dma_to_tcm_mm:           0.5
-      dma_to_fetch_store_mm:   0.0     # DMA → fetch_store chaining (ADR-0014 D6)
-      fetch_store_to_tcm_bw_gbs: 512.0
-      fetch_store_to_tcm_mm:   0.0
-      fetch_store_to_gemm_mm:  0.0     # fetch → GEMM chaining (ADR-0014 D6)
-      fetch_store_to_math_mm:  0.0     # fetch → MATH chaining (ADR-0014 D6)
-      gemm_to_fetch_store_mm:  0.0     # GEMM → store chaining (ADR-0014 D6)
-      gemm_to_math_mm:         0.0     # GEMM → MATH epilogue chaining (ADR-0014 D6)
-      math_to_fetch_store_mm:  0.0     # MATH → store chaining (ADR-0014 D6)
-      fetch_store_to_dma_mm:   0.0     # store → DMA writeback chaining (ADR-0014 D6)
-      gemm_to_tcm_bw_gbs:      512.0
-      gemm_to_tcm_mm:          0.5
-      math_to_tcm_bw_gbs:      512.0
-      math_to_tcm_mm:          0.5
-      cpu_to_ipcq_mm:          0.5     # PE_CPU → PE_IPCQ (ADR-0023)
-      ipcq_to_dma_mm:          0.0     # PE_IPCQ → PE_DMA token forwarding (ADR-0023)
-      dma_to_ipcq_mm:          0.0     # PE_DMA → PE_IPCQ metadata arrival (ADR-0023)
-
-  memory_map:
-    hbm_total_gb_per_cube: 48
-    hbm_slices_per_cube: 8
-    hbm_total_bw_gbs: 1024.0
-    hbm_mapping_mode: n_to_one        # one_to_one | n_to_one (ADR-0017 D8)
-    hbm_pseudo_channels: 64           # total pseudo channels per cube
-    hbm_channels_per_pe: 8            # = pseudo_channels / pes_per_cube
-    hbm_channel_bw_gbs: 32.0          # per-channel bandwidth (GB/s)
-
-  components:
-    noc_router: { kind: noc_router, impl: builtin.forwarding, attrs: { overhead_ns: 2.0 } }
-    m_cpu:      { kind: m_cpu,   impl: builtin.m_cpu,   attrs: { overhead_ns: 5.0 } }
-    hbm_ctrl:   { kind: hbm_ctrl, impl: builtin.hbm_ctrl, attrs: { capacity: 1, efficiency: 1.0, num_pcs: 8, burst_bytes: 256, switch_penalty_ns: 0.0 } }
-    sram:       { kind: sram,     impl: builtin.sram,     attrs: { size_mb: 32, overhead_ns: 2.0 } }
-
-  # Physical placement of non-PE components (mm coordinates)
-  placement:
-    m_cpu:  { pos_mm: [7.5, 3.0] }   # top center, below UCIe-N
-    sram:   { pos_mm: [1.5, 9.0] }   # left side, below HBM zone
-
-  ucie:
-    decompose: true
-    ports: [N, S, E, W]
-    overhead_ns: 8.0
-    n_connections: 4               # independent NOC↔UCIe connections per port
-    per_connection_bw_gbs: 128.0   # BW per connection; 4 × 128 = 512 GB/s = UCIe PHY BW
-
-  links:
-    # Router mesh links (ADR-0017 D5)
-    router_link_bw_gbs: 256.0        # inter-router XY mesh link BW
-    router_overhead_ns: 2.0          # per-router switching overhead
-    pe_to_router_bw_gbs: 256.0      # PE_DMA ↔ router (= N × channel_bw)
-    hbm_to_router_bw_gbs: 256.0     # HBM_CTRL ↔ router (= N × channel_bw)
-    sram_to_router_bw_gbs: 128.0    # SRAM ↔ router
-    m_cpu_to_router_mm: 0.0         # M_CPU ↔ router distance
-    pe_dma_to_noc_bw_gbs: 256.0     # PE → router BW (= HBM slice BW, no bottleneck)
-    noc_to_pe_cpu_mm: 0.0           # router → PE_CPU distance (command path)
-
-visualization:
-  emit_views: [system, sip, cube]
-  sip_ids: [0]
-  cubes: [0, 9, 15]