paper(§5): add roofline, capacity planning, parallelism selection subsections

Three new subsections at the top of §5, motivated by the deployment-sizing
questions the analytical visualization tool exposes:

- 5.1 Roofline Analysis — B*, L*, short vs long context batch effect
  (measured: 19× per-token cost reduction at short context, 28% at long)
- 5.2 Capacity Planning — 3-axis sizing (weights / KV / SLO), SLO targets,
  regime rules, deployment templates, per-axis playbook
- 5.3 Parallelism Selection — TP × CP × PP × DP × EP comparison,
  symptom-driven axis selection, add/stop criteria, misconceptions

Restructure:
- 05-gqa.tex trimmed to section header + intro
- Existing 5.4-5.7 content (placement, short/long ctx, composite) moved
  to 05x-fused-kernel.tex
- Summary extracted to 05z-summary.tex, cross-refs updated
- main.tex \input order sets the requested subsection sequence
- lmodern loaded to satisfy microtype font-expansion

Two new figures generated from tests/analytical_visualization/chip_roofline:
- roofline_short_context.png (S_kv=8K, batch wins)
- roofline_long_context.png (S_kv=1M, batch stalls at KV floor)
- Generator: tests/analytical_visualization/_gen_roofline_paper_figs.py

4 tables (regime rules, deployment templates, playbook, parallelism
criteria) use table* so they span both columns without overflow.

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

docs/adr-proposed/ADR-0070-algo-gqa-short-context-attention.md

@@ -0,0 +1,243 @@
+# ADR-0070: GQA Short-Context Attention — Unified A1/A2/A4/B Mapping for Prefill and Decode
+
+## Status
+
+Proposed — short-context (single tile to a handful of KV tiles) GQA
+attention benchmark covering **both prefill and decode** in **four
+mapping modes** (A1/A2/A4/B) and **three composite tiers** (without
+composite, GEMM-only composite, composite + softmax_merge fused).
+Compute-only attention — weight DMA, ITL batching, and end-to-end
+LLM driver scope are out of scope (separate ADRs).
+
+## Context
+
+A GQA layer at the LLaMA-3.1-70B headline shape (`h_q=64`, `h_kv=8`,
+`d_head=128`, GQA group `G = h_q/h_kv = 8`) on a 4×4 cube SIP can be
+mapped to AHBM in several ways that trade KV-cache memory, HBM read
+volume per cube, and intra-cube IPCQ traffic against PE parallelism.
+
+> **Bench-vs-headline shape.** The kernels and tests use a
+> `d_head = 64` proxy (`tests/attention/test_gqa_short_context.py`
+> `D_HEAD = 64`) so the per-tile working set fits comfortably in
+> scratch across all four modes. The mapping decisions are
+> `d_head`-independent — only the per-PE GEMM tile bytes scale.
+
+The four candidate mappings (per *GQA-mapping-on-AHBM-short-context* note §1):
+
+```
+Mode  kv_per_cube  C        group_size  Reduce / broadcast topology
+----  -----------  -------  ----------  ----------------------------------
+A1            1    h_kv      P (=8)     row chain + col bridge (2×4 mesh)
+A2            2    h_kv/2    P/2 (=4)   row chain only
+A4            4    h_kv/4    P/4 (=2)   single intra_W / intra_E hop
+B             8    1         1          no broadcast / no reduce (single PE)
+```
+
+Within each cube the 8 PEs split into `kv_per_cube` groups of
+`group_size = P/kv_per_cube` PEs; each group owns one KV head.
+
+The kernel must be able to compare these mappings at multiple context
+lengths (8K..64K KV tokens) and across the three composite tiers so
+the team can size the per-cube HBM / IPCQ / GEMM tradeoff empirically.
+
+## Decision
+
+### 1. Mapping (unified A1/A2/A4/B)
+
+Two unified kernels select mode at launch via `kv_per_cube ∈ {1,2,4,8}`:
+
+- **Prefill** (`gqa_attention_prefill_short_kernel`):
+  - Q-tile split — `T_q` rows floor-balanced across the group's
+    `group_size` PEs.
+  - FA2 head fusion — `G` Q heads fused into the M dimension of one
+    batched GEMM per PE per tile.
+  - **IPCQ KV broadcast** — group root (`pe_in_group == 0`) loads K/V
+    from HBM and IPCQ-broadcasts each tile to the rest of its group;
+    q-tiles are independent so no intra-group reduce.
+
+- **Decode** (`gqa_attention_decode_short_kernel`):
+  - Sequence-shard — each PE owns `S_local = S_kv/group_size` tokens
+    of the group's KV head.
+  - FA2 head fusion — same as prefill.
+  - **IPCQ chain reduce** — per-PE partial `(m, ℓ, O)` chain-reduce
+    up to group root (PE 0), which normalizes and stores.
+
+Geometry (2×4 mesh):
+
+```
+group_size = 8  : 2 row × 4 col   (A1)
+group_size = 4  : 1 row × 4 col   (A2)
+group_size = 2  : 1 row × 2 col   (A4)
+group_size = 1  : single PE       (B)
+```
+
+### 2. Shard addressing (ADR-0011 D-VA1 contract)
+
+Deploy is per-(sip, cube, pe) but `tl.load` receives a single global
+VA per tensor. The kernel computes its own shard base offset from
+`tl.program_id(axis=0)` (PE id) and `tl.program_id(axis=1)` (cube id):
+
+```python
+# Decode shape (pe=row_wise): K is split across PEs by sequence shard.
+cube_K_base       = cube_id        * kv_per_cube * K_HEAD_BYTES
+head_K_base       = group_id_in_cube * K_HEAD_BYTES
+pe_K_seq_offset   = pe_in_group     * n_tiles_per_pe * K_TILE_BYTES
+k_shard_base      = k_ptr + cube_K_base + head_K_base + pe_K_seq_offset
+```
+
+For **prefill** the K/V dp is `pe=replicate` (group root reads and
+broadcasts) so the `pe_in_group` term drops out:
+`k_head_shard_base = k_ptr + cube_K_base + head_K_base`.
+
+Skipping `cube_id` collapses all cubes onto cube 0's HBM region
+(observed pre-fix as an 11.5× per-cube DMA imbalance).
+
+### 3. Three composite tiers
+
+The same mapping is exercised against three GEMM/MATH dispatch styles
+so the contribution of the composite API vs the recipe-driven fusion
+can be isolated:
+
+```
+(1) without composite                  primitives only (tl.dot, tl.exp, …)
+(2) with composite (GEMM-only)         tl.composite(op="gemm")
+(3) with composite + softmax_merge     tl.composite(prologue=[softmax_merge],
+                                                    op="gemm", out=O,
+                                                    epilogue=[add])
+```
+
+File layout (`src/kernbench/benches/gqa_helpers/short_ctx/`):
+
+```
+_gqa_attention_{prefill,decode}_short.py                   (1)
+_gqa_attention_{prefill,decode}_short_composite.py         (2)
+_gqa_attention_{prefill,decode}_short_composite_fused.py   (3)
+```
+
+Tier (2) is **GEMM-only by definition** — no recipe-driven fusion,
+so both prefill's and decode's P·V stay `tl.dot`. Recipe-driven
+fusion is exactly what tier (3) adds.
+
+Tier (3) on **multi-cube modes (A1/A2/A4) of prefill** relies on the
+D4 supplement: any composite operand that is an IPCQ recv'd slot
+(non-root PE's K_T in Q·Kᵀ, V in P·V) is pinned and read in place
+instead of being DMA-streamed from HBM. Without the supplement every
+recv slot fed to a composite — both Q·Kᵀ's `b=K_T` and P·V's `b=V`
+— PageFaults on PA decode (PE scratch addresses overflow the 51-bit
+PA range).
+
+### 4. Caller contract — `_validate_config`
+
+Each kernel module exposes a single `_validate_config(...)` helper
+called by the bench wrapper before launch. The kernel itself is lean
+(no inline `if` guards): the contract block is enforced caller-side,
+sim cost zero, but catches every silent-shard-corruption foot-gun
+(`kv_per_cube=3 → group_size=2` via integer division, non-integer
+GQA group, mismatched cube count, partial-tile `S_kv`, …) before any
+address arithmetic runs.
+
+### 5. Tensor layouts (host-side, mode-invariant byte totals)
+
+- **Q**: `(kv_per_cube·T_q, h_q·d_head/kv_per_cube)`
+  `dp=(cube=column_wise, pe=replicate)` over `C = h_kv/kv_per_cube`.
+- **K**: `(h_kv·S_kv·d_head/TILE_S_KV, TILE_S_KV)` tile-major.
+  Decode: `dp=(cube=row_wise, pe=row_wise)`.
+  Prefill: `dp=(cube=row_wise, pe=replicate)` (broadcast model).
+- **V**: `(h_kv·S_kv, d_head)` native. Same dp as K.
+- **O**: same dp as Q.
+
+Caller pre-scales Q by `1/√d_head`.
+
+## Verification
+
+### Smoke + regression (`tests/attention/test_gqa_short_context.py`)
+
+- Prefill / decode smoke (all 4 modes, single-tile S_kv).
+- Multi-tile coverage (all 4 modes, S_kv chosen so each PE owns
+  ≥2 tiles).
+- Op-count invariants — prefill `dma_writes = group_size·kv_per_cube·C`,
+  `dma_reads = P·C + 2·kv_per_cube·C·n_tiles`; decode
+  `dma_writes = kv_per_cube·C`.
+- IPCQ topology — short-context kernel never emits inter-CUBE E/W
+  IPCQ.
+- **ADR-0011 D-VA1 regression** — `per_cube_disjoint_src_addrs`,
+  `per_cube_dma_balanced` (max/min DMA busy ratio < 1.1×) per mode
+  per phase.
+
+### Composite smoke
+
+- `test_{prefill,decode}_composite_smoke[A1/A2/A4/B]` — tier (2),
+  4 modes × 2 phases.
+- `test_{prefill,decode}_composite_fused_smoke[A1/A2/A4/B]` — tier
+  (3), 4 modes × 2 phases.
+- `test_decode_composite_fused_multitile[A1/A2/A4/B]` —
+  `n_tiles_per_pe == 2` per mode so the recipe-fused tile loop
+  actually executes.
+
+### Mode × context sweep (`tests/attention/test_gqa_short_context_sweep_*.py`)
+
+Six sweep files (3 tiers × {prefill, decode}) measure each
+`(mode, S_kv)` cell and dump the same metric set to CSV:
+
+```
+wall_us, n_pe, gemm_util, math_count, math_pipeline_us, hbm_bw_util,
+hbm_read_mb, hbm_write_kb, ipcq_kb, kv_cache_per_cube_mb
+```
+
+S_kv ∈ {8K, 16K, 32K, 64K}; outputs under `docs/sweeps/`.
+
+Composite/fused sweeps count both `gemm_f16` (non-pipeline) **and**
+`TileToken/GEMM` (pipeline composite path); fused additionally counts
+`TileToken/MATH`. Baseline sweeps count `gemm_f16` only.
+
+## Known limitations
+
+1. **Causal mask** — kernel is non-causal. Adding causal masking is
+   orthogonal to mapping and is scheduled as a separate ADR.
+2. **f16 online-softmax accumulator** — `(m, ℓ, O)` are f16 throughout.
+   Long-context numerical drift will need an f32 accumulator before
+   correctness-grade use.
+3. **End-to-end numeric validation** — current tests are op-count /
+   topology / per-cube-DMA invariants. No full-kernel
+   `np.allclose` against a torch/numpy reference yet.
+4. **Skinny-M GEMM underfill** — decode (`T_q=1`) gives `M = G = 8`,
+   well below `PE_SCHEDULER` supertile `TILE_M = 32`. The two GEMM
+   paths handle this differently:
+   - **Primitive `tl.dot`** dispatches the GEMM at the actual
+     `m = 8` (no padding); GEMM time reflects the real 8-row work.
+   - **`tl.composite(op="gemm")`** tiles at `TILE_M = 32`, padding
+     `M` 4× with zeros; GEMM time reflects the padded 32-row work.
+   Consequence: when sweep CSVs show tier (2)/(3) `gemm_util`
+   higher than tier (1), most of that gap is supertile padding
+   overhead, **not** extra useful work or fusion savings — read it as
+   "composite tile shape doesn't match decode-skinny shape" rather
+   than "composite/fusion is more compute-intensive." This is the
+   **intended decode shape** for this benchmark; lifting it requires
+   batched-M (multi-request inference, `B_sys > 1`), tracked as a
+   separate batched variant — see "Future work".
+5. **Multi-cube prefill composite without the ADR-0065 D4 supplement**
+   PageFaults. The supplement — IPCQ recv'd slots are pinned and read
+   in place as composite operands — ships alongside this ADR.
+
+## Future work
+
+- Batched-M variant (`B_sys ∈ {1,2,4,8,16}`) so composite/fused
+  pipeline overlap shows up in wall_clock rather than only in
+  per-engine utilization.
+- Long-context (S_kv ≥ 128K) sweep extension for the headline LLaMA
+  decode target.
+- Causal-mask + f32 accumulator promotion.
+- 3-variant comparison plots from the sweep CSVs (wall, gemm util,
+  hbm bw util, ipcq, kv cache) — generated under `docs/diagrams/`.
+
+## References
+
+- ADR-0011 — PhysAddr + VA contract (D-VA1: kernel-computed shard
+  offset).
+- ADR-0060 — GQA fused attention on AHBM (precursor of the unified
+  mapping).
+- ADR-0064 — CPU issue-cost model (composite supertile motivation).
+- ADR-0065 — Flat ops, composite, softmax_merge recipe (with the D4
+  supplement applied here).
+- ADR-0070 supersedes ADR-0060 §B.split.2 short-prefill clause.
+- *GQA-mapping-on-AHBM-short-context* note (research design rationale).

docs/report/1H-codesign-paper/Agentic_Runtime_Architecture.md

@@ -0,0 +1,848 @@
+# AHBM Agentic Runtime Architecture
+
+## Scope
+
+This note organizes the current design discussion for executing agentic workloads on a memory-centric AHBM-style architecture.
+
+The scope of this version is limited to:
+
+1. Motivation
+2. AHBM hardware assumptions
+3. Attention execution
+4. Attention execution policies
+5. Fan-in and join
+
+MoE execution and whole-model layer-aware scheduling are intentionally deferred to a later section.
+
+---
+
+# 1. Motivation
+
+Agentic workload patterns: loop, fan-out / fan-in. Loop 은 일반적인 LLM execution 과 다르지 않다. Sub-agent 가 만들어질 때 발생하는 Fan out / Fan in 이 주로 고려되어야 할 대상이다.
+
+## 1.1 Agentic fan-out
+
+An agentic workload often starts from one shared conversation or task context and then forks into several specialized sub-agents.
+
+```text
+Shared prefix
+   ├─ Agent A: Analyze performance
+   ├─ Agent B: Check correctness
+   └─ Agent C: Find alternatives
+```
+
+Each sub-agent sees:
+
+```text
+Shared prefix
++
+Private role or instruction
++
+Private generated continuation
+```
+
+The shared prefix may be long, while each private branch may initially contain only a small number of new tokens.
+
+This creates two important execution properties:
+
+1. All branches reuse the same prefix KV cache.
+2. New query tokens from multiple branches can potentially be batched.
+
+## 1.2 Why the workload is different from ordinary batching
+
+Ordinary batching groups unrelated requests that happen to arrive at similar times.
+
+Agentic fan-out is different because the branches are structurally related:
+
+```text
+Agent A context = Shared prefix + A suffix
+Agent B context = Shared prefix + B suffix
+Agent C context = Shared prefix + C suffix
+```
+
+The requests therefore have:
+
+- identical prefix KV,
+- different private suffix KV,
+- potentially synchronized execution points,
+- a later fan-in stage that combines their results.
+
+This structure creates opportunities that are not available in unrelated-request batching.
+
+## 1.3 Main optimization opportunity
+
+For attention, the key operation is:
+
+```text
+Q × Kᵀ
+```
+
+Multiple sub-agents can have different query rows while reading the same shared prefix KV.
+
+Instead of executing:
+
+```text
+Q_A × K_shared
+Q_B × K_shared
+Q_C × K_shared
+```
+
+independently, the runtime can combine the query rows:
+
+```text
+Q_combined =
+[
+  Q_A
+  Q_B
+  Q_C
+]
+```
+
+and execute:
+
+```text
+Q_combined × K_shared
+```
+
+The arithmetic is still row-independent, but the larger GEMM can improve utilization and amortize scheduling overhead.
+
+## 1.4 Main design questions
+
+The architecture must answer:
+
+- How should shared KV be distributed across CUBEs and PEs?
+- How should query rows from multiple sub-agents be grouped?
+- How should online softmax state be reduced across sequence shards?
+- Should Q be replicated or partitioned?
+- How should large fan-out results be joined back into the main agent?
+- Which data should be reused, recomputed, summarized, or materialized on demand?
+
+---
+
+# 2. AHBM Hardware Assumptions
+
+## 2.2 Sequence parallelism
+
+The sequence dimension of one KV head is distributed across the 32 PEs.
+
+```text
+KV head sequence
+   ├─ PE0  owns sequence shard 0
+   ├─ PE1  owns sequence shard 1
+   ├─ ...
+   └─ PE31 owns sequence shard 31
+```
+
+Each PE stores a different part of the sequence for the same KV head. A query row must attend to all 32 sequence shards, so every PE computes a partial attention result for its local KV shard.
+
+## 2.3 Q placement
+
+The baseline follows an AHBM-style replicated-Q execution model.
+
+Logically:
+
+```text
+The same Q rows are visible to all 32 PEs.
+```
+
+This does not require 32 independent physical copies. A possible implementation is:
+
+```text
+Chip-level Q source
+        ↓
+CUBE multicast
+        ↓
+Shared Q buffer within each CUBE
+        ↓
+8 PEs consume the same Q tile
+```
+
+Thus Q is logically replicated across PEs while physical traffic is reduced through multicast and shared buffering.
+
+## 2.4 KV placement
+
+KV remains stationary near the owning PE.
+
+```text
+PE0  → KV shard 0
+PE1  → KV shard 1
+...
+PE31 → KV shard 31
+```
+
+The baseline avoids:
+
+- KV remapping,
+- KV replication,
+- remote KV reads,
+- page-table reconstruction for every query group.
+
+## 2.5 GEMM engine assumptions
+
+Representative PE GEMM tile shapes include:
+
+```text
+16 × 16 × 16
+```
+
+and:
+
+```text
+8 × 64 × 8
+```
+
+For:
+
+```text
+8 sub-agents
+20 query tokens per sub-agent
+```
+
+the combined number of Q rows is:
+
+```text
+M = 8 × 20 = 160
+```
+
+If each PE owns 256 KV columns, a representative local GEMM is:
+
+```text
+Q_local: 160 × d
+K_local: d × 256
+```
+
+Both `M = 160` and `N = 256` align well with the assumed tile shapes.
+
+---
+
+# 3. Attention Execution
+
+## 3.1 Fan-out context structure
+
+Assume the parent agent has already prefetched a shared prefix.
+
+```text
+Shared prefix KV pages
+```
+
+The runtime forks the logical context:
+
+```text
+Branch A page table:
+[Shared prefix pages] + [A private pages]
+
+Branch B page table:
+[Shared prefix pages] + [B private pages]
+
+Branch C page table:
+[Shared prefix pages] + [C private pages]
+```
+
+The shared pages are referenced by multiple branches without being copied. Only private suffix pages are branch-specific.
+
+## 3.2 Sub-agent query batching
+
+Assume:
+
+```text
+8 sub-agents
+20 new tokens per sub-agent
+```
+
+The runtime forms:
+
+```text
+Q_A: 20 × d
+Q_B: 20 × d
+...
+Q_H: 20 × d
+```
+
+and concatenates them along the row dimension:
+
+```text
+Q_combined: 160 × d
+```
+
+The combined operation is:
+
+```text
+(160 × d) × (d × S)
+```
+
+where `S` is the total sequence length represented by all KV shards.
+
+Each query row remains logically independent. Batching changes the execution shape, not the attention semantics.
+
+## 3.3 Per-PE local attention
+
+Each PE owns only a local sequence shard.
+
+For PE `p`:
+
+```text
+Scores_p = Q_combined × K_pᵀ
+```
+
+The PE then computes local online-softmax state:
+
+```text
+m_p[row]
+l_p[row]
+o_p[row, :]
+```
+
+For 160 rows, each PE conceptually produces:
+
+```text
+m_p[160]
+l_p[160]
+o_p[160, d_v]
+```
+
+These may be processed in smaller row tiles for pipelining.
+
+## 3.4 Online softmax merge
+
+Each row has an independent online-softmax state.
+
+The reduction is always:
+
+```text
+same row index across different sequence shards
+```
+
+It is never:
+
+```text
+different query rows reduced together
+```
+
+Therefore 160 query rows do not imply 160 serialized communication rounds. The implementation exchanges vector or tiled payloads such as:
+
+```text
+m[160]
+l[160]
+o[160, d_v]
+```
+
+or:
+
+```text
+m[16]
+l[16]
+o[16, d_v]
+```
+
+The row states can be communicated and merged in parallel.
+
+## 3.5 Hierarchical reduction
+
+Because one KV head spans four CUBEs and 32 PEs, reduction is hierarchical.
+
+```text
+32 PE partial states
+        ↓
+8-PE reduction inside each CUBE
+        ↓
+4 CUBE-level states
+        ↓
+4-CUBE reduction
+        ↓
+Final attention outputs
+```
+
+The same online-softmax merge primitive is used at both levels.
+
+## 3.6 Prefill versus decode
+
+### Decode
+
+Decode typically has a very small number of Q rows.
+
+```text
+small Q
+   ↓ multicast
+32 PEs read local KV shards
+   ↓
+hierarchical reduction
+```
+
+The dominant concerns are usually KV bandwidth and reduction overhead.
+
+### Prefill
+
+Fan-out can make the query dimension much larger.
+
+```text
+Per-agent Q rows: 20
+Number of agents: 8
+Combined Q rows: 160
+```
+
+The larger `M` dimension can produce a much better GEMM shape. Agentic batching is therefore especially attractive for prefill or multi-token private suffix processing.
+
+## 3.7 Compute-cost clarification
+
+Replicating Q across 32 PEs does not multiply total attention FLOPs by 32. Each PE computes a different KV-column region.
+
+```text
+32 PEs × 160 rows × 256 local columns
+=
+160 rows × 8192 total columns
+```
+
+If Q is temporally tiled into four groups of 40 rows:
+
+```text
+4 × 32 PEs × 40 rows × 256 columns
+=
+160 rows × 8192 columns
+```
+
+The total arithmetic is identical. The policy changes Q traffic, reduction traffic, scheduling granularity, utilization, and buffering—not the mathematical amount of attention work.
+
+---
+
+# 4. Attention Execution Policies
+
+## 4.1 Policy A: Replicated Q, stationary KV
+
+Policy A keeps the existing KV placement unchanged.
+
+```text
+Q_combined
+    ↓ multicast
+PE0  computes against KV shard 0
+PE1  computes against KV shard 1
+...
+PE31 computes against KV shard 31
+```
+
+Each PE executes a local GEMM such as:
+
+```text
+(160 × d) × (d × 256)
+```
+
+### Advantages
+
+- No KV movement
+- No KV replication
+- No remote KV access
+- No page-table regrouping
+- Natural compatibility with sequence-parallel attention
+- Large local GEMM shapes
+- Simple hierarchical softmax reduction
+
+### Costs
+
+- Q must be distributed to all CUBEs
+- Every row requires a 32-way logical reduction
+- Large Q batches increase multicast and softmax-state traffic
+
+Policy A is the natural baseline for the assumed AHBM mapping.
+
+## 4.2 Policy B: Partitioned Q groups
+
+A possible alternative is to divide query rows among PE groups.
+
+```text
+Q group 0 → PE group 0
+Q group 1 → PE group 1
+...
+```
+
+However, every query row must still attend to the complete KV sequence. Because the 32 PEs already represent 32 sequence shards, spatially partitioning Q means each Q group must somehow access all KV shards.
+
+This requires one of the following:
+
+1. Regroup KV shards for each Q group.
+2. Read remote KV through symmetric memory.
+3. Replicate KV across Q-processing groups.
+4. Time-multiplex the same PEs over Q groups.
+
+The first three add memory-system complexity. The fourth is mainly temporal tiling and does not provide true spatial Q partitioning.
+
+## 4.3 Why Policy B is not automatically better
+
+The comparison is:
+
+```text
+Policy A cost
+=
+Q multicast
++
+hierarchical reduction
+```
+
+versus:
+
+```text
+Policy B cost
+=
+KV regrouping, replication, or remote access
++
+additional scheduling complexity
+```
+
+Policy B becomes attractive only if:
+
+```text
+Q distribution cost + reduction cost
+>
+remote or regrouped KV cost
+```
+
+Relevant factors include Q size, multicast bandwidth, reduction bandwidth, symmetric-memory bandwidth, remote-KV latency, KV replication capacity, page-table overhead, and GEMM utilization.
+
+## 4.4 Recommended baseline
+
+For:
+
+```text
+1 KV head = 4 CUBEs = 32 PEs
+```
+
+use:
+
+```text
+Policy A:
+replicated Q + stationary sequence-sharded KV
+```
+
+Policy B should be treated as an adaptive or future policy for cases where Q batches become extremely large, multicast becomes a bottleneck, reduction traffic dominates, or remote KV access becomes inexpensive.
+
+## 4.5 Runtime decision model
+
+A future runtime can estimate:
+
+```text
+T_A =
+T_Q_multicast
++
+T_local_GEMM
++
+T_hierarchical_reduce
+```
+
+and:
+
+```text
+T_B =
+T_Q_partition
++
+T_remote_or_replicated_KV
++
+T_local_GEMM
++
+T_group_reduce
++
+T_remap
+```
+
+The runtime selects the lower-cost policy for the current sequence length, agent count, Q size, KV-head mapping, bandwidth state, and interconnect congestion.
+
+---
+
+# 5. Fan-In and Join
+
+## 5.1 Fan-in problem
+
+After fan-out, each sub-agent produces a private result.
+
+```text
+Agent A → result A
+Agent B → result B
+Agent C → result C
+```
+
+The main agent must synthesize these results into a final continuation.
+
+## 5.2 Why branch KV cannot be concatenated
+
+Each branch has a different causal token history.
+
+```text
+A history:
+Shared prefix + A instruction + A reasoning
+
+B history:
+Shared prefix + B instruction + B reasoning
+
+C history:
+Shared prefix + C instruction + C reasoning
+```
+
+The K/V tensors of a token depend on token content, position, preceding causal context, and every transformer layer.
+
+Therefore:
+
+```text
+A private KV
++
+B private KV
++
+C private KV
+```
+
+does not form the KV cache of any valid single token sequence.
+
+## 5.3 Baseline text join
+
+The standard framework behavior is:
+
+```text
+Sub-agent result text
+        ↓
+Framework gathers and formats results
+        ↓
+Main-agent join prompt
+        ↓
+Continuation prefill
+        ↓
+New main-agent KV suffix
+```
+
+The framework normally aggregates the inputs, while the main LLM performs final synthesis.
+
+## 5.4 Main-agent KV after join
+
+The main-agent page table becomes:
+
+```text
+[Shared prefix KV pages]
++
+[Join-input KV pages]
++
+[Main continuation KV pages]
+```
+
+The private A/B/C pages remain separate and are not attached directly.
+
+## 5.5 Cost of full-text gather
+
+Assume:
+
+```text
+8 sub-agents
+1000 output tokens per agent
+```
+
+A full-text join creates:
+
+```text
+8000 join-input tokens
+```
+
+These tokens must pass through all transformer layers during continuation prefill. This increases prefill work, KV allocation, future decode KV reads, and context-window occupancy.
+
+## 5.6 Schema-constrained join
+
+A practical optimization is to constrain each sub-agent to a compact result schema.
+
+```json
+{
+  "claim": "memory_bandwidth_bottleneck",
+  "confidence": 0.91,
+  "evidence_ids": [17, 24]
+}
+```
+
+The framework can serialize it compactly:
+
+```text
+Finding: memory bandwidth bottleneck
+Confidence: 0.91
+Evidence: E17, E24
+```
+
+Example:
+
+```text
+Before:
+8 × 1000 = 8000 tokens
+
+After:
+8 × 50 = 400 tokens
+```
+
+The main LLM still performs final reasoning, but the continuation prefill is much shorter.
+
+## 5.7 Deduplication and aggregation
+
+Different sub-agents may produce overlapping findings.
+
+```json
+[
+  {"claim": "memory_bw", "confidence": 0.91},
+  {"claim": "memory_bw", "confidence": 0.87},
+  {"claim": "compute", "confidence": 0.42}
+]
+```
+
+The framework can combine duplicates:
+
+```text
+Primary finding:
+- Memory-bandwidth bottleneck
+- Supported by 2 agents
+- Maximum confidence: 0.91
+
+Alternative:
+- Compute bottleneck
+- Confidence: 0.42
+```
+
+This is preprocessing, not final reasoning.
+
+## 5.8 Pointer-based join
+
+Detailed evidence can remain outside the initial main-agent context.
+
+```text
+Agent A:
+- Finding: memory-bandwidth bottleneck
+- Evidence handle: E17
+
+Agent B:
+- Finding: reduction error
+- Evidence handle: E24
+```
+
+The main agent initially receives only summaries and handles. The framework retrieves and appends evidence only when requested.
+
+```text
+Effective input
+=
+summary tokens
++
+tokens for evidence actually used
+```
+
+The transformer does not directly dereference the handle; the framework resolves it through retrieval or a tool call.
+
+## 5.9 Hierarchical reduction
+
+For large fan-out width, local reducer agents can summarize groups of branches.
+
+```text
+16 sub-agents
+      ↓
+4 local reducers
+      ↓
+4 summaries
+      ↓
+Main agent
+```
+
+Benefits:
+
+- Reducers operate in parallel.
+- The final join prompt is shorter.
+- Main-agent KV growth is smaller.
+- Fan-in traffic is organized hierarchically.
+
+Costs:
+
+- Additional reducer inference
+- Possible loss of detail
+- Need for fallback evidence retrieval
+
+## 5.10 Latent-state join
+
+A more aggressive approach is to replace long text outputs with learned latent representations.
+
+```text
+Sub-agent hidden states
+        ↓
+Compression or projection
+        ↓
+Small latent-token set
+        ↓
+Main model
+```
+
+Example:
+
+```text
+1000 text tokens
+    ↓
+16 latent tokens
+```
+
+This could reduce join prefill and KV growth, but it requires training the main model to consume latent tokens, aligning branch representations, defining causal and positional semantics, and validating accuracy. It is a model-system co-design direction rather than a drop-in runtime optimization.
+
+## 5.11 KV and compute impact
+
+The shared parent prefix KV is already present. The primary optimization target is the newly created join suffix.
+
+```text
+New main KV size
+∝
+number of join-input tokens
+```
+
+Reducing join input reduces continuation-prefill work, new KV allocation, Q rows processed during join, later decode-time KV reads, and context-window consumption.
+
+A practical flow is:
+
+```text
+Sub-agent full outputs
+        ↓
+Schema-constrained results
+        ↓
+Deduplication and ranking
+        ↓
+Summaries + evidence handles
+        ↓
+Selective evidence materialization
+        ↓
+Main-agent continuation prefill
+```
+
+## 5.12 Recommended practical join design
+
+For an unchanged LLaMA-style model on AHBM:
+
+1. Reuse shared prefix KV through page-table references.
+2. Keep branch-private KV isolated.
+3. Require schema-constrained sub-agent outputs.
+4. Deduplicate repeated claims in the framework.
+5. Pass compact summaries, confidence values, and evidence handles.
+6. Retrieve detailed evidence only when requested.
+7. Introduce hierarchical reducer agents for wide fan-out.
+8. Keep the main LLM responsible for final synthesis.
+
+---
+
+# Current Baseline Summary
+
+```text
+Shared prefix prefill
+        ↓
+Shared KV page reuse
+        ↓
+Agentic fan-out
+        ↓
+Combine private Q rows
+        ↓
+Replicated-Q attention over stationary sequence-sharded KV
+        ↓
+Hierarchical online-softmax reduction
+        ↓
+Independent private branch continuations
+        ↓
+Schema/pointer-based fan-in
+        ↓
+Main-agent continuation prefill
+```
+
+Main design principles:
+
+- Reuse shared prefix KV without copying it.
+- Batch sub-agent Q rows when they read the same KV.
+- Keep KV stationary and multicast Q.
+- Reduce softmax state hierarchically by matching row index.
+- Do not directly merge branch-private KV.
+- Reduce fan-in cost by shortening and selectively materializing join inputs.

docs/report/1H-codesign-paper/main.tex

@@ -15,6 +15,7 @@
 \usepackage{caption}
 \usepackage{subcaption}
 \captionsetup{font=small,labelfont=bf}
+\usepackage{lmodern}  % scalable Latin Modern fonts (required by microtype expansion)
 \usepackage{microtype}
 \usepackage{tikz}
 \usetikzlibrary{arrows.meta,positioning,calc,fit}
@@ -28,7 +29,7 @@
 \date{
 \small
 AGI Computing Lab, System Technology Group\\
-2026 H1 Report
+2026 Q1-Q3 Report
 }
 
 \begin{document}
@@ -39,9 +40,16 @@ AGI Computing Lab, System Technology Group\\
 \input{sections/02-platform}
 \input{sections/03-gemm}
 \input{sections/04-allreduce}
-\input{sections/05-gqa}
+\input{sections/05-gqa}                  % section header + intro
-\input{sections/06-discussion}
+\input{sections/05a-roofline}            % 5.1 Roofline Analysis
-\input{sections/07-conclusion}
+\input{sections/05b-capacity-planning}   % 5.2 Capacity Planning
-\input{sections/08-future-work}
+\input{sections/05c-parallelism-selection} % 5.3 Parallelism Selection
+\input{sections/05x-fused-kernel}        % 5.4-5.7 Placement + short/long ctx + composite
+\input{sections/05z-summary}             % 5.8 Summary
+\input{sections/06-agentic}
+\input{sections/07-hw-spec-search}
+\input{sections/08-discussion}
+\input{sections/09-conclusion}
+\input{sections/10-future-work}
 
 \end{document}

docs/report/1H-codesign-paper/sections/04-allreduce.tex

@@ -145,9 +145,9 @@ with credit return, splitting the control plane into PE\_IPCQ and the
 data plane into PE\_DMA, with head updates riding the payload and tail
 updates riding a 16\,B side-channel credit (\S\ref{sec:allreduce}).
 
-\begin{figure}[t]
+\begin{figure*}[t]
 \centering
-\includegraphics[width=\linewidth]{ipcq_alternatives_architecture_stacked.png}
+\includegraphics[width=0.78\linewidth]{ipcq_alternatives_architecture_flow.png}
 \caption{Per-send data and control flow for the four PE-to-PE
 signalling mechanisms (sender\,$\rightarrow$\,NoC\,$\rightarrow$\,receiver).
 Doorbell and RDMA-CQ each issue two fabric transactions (payload then
@@ -158,7 +158,7 @@ the payload flit train and returns the tail credit on a side channel,
 so a send is one MMIO write and a receive is a flip-flop read. This is
 a \emph{design schematic}, not a measured comparison.}
 \label{fig:ipcq-arch}
-\end{figure}
+\end{figure*}
 
 \begin{figure}[t]
 \centering

docs/report/1H-codesign-paper/sections/05-gqa.tex

@@ -24,224 +24,3 @@ kernel that uses the composite command and PE\_IPCQ at the same time.
 Multi-head attention (MHA) was studied in prior work and serves here as
 the established baseline rather than being re-derived.
 
-\subsection{Data Placement Policy}
-\label{sec:gqa-placement}
-
-Long-context decode is bound by the KV cache, so the first-order design
-question is how to place that cache---and the running softmax state it
-feeds---across the two hardware axes the machine exposes: the CUBEs and,
-within each CUBE, the PEs (here $C{=}8$ CUBEs $\times$ $P{=}8$ PEs, for
-$C\!\cdot\!P{=}64$ attention engines over one KV-head group on the
-LLaMA-3.1-70B target). Each axis can \emph{replicate} the KV cache or
-\emph{shard} it, and a shard can run along the sequence dimension
-$S_{kv}$ or the head dimension $d_{\text{head}}$. The cross product is a
-small, enumerable taxonomy; six placements span its meaningful corners
-(Figure~\ref{fig:gqa-kv-sharding}).
-
-\begin{figure}[t]
-\centering
-\includegraphics[width=\linewidth]{gqa_long_ctx_6cases_kv_sharding_diagram.png}
-\caption{The six KV-placement strategies, drawn on the
-$S_{kv}\!\times\!d_{\text{head}}$ KV tensor (rows = sequence, columns =
-head dimension). Cube colour bands and dashed PE dividers show which
-axis each level shards. Cases~1--3 either replicate the cache or shard
-it on a single axis (8-way at most); Cases~4--6 reach a full 64-way
-split, three different ways: Case~4 splits $S_{kv}$ across CUBEs and
-$d_{\text{head}}$ across PEs, Case~5 the mirror, and Case~6~$\star$
-splits $S_{kv}$ on \emph{both} axes.}
-\label{fig:gqa-kv-sharding}
-\end{figure}
-
-Two quantities decide which placement is viable, and they pull against
-each other. The first is \textbf{per-PE KV memory}. With a per-PE HBM
-budget of \SI{6.0}{\giga\byte} and \SI{1.76}{\giga\byte} of attention
-weights resident, the KV headroom is \SI{4.24}{\giga\byte} per PE. At a
-production context of $S_{kv}{=}1\,\text{M}$ tokens the unsharded cache
-is \SI{40}{\giga\byte}/PE (Case~1), an 8-way shard is \SI{5}{\giga\byte}
-(Cases~2--3)---both \emph{over} the headroom---while only the 64-way
-placements bring it to \SI{640}{\mega\byte}/PE (Cases~4--6), comfortably
-inside budget. Memory alone therefore eliminates Cases~1--3 at long
-context. The second quantity is \textbf{communication per token}, and it
-is what separates the three survivors. Sharding $d_{\text{head}}$
-(Cases~4--5) makes each PE hold only a slice of every head, so the
-$Q\!\cdot\!K^{\top}$ score is \emph{partial} and must be all-reduced
-across the slice owners on every token---a reduction whose volume scales
-with $S_{kv}$ ($\sim$\SI{166}{\mega\byte}/token analytically, intra-CUBE
-on the NoC for Case~4, inter-CUBE on UCIe for Case~5). Case~6~$\star$
-instead shards $S_{kv}$ on both axes, so every PE computes a
-\emph{complete} score over its own token range and only the small running
-softmax state $(m,\ell,O)$ is merged across PEs
-($\sim$\SI{6.2}{\mega\byte}/token)---a $\sim$27$\times$ lighter collective
-than the $d_{\text{head}}$-split designs.
-
-\begin{figure}[t]
-\centering
-\includegraphics[width=\linewidth]{gqa_long_ctx_6cases_summary.png}
-\caption{Long-context placement analysis at $S_{kv}{=}1\,\text{M}$ tokens.
-\emph{Left:} the per-PE HBM budget---\SI{1.76}{\giga\byte} of attention
-weights leave \SI{4.24}{\giga\byte} of KV headroom (red line).
-\emph{Middle:} per-PE KV memory per case (log scale); only the 64-way
-placements (Cases~4--6, \SI{640}{\mega\byte}) clear the headroom, while
-the unsharded (\SI{40}{\giga\byte}) and 8-way (\SI{5}{\giga\byte}) cases
-overflow. \emph{Right:} analytical communication per token (log scale);
-the $d_{\text{head}}$-split Cases~4--5 pay a partial-score all-reduce
-($\sim$\SI{166}{\mega\byte}/token) that the both-axes-$S_{kv}$ split of
-Case~6~$\star$ avoids ($\sim$\SI{6.2}{\mega\byte}/token, merging only the
-softmax state). On these two axes Case~6 (marked $\star$ in the figure)
-is the single placement that lands both inside the memory budget and at
-low per-token communication; whether that combination is the right one to
-pick is a regime-specific question taken up next.}
-\label{fig:gqa-budget}
-\end{figure}
-
-These two costs---per-PE KV memory and per-token communication---are
-intrinsic properties of each placement, fixed by how it shards the cache
-and independent of the workload regime. They do not by themselves name a
-winner: a short prompt where the whole cache fits on one PE values low
-communication and tolerates replication, whereas a million-token decode
-is bound by the memory wall and will pay communication to escape it.
-Which placement is appropriate is therefore a per-regime question, which
-the short- and long-context subsections that follow answer by running the
-options on the simulator and reading off latency, traffic, and the
-redundant compute each one induces.
-
-\subsection{Inference with Short-Context Length}
-\label{sec:gqa-short}
-
-The fused GQA kernel issues its matrix products as scheduler-managed
-composite commands and keeps the online-softmax merge and the cross-device
-KV reduction inside the kernel, on PE\_IPCQ. Two kernel families cover the
-two phases. The \emph{prefill} kernel is head-parallel and rotates the KV
-shards around an inter-CUBE ring (``Ring KV''). The \emph{decode} kernel
-is head-replicated with a statically sharded KV cache and reduces partial
-attention outputs through an M-fold intra-CUBE chain and, for multiple
-users, a two-level reduce-to-root. Two further primitives make long
-context practical: a \emph{lazy load} that issues the KV \textsf{DMA\_READ}
-and returns immediately, auto-waiting only at first use so that KV load
-overlaps score computation; and per-tile \emph{scratch recycling} that
-keeps the running softmax accumulators ($m,\ell,O$) in a persistent arena
-while freeing per-tile temporaries, so the kernel fits the
-\SI{1}{\mebi\byte} scratch budget across many tiles. A further refinement
-that restructures the decode step into two stateful composites (a named
-\textsf{softmax\_merge} recipe) is designed but not yet wired into the
-measured path; results below reflect the implemented kernel only.
-
-% TODO: CUBE <-> KV-head mapping diagram for the short-context regime
-% (h_kv=8 KV heads -> 8 CUBEs, 1:1; intra-CUBE PE usage).
-% Bench code: src/kernbench/benches/gqa_helpers/short_ctx/
-
-% TODO: prefill performance figure (latency, stage breakdown).
-% TODO: decode performance figure (latency, stage breakdown).
-%       Bench output for short_ctx to be generated.
-
-\subsection{Inference with Long-Context Length}
-\label{sec:gqa-long}
-
-% TODO: prefill long-context kernel implementation description
-% (Sequence-Parallel partition of S_kv, per-case mechanics).
-% Bench code: src/kernbench/benches/gqa_helpers/long_ctx/
-
-% TODO: prefill long-context performance figure.
-
-Long-context decode is the regime where the KV cache, not attention
-compute, sets serving cost, so the placement question of
-\S\ref{sec:gqa-placement} becomes decisive here. To pick the right
-placement for this regime we run each of the six options as a fused
-decode kernel on the simulator (one decode step on the LLaMA-3.1-70B
-single-KV-head-group target, $C{=}8$ CUBEs $\times$ $P{=}8$ PEs) at a
-tractable $S_{kv}{=}8192$, and read off end-to-end latency, on-device op
-traffic, and the redundant compute each one induces. The swept context is
-small enough that all six fit in memory at $S_{kv}{=}8192$; the
-placements the long-context memory budget rules out (Cases~1--3) are
-drawn in red, run here only to expose their issue and communication
-structure.
-
-\begin{figure}[t]
-\centering
-\includegraphics[width=\linewidth]{gqa_decode_long_ctx_6cases_latency.png}
-\caption{Measured end-to-end decode latency per placement
-($S_{kv}{=}8192$; red = ruled out by the long-context memory budget,
-blue = the predicted Pareto choice). The fastest raw latency belongs to
-Case~3 (\SI{17.8}{\micro\second})---but Case~3 replicates the full KV
-cache into every CUBE, so it overflows the per-PE budget at production
-context and wastes 8$\times$ the compute
-(Figure~\ref{fig:gqa-6cases-par}). Among the placements that actually fit
-1\,M-token memory (Cases~4--6), Case~6~$\star$ is the fastest
-(\SI{30.6}{\micro\second}, versus \SI{31.4}{} and \SI{34.5}{\micro\second}
-for the $d_{\text{head}}$-split Cases~4 and~5)---making it the placement
-of choice for long-context decode.}
-\label{fig:gqa-6cases-lat}
-\end{figure}
-
-\begin{figure}[t]
-\centering
-\includegraphics[width=\linewidth]{gqa_decode_long_ctx_6cases_parallelism.png}
-\caption{Redundant compute per placement, measured as active-PE
-$\times$ $S_{\text{local}}$ (PE-tokens; lower means less wasted work).
-The minimum is \num{8192} PE-tokens---one pass over the sequence.
-Case~3 inflates this 8$\times$ to \num{65536} by replicating the KV
-cache across all eight CUBEs so every CUBE redundantly re-attends the
-whole sequence; the $d_{\text{head}}$-split Cases~4--5 likewise carry
-\num{65536} because each token is processed across eight head slices.
-Case~6~$\star$ achieves the full 64-way split at the minimal
-\num{8192} PE-tokens---fully parallel, no replication.}
-\label{fig:gqa-6cases-par}
-\end{figure}
-
-\begin{figure}[t]
-\centering
-\includegraphics[width=\linewidth]{gqa_decode_long_ctx_6cases_traffic.png}
-\caption{Measured on-device op traffic per placement. The unsharded
-Case~1 issues no IPCQ copies (each PE has the full cache, nothing to
-reduce); the single-axis Case~3 charges 168. Among the 64-way splits,
-the $d_{\text{head}}$-split Cases~4--5 charge the most---280 IPCQ copies
-and 8 DMA writes each, the partial-score all-reduce that head-slicing
-forces---while Case~6~$\star$ needs only 189 IPCQ copies and a single DMA
-write, because it merges just the running softmax state $(m,\ell,O)$
-rather than partial scores. This is the on-device collective traffic that
-PE\_IPCQ and the torus links of \S\ref{sec:allreduce} are provisioned to
-absorb at link speed.}
-\label{fig:gqa-6cases-traffic}
-\end{figure}
-
-The measurements select the right placement for this regime. The fastest
-raw latency (Case~3, \SI{17.8}{\micro\second}) comes from replicating the
-full KV cache into every CUBE---which is exactly the placement the
-long-context memory budget forbids, and which the parallelism panel shows
-wastes 8$\times$ the compute. Restricting attention to the placements that
-fit production-context memory (the 64-way splits, Cases~4--6), the choice
-is Case~6~$\star$: it is the fastest of the three
-(\SI{30.6}{\micro\second}), and the op-count panel shows why---its
-softmax-state-only reduction charges 189 IPCQ copies and one DMA write
-against the 280 copies and 8 DMA writes the $d_{\text{head}}$-split
-Cases~4--5 pay for their partial-score all-reduce. For long-context
-decode, then, the appropriate data placement is the both-axes sequence
-shard (Case~6): it is the cheapest-communicating member of the only
-memory-feasible family, and the cross-PE softmax reduction it does pay is
-precisely the traffic the communication-side codesign of this report is
-built to move quickly.
-
-\subsection{Comprehensive Analysis}
-\label{sec:gqa-analysis}
-
-These panels are the clearest statement of the codesign thesis in the
-report. Because the composite command keeps GEMM issue cheap and the MAC
-array barely occupied, the fused attention kernel's latency is set almost
-entirely by data movement: streaming the KV cache and reducing partials
-across devices. That is precisely the cost that the communication-side
-work targets---PE\_IPCQ for the on-device reduction, the lazy load for
-load/compute overlap, fast TCM staging and torus links for the reduction
-itself. In other words, the two enablers are not independent features that
-happen to appear in the same kernel; the GEMM optimization is what
-\emph{exposes} the data-movement bottleneck (by removing the compute and
-issue overhead that would otherwise hide it), and the communication
-optimization is what \emph{attacks} it. For an attention-dominated decoder
-the meaningful hardware investments are therefore the ones that move data
-faster and reduce it on-device---not additional MAC throughput, which this
-workload cannot use.
-
-% TODO: cross-regime DP (data parallelism) applicability:
-%   - Does Case-4 long-context placement compose with batch-level DP
-%     without further changes?
-%   - Does the short-context placement compose the same way?
-%   - Implications for multi-user serving (single vs. mixed regimes).

docs/report/1H-codesign-paper/sections/05a-roofline.tex

@@ -0,0 +1,85 @@
+\subsection{Roofline Analysis: Batch and Context Regimes}
+\label{sec:roofline}
+
+Before deploying any sharding strategy, the workload's position on the
+decode roofline sets what \emph{can} be improved and what \emph{cannot}.
+Two quantities matter: the machine's \textbf{critical batch}
+$B^\ast = C \cdot b / (2 W)$ (where $C$ is per-PE FLOPs, $W$ is per-PE
+HBM bandwidth, $b$ is bytes per parameter), and its \textbf{balance
+context} $L^\ast = 2 N / (\mathrm{AI} \cdot \mathrm{kv\_bpt})$ (where
+$N$ is active parameters and $\mathrm{AI} = C/W$ is arithmetic
+intensity). Above $B^\ast$ the deployment leaves the memory-bound
+regime; above $L^\ast$ per-user KV streaming dominates and no amount
+of batching hides it. For Llama-3.1-70B on the default AHBM machine
+($C \!=\!  \SI{8}{TFLOPS}$/PE, $W\!=\!  \SI{256}{GB/s}$/PE), we get
+$B^\ast \!\approx\! 31$ and $L^\ast \!\approx\! \SI{13.4}{K}$ tokens.
+
+\paragraph{Short-context regime ($S_{kv} \!<\! L^\ast$).}
+Figure~\ref{fig:roofline-short} shows one decode step at
+$S_{kv}\!=\!  \SI{8}{K}$. The step-latency panel (left) makes it
+obvious that weight fetch dominates at low $B$: at $B\!=\!1$, one
+step spends $\sim\!\SI{535}{ms}$ streaming weights and only
+$\SI{28}{ms}$ on per-sequence compute and KV combined. The
+cost-per-token panel (right) is where the batch story lives: dividing
+step latency by $B$ shrinks the weight term as $1/B$, so per-token
+cost drops from $\SI{562}{ms}$ at $B\!=\!1$ to $\SI{29.7}{ms}$ at
+$B\!=\!256$ --- a $19\times$ reduction. This is the memory-bound-to-
+compute-bound crossover: past $B^\ast$, weight cost is fully
+amortized and per-token time asymptotes to the compute + KV floor
+($\sim\!\SI{27}{ms}$). \textbf{At short context, batching directly
+lowers cost per token.}
+
+\begin{figure*}[h]
+\centering
+\includegraphics[width=\textwidth]{roofline_short_context.png}
+\caption{Roofline decomposition for Llama-3.1-70B decode at
+$S_{kv}\!=\!  \SI{8}{K}$ (short, well below $L^\ast\!\approx\!  \SI{13.4}{K}$).
+\textbf{Left}: step latency vs.\ batch $B$; weight fetch is flat
+(one HBM sweep per step), compute and KV grow linearly with $B$.
+\textbf{Right}: per-token cost ($=$ step $\div B$); the $1/B$
+weight-fetch term amortizes rapidly, dropping total per-token time
+from \SI{562}{ms} at $B\!=\!1$ to \SI{29.7}{ms} at $B\!=\!256$
+--- a $19\times$ throughput gain from batching alone.}
+\label{fig:roofline-short}
+\end{figure*}
+
+\paragraph{Long-context regime ($S_{kv} \!\gg\! L^\ast$).}
+Figure~\ref{fig:roofline-long} shows the same decomposition at
+$S_{kv}\!=\!  \SI{1}{M}$ ($\sim\!78 \times L^\ast$). The step-latency
+panel shows KV fetch has swelled by two orders of magnitude:
+per-token KV cost is now $\sim\!\SI{1342}{ms}$, dwarfing both
+weight ($\SI{535}{ms}$ at $B\!=\!1$) and compute ($\SI{17}{ms}$).
+The cost-per-token panel is the punchline: increasing $B$ still
+shrinks the weight term but leaves the giant KV floor untouched,
+because \textbf{KV fetch is per-sequence} --- adding another user
+adds a full extra copy of their KV read. Total per-token cost falls
+only from \SI{1894}{ms} at $B\!=\!1$ to \SI{1361}{ms} at $B\!=\!256$
+--- a mere $28\%$ reduction despite $256\times$ the batch.
+\textbf{At long context, batching stops paying because per-user KV
+streaming dominates.}
+
+\begin{figure*}[h]
+\centering
+\includegraphics[width=\textwidth]{roofline_long_context.png}
+\caption{Same decomposition at $S_{kv}\!=\!  \SI{1}{M}$ (long,
+$\sim\!78\times L^\ast$). KV fetch has grown to \SI{1342}{ms}/token
+and is now the dominant term at every $B$. Batching only shrinks
+the weight component; the KV floor is unmovable because each new
+user brings a full per-sequence KV read. Total per-token cost falls
+just $28\%$ from $B\!=\!1$ to $B\!=\!256$, versus $19\times$ at
+short context.}
+\label{fig:roofline-long}
+\end{figure*}
+
+\paragraph{Implication for deployment.} The two regimes call for
+opposite strategies. In the short-context regime, the operator packs
+$B$ as high as HBM allows to sit on the compute floor --- this is
+where cost-per-token is minimized and hardware utilization is
+highest. In the long-context regime, per-user KV is the binding
+resource; batching offers little benefit, so the operator instead
+shrinks $\mathrm{kv\_bpt}$ (GQA / MQA / MLA, INT4 KV, sparse
+attention) and shards the sequence dimension itself (CP), routing
+long-context requests to a dedicated pool with more chips per user.
+The next two subsections (\S\ref{sec:capacity-planning},
+\S\ref{sec:parallelism-selection}) turn these regime observations
+into concrete sizing and sharding rules.

docs/report/1H-codesign-paper/sections/05b-capacity-planning.tex

@@ -0,0 +1,118 @@
+\subsection{Capacity Planning: LLM Serving on AHBM}
+\label{sec:capacity-planning}
+
+Once a fused kernel is chosen, the deployment question is orthogonal:
+\emph{how many cubes (or SIPs) does a given model, context length, and
+latency SLO require?} Hyperscalers decide this along three axes that
+must all be satisfied simultaneously. The total PE count is
+$\max(A, B, C) \times N_{\text{replicas}}$, where $A$ is the capacity
+floor, $B$ the KV-cache headroom, and $C$ the throughput SLO
+(Table~\ref{tab:capacity-axes}).
+
+\begin{table}[h]
+\centering
+\small
+\caption{Three-axis sizing. $N$ = active params, $b$ = bytes per param,
+$\mathrm{HBM}_{\mathrm{PE}}$ = per-PE HBM budget, $S_{kv}$ = context
+length, $\mathrm{kv\_bpt}$ = KV cache bytes per token per user
+($=2 \cdot h_{kv} \cdot d_{\text{head}} \cdot b \cdot L$).}
+\label{tab:capacity-axes}
+\begin{tabular}{@{}l l@{}}
+\toprule
+Axis & Formula \\
+\midrule
+A. Capacity floor & $\lceil N b / \mathrm{HBM}_{\mathrm{PE}} \rceil$ \\
+B. KV headroom & $\lceil (N b + u \cdot S_{kv} \cdot \mathrm{kv\_bpt}) / \mathrm{HBM}_{\mathrm{PE}} \rceil$ \\
+C. Throughput SLO & $N_{\text{replicas}} = \lceil n_{\text{users}} / B_{\mathrm{SLO}} \rceil$ \\
+\bottomrule
+\end{tabular}
+\end{table}
+
+\paragraph{SLO targets.} The per-token latency budget $B_{\mathrm{SLO}}$
+follows the use case. TTFT (Time To First Token) is dominated by
+prefill; TPOT (Time Per Output Token) by one decode step
+(Table~\ref{tab:slo-targets}).
+
+\begin{table}[h]
+\centering
+\small
+\caption{Typical SLO targets. Voice needs a tight per-token cadence;
+chat balances both; batch is priced on throughput not latency.}
+\label{tab:slo-targets}
+\begin{tabular}{@{}l l l@{}}
+\toprule
+Workload & TTFT & TPOT \\
+\midrule
+Voice / real-time & 200--300 ms & 10--25 ms \\
+Interactive chat & 300 ms -- 1 s & 20--50 ms \\
+Batch / offline & seconds -- minutes & N/A \\
+\bottomrule
+\end{tabular}
+\end{table}
+
+\paragraph{Regime-aware rules of thumb.} Context length relative to the
+machine's balance point $L^*=2N/(\mathrm{AI}\cdot\mathrm{kv\_bpt})$
+determines which strategy applies. Below $L^*$ the deployment is
+compute-bound and batch scales to $\sim\!2B^*$; above, KV streaming
+dominates and $B$ must shrink (Table~\ref{tab:regime-rules}).
+
+\begin{table*}[t]
+\centering
+\small
+\caption{Deployment strategy by context regime.}
+\label{tab:regime-rules}
+\begin{tabular}{@{}l l l@{}}
+\toprule
+Regime & Batch & Cost/token \\
+\midrule
+Short ($S_{kv} < L^*$) & $B \approx 2 B^*$ & low \\
+Long ($S_{kv} > L^*$) & smaller $B$, more CP & higher \\
+Extreme ($S_{kv} \gg L^*$) & dedicated pool, disagg.\ prefill & much higher \\
+\bottomrule
+\end{tabular}
+\end{table*}
+
+\paragraph{Sample deployment templates.} A single base model is
+typically routed to one of several sharding tiers depending on the
+request's expected context length; the API gateway dispatches to the
+appropriate replica pool (Table~\ref{tab:deploy-templates}).
+
+\begin{table*}[t]
+\centering
+\small
+\caption{Sample deployment tiers for one base model. TP is fixed to
+match $h_{kv}$; CP grows with context; $B$ shrinks to hold TPOT.}
+\label{tab:deploy-templates}
+\begin{tabular}{@{}l c c c c@{}}
+\toprule
+Tier & CP$\times$TP & Cubes/repl. & Max $S_{kv}$ & Typical $B$ \\
+\midrule
+Small & $1 \times 8$ & 1 & 32k & 64 \\
+Medium & $4 \times 8$ & 4 & 128k & 32 \\
+Large & $32 \times 8$ & 32 & 1M & 4 \\
+\bottomrule
+\end{tabular}
+\end{table*}
+
+\paragraph{What to do when each axis binds.} The dominant axis dictates
+the first lever to pull (Table~\ref{tab:axis-playbook}). Axes $A$ and
+$C$ scale nearly linearly with hardware; axis $B$ is the one where
+algorithmic techniques (GQA, MLA, INT4 KV, sparse attention) buy the
+most because they attack $\mathrm{kv\_bpt}$ directly rather than adding
+chips.
+
+\begin{table*}[t]
+\centering
+\small
+\caption{Playbook per binding axis.}
+\label{tab:axis-playbook}
+\begin{tabular}{@{}l l l@{}}
+\toprule
+Binding & First lever & Second lever \\
+\midrule
+A. Capacity & $\uparrow$ TP / PP & bigger-HBM chip; FP8 / INT4 weights \\
+B. KV & $\uparrow$ CP; $\downarrow B$ & GQA / MQA / MLA; INT4 KV; sparse attn \\
+C. Throughput & $\uparrow$ replicas (DP) & loosen SLO; smaller model; spec.\ decoding \\
+\bottomrule
+\end{tabular}
+\end{table*}

docs/report/1H-codesign-paper/sections/05c-parallelism-selection.tex

@@ -0,0 +1,91 @@
+\subsection{Parallelism Selection: TP $\times$ CP $\times$ PP $\times$ DP $\times$ EP}
+\label{sec:parallelism-selection}
+
+Given a capacity-feasible layout (\S\ref{sec:capacity-planning}), the
+remaining decision is \emph{how to shard}: which axes to enable and at
+what degrees. Memory is the feasibility filter; once you fit, the
+design problem is minimising \emph{exposed} communication (communication
+that could not be overlapped with compute). The core principle is to
+\textbf{rank techniques by how much they communicate per unit of
+compute, and assign the chattiest ones to the fastest links.}
+
+\begin{table*}[h]
+\centering
+\small
+\caption{Parallelism axes at a glance. TP AllReduce volume does not
+shrink with degree — only compute does — which sets its practical
+ceiling near a fast-link (NVLink / intra-SIP) domain.}
+\label{tab:parallelism-compare}
+\begin{tabular}{@{}l l l l l@{}}
+\toprule
+Axis & Shards & Comm pattern & Frequency & Hard ceiling \\
+\midrule
+DP & optimizer state (w/ ZeRO) & AllReduce (grads) & 1$\times$/step & critical batch \\
+TP & weights + acts + KV heads & AllReduce & 2$\times$/layer & fast-link domain; $\lesssim h_{kv}$ \\
+PP & weights + KV by layer & P2P activation hand-off & per stage boundary & $n_{\text{layers}}$ \\
+CP & activations + KV by position & Ring P2P per hop & per layer (overlappable) & $\mathrm{seq\_len} / \mathrm{block}$ \\
+EP & expert weights (MoE only) & All-to-all & 2$\times$/MoE layer & $n_{\text{experts}}$ \\
+\bottomrule
+\end{tabular}
+\end{table*}
+
+\paragraph{Symptom-driven axis selection.} The dominant problem dictates
+the first knob to try (Table~\ref{tab:parallelism-problem}).
+
+\begin{table}[h]
+\centering
+\small
+\caption{Which axis to reach for first, by dominant problem.}
+\label{tab:parallelism-problem}
+\begin{tabular}{@{}l l@{}}
+\toprule
+Dominant problem & First axis \\
+\midrule
+Weight memory doesn't fit & TP within a fast domain \\
+KV of ONE long sequence & CP \\
+KV for MANY sequences & DP replicas \\
+Single-request TPOT & TP ($\lesssim 8$) \\
+Aggregate throughput & Outer DP \\
+Model spans multiple nodes & TP intra-node, PP inter-node \\
+MoE expert weight memory & EP \\
+\bottomrule
+\end{tabular}
+\end{table}
+
+\paragraph{When to add, when to stop.} Every axis has a natural
+saturation point past which further degree hurts more than it helps
+(Table~\ref{tab:parallelism-criteria}). Sizing decisions are almost
+always constrained by two axes simultaneously (e.g.\ TP by NVLink
+domain and $h_{kv}$; CP by ring-hop hiding and per-rank block size).
+
+\begin{table*}[t]
+\centering
+\small
+\caption{Add / stop criteria per axis.}
+\label{tab:parallelism-criteria}
+\begin{tabular}{@{}l l l@{}}
+\toprule
+Axis & Add when & Stop when \\
+\midrule
+DP & more independent requests & global batch $>$ critical \\
+TP & weights need sharding or TPOT tight & collectives dominate; local GEMMs shrink \\
+PP & depth too large; TP would cross slow links & bubble $>$ $\sim$10\% \\
+CP & one sequence needs position sharding & ring hops can't overlap compute \\
+EP & MoE expert memory & expert GEMMs too small; routing imbalance \\
+\bottomrule
+\end{tabular}
+\end{table*}
+
+\paragraph{Common misconceptions.} Three widely repeated rules survive
+contact with real workloads only partially:
+\emph{(i)} memory is not just a feasibility filter but a continuous
+performance variable --- more free HBM enables larger $B$, more prefix
+cache, higher throughput even after weights fit;
+\emph{(ii)} the ``TP $\leq h_{kv}$'' ceiling is an efficiency
+preference, not a correctness constraint --- vLLM and others support
+KV-head replication and head-dim splitting for TP $>$ $h_{kv}$;
+\emph{(iii)} ``always start at TP=8'' is disproven by public
+disclosures (DeepSeek-V3 inference uses TP=4, some MLPs at TP=1;
+DeepSeek-V3 training uses PP=16 + EP=64 with no TP), which shows the
+right starting point is the smallest TP that fits memory and meets
+latency, followed by measurement.

docs/report/1H-codesign-paper/sections/05x-fused-kernel.tex

@@ -0,0 +1,515 @@
+\subsection{Data Placement Policy}
+\label{sec:gqa-placement}
+
+Long-context decode is bound by the KV cache, so the first-order design
+question is how to place that cache---and the running softmax state it
+feeds---across the two hardware axes the machine exposes: the CUBEs and,
+within each CUBE, the PEs (here $C{=}8$ CUBEs $\times$ $P{=}8$ PEs, for
+$C\!\cdot\!P{=}64$ attention engines over one KV-head group on the
+LLaMA-3.1-70B target). Each axis can \emph{replicate} the KV cache or
+\emph{shard} it, and a shard can run along the sequence dimension
+$S_{kv}$ or the head dimension $d_{\text{head}}$. The cross product is a
+small, enumerable taxonomy; six placements span its meaningful corners
+(Figure~\ref{fig:gqa-kv-sharding}).
+
+\begin{figure}[t]
+\centering
+\includegraphics[width=\linewidth]{gqa_long_ctx_6cases_kv_sharding_diagram.png}
+\caption{The six KV-placement strategies, drawn on the
+$S_{kv}\!\times\!d_{\text{head}}$ KV tensor (rows = sequence, columns =
+head dimension). Cube colour bands and dashed PE dividers show which
+axis each level shards. Cases~1--3 either replicate the cache or shard
+it on a single axis (8-way at most); Cases~4--6 reach a full 64-way
+split, three different ways: Case~4 splits $S_{kv}$ across CUBEs and
+$d_{\text{head}}$ across PEs, Case~5 the mirror, and Case~6~$\star$
+splits $S_{kv}$ on \emph{both} axes.}
+\label{fig:gqa-kv-sharding}
+\end{figure}
+
+Two quantities decide which placement is viable, and they pull against
+each other. The first is \textbf{per-PE KV memory}. With a per-PE HBM
+budget of \SI{6.0}{\giga\byte} and \SI{1.76}{\giga\byte} of attention
+weights resident, the KV headroom is \SI{4.24}{\giga\byte} per PE. At a
+production context of $S_{kv}{=}1\,\text{M}$ tokens the unsharded cache
+is \SI{40}{\giga\byte}/PE (Case~1), an 8-way shard is \SI{5}{\giga\byte}
+(Cases~2--3)---both \emph{over} the headroom---while only the 64-way
+placements bring it to \SI{640}{\mega\byte}/PE (Cases~4--6), comfortably
+inside budget. Memory alone therefore eliminates Cases~1--3 at long
+context. The second quantity is \textbf{communication per token}, and it
+is what separates the three survivors. Sharding $d_{\text{head}}$
+(Cases~4--5) makes each PE hold only a slice of every head, so the
+$Q\!\cdot\!K^{\top}$ score is \emph{partial} and must be all-reduced
+across the slice owners on every token---a reduction whose volume scales
+with $S_{kv}$ ($\sim$\SI{166}{\mega\byte}/token analytically, intra-CUBE
+on the NoC for Case~4, inter-CUBE on UCIe for Case~5). Case~6~$\star$
+instead shards $S_{kv}$ on both axes, so every PE computes a
+\emph{complete} score over its own token range and only the small running
+softmax state $(m,\ell,O)$ is merged across PEs
+($\sim$\SI{6.2}{\mega\byte}/token)---a $\sim$27$\times$ lighter collective
+than the $d_{\text{head}}$-split designs.
+
+\begin{figure}[t]
+\centering
+\includegraphics[width=\linewidth]{gqa_long_ctx_6cases_summary.png}
+\caption{Long-context placement analysis at $S_{kv}{=}1\,\text{M}$ tokens.
+\emph{Left:} the per-PE HBM budget---\SI{1.76}{\giga\byte} of attention
+weights leave \SI{4.24}{\giga\byte} of KV headroom (red line).
+\emph{Middle:} per-PE KV memory per case (log scale); only the 64-way
+placements (Cases~4--6, \SI{640}{\mega\byte}) clear the headroom, while
+the unsharded (\SI{40}{\giga\byte}) and 8-way (\SI{5}{\giga\byte}) cases
+overflow. \emph{Right:} analytical communication per token (log scale);
+the $d_{\text{head}}$-split Cases~4--5 pay a partial-score all-reduce
+($\sim$\SI{166}{\mega\byte}/token) that the both-axes-$S_{kv}$ split of
+Case~6~$\star$ avoids ($\sim$\SI{6.2}{\mega\byte}/token, merging only the
+softmax state). On these two axes Case~6 (marked $\star$ in the figure)
+is the single placement that lands both inside the memory budget and at
+low per-token communication; whether that combination is the right one to
+pick is a regime-specific question taken up next.}
+\label{fig:gqa-budget}
+\end{figure}
+
+These two costs---per-PE KV memory and per-token communication---are
+intrinsic properties of each placement, fixed by how it shards the cache
+and independent of the workload regime. They do not by themselves name a
+winner: a short prompt where the whole cache fits on one PE values low
+communication and tolerates replication, whereas a million-token decode
+is bound by the memory wall and will pay communication to escape it.
+Which placement is appropriate is therefore a per-regime question, which
+the short- and long-context subsections that follow answer by running the
+options on the simulator and reading off latency, traffic, and the
+redundant compute each one induces.
+
+\subsection{Inference with Short-Context Length}
+\label{sec:gqa-short}
+
+The fused GQA kernel issues its matrix products as scheduler-managed
+composite commands and keeps the online-softmax merge and the cross-device
+KV reduction inside the kernel, on PE\_IPCQ. Two kernel families cover the
+two phases. The \emph{prefill} kernel splits the query tile across the PEs
+of a group and broadcasts each KV tile from the group's root PE over
+intra-CUBE IPCQ, so the HBM K/V read is paid once per group instead of
+once per PE. The \emph{decode} kernel sequence-shards the KV cache across
+the same group of PEs, runs per-PE local attention, then chain-reduces the
+partial $(m,\ell,O)$ triples back up to the group root. Two further
+primitives make long context practical: a \emph{lazy load} that issues the
+KV \textsf{DMA\_READ} and returns immediately, auto-waiting only at first
+use so KV load overlaps score computation; and per-tile \emph{scratch
+recycling} that keeps the running accumulators in a persistent arena while
+freeing per-tile temporaries, so the kernel fits the
+\SI{1}{\mebi\byte} scratch budget across many tiles. A further refinement
+restructures the decode step so the per-tile matrix products and the
+online-softmax merge are issued as \emph{composite} commands rather than
+hand-tiled primitives; \S\ref{sec:gqa-composite} measures it against the
+primitive baseline at long context.
+
+This is a different decomposition from the six-case placement taxonomy of
+\S\ref{sec:gqa-placement}: that taxonomy spreads a single KV-head group
+across all 64 PEs and is the right lens for long context
+(\S\ref{sec:gqa-long}), whereas short context---where the whole cache fits
+comfortably---turns instead on how the eight KV heads are distributed
+across the eight CUBEs. The design question for short context is therefore
+not whether to fuse softmax---that is settled---but how to distribute the
+eight KV heads across the eight CUBEs of a SIP. Four mappings cover the spectrum, named
+by KV-head count per CUBE: \textsf{1-kv-per-cube} dedicates one whole
+CUBE to each head and uses all eight PEs of that CUBE on the head's
+sequence shard; \textsf{2-kv-per-cube} packs two heads per CUBE with
+group size four; \textsf{4-kv-per-cube} packs four heads with group size
+two; \textsf{8-kv-per-cube} loads all heads onto a single CUBE with one
+PE per head. The mapping fixes a trade-off: as KV heads per CUBE grows, the per-CUBE KV
+footprint grows linearly and the mapping uses proportionally fewer CUBEs
+(eight down to one), but the intra-CUBE IPCQ broadcast/reduce cost shrinks
+to zero (8-kv-per-cube uses a single PE per head and has no group
+communication).
+
+A second axis is how each GEMM tile is issued. We isolate this with three
+\emph{composite tiers} applied to the same mapping: \textsf{without
+composite} uses primitive \textsf{tl.dot} and unfused softmax;
+\textsf{with composite (GEMM-only)} issues the GEMMs as
+\textsf{tl.composite(op="gemm")} commands so the scheduler can pipeline
+the tile stages but leaves softmax as primitives; \textsf{with composite
++ softmax\_merge} additionally folds the softmax fold into the $P\!\cdot\!V$
+composite through a named \textsf{softmax\_merge} prologue recipe.
+
+We sweep all four mappings $\times$ three composite tiers $\times$
+$S_{kv}\in\{8\text{K},16\text{K},32\text{K},64\text{K}\}$ for both
+phases (prefill uses $T_q{=}8$ sliced tiles; decode uses $T_q{=}1$). The
+headline latency comparison is in Figure~\ref{fig:gqa-short-wall-baseline}.
+The mappings separate cleanly, and in proportion to how many PEs each
+dedicates to a head: \textsf{1-kv-per-cube} spreads one head across all
+eight PEs of its CUBE, \textsf{8-kv-per-cube} gives each head a single PE,
+and the four mappings form a $\{64,32,16,8\}$-active-PE ladder. Decode
+latency tracks that ladder inversely---at $S_{kv}{=}64$K the
+\textsf{8-kv-per-cube}/\textsf{1-kv-per-cube} ratio is $7.4\times$
+($\SI{87.0}{}$ vs.\ $\SI{11.8}{\micro\second}$), and even at $8$K it is
+$4.8\times$ ($\SI{10.8}{}$ vs.\ $\SI{2.2}{\micro\second}$); prefill shows
+the same ordering more gently ($2.3\times$ at $64$K). The reason is that
+decode here is \emph{bandwidth-bound}: each active PE streams its KV shard
+at \SI{46}{}--\SI{76}{\percent} of the \SI{256}{\giga\byte\per\second}
+per-PE ceiling (rising with context), so a mapping's speed is set by how
+many PEs it puts to work---more PEs, more aggregate HBM bandwidth, lower
+latency. The $M{=}G{=}8$ score GEMM is skinny and leaves the MAC array
+lightly loaded throughout; the lever here is data movement, not compute.
+
+\begin{figure}[t]
+\centering
+\includegraphics[width=\linewidth]{gqa_short_context/wall_variant1_mode_compare.png}
+\caption{Wall-clock latency of the four short-context mappings, baseline
+kernel (no composite; log scale). The mappings separate along the
+$\{64,32,16,8\}$-active-PE ladder set by heads-per-CUBE:
+\textsf{1-kv-per-cube} (64 PEs) is fastest and \textsf{8-kv-per-cube}
+(8 PEs) slowest, by $4.8\times$ at $8$K growing to $7.4\times$ at $64$K in
+decode ($2.3\times$ at $64$K in prefill). Decode is bandwidth-bound, so
+latency scales inversely with the number of active PEs. Composite tiers
+track the baseline at this skinny shape
+(Figure~\ref{fig:gqa-short-a1-variant}).}
+\label{fig:gqa-short-wall-baseline}
+\end{figure}
+
+Figure~\ref{fig:gqa-short-tradeoff} shows what the mappings trade for that
+latency. The differentiator is \emph{density}: \textsf{8-kv-per-cube}
+concentrates all eight heads onto one CUBE, so its per-CUBE KV footprint is
+$8\times$ that of \textsf{1-kv-per-cube} (one CUBE per head), and it
+occupies a single one of the SIP's sixteen CUBEs against
+\textsf{1-kv-per-cube}'s eight. Per-PE HBM \emph{utilization}, by contrast,
+is similar across mappings---every active PE runs near the same
+\SI{46}{}--\SI{76}{\percent} of the per-PE ceiling---so the latency ladder
+comes from the \emph{count} of active PEs, not from any per-PE
+concentration. IPCQ traffic runs opposite to density:
+\textsf{1-kv-per-cube} pays the eight-PE chain-reduce while
+\textsf{8-kv-per-cube} (one PE per head) pays none, but absolute IPCQ volume
+stays under \SI{120}{\kibi\byte} per run, two orders of magnitude below the
+HBM traffic. The result is a genuine trade-off rather than a single winner:
+\textsf{1-kv-per-cube} minimizes per-request latency by spending the most
+hardware (eight CUBEs, 64 PEs) on one request, while denser mappings leave
+CUBEs free for other requests---the batched-serving axis taken up below.
+
+\begin{figure}[t]
+\centering
+\includegraphics[width=\linewidth]{gqa_short_context/per_cube_tradeoff_mode_compare.png}
+\caption{Per-CUBE trade-off across the four mappings (baseline kernel;
+log scale). Top row: per-CUBE KV footprint---\textsf{8-kv-per-cube} packs
+all eight heads onto one CUBE, an $8\times$ larger per-CUBE KV than
+\textsf{1-kv-per-cube} and one CUBE used against eight (the $8\times$ ratio
+is context-invariant; the absolute footprint grows with $S_{kv}$). Bottom
+row: IPCQ traffic per run---\textsf{1-kv-per-cube} pays the eight-PE
+chain-reduce while \textsf{8-kv-per-cube} (single PE per head) pays zero.
+The density (top) is what sets how many concurrent requests a SIP can hold.}
+\label{fig:gqa-short-tradeoff}
+\end{figure}
+
+The composite ablation in Figure~\ref{fig:gqa-short-a1-variant} shows the
+three tiers within \textsf{1-kv-per-cube}: wall-clock is nearly flat, with
+the GEMM-only composite tier a few percent \emph{slower}
+(\SI{14.0}{} vs.\ \SI{11.8}{\micro\second} at $64$K decode) and the
+\textsf{softmax\_merge} tier matching the baseline. This is the expected
+outcome for the decode-skinny shape---$M{=}G{=}8$ is well below the
+scheduler's $\textsf{TILE\_M}{=}32$ supertile (\S\ref{sec:gemm}), so the
+composite path pads $M$ by $4\times$ with zeros and the fusion has no slack
+to win back; the padding is pure overhead here. It shows up in
+Figure~\ref{fig:gqa-short-gemm-util} as inflated GEMM-engine busy time on
+the composite tiers---the padded-$M$ GEMM is charged in full against a
+sub-\SI{15}{\micro\second} decode wall, driving the accounted utilization
+well above the baseline's \SI{12}{}--\SI{18}{\percent}. Fusion benefit
+requires lifting $M$, either by the long-context Q-tile split
+(\S\ref{sec:gqa-composite}) or by batching multiple users---which is also
+the lever that converts a mapping's density into throughput.
+
+\paragraph{Latency versus batched throughput.} The latency ladder and the
+density trade-off pull in opposite directions once a SIP serves more than
+one request. Decode users are independent---each attends its own KV
+cache---and these mappings generate no cross-CUBE traffic, so a SIP runs
+several users at once on disjoint CUBE groups, up to $16/C$ users: two for
+\textsf{1-kv-per-cube}, sixteen for \textsf{8-kv-per-cube}. We measure this
+directly, launching $B$ concurrent users and reading aggregate latency off
+the shared clock (Figure~\ref{fig:gqa-batch}). The overlap is nearly
+perfect: aggregate latency stays within \SI{3}{}--\SI{5}{\percent} of the
+single-user latency all the way to a full SIP (\textsf{8-kv-per-cube} at
+$16$ users is \SI{11.4}{} vs.\ \SI{10.8}{\micro\second}), so that small
+residual is the only cross-user interference the shared fabric adds.
+Aggregate throughput at $8$K therefore rises almost linearly with $B$, from
+\SI{0.86}{}~requests\,/\,\si{\micro\second} (\textsf{1-kv-per-cube}, capped
+at two users) to \SI{1.41}{} (\textsf{8-kv-per-cube}, sixteen users): the
+\emph{dense} mapping wins on throughput even though it is the \emph{slowest}
+per request.
+
+The reason is worth spelling out, because both mappings fill the same
+hardware: at capacity \textsf{1-kv-per-cube} runs two users
+$\times$ 64 PEs and \textsf{8-kv-per-cube} sixteen users $\times$ 8 PEs, so
+each puts all 128 PEs of the SIP to work. What differs is per-PE
+\emph{efficiency}. Splitting one head across eight PEs
+(\textsf{1-kv-per-cube}) leaves each PE only $S_{kv}/8$ tokens---a single
+tile---so its fixed overheads (pipeline fill on the lazy KV load, and the
+eight-PE online-softmax chain-reduce) dominate the little streaming work it
+does, and it sustains just \SI{46}{\percent} of the per-PE HBM ceiling.
+Giving each head a whole PE (\textsf{8-kv-per-cube}) streams the full
+$S_{kv}$-token cache contiguously with no reduce, amortizing those overheads
+over $8\times$ more work and reaching \SI{76}{\percent}. The measured
+throughput ratio ($1.41/0.86 = 1.6\times$) matches the utilization ratio
+($0.76/0.46 = 1.7\times$) almost exactly: because decode is
+bandwidth-bound, aggregate throughput is set by total HBM efficiency, not
+by PE count. \textsf{1-kv-per-cube} is really spending hardware on
+\emph{latency}---eight PEs per head buy only a $4.8\times$ speedup, a $60\%$
+strong-scaling efficiency---and that same $40\%$ overhead is what caps its
+throughput once the SIP is full. The design conclusion is regime-dependent: \textsf{1-kv-per-cube} is the
+latency-optimal choice for single-stream or low-batch decode, while denser
+mappings are the throughput-optimal choice for high-batch short-context
+serving. At long context the per-request latencies already scale as $1/C$,
+so the same accounting predicts the throughputs converge; a measured
+long-context batch sweep is 2H work (\S\ref{sec:future}).
+
+\begin{figure}[t]
+\centering
+\includegraphics[width=\linewidth]{gqa_short_context/batch_scaling.png}
+\caption{Batch scaling at $S_{kv}{=}8$K: aggregate throughput (requests per
+\si{\micro\second}) versus the number of concurrent decode users on one SIP,
+each user on a disjoint CUBE group. Each mapping scales almost linearly with
+$B$ up to its SIP capacity $16/C$---\textsf{1-kv-per-cube} saturates at two
+users, \textsf{8-kv-per-cube} at sixteen. Solid segments are measured
+concurrent runs; the dashed extensions are the saturated-throughput ceiling
+beyond capacity, where further users run in waves at that sustained rate
+(so the SIP is full, not idle). The dense mapping reaches the highest
+ceiling; the spread mapping is latency-optimal but capacity-limited.
+Concurrent users overlap to within \SI{3}{}--\SI{5}{\percent} of the
+single-user latency.}
+\label{fig:gqa-batch}
+\end{figure}
+
+\begin{figure}[t]
+\centering
+\includegraphics[width=\linewidth]{gqa_short_context/wall_a1_variant_compare.png}
+\caption{Composite-tier ablation on \textsf{1-kv-per-cube} (log scale):
+wall-clock is nearly flat across tiers---the GEMM-only composite runs a few
+percent slower (padding overhead) and \textsf{softmax\_merge} matches the
+baseline. With $M{=}G{=}8 < \textsf{TILE\_M}{=}32$ the composite scheduler
+pads $M$ by $4\times$, leaving no fusion slack to recover at this shape.
+Fusion pays off only when $M$ fills the supertile---long-context Q-tile
+splitting or batched-$M$ inference.}
+\label{fig:gqa-short-a1-variant}
+\end{figure}
+
+\begin{figure}[t]
+\centering
+\includegraphics[width=\linewidth]{gqa_short_context/gemm_util_a1_variant_ablation.png}
+\caption{GEMM-engine busy fraction on \textsf{1-kv-per-cube} across the
+three composite tiers. The baseline runs at \SI{12}{}--\SI{18}{\percent};
+the composite tiers read much higher---up to and past
+\SI{100}{\percent} of the short decode wall---because the padded
+$8\!\to\!32$ $M$ dimension is charged as GEMM time, bookkeeping overhead
+rather than useful work. The engine is never the bottleneck at this shape;
+decode is KV-bandwidth-bound.}
+\label{fig:gqa-short-gemm-util}
+\end{figure}
+
+\subsection{Inference with Long-Context Length}
+\label{sec:gqa-long}
+
+% TODO: prefill long-context kernel implementation description
+% (Sequence-Parallel partition of S_kv, per-case mechanics).
+% Bench code: src/kernbench/benches/gqa_helpers/long_ctx/
+
+% TODO: prefill long-context performance figure.
+
+Long-context decode is the regime where the KV cache, not attention
+compute, sets serving cost, so the placement question of
+\S\ref{sec:gqa-placement} becomes decisive here. To pick the right
+placement for this regime we run each of the six options as a fused
+decode kernel on the simulator (one decode step on the LLaMA-3.1-70B
+single-KV-head-group target, $C{=}8$ CUBEs $\times$ $P{=}8$ PEs) at a
+tractable $S_{kv}{=}8192$, and read off end-to-end latency, on-device op
+traffic, and the redundant compute each one induces. The swept context is
+small enough that all six fit in memory at $S_{kv}{=}8192$; the
+placements the long-context memory budget rules out (Cases~1--3) are
+drawn in red, run here only to expose their issue and communication
+structure.
+
+\begin{figure}[t]
+\centering
+\includegraphics[width=\linewidth]{gqa_decode_long_ctx_6cases_latency.png}
+\caption{Measured end-to-end decode latency per placement
+($S_{kv}{=}8192$; red = ruled out by the long-context memory budget,
+blue = the predicted Pareto choice). The fastest raw latency belongs to
+Case~3 (\SI{17.8}{\micro\second})---but Case~3 replicates the full KV
+cache into every CUBE, so it overflows the per-PE budget at production
+context and wastes 8$\times$ the compute
+(Figure~\ref{fig:gqa-6cases-par}). Among the placements that actually fit
+1\,M-token memory (Cases~4--6), Case~6~$\star$ is the fastest
+(\SI{30.6}{\micro\second}, versus \SI{31.4}{} and \SI{34.5}{\micro\second}
+for the $d_{\text{head}}$-split Cases~4 and~5)---making it the placement
+of choice for long-context decode.}
+\label{fig:gqa-6cases-lat}
+\end{figure}
+
+\begin{figure}[t]
+\centering
+\includegraphics[width=\linewidth]{gqa_decode_long_ctx_6cases_parallelism.png}
+\caption{Redundant compute per placement, measured as active-PE
+$\times$ $S_{\text{local}}$ (PE-tokens; lower means less wasted work).
+The minimum is \num{8192} PE-tokens---one pass over the sequence.
+Case~3 inflates this 8$\times$ to \num{65536} by replicating the KV
+cache across all eight CUBEs so every CUBE redundantly re-attends the
+whole sequence; the $d_{\text{head}}$-split Cases~4--5 likewise carry
+\num{65536} because each token is processed across eight head slices.
+Case~6~$\star$ achieves the full 64-way split at the minimal
+\num{8192} PE-tokens---fully parallel, no replication.}
+\label{fig:gqa-6cases-par}
+\end{figure}
+
+\begin{figure}[t]
+\centering
+\includegraphics[width=\linewidth]{gqa_decode_long_ctx_6cases_traffic.png}
+\caption{Measured on-device op traffic per placement. The unsharded
+Case~1 issues no IPCQ copies (each PE has the full cache, nothing to
+reduce); the single-axis Case~3 charges 168. Among the 64-way splits,
+the $d_{\text{head}}$-split Cases~4--5 charge the most---280 IPCQ copies
+and 8 DMA writes each, the partial-score all-reduce that head-slicing
+forces---while Case~6~$\star$ needs only 189 IPCQ copies and a single DMA
+write, because it merges just the running softmax state $(m,\ell,O)$
+rather than partial scores. This is the on-device collective traffic that
+PE\_IPCQ and the torus links of \S\ref{sec:allreduce} are provisioned to
+absorb at link speed.}
+\label{fig:gqa-6cases-traffic}
+\end{figure}
+
+Taken together the three panels select the placement. The only
+memory-feasible family at production context is the 64-way splits
+(Cases~4--6), and within it Case~6~$\star$ is both the fastest and the
+lightest-communicating, because it merges only the running softmax state
+rather than the partial scores that the $d_{\text{head}}$-split
+Cases~4--5 must all-reduce. Case~3's lower raw latency is beside the
+point---it replicates the full cache into every CUBE, overflowing the
+per-PE budget and wasting $8\times$ the compute. For long-context decode
+the placement of choice is therefore the both-axes sequence shard, and the
+cross-PE softmax reduction it does pay is precisely the traffic the
+communication-side codesign of this report is built to move quickly.
+
+\subsection{Use of Composite Commands}
+\label{sec:gqa-composite}
+
+The decode kernel of \S\ref{sec:gqa-long} issues its local attention as
+primitive operations: it walks each PE's $S_{\text{local}}$ token slice in
+\SI{1024}{}-token tiles, and for every tile issues a $Q\!\cdot\!K^{\top}$
+\textsf{dot}, the online-softmax primitives, and a $P\!\cdot\!V$
+\textsf{dot}, merging the running $(m,\ell,O)$ state by hand. The number
+of PE\_CPU commands this costs grows with the context. At a production
+context of $S_{kv}{=}1\,\text{M}$ tokens the Case-6 64-way split gives
+each PE $S_{\text{local}}{=}16384$ tokens, so its local attention is a
+$Q\!\cdot\!K^{\top}$ of $(8,128)\!\cdot\!(128,16384)$ and a
+$P\!\cdot\!V$ of $(8,16384)\!\cdot\!(16384,128)$---sixteen hand-issued
+tiles, each a fresh batch of CPU commands.
+
+The composite command lets the kernel hand that tiling to PE\_SCHEDULER.
+We compare three command forms of the \emph{same} Case-6 kernel---identical
+placement and identical $(m,\ell,O)$ reduce, differing only in how the
+local attention is issued:
+\begin{itemize}
+  \item \textbf{primitive}---the hand-tiled \textsf{dot}/softmax kernel
+    above (the baseline of \S\ref{sec:gqa-long}).
+  \item \textbf{composite}---each matrix product is one coarse
+    \textsf{composite} GEMM over the \emph{whole} $S_{\text{local}}$, with
+    $K$ and $V$ passed as HBM references so PE\_SCHEDULER streams and
+    tiles them on the fixed $32\!\times\!64\!\times\!32$ MAC tile; the
+    softmax stays primitive.
+  \item \textbf{composite\,+\,softmax\_merge}---additionally folds the
+    online-softmax merge and $P\!\cdot\!V$ into a single stateful
+    \textsf{softmax\_merge} recipe composite.
+\end{itemize}
+
+\begin{figure}[t]
+\centering
+\includegraphics[width=\linewidth]{gqa_decode_long_ctx_composite.png}
+\caption{Three command forms of the Case-6 decode kernel, swept over
+context length ($S_{\text{local}}{=}S_{kv}/64$ per PE). \emph{Right:}
+PE\_CPU commands issued. The hand-tiled primitive kernel rises
+$O(n_{\text{tiles}})$---from 96 commands at one tile to 426 at the
+1\,M-token, sixteen-tile production point---while both composite forms
+issue a context-\emph{independent} $O(1)$ count (94 and 98) that
+\emph{saturates}: one coarse descriptor offloads the entire per-tile
+fan-out. \emph{Left:} the consequence for wall-clock latency is none---all
+three land on the same curve (\SI{30.6}{}, \SI{231}{},
+\SI{461}{\micro\second} at 8\,K\,/\,64\,K\,/\,128\,K), because decode is
+bound by streaming the KV cache, not by issue. Command-count is measured
+at emit time (exact, to 1\,M); latency on the data-mode engine over the
+tractable range.}
+\label{fig:gqa-composite}
+\end{figure}
+
+Figure~\ref{fig:gqa-composite} reads off the two quantities that matter,
+and they point in opposite directions. The PE\_CPU command count (right)
+collapses from a context-growing $O(n_{\text{tiles}})$ to a flat $O(1)$:
+at 1\,M tokens the composite form issues \num{94} commands against the
+primitive kernel's \num{426}, a $4.5\times$ reduction
+($4.2\times$ in modeled dispatch cost), and---crucially---that number no
+longer grows with context. The wall-clock latency (left), by contrast,
+is unchanged across all three forms: decode is bound by streaming the KV
+cache out of HBM, so the command form does not move the critical path.
+
+That juxtaposition is the point. The composite command is not a latency
+optimization for this memory-bound decode; it is a \emph{CPU-issue}
+optimization. Its value is removing the per-tile dispatch work that would
+otherwise grow without bound as context grows, freeing PE\_CPU to run
+ahead and keep the engines fed---which is exactly what lets the
+data-movement cost analyzed next show through as the true bottleneck
+rather than being masked by issue overhead. The \textsf{softmax\_merge}
+recipe folds the online merge into the same descriptor; on this
+memory-bound path its marginal cost over the plain GEMM composite is small
+(98 vs.\ 94 commands), and like the plain composite it keeps the issued
+count flat as context scales.
+
+\paragraph{The compute-bound mirror: prefill.} Decode's verdict---command
+form is latency-neutral---is a property of its regime, not of the
+composite command. A decode step has $T_q{=}1$, so its score and context
+products are skinny ($M{=}G\,T_q{=}8$): the MAC array is barely fed and
+the kernel is bound by streaming the KV cache. Prefill is the opposite
+corner. It processes a block of query positions at once, so $M{=}G\,T_q$
+is large and tile-filling, the GEMMs carry real arithmetic intensity
+($\sim$$M$ flops/byte, well above the roofline ridge), and the kernel is
+\emph{compute-bound}. This is the regime the composite command was built
+for (\S\ref{sec:gemm}): it streams the per-HW-tile
+DMA$\rightleftarrows$compute pipeline so the MAC array stays fed, whereas
+the primitive kernel's blocking \textsf{tl.dot} serializes each tile's
+load and compute and starves the array between tiles. We run the same
+three command forms on a single-rank compute-bound prefill (FlashAttention
+$Q$-block $\times$ $S_{kv}$-tile, online softmax) and sweep the context
+length (Figure~\ref{fig:gqa-prefill-cb}).
+
+\begin{figure}[t]
+\centering
+\includegraphics[width=\linewidth]{gqa_prefill_compute_bound.png}
+\caption{Compute-bound prefill, three command forms, swept over context
+length ($M{=}8\,T_q$ tile-filling). \emph{Left:} end-to-end latency.
+\emph{Right:} MAC utilization (achieved $\div$ the
+\SI{8}{\tera\flop\per\second} per-PE peak). The hand-tiled primitive sits
+flat at $\sim$\SI{68}{\percent}---its serial load$\to$dot path leaves the
+MAC array idle between tiles regardless of context. The composite forms
+climb with context (\SI{67}{}$\to$\SI{80}{\percent} plain,
+\SI{67}{}$\to$\SI{83}{\percent} with the recipe) because a deeper $P\!\cdot
+\!V$ reduction gives more HW tiles to pipeline, and they convert that into
+wall-clock: at \num{1024} the recipe form is \SI{646.9}{} vs.\
+\SI{794.1}{\micro\second} (\SI{19}{\percent} faster). The margin
+\emph{grows} with context---the compute-bound mirror of the GEMM result of
+\S\ref{sec:gemm}.}
+\label{fig:gqa-prefill-cb}
+\end{figure}
+
+The two studies together state the composite command's value precisely. It
+has two distinct benefits, and which one matters is set by the workload's
+roofline position. The first is \emph{host-issue offload}: one macro
+command in place of $O(n_{\text{tiles}})$ fine ones, which removes
+PE\_CPU dispatch work and is regime-independent (it shows in the decode
+command count). The second is \emph{MAC-array feeding}: the
+scheduler-internal per-tile DMA$\rightleftarrows$compute pipeline, which
+only converts to latency when the workload is compute-bound enough to have
+a MAC array worth keeping busy (it shows in the prefill utilization). A
+memory-bound decode exercises only the first; a compute-bound prefill
+exercises both. The composite command is the single mechanism that
+delivers each where it applies.
+
+% Summary moved to sections/05z-summary.tex so the three new
+% subsections (roofline, capacity planning, parallelism selection)
+% can appear between the technical results and the closing summary.
+
+% TODO: cross-regime DP (data parallelism) applicability:
+%   - Does Case-4 long-context placement compose with batch-level DP
+%     without further changes?
+%   - Does the short-context placement compose the same way?
+%   - Implications for multi-user serving (single vs. mixed regimes).

docs/report/1H-codesign-paper/sections/05z-summary.tex

@@ -0,0 +1,22 @@
+\subsection{Summary}
+\label{sec:gqa-analysis}
+
+The section's results line up into one picture of the fused kernel.
+Placement is decided by memory first and communication second: only the
+64-way splits fit production context, and among them the both-axes
+sequence shard (Case~6) minimizes the per-token collective by merging
+softmax state rather than partial scores. Command form is decided by
+roofline position: the composite command removes the per-tile issue work
+in both regimes, but converts to wall-clock only in compute-bound prefill,
+while memory-bound decode stays bound by KV streaming regardless of command
+form. The thread connecting the two is that, once composite issue makes
+GEMM cheap and leaves the MAC array idle, the fused kernel's latency is set
+almost entirely by data movement---streaming the KV cache and reducing
+partials across PEs---which is exactly the cost the communication-side
+work (on-device reduction, lazy load/compute overlap, fast TCM staging and
+torus links) is built to attack. Roofline framing
+(\S\ref{sec:roofline}) turns this observation into a sizing and sharding
+programme in the next two subsections: capacity planning
+(\S\ref{sec:capacity-planning}) and parallelism selection
+(\S\ref{sec:parallelism-selection}). What this implies for hardware
+investment across the report as a whole is taken up in the discussion.

docs/report/1H-codesign-paper/sections/06-agentic.tex

@@ -0,0 +1,309 @@
+\section{Supporting Agentic Workloads}
+\label{sec:agentic}
+
+The fused GQA kernel of \S\ref{sec:gqa} solved the efficient execution of a
+\emph{single} logical attention stream: one query stream against one KV
+cache, tiled and merged with an online softmax. Agentic inference
+introduces a different execution model, in which one request dynamically
+\emph{forks} into several cooperating reasoning branches and later
+\emph{joins} them. The purpose of this section is to show that this
+multi-stream model needs no new attention algorithm---the expensive
+hardware machinery built for \S\ref{sec:gqa} is reused unchanged---and to
+work out the resulting design across three implementation layers. The
+central claim, stated once here and defended throughout, is:
+
+\begin{quote}
+\emph{The proposed agentic execution does not change the semantics of
+transformer attention; it changes only how independent query rows are
+scheduled over a shared KV cache.}
+\end{quote}
+
+\noindent Attention stays identical; only the scheduling differs. What
+follows is a \emph{design} built on the measured \S\ref{sec:gqa} kernel; it
+is not yet a KernBench measurement, and points that go beyond the
+implemented path are flagged as such.
+
+\subsection{Motivation: from one attention stream to many}
+\label{sec:agentic-why}
+
+Agentic execution has two recurring shapes: a \emph{loop} (an agent that
+repeatedly generates, calls a tool, and continues) and a \emph{fan-out /
+fan-in} (a parent agent forks several specialised sub-agents and later
+synthesises their results). The loop is ordinary autoregressive decoding
+and needs nothing new. The fan-out is the interesting case, because the
+branches are \emph{structurally related} rather than independent:
+\[
+\begin{aligned}
+  \text{context}_i \;=\;& \underbrace{\text{shared prefix}}_{\text{common}} \\
+  &+\; \underbrace{\text{private role}_i + \text{private suffix}_i}_{\text{branch-specific}} .
+\end{aligned}
+\]
+Ordinary batching groups unrelated requests that merely arrive together;
+agentic fan-out groups requests that share an identical prefix KV cache and
+differ only in a short private suffix. That structure creates two levers
+that unrelated-request batching cannot pull:
+
+\begin{enumerate}
+  \item \textbf{Shared-prefix KV reuse.} All branches attend to the same
+  prefix KV, so it is read (and stored) once, not once per branch.
+  \item \textbf{Query-row batching.} The new query rows of many sub-agents
+  read the \emph{same} shared KV, so they can be concatenated into one
+  taller GEMM.
+\end{enumerate}
+
+\noindent Both levers land directly on the \S\ref{sec:gqa} design, which
+already multicasts a (small) query against a stationary KV cache and merges
+the result hierarchically. Fan-out simply makes the query taller, and
+fan-in adds a join step; neither touches the attention math.
+
+\subsection{Architecture: three execution layers}
+\label{sec:agentic-arch}
+
+Before the mechanics, we fix the layering, because the recurring reader
+question is ``whose job is this?''. Agentic execution divides cleanly along
+the request-flow layers already used in this report: the \textbf{agentic
+framework} decides \emph{what} to run and how to combine it, the
+\textbf{runtime} decides \emph{how} to map it onto the hardware without
+copying shared state, and the \textbf{kernel} does the actual math and
+reduction on the PEs.
+
+\[
+\text{Framework} \;\longrightarrow\; \text{Runtime}
+\;\longrightarrow\; \text{Kernel}
+\]
+
+\noindent Keeping the split this way preserves the topology-agnostic
+runtime boundary: the framework never sees the SIP/CUBE/PE hierarchy, and
+the kernel never sees the agent tree. The remainder of the section walks
+the fan-out $\rightarrow$ runtime $\rightarrow$ kernel $\rightarrow$ fan-in
+path through exactly these three layers.
+Table~\ref{tab:agentic-levels} states each layer's responsibilities; the
+kernel row is unchanged from \S\ref{sec:gqa}.
+
+\begin{table*}[t]
+\centering
+\caption{Division of responsibilities for agentic attention. Only the
+framework and runtime rows are new work; the kernel is the
+\S\ref{sec:gqa} fused GQA kernel, unchanged in its math and reduction.}
+\label{tab:agentic-levels}
+\small
+\begin{tabular}{@{}p{0.16\textwidth}p{0.40\textwidth}p{0.36\textwidth}@{}}
+\toprule
+\textbf{Level} & \textbf{Responsibilities} & \textbf{Explicitly not its job} \\
+\midrule
+\textbf{Agentic framework}
+& Manage the agent tree: fork sub-agents and assign roles; identify the
+shared prefix; mark which branches are co-schedulable (same shared prefix).
+Fan-in: enforce schema-constrained results, deduplicate/rank findings,
+hold evidence behind pointers, insert hierarchical reducers, and drive the
+main agent's final synthesis.
+& No topology, routing, KV placement, or scheduling. Sees agents and text,
+not CUBEs or PEs. \\
+\addlinespace
+\textbf{Runtime}
+& Fork the logical context by page table (shared-prefix pages $+$ private
+suffix pages) with \emph{no copy} of shared KV. Concatenate co-scheduled
+sub-agent query rows into $Q_{\text{cmb}}$. Select the execution policy
+(Stationary-KV vs.\ Distributed-Q, \S\ref{sec:agentic-choice}) from the
+cost model. Launch the
+fused GQA kernel (composite command $+$ PE\_IPCQ) and hand the batched work
+and policy to the simulation engine.
+& No attention math and no per-hop routing (delegated to the engine and
+policy); no agent-level semantics. \\
+\addlinespace
+\textbf{Kernel}
+& Multicast $Q_{\text{cmb}}$ to all PEs; per-PE local attention over the
+stationary KV shard via composite-command $Q\!\cdot\!K^{\top}$ and
+$P\!\cdot\!V$; produce per-row online-softmax state; merge that state
+hierarchically over PE\_IPCQ by matching row index (row-parallel, not
+serialised).
+& No knowledge of agents, page tables, or policy selection. Identical to
+\S\ref{sec:gqa}; a taller $Q$ is the only difference. \\
+\bottomrule
+\end{tabular}
+\end{table*}
+
+\subsection{Stationary-KV Execution Policy}
+\label{sec:agentic-policyA}
+
+The core attention operation is $Q\!\cdot\!K^{\top}$, and under fan-out
+multiple sub-agents contribute different query rows while reading a common
+$K$. Rather than issuing $Q_A\!\cdot\!K_{\text{shared}}$,
+$Q_B\!\cdot\!K_{\text{shared}}$, \dots\ as separate small GEMMs, the runtime
+concatenates the rows,
+\[
+  Q_{\text{cmb}} = \begin{bmatrix} Q_A \\ Q_B \\ \vdots \end{bmatrix},
+  \qquad
+  Q_{\text{cmb}}\!\cdot\!K_{\text{shared}},
+\]
+and issues one taller GEMM. The arithmetic is unchanged---each row is still
+independent---but the $M$ dimension grows. For a representative fan-out of
+$8$ sub-agents at $20$ new tokens each, $M = 8\times20 = 160$; with a per-PE
+KV shard of $256$ sequence positions this is a $160\times d$ by $d\times256$
+local product---an $M{=}160$, $N{=}256$ shape that sits well above the
+per-command issue overhead the composite command (\S\ref{sec:gemm}) is built
+to amortise. Fan-out is therefore especially valuable during \emph{prefill}
+of the private suffixes, where the combined $M$ is large; during decode the
+query stays short and the workload remains, as \S\ref{sec:gqa} found,
+KV-bandwidth bound.
+
+The Stationary-KV policy maps onto the \S\ref{sec:gqa} placement
+essentially unchanged: logically replicated $Q$ over a stationary,
+sequence-sharded KV cache. One
+KV head spans four CUBEs and 32 PEs, with each PE owning a different KV
+\emph{sequence} shard $K_p[S_p,d]$, $V_p[S_p,d_v]$. The combined $Q$ is
+multicast to all of them; each PE forms a complete local score
+$Q_{\text{cmb}}\!\cdot\!K_p^{\top}$ over its own shard, and the per-row
+online-softmax state is reduced hierarchically---8-PE merge inside each
+CUBE, then a 4-CUBE merge---using the same PE\_IPCQ merge primitive from
+\S\ref{sec:allreduce}.
+
+\paragraph{The reduction algorithm does not change.} This is the point to
+emphasise. Agentic batching does \emph{not} require a new reduction
+algorithm: the hierarchical online-softmax reduction of \S\ref{sec:gqa}
+remains exactly as-is, and only the query batch becomes larger. The
+reduction is always \emph{same row index across sequence shards}, never
+\emph{different query rows against each other}, so $M$ batched rows do
+\emph{not} become $M$ serialised communication rounds; the $(m,\ell,O)$ row
+states are exchanged as vectors/tiles and merged in parallel. Because the
+merge is row-indexed, a taller $Q$ widens each payload but adds no rounds.
+This is what makes agentic support a reuse of \S\ref{sec:gqa} rather than a
+redesign.
+
+\paragraph{Logical vs.\ physical $Q$ replication.} ``Replicated $Q$'' is a
+\emph{logical} statement. Because the shard axis is the KV \emph{sequence}
+dimension $S$, every PE must form the full score $Q\!\cdot\!K_p^{\top}$
+against its local $K_p$ and therefore needs $Q$'s \emph{entire} hidden
+dimension $d$; what is partitioned across PEs is $K$/$V$ along $S$, never
+$Q$ along its columns. Splitting $Q$ (and $K$) on the hidden dimension
+would instead make each PE's product \emph{partial} and force a pre-softmax
+hidden-dimension reduction ($QK^{\top}=\sum_i Q_iK_i^{\top}$)---that is
+tensor-/head-parallel attention, a different structure from the
+sequence-parallel one assumed here, and one that cannot coexist with using
+the PE axis for sequence shards. Logical replication also does not mean 32
+physical copies: $Q$ can be multicast once into a CUBE-local shared buffer
+(shared SRAM) that all PEs in the CUBE read, and a large $Q$ can further be
+\emph{row}-tiled in time ($Q[0{:}16,:],\,Q[16{:}32,:],\dots$)---row tiling
+splits the $M$ dimension, not the hidden-dimension columns. In short:
+the Stationary-KV policy uses logically replicated $Q$ across
+sequence-parallel PEs while $K$ and $V$ are partitioned along the sequence
+dimension; $Q$ may be
+temporally row-tiled or physically shared through multicast buffers, but it
+is not partitioned along the hidden-dimension columns.
+
+\paragraph{Replication is not $32\times$ the compute work (attention
+FLOPs).} Multicasting $Q$ to 32 PEs does not multiply attention FLOPs,
+because each PE computes against a different KV sequence shard rather than
+the same one. Let the KV sequence length be $S$; with sequence parallelism
+over 32 PEs, each PE owns $S/32$ positions. The score GEMM
+$Q[M,d]\!\cdot\!K^{\top}[d,S]$ costs $\propto M\,d\,S$, so each PE performs
+$M\,d\,(S/32)$ and the 32 shards sum to
+\[
+  32 \cdot M\,d\,\tfrac{S}{32} \;=\; M\,d\,S,
+\]
+identical to attention over one undivided sequence. Replication therefore
+changes $Q$ distribution, reduction traffic, buffering, and scheduling---not
+the total attention FLOPs.
+
+\subsection{Distributed-Q Execution Policy (alternative)}
+\label{sec:agentic-policyB}
+
+The natural alternative is to partition (distribute) the query rows across
+PE groups (\emph{the Distributed-Q policy}) rather than replicate them. It
+is not automatically
+better. Because the 32 PEs already shard the KV \emph{sequence}, every query
+row must still attend to \emph{all} shards; partitioning $Q$ across PE
+groups therefore forces each group to reach every KV shard, which requires
+one of: regrouping KV shards per $Q$ group, replicating KV across groups, or
+reading remote KV through symmetric memory. Each of these adds
+memory-system complexity that the Stationary-KV policy avoids entirely.
+Time-multiplexing the same PEs over $Q$ groups is the fourth option, but
+that is temporal tiling---already available under the Stationary-KV policy
+as row tiling---not true spatial $Q$ partitioning. The Distributed-Q policy
+is thus a proposed adaptive extension, not the
+baseline, and is only worth its complexity when $Q$ grows large enough that
+multicast and reduction traffic dominate remote/regrouped-KV cost.
+
+\subsection{Why the Stationary-KV policy is the baseline}
+\label{sec:agentic-choice}
+
+The choice reduces to a cost comparison the runtime can estimate. The
+Stationary-KV policy pays for $Q$ multicast and a hierarchical reduction;
+the Distributed-Q policy pays for moving or replicating KV plus extra
+scheduling:
+\[
+  T_{\text{SK}} = T_{Q\text{-mcast}} + T_{\text{local GEMM}} + T_{\text{hier.\ reduce}},
+\]
+\[
+\begin{aligned}
+  T_{\text{DQ}} = {}& T_{Q\text{-part}} + T_{\text{remote/repl.\ KV}} + T_{\text{local GEMM}} \\
+        &+ T_{\text{grp.\ reduce}} + T_{\text{remap}}.
+\end{aligned}
+\]
+For the assumed mapping (one KV head $=$ 4 CUBEs $=$ 32 PEs), the KV cache
+is large and stationary while $Q$ is comparatively small, so $T_{\text{SK}}$
+is the lower cost and the Stationary-KV policy is the recommended baseline.
+A future runtime can compute both estimates per launch---from the current
+agent count, $Q$ size, KV-head mapping, and interconnect state---and switch
+to the Distributed-Q policy only in the regime where a very large $Q$ batch
+makes multicast and reduction traffic outweigh the cost of remote or
+regrouped KV. Until that regime is measured, the Distributed-Q policy
+remains a designed, not-yet-implemented option.
+
+\subsection{Fan-in: joining sub-agent branches}
+\label{sec:agentic-fanin}
+
+Fan-out is only half of the pattern; after it, each branch produces a
+private continuation and the parent must synthesise them. We treat fan-in
+as a runtime/framework design problem with a clear optimization ladder.
+
+\paragraph{Problem.} The branch KV caches \emph{cannot} be concatenated.
+Each token's K/V depends on its full causal history, so stacking several
+branches' private KV does not form the cache of any single valid
+sequence. Join must therefore happen at the token/text level, and its cost
+is dominated by the number of join-input tokens, because the new main-agent
+KV grows in proportion to them.
+
+\paragraph{Baseline.} A naive full-text gather concatenates every branch's
+raw output: $8$ agents $\times\ 1000$ tokens $=\ 8000$ tokens pushed through
+every layer during continuation prefill---inflating prefill work, KV
+allocation, and later decode-time KV reads. This is the cost the ladder
+below drives down.
+
+\paragraph{Optimization 1 --- schema-constrained results.} Constrain each
+sub-agent to emit a compact structured result (claim, confidence, evidence
+handles) instead of free text, cutting join input by an order of magnitude
+($8\times1000 \to 8\times50$).
+
+\paragraph{Optimization 2 --- deduplication and ranking.} Overlapping
+findings across branches are merged and ranked in the framework before they
+reach the main context. This is preprocessing, not reasoning, and shrinks
+the input further without involving the main model.
+
+\paragraph{Optimization 3 --- pointer-based evidence.} Detailed evidence
+stays outside the main context behind handles, materialised only when the
+main agent actually requests it, so the effective input is summaries plus
+only the evidence truly used.
+
+\paragraph{Optimization 4 --- hierarchical reducers.} For wide fan-out,
+intermediate reducer agents summarise groups of branches, shrinking the
+final join prompt and organising fan-in traffic hierarchically---mirroring
+how the hierarchical online-softmax merge organises attention reduction.
+
+\paragraph{Future --- latent-state join.} A more aggressive step replaces
+text outputs with a few learned latent tokens, further cutting join prefill
+and KV growth. This requires training the main model to consume latent
+tokens and aligning branch representations; it is a model--system co-design
+direction, not a drop-in runtime optimization, and is out of scope here.
+
+\medskip
+\noindent Throughout, the shared prefix KV is reused by page-table
+reference and only the join suffix is new, so shortening join input is the
+primary lever on continuation-prefill cost, KV growth, and later
+decode-time KV reads. Taken together, \S\ref{sec:agentic-policyA}--\ref{sec:agentic-fanin}
+show why this workload is a natural extension of the 1H design rather than a
+new one: the decisive hardware levers of this report---cheap composite
+issue and a fast on-device reduction path---are exactly what make agentic
+fan-out efficient, and no part of the attention math is altered to get
+there. Quantifying the fan-out speedup and the fan-in join savings on
+KernBench is left as measured 2H work.

docs/report/1H-codesign-paper/sections/07-hw-spec-search.tex

@@ -0,0 +1,67 @@
+\section{Hardware Performance-Spec Search for GQA}
+\label{sec:hwspec}
+
+\emph{Work in progress --- this section states the study's intent and
+method; experimental data, figures, and conclusions are not yet
+available and will be added in a later revision.}
+
+The studies so far fixed the modeled hardware (\S\ref{sec:hw}) and
+varied the \emph{software}: the composite command, the reduction path, and
+the KV placement. This section inverts the question. Given the fused GQA
+kernel of \S\ref{sec:gqa} as the target workload, \emph{which hardware
+specification best serves it}, and where does spending more silicon stop
+paying off? The discussion of \S\ref{sec:discussion} already gives a
+strong prior---attention is data-movement bound, so raw MAC throughput is
+not the binding constraint---but that conclusion was drawn at one operating
+point. A systematic sweep is needed to turn it into a defensible
+performance-spec recommendation across context lengths, agent counts, and
+sharding regimes.
+
+\subsection{Sweep axes}
+\label{sec:hwspec-axes}
+
+Four hardware knobs are varied, spanning the compute side and both levels
+of the interconnect so that the compute/communication balance can be
+located rather than assumed. The KernBench cost model already exposes each
+as a parameter, so the sweep reuses the same deterministic engine as the
+rest of the report; no production change is required.
+
+\begin{table}[h]
+\centering
+\caption{Hardware knobs varied in the performance-spec search. Ranges and
+step counts are to be finalised with the experiment.}
+\label{tab:hwspec-axes}
+\small
+\begin{tabular}{@{}p{0.30\textwidth}p{0.14\textwidth}@{}}
+\toprule
+\textbf{Knob} & \textbf{Axis} \\
+\midrule
+GEMM throughput (TFLOPS) & compute \\
+MATH-engine ALU count & compute \\
+CUBE-to-CUBE (die-to-die) bandwidth & interconnect \\
+SIP-to-SIP (card-to-card) bandwidth & interconnect \\
+\bottomrule
+\end{tabular}
+\end{table}
+
+The first two axes scale the on-PE compute engines; the latter two scale
+the two inter-device links that carry the KV reduction and cross-device
+softmax merge. Sweeping them jointly---rather than one at a time---is what
+exposes the interactions: for example, whether added die-to-die bandwidth
+only helps once card-to-card bandwidth is also raised, or whether GEMM
+throughput is genuinely inert for attention once issue overhead is removed.
+
+\subsection{Method and target output}
+\label{sec:hwspec-method}
+
+The intended procedure is a joint sweep over the four axes, running the
+fused GQA kernel at representative decode and prefill configurations
+(including the agentic fan-out shapes of \S\ref{sec:agentic}), and reading
+latency and per-engine busy time from the engine's completion timestamps
+and op log---the same measurement path used in \S\ref{sec:gqa}. The target
+deliverables are: (i) the sensitivity of GQA latency to each knob in
+isolation, (ii) the Pareto frontier of latency against a simple
+area/cost proxy, and (iii) a recommended balanced specification---the knob
+combination past which further investment does not move GQA latency. These
+results are the natural quantitative counterpart to the qualitative
+ranking in \S\ref{sec:discussion}, and completing them is a 2H objective.

docs/report/1H-codesign-paper/sections/10-future-work.tex

@@ -1,4 +1,4 @@
-\section{Future Work --- 2H}
+\section{Future Work}
 \label{sec:future}
 
 The 1H work covered attention end to end. The natural next step is to
@@ -30,11 +30,22 @@ vs.\ sequence parallelism---under a single, software-stack-independent
 model, so the interconnect and memory implications of each choice are
 measured rather than assumed.
 
-\paragraph{Agentic workloads.} Agentic inference interleaves many
+\paragraph{Agentic workloads: from design to implementation.}
-short, bursty decode requests with tool use and long shared contexts,
+Section~\ref{sec:agentic} established \emph{how} the fused GQA design
-which stresses the system differently from a single long generation:
+extends to agentic fan-out/fan-in and \emph{which} responsibilities fall to
-context reuse across requests, dynamic batching, and uneven expert load all
+the framework, the runtime, and the kernel. The 2H step is no longer to
-change how compute and data should be dispersed. Characterizing how total
+analyse the workload but to \emph{build} that support: the detailed design
-compute and data movement distribute across the SIP/CUBE/PE hierarchy under
+and implementation of all three levels. This is necessary work rather than
-such workloads---and which of the 1H hardware levers still dominate when the
+optional, because today's agentic frameworks are effectively all
-workload is this irregular---is the broader 2H agenda.
+closed-source---there is no open substrate that exposes shared-prefix KV
+reuse, cross-agent query batching, and schema/pointer-based fan-in down to
+the hardware. Concretely, 2H targets: a \textbf{framework} layer that
+manages the agent tree, identifies co-schedulable branches, and enforces
+compact fan-in; a \textbf{runtime} layer that forks logical contexts by
+page table without copying shared KV, batches sub-agent query rows, and
+selects the execution policy; and a \textbf{kernel} layer that runs the
+batched replicated-Q attention and hierarchical online-softmax merge over
+PE\_IPCQ. Bringing these up on KernBench turns the
+Section~\ref{sec:agentic} design into a measured, end-to-end agentic path
+and lets us confirm which of the 1H hardware levers still dominate once the
+full framework/runtime/kernel stack is in place.

docs/report/1H-codesign-paper/toc.md

@@ -45,17 +45,28 @@
 6. **Fused Grouped Query Attention** (composite + PE_IPCQ) — `sections/05-gqa.tex`
    necessity · design · results · analysis. (+ GQA seq/head/user configs)
 
-7. **Discussion** — `sections/06-discussion.tex`
+7. **Supporting Agentic Workloads** — `sections/06-agentic.tex`
+   How the §6 GQA design extends to agentic fan-out/fan-in: shared-prefix KV
+   reuse, batched replicated-Q attention, schema/pointer-based fan-in.
+   Implementation split across three levels (agentic framework / runtime /
+   kernel) with each component's responsibilities. Design, not yet measured.
+
+8. **Hardware Performance-Spec Search for GQA** — `sections/07-hw-spec-search.tex`
+   *Work in progress.* Joint sweep of GEMM TFLOPS, MATH-engine ALU count,
+   CUBE↔CUBE (die-to-die) BW, SIP↔SIP (card-to-card) BW to find the balanced
+   spec for GQA. Intent + method only; data/figures/conclusions deferred.
+
+9. **Discussion** — `sections/08-discussion.tex`
    Which HW changes are meaningful, and under what regimes (cross-cutting).
 
-8. **Conclusion** — `sections/07-conclusion.tex`
+10. **Conclusion** — `sections/09-conclusion.tex`
-   The codesign thesis, stated plainly, supported by the measured results.
+    The codesign thesis, stated plainly, supported by the measured results.
 
-9. **Future Work — 2H** — `sections/08-future-work.tex`
+11. **Future Work — 2H** — `sections/10-future-work.tex`
-   Add the FFN/MoE layer toward full LLM decoding; how compute & data should
+    Add the FFN/MoE layer toward full LLM decoding; how compute & data should
-   be distributed for agentic / MoE workloads.
+    be distributed for agentic / MoE workloads.
 
-10. **References** *(optional)* — external literature only (FlashAttention,
+12. **References** *(optional)* — external literature only (FlashAttention,
     Megatron-LM, GPT-3, Llama 3). No ADR/SPEC entries.
 
 ## Per-section structure (§4/§5/§6)

docs/sweeps/decode_batch_sweep.csv

@@ -0,0 +1,12 @@
+mode,kv_per_cube,C,S_kv,B,capacity,agg_latency_us,mean_user_latency_us,throughput_users_per_us
+A1,1,8,8192,1,2,2.2471645000005376,2.2471645000005376,0.44500524994932983
+A1,1,8,8192,2,2,2.3131845000012543,2.2471645000012357,0.864608940617973
+A2,2,4,8192,1,4,3.3555127500005475,3.3555127500005475,0.2980170467240325
+A2,2,4,8192,2,4,3.3935227500010514,3.355512750001042,0.5893580645656141
+A2,2,4,8192,4,4,3.4695427500010703,3.3555127500011586,1.1528896711241752
+A4,4,2,8192,1,8,5.586135500000557,5.586135500000557,0.17901463364071643
+A4,4,2,8192,2,8,5.602295500000823,5.586135500000673,0.35699652044411195
+A4,4,2,8192,8,8,5.783285499986028,5.586135499984957,1.3832967436968704
+B,8,1,8192,1,16,10.816091000000947,10.816091000000947,0.09245484343649775
+B,8,1,8192,2,16,10.848411000001535,10.816091000001236,0.18435879687815265
+B,8,1,8192,16,16,11.38492099986691,10.816090999974403,1.4053676788962384

docs/sweeps/plot_scripts/plot_batch_scaling.py

@@ -0,0 +1,69 @@
+"""Batch-scaling figure for concurrent decode users on one SIP.
+
+Reads ``docs/sweeps/decode_batch_sweep.csv`` (from
+``tests/attention/test_gqa_decode_batch_sweep.py``) and plots aggregate
+throughput (requests / us) versus batch size B for the four KV mappings,
+one panel per context length. Each mapping's curve runs up to its SIP
+capacity 16/C (A1=2 ... B=16 users).
+"""
+from pathlib import Path
+
+import csv
+
+import matplotlib.pyplot as plt
+
+ROOT = Path(__file__).resolve().parents[3]
+CSV = ROOT / "docs" / "sweeps" / "decode_batch_sweep.csv"
+OUT = (ROOT / "docs" / "report" / "1H-codesign-paper"
+       / "figures" / "gqa_short_context" / "batch_scaling.png")
+
+MODE_LABEL = {"A1": "1-kv-per-cube", "A2": "2-kv-per-cube",
+              "A4": "4-kv-per-cube", "B": "8-kv-per-cube"}
+MODE_COLOR = {"A1": "#1f77b4", "A2": "#2ca02c",
+              "A4": "#ffd000", "B": "#d62728"}
+MODES = ["A1", "A2", "A4", "B"]
+
+
+def main():
+    rows = list(csv.DictReader(CSV.open()))
+    contexts = sorted({int(r["S_kv"]) for r in rows})
+    fig, axes = plt.subplots(1, len(contexts), figsize=(6.5 * len(contexts), 4.5),
+                             squeeze=False)
+    xmax = max(int(r["B"]) for r in rows)  # largest SIP capacity (= 16)
+    for col, S in enumerate(contexts):
+        ax = axes[0][col]
+        for m in MODES:
+            pts = sorted(
+                ((int(r["B"]), float(r["throughput_users_per_us"]))
+                 for r in rows if r["mode"] == m and int(r["S_kv"]) == S),
+                key=lambda t: t[0],
+            )
+            if not pts:
+                continue
+            xs, ys = zip(*pts)
+            # Measured: concurrent users up to the SIP capacity 16/C.
+            ax.plot(xs, ys, marker="o", color=MODE_COLOR[m],
+                    label=MODE_LABEL[m])
+            # Beyond capacity the SIP is full: extra users run in waves at
+            # the saturated rate, so throughput plateaus. Draw that ceiling
+            # as a dashed extension (projected, not concurrently measured).
+            if xs[-1] < xmax:
+                ax.plot([xs[-1], xmax], [ys[-1], ys[-1]],
+                        linestyle="--", color=MODE_COLOR[m], alpha=0.55)
+            ax.annotate(f"{ys[-1]:.2f}", (xs[-1], ys[-1]),
+                        textcoords="offset points", xytext=(4, 4), fontsize=8)
+        ax.set_title(f"S_kv = {S // 1024}K")
+        ax.set_xlabel("batch size B (concurrent users)")
+        ax.set_ylabel("throughput (requests / µs)")
+        ax.grid(True, alpha=0.3)
+        ax.legend(fontsize=9)
+    fig.suptitle("Batch scaling: aggregate throughput vs concurrent users "
+                 "(one SIP)", fontsize=12, fontweight="bold")
+    fig.tight_layout()
+    OUT.parent.mkdir(parents=True, exist_ok=True)
+    fig.savefig(OUT, dpi=140, bbox_inches="tight")
+    print(f"  ✓ {OUT}")
+
+
+if __name__ == "__main__":
+    main()

docs/sweeps/plot_scripts/plot_short_context.py

@@ -0,0 +1,231 @@
+"""Generate 6 comparison figures for the short-context GQA attention sweep.
+
+Sweep matrix: 3 variants × 2 phases × 4 modes × 4 contexts (8K/16K/32K/64K).
+
+Figures
+  Wall-clock latency:
+    1. wall_variant1_mode_compare.png  — (1) without composite
+    2. wall_variant2_mode_compare.png  — (2) with composite (GEMM-only)
+    3. wall_variant3_mode_compare.png  — (3) with composite + softmax_merge
+    4. wall_a1_variant_compare.png     — 1-kv-per-cube: variant comparison
+
+  Mode trade-off + composite ablation:
+    5. per_cube_tradeoff_mode_compare.png      — HBM BW + IPCQ (variant-invariant)
+    6. gemm_util_a1_variant_ablation.png       — 1-kv-per-cube GEMM util
+
+Output: docs/report/1H-codesign-paper/figures/gqa_short_context/
+"""
+from __future__ import annotations
+
+import csv
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+ROOT = Path(__file__).resolve().parents[3]
+SWEEPS = ROOT / "docs" / "sweeps"
+OUT = (ROOT / "docs" / "report" / "1H-codesign-paper"
+       / "figures" / "gqa_short_context")
+OUT.mkdir(parents=True, exist_ok=True)
+
+CONTEXTS = [8192, 16384, 32768, 65536]
+CONTEXT_LABELS = ["8K", "16K", "32K", "64K"]
+MODES = ["A1", "A2", "A4", "B"]
+MODE_LABEL = {
+    "A1": "1-kv-per-cube",
+    "A2": "2-kv-per-cube",
+    "A4": "4-kv-per-cube",
+    "B":  "8-kv-per-cube",
+}
+MODE_COLOR = {
+    MODE_LABEL["A1"]: "#1f77b4",
+    MODE_LABEL["A2"]: "#2ca02c",
+    MODE_LABEL["A4"]: "#ffd000",
+    MODE_LABEL["B"]:  "#d62728",
+}
+VARIANTS = [
+    ("baseline",        "(1) without composite"),
+    ("composite",       "(2) with composite (GEMM-only)"),
+    ("composite_fused", "(3) with composite + softmax_merge"),
+]
+VARIANT_COLOR = {
+    "baseline":        "#1f77b4",
+    "composite":       "#2ca02c",
+    "composite_fused": "#d62728",
+}
+
+CSV_MAP = {
+    ("decode",  "baseline"):         SWEEPS / "short_context_decode_sweep.csv",
+    ("decode",  "composite"):        SWEEPS / "short_context_decode_composite_sweep.csv",
+    ("decode",  "composite_fused"):  SWEEPS / "short_context_decode_composite_fused_sweep.csv",
+    ("prefill", "baseline"):         SWEEPS / "short_context_prefill_sweep.csv",
+    ("prefill", "composite"):        SWEEPS / "short_context_prefill_composite_sweep.csv",
+    ("prefill", "composite_fused"):  SWEEPS / "short_context_prefill_composite_fused_sweep.csv",
+}
+
+
+def load(phase, variant):
+    """Return {(mode, S_kv): row dict}."""
+    rows = {}
+    with CSV_MAP[(phase, variant)].open() as f:
+        for r in csv.DictReader(f):
+            rows[(r["mode"], int(r["S_kv"]))] = r
+    return rows
+
+
+def grouped_bars(ax, *, x_labels, groups, group_color, ylabel, title,
+                 y_log=False, value_fmt="{:.3g}"):
+    """Grouped bars: x = x_labels, groups = list of (label, values)."""
+    n = len(groups)
+    nx = len(x_labels)
+    width = 0.8 / n
+    x = np.arange(nx)
+    for i, (label, values) in enumerate(groups):
+        offset = (i - (n - 1) / 2) * width
+        bars = ax.bar(x + offset, values, width, label=label,
+                      color=group_color[label], edgecolor="black",
+                      linewidth=0.4)
+        for bar, v in zip(bars, values):
+            if v <= 0 and y_log:
+                continue
+            ax.text(bar.get_x() + bar.get_width() / 2, bar.get_height(),
+                    value_fmt.format(v), ha="center", va="bottom",
+                    fontsize=7)
+    ax.set_xticks(x)
+    ax.set_xticklabels(x_labels)
+    ax.set_xlabel("S_kv")
+    ax.set_ylabel(ylabel)
+    ax.set_title(title)
+    if y_log:
+        ax.set_yscale("log")
+    ax.grid(True, axis="y", alpha=0.3, which="both")
+    ax.legend(fontsize=8, loc="upper left")
+
+
+# ── 1-3. Wall mode comparison per variant ──────────────────────────
+
+def fig_wall_mode_compare(variant_key, suptitle, fname):
+    fig, axes = plt.subplots(1, 2, figsize=(13, 4.8))
+    for ax, phase in zip(axes, ("prefill", "decode")):
+        rows = load(phase, variant_key)
+        groups = [
+            (MODE_LABEL[m], [float(rows[(m, s)]["wall_us"]) for s in CONTEXTS])
+            for m in MODES
+        ]
+        grouped_bars(ax, x_labels=CONTEXT_LABELS, groups=groups,
+                     group_color=MODE_COLOR,
+                     ylabel="Wall-clock latency (μs)",
+                     title=phase.capitalize(),
+                     y_log=True, value_fmt="{:.0f}")
+    fig.suptitle(suptitle, fontsize=12, fontweight="bold")
+    fig.tight_layout()
+    out = OUT / fname
+    fig.savefig(out, dpi=140, bbox_inches="tight")
+    plt.close(fig)
+    print(f"  ✓ {out}")
+
+
+# ── 4. Wall 1-kv-per-cube variant comparison ───────────────────────
+
+def fig_wall_a1_variant_compare(fname):
+    fig, axes = plt.subplots(1, 2, figsize=(13, 4.8))
+    for ax, phase in zip(axes, ("prefill", "decode")):
+        groups = []
+        for variant_key, _ in VARIANTS:
+            rows = load(phase, variant_key)
+            values = [float(rows[("A1", s)]["wall_us"]) for s in CONTEXTS]
+            groups.append((variant_key, values))
+        grouped_bars(ax, x_labels=CONTEXT_LABELS, groups=groups,
+                     group_color=VARIANT_COLOR,
+                     ylabel="Wall-clock latency (μs)",
+                     title=phase.capitalize(),
+                     y_log=True, value_fmt="{:.0f}")
+        ax.legend([v[1] for v in VARIANTS], fontsize=8, loc="upper left")
+    fig.suptitle("1-kv-per-cube: variant comparison",
+                 fontsize=12, fontweight="bold")
+    fig.tight_layout()
+    out = OUT / fname
+    fig.savefig(out, dpi=140, bbox_inches="tight")
+    plt.close(fig)
+    print(f"  ✓ {out}")
+
+
+# ── 5. Per-cube trade-off (HBM BW + IPCQ, mode compare) ────────────
+
+def fig_per_cube_tradeoff(fname):
+    fig, axes = plt.subplots(2, 2, figsize=(13, 9))
+    metrics = [
+        ("kv_cache_per_cube_mb", "Per-cube KV footprint (MB)", True, "{:.0f}"),
+        ("ipcq_kb",              "IPCQ traffic (KB)",          True, "{:.1f}"),
+    ]
+    for row, (key, ylabel, log, fmt) in enumerate(metrics):
+        for col, phase in enumerate(("prefill", "decode")):
+            ax = axes[row][col]
+            rows = load(phase, "baseline")
+            groups = [
+                (MODE_LABEL[m], [float(rows[(m, s)][key]) for s in CONTEXTS])
+                for m in MODES
+            ]
+            grouped_bars(ax, x_labels=CONTEXT_LABELS, groups=groups,
+                         group_color=MODE_COLOR, ylabel=ylabel,
+                         title=phase.capitalize(),
+                         y_log=log, value_fmt=fmt)
+    fig.suptitle("Per-cube trade-off: KV footprint + IPCQ mode comparison",
+                 fontsize=12, fontweight="bold")
+    fig.tight_layout()
+    out = OUT / fname
+    fig.savefig(out, dpi=140, bbox_inches="tight")
+    plt.close(fig)
+    print(f"  ✓ {out}")
+
+
+# ── 6. GEMM util ablation (1-kv-per-cube, variant compare) ─────────
+
+def fig_gemm_util_a1_ablation(fname):
+    fig, axes = plt.subplots(1, 2, figsize=(13, 5.4))
+    for ax, phase in zip(axes, ("prefill", "decode")):
+        groups = []
+        for variant_key, _ in VARIANTS:
+            rows = load(phase, variant_key)
+            values = [float(rows[("A1", s)]["gemm_util"]) for s in CONTEXTS]
+            groups.append((variant_key, values))
+        grouped_bars(ax, x_labels=CONTEXT_LABELS, groups=groups,
+                     group_color=VARIANT_COLOR,
+                     ylabel="GEMM engine utilization (per PE)",
+                     title=phase.capitalize(),
+                     y_log=False, value_fmt="{:.4f}")
+        ax.legend([v[1] for v in VARIANTS], fontsize=8, loc="upper left")
+    fig.suptitle("1-kv-per-cube: GEMM utilization variant comparison",
+                 fontsize=12, fontweight="bold")
+    fig.text(
+        0.5, -0.04,
+        "⚠ Higher gemm_util on composite/fused is largely supertile "
+        "padding overhead (M=G=8 padded to TILE_M=32), not pure fusion "
+        "gain. See ADR-0070 limitation #4.",
+        ha="center", fontsize=9, style="italic", color="#555")
+    fig.tight_layout()
+    out = OUT / fname
+    fig.savefig(out, dpi=140, bbox_inches="tight")
+    plt.close(fig)
+    print(f"  ✓ {out}")
+
+
+if __name__ == "__main__":
+    print(f"Output: {OUT}")
+    fig_wall_mode_compare(
+        "baseline",
+        "(1) Without composite: latency mode comparison",
+        "wall_variant1_mode_compare.png")
+    fig_wall_mode_compare(
+        "composite",
+        "(2) With composite (GEMM-only): latency mode comparison",
+        "wall_variant2_mode_compare.png")
+    fig_wall_mode_compare(
+        "composite_fused",
+        "(3) With composite + softmax_merge: latency mode comparison",
+        "wall_variant3_mode_compare.png")
+    fig_wall_a1_variant_compare("wall_a1_variant_compare.png")
+    fig_per_cube_tradeoff("per_cube_tradeoff_mode_compare.png")
+    fig_gemm_util_a1_ablation("gemm_util_a1_variant_ablation.png")
+    print("\nDone.")

docs/sweeps/short_context_decode_composite_fused_sweep.csv

@@ -0,0 +1,17 @@
+variant,mode,kv_per_cube,C,S_kv,wall_us,n_pe,gemm_util,math_count,math_pipeline_us,hbm_bw_util,hbm_read_mb,hbm_write_kb,ipcq_kb,kv_cache_per_cube_mb
+composite_fused,A1,1,8,8192,2.2471645000005376,64,0.11665545624270428,656,0.0,0.457688,16.0625,8.0,57.75,2.0
+composite_fused,A1,1,8,16384,4.090972500001895,64,0.4907507933676371,1296,13.903615999992937,0.501714,32.0625,8.0,57.75,4.0
+composite_fused,A1,1,8,32768,6.65601250000007,64,0.8261186408347434,2576,41.71084799996763,0.61606,64.0625,8.0,57.75,8.0
+composite_fused,A1,1,8,65536,11.786092499984429,64,1.05893212704674,5136,97.32531200143696,0.695438,128.0625,8.0,57.75,16.0
+composite_fused,A2,2,4,8192,3.8395887500001407,32,0.5228809986460942,624,6.951807999995537,0.534693,16.03125,8.0,24.75,4.0
+composite_fused,A2,2,4,16384,6.404628750001605,32,0.858544064722257,1264,20.855423999989405,0.640318,32.03125,8.0,24.75,8.0
+composite_fused,A2,2,4,32768,11.534708749999874,32,1.0820101547616419,2544,48.66265599996224,0.710638,64.03125,8.0,24.75,16.0
+composite_fused,A2,2,4,65536,21.794868749984772,32,1.21334541192213,5104,104.27712000153959,0.751966,128.03125,8.0,24.75,32.0
+composite_fused,A4,4,2,8192,5.915787499999919,16,0.9294884239792724,608,10.427711999993306,0.693399,16.015625,8.0,8.25,8.0
+composite_fused,A4,4,2,16384,11.045867500001041,16,1.1298951395303998,1248,24.33132799998764,0.742178,32.015625,8.0,8.25,16.0
+composite_fused,A4,4,2,32768,21.306027499999384,16,1.2411841672219757,2528,52.13855999995954,0.769266,64.015625,8.0,8.25,32.0
+composite_fused,A4,4,2,65536,41.82634749998455,16,1.2999645260103314,5088,107.7530240015909,0.783573,128.015625,8.0,8.25,64.0
+composite_fused,B,8,1,8192,10.856904999999854,8,1.149560763397397,600,12.16566399999219,0.75528,16.0078125,8.0,0.0,16.0
+composite_fused,B,8,1,16384,21.117065000001226,8,1.2522906947682175,1240,26.069279999986758,0.776244,32.0078125,8.0,0.0,32.0
+composite_fused,B,8,1,32768,41.637384999999774,8,1.3058641410537761,2520,53.8765119999582,0.787177,64.0078125,8.0,0.0,64.0
+composite_fused,B,8,1,65536,82.67802499997313,8,1.33323087973183,5080,109.49097600161657,0.792762,128.0078125,8.0,0.0,128.0

docs/sweeps/short_context_decode_composite_sweep.csv

@@ -0,0 +1,17 @@
+variant,mode,kv_per_cube,C,S_kv,wall_us,n_pe,gemm_util,math_count,math_pipeline_us,hbm_bw_util,hbm_read_mb,hbm_write_kb,ipcq_kb,kv_cache_per_cube_mb
+composite,A1,1,8,8192,2.532214500000584,64,0.33402541530150626,656,0.0,0.406166,16.0625,8.0,57.75,2.0
+composite,A1,1,8,16384,4.176996500001463,64,0.40499148132408613,1360,0.0,0.491382,32.0625,8.0,57.75,4.0
+composite,A1,1,8,32768,7.466560500002,64,0.4531264428807111,2768,0.0,0.549182,64.0625,8.0,57.75,8.0
+composite,A1,1,8,65536,14.045688499950572,64,0.4817558071541797,5584,0.0,0.58356,128.0625,8.0,57.75,16.0
+composite,A2,2,4,8192,3.9256127500006404,32,0.4309258471789165,656,0.0,0.522976,16.03125,8.0,24.75,4.0
+composite,A2,2,4,16384,7.215176750001498,32,0.4689138072801237,1360,0.0,0.568385,32.03125,8.0,24.75,8.0
+composite,A2,2,4,32768,13.794304750002688,32,0.4905351971318266,2768,0.0,0.594231,64.03125,8.0,24.75,16.0
+composite,A2,2,4,65536,26.95256074991636,32,0.5021112511805295,5584,0.0,0.608068,128.03125,8.0,24.75,32.0
+composite,A4,4,2,8192,6.726335500000743,16,0.5029924540606681,656,0.0,0.609842,16.015625,8.0,8.25,8.0
+composite,A4,4,2,16384,13.305463500001585,16,0.5085574057667106,1360,0.0,0.616138,32.015625,8.0,8.25,16.0
+composite,A4,4,2,32768,26.463719500003965,16,0.5113863151276257,2768,0.0,0.619338,64.015625,8.0,8.25,32.0
+composite,A4,4,2,65536,52.78023149984703,16,0.5128126048745716,5584,0.0,0.620952,128.015625,8.0,8.25,64.0
+composite,B,8,1,8192,13.096491000001318,8,0.5166721375947838,656,0.0,0.626122,16.0078125,8.0,0.0,16.0
+composite,B,8,1,16384,26.254747000003057,8,0.5154566524738308,1360,0.0,0.624344,32.0078125,8.0,0.0,32.0
+composite,B,8,1,32768,52.57125900000788,8,0.5148510519664782,2768,0.0,0.623459,64.0078125,8.0,0.0,64.0
+composite,B,8,1,65536,105.20428299969761,8,0.514548785079354,5584,0.0,0.623016,128.0078125,8.0,0.0,128.0

docs/sweeps/short_context_decode_sweep.4ctx.csv

@@ -0,0 +1,17 @@
+mode,kv_per_cube,C,S_kv,wall_us,n_pe,gemm_util,math_count,hbm_bw_util,hbm_read_mb,hbm_write_kb,ipcq_kb,kv_cache_per_cube_mb
+A1,1,8,8192,149.2685345000007,64,0.001756190618994587,656,0.013780533230866485,16.0625,8.0,57.75,2.0
+A1,1,8,16384,282.2252665000025,64,0.0018576933472394743,1360,0.014545118695104373,32.0625,8.0,57.75,4.0
+A1,1,8,32768,548.138730500006,64,0.0019129755692392889,2768,0.014961540835691614,64.0625,8.0,57.75,8.0
+A1,1,8,65536,1080.073058500367,64,0.0019416760593127277,5584,0.01517767698303756,128.0625,8.0,57.75,16.0
+A2,2,4,8192,142.16883275000066,32,0.0036877843748065785,656,0.028881154333033524,16.03125,8.0,24.75,4.0
+A2,2,4,16384,277.0102967500018,32,0.0037853322143696724,1360,0.029609007665885423,32.03125,8.0,24.75,8.0
+A2,2,4,32768,546.8006247504073,32,0.0038353138330044275,2768,0.02998167752182655,64.03125,8.0,24.75,16.0
+A2,2,4,65536,1087.0962007518285,32,0.003858263874988177,5584,0.03015188534127058,128.03125,8.0,24.75,32.0
+A4,4,2,8192,142.21417550000066,16,0.007373217165681769,656,0.05768763888097753,16.015625,8.0,8.25,8.0
+A4,4,2,16384,280.8788035001758,16,0.007466394665122732,1360,0.058373931374247456,32.015625,8.0,8.25,16.0
+A4,4,2,32768,558.5655195010398,16,0.007509063580845671,2768,0.058686042828569145,64.015625,8.0,8.25,32.0
+A4,4,2,65536,1114.319951501857,16,0.007528006645388841,5584,0.05882332081702009,128.015625,8.0,8.25,64.0
+B,8,1,8192,147.93485600006025,8,0.014176185766516818,656,0.11085960701508589,16.0078125,8.0,0.0,16.0
+B,8,1,16384,294.3171420004266,8,0.014250967413897551,1360,0.1113900460475132,32.0078125,8.0,0.0,32.0
+B,8,1,32768,587.2722140010417,8,0.01428401991446495,2768,0.11162115018757507,64.0078125,8.0,0.0,64.0
+B,8,1,65536,1173.1823580017838,8,0.014300603725891686,5584,0.11173710472707321,128.0078125,8.0,0.0,128.0

docs/sweeps/short_context_decode_sweep.csv

@@ -0,0 +1,17 @@
+mode,kv_per_cube,C,S_kv,wall_us,n_pe,gemm_util,math_count,hbm_bw_util,hbm_read_mb,hbm_write_kb,ipcq_kb,kv_cache_per_cube_mb
+A1,1,8,8192,2.2471645000005376,64,0.11665545624270428,656,0.457688,16.0625,8.0,57.75,2.0
+A1,1,8,16384,3.6068965000013704,64,0.14535709577464812,1360,0.569049,32.0625,8.0,57.75,4.0
+A1,1,8,32768,6.326360500001814,64,0.1657471147905734,2768,0.648161,64.0625,8.0,57.75,8.0
+A1,1,8,65536,11.765288499959512,64,0.17824909265982342,5584,0.696668,128.0625,8.0,57.75,16.0
+A2,2,4,8192,3.3555127500005475,32,0.15624676139283236,656,0.611829,16.03125,8.0,24.75,4.0
+A2,2,4,16384,6.074976750001311,32,0.17260576347056114,1360,0.675064,32.03125,8.0,24.75,8.0
+A2,2,4,32768,11.513904750002315,32,0.18214081543451005,2768,0.711922,64.03125,8.0,24.75,16.0
+A2,2,4,65536,22.391760749934242,32,0.18731461303283142,5584,0.731921,128.03125,8.0,24.75,32.0
+A4,4,2,8192,5.586135500000557,16,0.1877104484844272,656,0.734318,16.015625,8.0,8.25,8.0
+A4,4,2,16384,11.025063500001211,16,0.19021677290108474,1360,0.743578,32.015625,8.0,8.25,16.0
+A4,4,2,32768,21.902919500003218,16,0.19149520227204342,2768,0.748302,64.015625,8.0,8.25,32.0
+A4,4,2,65536,43.65863149988279,16,0.19214088284049563,5584,0.750688,128.015625,8.0,8.25,64.0
+B,8,1,8192,10.816091000000947,8,0.1938918598225168,656,0.75813,16.0078125,8.0,0.0,16.0
+B,8,1,16384,21.693947000002314,8,0.1933398288471476,1360,0.755602,32.0078125,8.0,0.0,32.0
+B,8,1,32768,43.449659000006385,8,0.19306499045255035,2768,0.754344,64.0078125,8.0,0.0,64.0
+B,8,1,65536,86.96108299976913,8,0.19292786406575965,5584,0.753716,128.0078125,8.0,0.0,128.0

docs/sweeps/short_context_prefill_composite_fused_sweep.csv

@@ -0,0 +1,17 @@
+variant,mode,kv_per_cube,C,S_kv,wall_us,n_pe,gemm_util,math_count,math_pipeline_us,hbm_bw_util,hbm_read_mb,hbm_write_kb,ipcq_kb,kv_cache_per_cube_mb
+composite_fused,A1,1,8,8192,35.53158299944201,64,0.3512557265116541,4800,97.32531200143696,0.0290446,16.0625,64.0,114688.0,2.0
+composite_fused,A1,1,8,16384,63.02314799463935,64,0.4196030322618738,9920,208.55424000307917,0.0326229,32.0625,64.0,229376.0,4.0
+composite_fused,A1,1,8,32768,118.00299802135117,64,0.4607744625054221,20160,431.01209600543973,0.0347788,64.0625,64.0,458752.0,8.0
+composite_fused,A1,1,8,65536,227.95666800683736,64,0.4835519703723695,40640,875.9278079059123,0.0359717,128.0625,64.0,917504.0,16.0
+composite_fused,A2,2,4,8192,35.27508350025257,32,0.5747630902258046,2400,96.40870400077105,0.0585116,16.0625,64.0,49152.0,4.0
+composite_fused,A2,2,4,16384,65.58160349091561,32,0.653336388872743,4960,206.59008000165224,0.0627005,32.0625,64.0,98304.0,8.0
+composite_fused,A2,2,4,32768,126.19361352364719,32,0.69726913692688,10080,426.9528320015669,0.0650429,64.0625,64.0,196608.0,16.0
+composite_fused,A2,2,4,65536,247.42072338639758,32,0.7205501525262498,20320,867.678335950613,0.0662839,128.0625,64.0,393216.0,32.0
+composite_fused,A4,4,2,8192,45.62825600103196,16,0.7872118539828873,1200,92.28038400042057,0.0904703,16.0625,64.0,16384.0,8.0
+composite_fused,A4,4,2,16384,87.81637548340065,16,0.8628360665552867,2480,197.74368000090124,0.09365,32.0625,64.0,32768.0,16.0
+composite_fused,A4,4,2,32768,172.19397554247547,16,0.9029073142513345,5040,408.6702720000148,0.0953343,64.0625,64.0,65536.0,32.0
+composite_fused,A4,4,2,65536,340.9491752808783,16,0.9235491703853554,10160,830.5234559737444,0.096202,128.0625,64.0,131072.0,64.0
+composite_fused,B,8,1,8192,68.33786600000411,8,2.036048828365075,712,121.6138560003899,0.120811,16.0625,64.0,0.0,16.0
+composite_fused,B,8,1,16384,133.87386600000133,8,2.209234414765302,1480,260.6011200008355,0.122862,32.0625,64.0,0.0,32.0
+composite_fused,B,8,1,32768,264.9458659999771,8,2.2985744263323977,3016,538.575647996068,0.12392,64.0625,64.0,0.0,64.0
+composite_fused,B,8,1,65536,527.0898659999762,8,2.3439567929361727,6088,1094.5247039788662,0.124457,128.0625,64.0,0.0,128.0

docs/sweeps/short_context_prefill_composite_sweep.csv

@@ -0,0 +1,17 @@
+variant,mode,kv_per_cube,C,S_kv,wall_us,n_pe,gemm_util,math_count,math_pipeline_us,hbm_bw_util,hbm_read_mb,hbm_write_kb,ipcq_kb,kv_cache_per_cube_mb
+composite,A1,1,8,8192,35.312546999435405,64,0.19162004939598534,5248,0.0,0.0292247,16.0625,64.0,114688.0,2.0
+composite,A1,1,8,16384,62.80411199463298,64,0.21548245124172694,10880,0.0,0.0327367,32.0625,64.0,229376.0,4.0
+composite,A1,1,8,32768,117.78396202135924,64,0.2297967187462748,22144,0.0,0.0348435,64.0625,64.0,458752.0,8.0
+composite,A1,1,8,65536,227.7376320068799,64,0.23769780826259174,44672,0.0,0.0360063,128.0625,64.0,917504.0,16.0
+composite,A2,2,4,8192,35.873466499721395,32,0.21785371648964683,2624,0.0,0.0575356,16.0625,64.0,49152.0,4.0
+composite,A2,2,4,16384,67.67155399355386,32,0.23097350479268527,5440,0.0,0.0607641,32.0625,64.0,98304.0,8.0
+composite,A2,2,4,32768,129.35661752446276,32,0.24166271963282046,11072,0.0,0.0634525,64.0625,64.0,196608.0,16.0
+composite,A2,2,4,65536,254.34501247150266,32,0.24581313146006306,22336,0.0,0.0644793,128.0625,64.0,393216.0,32.0
+composite,A4,4,2,8192,46.88723350110906,16,0.2114076532174909,1312,0.0,0.088041,16.0625,64.0,16384.0,8.0
+composite,A4,4,2,16384,89.10037348339473,16,0.22249783278807841,2720,0.0,0.0923004,32.0625,64.0,32768.0,16.0
+composite,A4,4,2,32768,173.47797354247328,16,0.22855512537871095,5536,0.0,0.0946287,64.0625,64.0,65536.0,32.0
+composite,A4,4,2,65536,342.2399932809714,16,0.23170453934016594,11168,0.0,0.0958392,128.0625,64.0,131072.0,64.0
+composite,B,8,1,8192,68.73804599999497,8,0.3341725483628459,656,0.0,0.120108,16.0625,64.0,0.0,16.0
+composite,B,8,1,16384,134.27404599999218,8,0.3421415930417734,1360,0.0,0.122496,32.0625,64.0,0.0,32.0
+composite,B,8,1,32768,265.34604599999824,8,0.3462703641501576,2768,0.0,0.123733,64.0625,64.0,0.0,64.0
+composite,B,8,1,65536,527.4900460000299,8,0.3483723443539211,5584,0.0,0.124363,128.0625,64.0,0.0,128.0

docs/sweeps/short_context_prefill_sweep.4ctx.csv

@@ -0,0 +1,17 @@
+mode,kv_per_cube,C,S_kv,wall_us,n_pe,gemm_util,math_count,hbm_bw_util,hbm_read_mb,hbm_write_kb,ipcq_kb,kv_cache_per_cube_mb
+A1,1,8,8192,1103.359903999566,64,0.0019006962210579498,5248,0.001870649814732448,16.0625,64.0,114688.0,2.0
+A1,1,8,16384,2182.457823968016,64,0.0019218259129421755,10880,0.0018841143021603984,32.0625,64.0,229376.0,4.0
+A1,1,8,32768,4343.2822238495155,64,0.0019313983221991736,22144,0.0018898150239762977,64.0625,64.0,458752.0,8.0
+A1,1,8,65536,8666.440783650422,64,0.001935883071136998,44672,0.0018923570136126979,128.0625,64.0,917504.0,16.0
+A2,2,4,8192,1097.2535885006619,32,0.003822547535007557,2624,0.0037621203004135906,16.0625,64.0,49152.0,4.0
+A2,2,4,16384,2179.78425743804,32,0.003848366172648072,5440,0.003772850442394648,32.0625,64.0,98304.0,8.0
+A2,2,4,32768,4344.834905351148,32,0.003861416225357031,11072,0.003778279349528763,64.0625,64.0,196608.0,16.0
+A2,2,4,65536,8674.946611121853,32,0.003867969856656271,22336,0.003781003096658628,128.0625,64.0,393216.0,32.0
+A4,4,2,8192,1103.675383502932,16,0.007600611670231015,1312,0.007480460399321815,16.0625,64.0,16384.0,8.0
+A4,4,2,16384,2197.0616234262543,16,0.00763620638634706,2720,0.007486362614786297,32.0625,64.0,32768.0,16.0
+A4,4,2,32768,4383.834103368879,16,0.007654129058897454,5536,0.007489334501679554,64.0625,64.0,65536.0,32.0
+A4,4,2,65536,8757.379062883949,16,0.007663121981838932,11168,0.007490825682997995,128.0625,64.0,131072.0,64.0
+B,8,1,8192,1125.9457520018113,8,0.014900554462921552,656,0.01466500492643045,16.0625,64.0,0.0,16.0
+B,8,1,16384,2242.543191943168,8,0.014962669223294344,1360,0.014669059716747552,32.0625,64.0,0.0,32.0
+B,8,1,32768,4475.650071825993,8,0.014994216018466134,2768,0.014671388277951352,64.0625,64.0,0.0,64.0
+B,8,1,65536,8942.79183160376,8,0.015008481750139946,5584,0.014671033662702533,128.0625,64.0,0.0,128.0

docs/sweeps/short_context_prefill_sweep.csv

@@ -0,0 +1,17 @@
+mode,kv_per_cube,C,S_kv,wall_us,n_pe,gemm_util,math_count,hbm_bw_util,hbm_read_mb,hbm_write_kb,ipcq_kb,kv_cache_per_cube_mb
+A1,1,8,8192,34.765078999497696,64,0.06032352177393696,5248,0.029685,16.0625,64.0,114688.0,2.0
+A1,1,8,16384,62.2566439945763,64,0.06737118692698127,10880,0.0330246,32.0625,64.0,229376.0,4.0
+A1,1,8,32768,117.23649402130023,64,0.07155287327558905,22144,0.0350062,64.0625,64.0,458752.0,8.0
+A1,1,8,65536,227.190164006345,64,0.07384657726472504,44672,0.0360931,128.0625,64.0,917504.0,16.0
+A2,2,4,8192,34.98345650027832,32,0.11989392757596813,2624,0.0589993,16.0625,64.0,49152.0,4.0
+A2,2,4,16384,65.32612049105158,32,0.12841123790818756,5440,0.0629457,32.0625,64.0,98304.0,8.0
+A2,2,4,32768,126.01144852332584,32,0.13314041062637966,11072,0.0651369,64.0625,64.0,196608.0,16.0
+A2,2,4,65536,247.3821043861434,32,0.1356380732680083,22336,0.0662942,128.0625,64.0,393216.0,32.0
+A4,4,2,8192,46.571509501108665,16,0.18012317165281877,1312,0.0886379,16.0625,64.0,16384.0,8.0
+A4,4,2,16384,88.98265048367577,16,0.18854479956273557,2720,0.0924225,32.0625,64.0,32768.0,16.0
+A4,4,2,32768,173.195009542468,16,0.19373786859461295,5536,0.0947833,64.0625,64.0,65536.0,32.0
+A4,4,2,65536,341.9424492809754,16,0.19625777419912674,11168,0.0959226,128.0625,64.0,131072.0,64.0
+B,8,1,8192,68.50162199999485,8,0.24491706196386726,656,0.120523,16.0625,64.0,0.0,16.0
+B,8,1,16384,134.03762199999207,8,0.25033592434218427,1360,0.122712,32.0625,64.0,0.0,32.0
+B,8,1,32768,265.1096219999986,8,0.25313628186615866,2768,0.123843,64.0625,64.0,0.0,64.0
+B,8,1,65536,527.2536220000293,8,0.2545600872134325,5584,0.124418,128.0625,64.0,0.0,128.0

scratch_dot_tiling_experiment.py

@@ -0,0 +1,204 @@
+#!/usr/bin/env python3
+"""SCRATCH EXPERIMENT (not production; do not commit).
+
+Question: does charging the *primitive* decode kernel for per-HW-tile
+(16x16x16) CPU dispatch flip the "composite gives no decode-latency
+benefit" conclusion?
+
+We monkeypatch TLContext.dot so that, in the primitive kernel, every
+tl.dot whose (M,K,N) exceeds the HW GEMM tile (mac_m/mac_k/mac_n) is
+split by the CPU into ceil(M/mac_m)*ceil(K/mac_k)*ceil(N/mac_n)
+HW-tile-sized GemmCmds.  Each tile GemmCmd is emitted through the normal
+_emit() path, so it (a) charges PE_CPU dispatch overhead via
+_charge_dispatch (PeCpuOverheadCmd), and (b) blocks like a normal
+single-op GemmCmd on PE_GEMM at the cycle-accurate ceil-product latency.
+
+We also inject mac_m/mac_k/mac_n into every pe_gemm topology node so BOTH
+the primitive-tiled and the composite variants run on the *same*
+cycle-accurate engine (fair comparison).  Composite is left untouched:
+the CPU emits ONE CompositeCmd, and PE_SCHEDULER tiles internally (no
+per-HW-tile CPU dispatch).
+
+Data correctness:
+  Inputs are ctx.zeros (q/k/v), so every matmul result is zeros and the
+  DataExecutor replay is trivial.  To keep replay numerically correct
+  regardless, exactly ONE emitted tile per dot carries the *real* full
+  operands+output handles (so the DataExecutor computes the true (M,N)
+  result via the recorded handle shapes), while its timing fields
+  (m,k,n) are the HW-tile size so the engine charges exactly one tile of
+  cycle time.  The remaining n_tiles-1 emitted tiles are timing-only
+  GemmCmds (16x16x16) writing to throwaway scratch.  Net: n_tiles tiles
+  of engine time + n_tiles dispatch charges, and a correct final output.
+
+Engine mode: enable_data=True (same as the production sweep's
+_engine_latency_ns), op_log end-to-end latency.
+"""
+from __future__ import annotations
+
+import json
+from math import ceil
+from pathlib import Path
+
+# --- HW GEMM tile under test -------------------------------------------------
+MAC_M = 16
+MAC_K = 16
+MAC_N = 16
+# Legacy alt to also try: (8, 16, 32)
+
+S_KV_LATENCY = (8192, 32_768, 65_536, 131_072)
+
+ROOT = Path(__file__).resolve().parent
+SWEEP_JSON = (
+    ROOT / "src" / "kernbench" / "benches" / "1H_milestone_output"
+    / "gqa" / "long_ctx" / "sweep_decode_composite.json"
+)
+
+
+# --- mac-dim topology override -----------------------------------------------
+def _topo_with_mac(mac_m: int, mac_k: int, mac_n: int):
+    """Compiled topology with mac dims injected into every pe_gemm node."""
+    from kernbench.topology.builder import resolve_topology
+
+    handle = resolve_topology("topology.yaml")
+    g = handle.topology_obj
+    n = 0
+    for node in g.nodes.values():
+        if node.kind == "pe_gemm":
+            node.attrs["mac_m"] = mac_m
+            node.attrs["mac_k"] = mac_k
+            node.attrs["mac_n"] = mac_n
+            n += 1
+    print(f"    injected mac=({mac_m},{mac_k},{mac_n}) into {n} pe_gemm nodes")
+    return handle
+
+
+# --- tiling monkeypatch for TLContext.dot ------------------------------------
+def _make_tiled_dot(orig_dot, mac_m: int, mac_k: int, mac_n: int):
+    from kernbench.common.pe_commands import GemmCmd
+
+    def tiled_dot(self, a, b):
+        if len(a.shape) < 2 or len(b.shape) < 2:
+            return orig_dot(self, a, b)
+        m, k = a.shape[-2], a.shape[-1]
+        k2, n = b.shape[-2], b.shape[-1]
+        if k != k2:
+            raise ValueError(f"dot shape mismatch: a.K={k} != b.K={k2}")
+
+        n_tiles = ceil(m / mac_m) * ceil(k / mac_k) * ceil(n / mac_n)
+
+        out_shape = (*a.shape[:-2], m, n)
+        out = self._make_compute_out(shape=out_shape, dtype=a.dtype)
+        self._await_pending(a, b)
+
+        if n_tiles <= 1:
+            self._emit(GemmCmd(a=a, b=b, out=out, m=m, k=k, n=n))
+            return out
+
+        # One real-data tile: full handles (so DataExecutor computes the
+        # true result), but timing fields = HW tile (one tile of cycles).
+        self._emit(GemmCmd(a=a, b=b, out=out, m=mac_m, k=mac_k, n=mac_n))
+        # Remaining timing-only tiles: throwaway scratch, 16x16x16.
+        scratch = self._make_compute_out(shape=(mac_m, mac_n), dtype=a.dtype)
+        for _ in range(n_tiles - 1):
+            self._emit(GemmCmd(a=a, b=b, out=scratch,
+                               m=mac_m, k=mac_k, n=mac_n))
+        return out
+
+    return tiled_dot
+
+
+# --- latency runner (replicates sweep's _engine_latency_ns) ------------------
+def _engine_latency_ns(variant: str, S_kv: int, topo) -> float:
+    from kernbench.benches.gqa_helpers.long_ctx.gqa_decode_long_ctx_composite import (  # noqa: E501
+        _end_to_end_ns, _run_panel_fn,
+    )
+    from kernbench.runtime_api.bench_runner import run_bench
+    from kernbench.runtime_api.types import resolve_device
+    from kernbench.sim_engine.engine import GraphEngine
+
+    result = run_bench(
+        topology=topo, bench_fn=_run_panel_fn(variant, S_kv),
+        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"{variant}@{S_kv} failed: {result.completion}"
+        )
+    return _end_to_end_ns(result.engine.op_log)
+
+
+def _emit_dispatch(variant: str, S_kv: int) -> int:
+    """PE_CPU command count at the center rank (cube 6, pe 0)."""
+    from kernbench.benches.gqa_helpers.long_ctx.gqa_decode_long_ctx_composite import (  # noqa: E501
+        _emit_dispatch as prod_emit,
+    )
+    return prod_emit(variant, S_kv)[0]
+
+
+def main() -> None:
+    import kernbench.triton_emu.tl_context as tlc
+
+    # Baseline (A): primitive UNTILED latencies from the production sweep
+    # (mac=0 / TFLOPS model). Read straight off the committed sweep JSON.
+    sweep = json.loads(SWEEP_JSON.read_text())
+    base_A = {}
+    for r in sweep["rows"]:
+        if r["variant"] == "primitive" and r["latency_ns"] is not None:
+            base_A[r["S_kv"]] = r["latency_ns"]
+
+    print(f"== mac tile = ({MAC_M},{MAC_K},{MAC_N}) ==")
+
+    # --- command-count sanity (emit-time, mac-independent) ---------------
+    orig_dot = tlc.TLContext.dot
+    print("\n[dispatch counts @ S_kv=131072]")
+    n_prim_untiled = _emit_dispatch("primitive", 131072)
+    tlc.TLContext.dot = _make_tiled_dot(orig_dot, MAC_M, MAC_K, MAC_N)
+    try:
+        n_prim_tiled = _emit_dispatch("primitive", 131072)
+        n_comp = None
+    finally:
+        tlc.TLContext.dot = orig_dot
+    n_comp = _emit_dispatch("composite", 131072)
+    print(f"  primitive UNTILED PE_CPU cmds : {n_prim_untiled}")
+    print(f"  primitive TILED   PE_CPU cmds : {n_prim_tiled}  "
+          f"(x{n_prim_tiled / max(n_prim_untiled,1):.0f})")
+    print(f"  composite         PE_CPU cmds : {n_comp}")
+
+    # --- latency sweep ----------------------------------------------------
+    rows = []
+    topo = _topo_with_mac(MAC_M, MAC_K, MAC_N)
+
+    for S_kv in S_KV_LATENCY:
+        A = base_A.get(S_kv)
+
+        # (C) composite on the mac engine (untouched dot path)
+        C = _engine_latency_ns("composite", S_kv, topo)
+
+        # (B) primitive TILED on the mac engine
+        tlc.TLContext.dot = _make_tiled_dot(orig_dot, MAC_M, MAC_K, MAC_N)
+        try:
+            B = _engine_latency_ns("primitive", S_kv, topo)
+        finally:
+            tlc.TLContext.dot = orig_dot
+
+        gap_pct = (B - C) / C * 100.0 if C else float("nan")
+        rows.append((S_kv, A, B, C, gap_pct))
+        print(f"  S_kv={S_kv:>7}: A(untiled)={A!s:>12}  "
+              f"B(tiled)={B:12.2f}  C(comp)={C:12.2f}  (B-C)/C={gap_pct:+6.1f}%")
+
+    # --- final table ------------------------------------------------------
+    print("\n==================== RESULT TABLE ====================")
+    print(f"{'S_kv':>8} | {'A untiled(ns)':>14} | {'B tiled(ns)':>14} | "
+          f"{'C comp(ns)':>14} | {'(B-C)/C':>9}")
+    print("-" * 72)
+    for S_kv, A, B, C, gap in rows:
+        a_s = f"{A:.2f}" if A is not None else "n/a"
+        print(f"{S_kv:>8} | {a_s:>14} | {B:>14.2f} | {C:>14.2f} | "
+              f"{gap:>+8.1f}%")
+
+
+if __name__ == "__main__":
+    main()

scripts/paper/paper_plot_gqa_decode_composite_breakdown.py

@@ -0,0 +1,131 @@
+"""Latency breakdown bar chart for the Case-6 composite-command decode study.
+
+Reads sweep_decode_composite.json and writes a single-figure stacked bar
+chart comparing three variants — primitive hand-tiled (16×16×16),
+composite GEMM, composite + softmax_merge — at the S_kv = 131 072 point:
+
+  bottom stack: PE_CPU dispatch time (from pe_cpu_dispatch_cycles, ns at
+                1 GHz — ADR-0064 Rev2 D3)
+  top stack:    engine time (latency_ns − dispatch cycles) — DMA / GEMM /
+                MATH / IPCQ work the engine flushes on the critical path
+
+The dispatch/engine split is a first-order breakdown; in reality the
+two paths overlap partially. The note on the plot calls this out.
+
+Run (after the composite bench sweep):
+    python scripts/paper/paper_plot_gqa_decode_composite_breakdown.py
+"""
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt  # noqa: E402
+
+_REPO_ROOT = Path(__file__).resolve().parents[2]
+_FIG_DIR = (
+    _REPO_ROOT / "src" / "kernbench" / "benches"
+    / "1H_milestone_output" / "gqa" / "long_ctx"
+)
+_SWEEP_JSON = _FIG_DIR / "sweep_decode_composite.json"
+_PAPER_FIG_DIR = (
+    _REPO_ROOT / "docs" / "report" / "1H-codesign-paper" / "figures"
+)
+
+# The three variants to compare (coarse `primitive` intentionally dropped).
+_ORDER = ("primitive_tiled", "composite", "composite_extended")
+_LABELS = {
+    "primitive_tiled":    "primitive hand-tiled\n(16×16×16)",
+    "composite":          "composite\nGEMM",
+    "composite_extended": "composite +\nsoftmax_merge",
+}
+_S_KV_TARGET = 131_072
+
+
+def main() -> None:
+    sweep = json.loads(_SWEEP_JSON.read_text())
+    rows = {(r["variant"], r["S_kv"]): r for r in sweep["rows"]}
+
+    dispatch_us = []
+    engine_us = []
+    totals_us = []
+    for v in _ORDER:
+        r = rows[(v, _S_KV_TARGET)]
+        disp_ns = r["pe_cpu_dispatch_cycles"]
+        total_ns = r["latency_ns"]
+        eng_ns = max(0.0, total_ns - disp_ns)
+        dispatch_us.append(disp_ns / 1e3)
+        engine_us.append(eng_ns / 1e3)
+        totals_us.append(total_ns / 1e3)
+
+    xs = list(range(len(_ORDER)))
+    labels = [_LABELS[v] for v in _ORDER]
+
+    fig, ax = plt.subplots(figsize=(7.5, 5.0))
+
+    bars_eng = ax.bar(
+        xs, engine_us,
+        color="#4f8a4f", edgecolor="#2a4a2a", label="engine (DMA + GEMM + MATH + IPCQ)",
+    )
+    bars_disp = ax.bar(
+        xs, dispatch_us, bottom=engine_us,
+        color="#c0504d", edgecolor="#5a2624", label="PE_CPU dispatch",
+    )
+
+    # Segment value labels — engine at mid, dispatch at mid of its stack.
+    for i, (eng, disp, tot) in enumerate(zip(engine_us, dispatch_us, totals_us)):
+        ax.text(i, eng / 2, f"{eng:.1f} µs",
+                ha="center", va="center", fontsize=9, color="white")
+        if disp / max(totals_us) > 0.04:  # only label if visible
+            ax.text(i, eng + disp / 2, f"{disp:.1f} µs",
+                    ha="center", va="center", fontsize=9, color="white")
+        else:
+            ax.annotate(f"{disp:.2f} µs",
+                        xy=(i, tot), xytext=(0, 6),
+                        textcoords="offset points",
+                        ha="center", va="bottom", fontsize=8,
+                        color="#5a2624")
+        offset_pts = 14 if disp / max(totals_us) > 0.04 else 20
+        ax.annotate(f"total {tot:.1f}",
+                    xy=(i, tot), xytext=(0, offset_pts),
+                    textcoords="offset points",
+                    ha="center", va="bottom", fontsize=9,
+                    color="#333", fontweight="bold")
+
+    ax.set_xticks(xs)
+    ax.set_xticklabels(labels, fontsize=10)
+    ax.set_ylabel("time (µs)")
+    ax.set_title(
+        f"Case-6 decode latency breakdown at $S_{{kv}}=${_S_KV_TARGET // 1024}K\n"
+        "(engine dominates all three — primitive hand-tiled adds "
+        f"{dispatch_us[0]:.0f} µs PE_CPU dispatch overhead)",
+        fontsize=11,
+    )
+    ax.legend(loc="upper right", fontsize=9)
+    ax.grid(True, axis="y", ls=":", alpha=0.5)
+    ax.set_ylim(0, max(totals_us) * 1.15)
+
+    fig.text(
+        0.5, 0.02,
+        "First-order breakdown: dispatch and engine paths overlap partially on the real "
+        "critical path;\ntreat the split as an upper bound on dispatch's contribution.",
+        ha="center", fontsize=8, color="#666",
+    )
+    fig.tight_layout(rect=(0, 0.06, 1, 1))
+
+    out = _FIG_DIR / "gqa_decode_long_ctx_composite_breakdown.png"
+    fig.savefig(out, dpi=150)
+    plt.close(fig)
+    print(f"wrote {out}")
+
+    if _PAPER_FIG_DIR.is_dir():
+        dst = _PAPER_FIG_DIR / out.name
+        dst.write_bytes(out.read_bytes())
+        print(f"copied {dst}")
+
+
+if __name__ == "__main__":
+    main()

scripts/paper/paper_plot_gqa_decode_long_ctx_composite.py

@@ -0,0 +1,147 @@
+"""Comparative figure for the Case-6 composite-command decode study.
+
+Reads sweep_decode_composite.json (emitted by milestone-1h-gqa, sweep
+``composite``) and writes one two-panel PNG into the bench-output dir:
+
+  gqa_decode_long_ctx_composite.png
+    Left  — end-to-end decode latency (µs) vs context length, per command
+            form (primitive / composite / composite_extended).
+    Right — PE_CPU command count vs context length: the hand-tiled
+            primitive kernel issues O(n_tiles) commands (rises with
+            context), while the coarse composite forms issue O(1) and
+            *saturate* — PE_SCHEDULER absorbs the per-tile fan-out.
+
+The x-axis is the global context length S_kv; each PE owns
+S_local = S_kv/(C·P=64) tokens, so at the 1M production point every PE
+runs Q·Kᵀ of (G·T_q, d_head)·(d_head, 16384) and P·V of
+(G·T_q, 16384)·(16384, d_head).
+
+Run (after the bench):
+    GQA_1H_RUN=1 GQA_1H_SWEEPS=composite python -m kernbench.cli.main run \\
+        --bench milestone-1h-gqa --topology topology.yaml
+    python scripts/paper/paper_plot_gqa_decode_long_ctx_composite.py
+"""
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt  # noqa: E402
+
+_REPO_ROOT = Path(__file__).resolve().parents[2]
+_FIG_DIR = (
+    _REPO_ROOT / "src" / "kernbench" / "benches"
+    / "1H_milestone_output" / "gqa" / "long_ctx"
+)
+_SWEEP_JSON = _FIG_DIR / "sweep_decode_composite.json"
+_PAPER_FIG_DIR = (
+    _REPO_ROOT / "docs" / "report" / "1H-codesign-paper" / "figures"
+)
+
+_N_RANKS = 64  # C·P for the Case-6 64-way split.
+
+# variant key → (display label, colour, marker)
+_VARIANT_STYLE = {
+    "primitive_tiled":    ("primitive hand-tiled (16×16×16)", "#8b5a2b", "D"),
+    "primitive":          ("primitive (tl.dot, hand-tiled)",  "#c0504d", "o"),
+    "composite":          ("composite GEMM",                  "#3b6ea5", "s"),
+    "composite_extended": ("composite + softmax_merge",       "#4f8a4f", "^"),
+}
+_ORDER = ("primitive_tiled", "composite", "composite_extended")
+
+
+def _load() -> dict:
+    return json.loads(_SWEEP_JSON.read_text())
+
+
+def _series(rows: list[dict], variant: str, key: str):
+    """Sorted (S_kv, value) series for a variant, skipping null values
+    (latency is only measured over the tractable S_kv subset)."""
+    pts = sorted(
+        ((r["S_kv"], r[key]) for r in rows
+         if r["variant"] == variant and r.get(key) is not None),
+        key=lambda t: t[0],
+    )
+    return [p[0] for p in pts], [p[1] for p in pts]
+
+
+def _xticklabels(s_kvs: list[int]) -> list[str]:
+    out = []
+    for s in s_kvs:
+        if s >= 1 << 20:
+            out.append(f"{s // (1 << 20)}M")
+        else:
+            out.append(f"{s // 1024}K")
+    return out
+
+
+def main() -> None:
+    sweep = _load()
+    rows = sweep["rows"]
+    s_kv_op = sweep["s_kv_opcount"]
+    s_kv_lat = sweep["s_kv_latency"]
+
+    fig, (ax_lat, ax_cmd) = plt.subplots(1, 2, figsize=(13.0, 4.8))
+
+    for v in _ORDER:
+        label, color, marker = _VARIANT_STYLE[v]
+        xs, lat = _series(rows, v, "latency_ns")
+        ax_lat.plot(xs, [y / 1e3 for y in lat], marker=marker,
+                    color=color, label=label, lw=2)
+        xs, cmds = _series(rows, v, "pe_cpu_cmd_count")
+        ax_cmd.plot(xs, cmds, marker=marker, color=color, label=label, lw=2)
+
+    ax_lat.set_xticks(s_kv_lat)
+    ax_lat.set_xticklabels(_xticklabels(s_kv_lat))
+    ax_cmd.set_xticks(s_kv_op)
+    ax_cmd.set_xticklabels(_xticklabels(s_kv_op), fontsize=8)
+    for ax in (ax_lat, ax_cmd):
+        ax.set_xscale("log", base=2)
+        ax.set_xlabel(
+            r"context length $S_{kv}$ "
+            r"($S_{\mathrm{local}}=S_{kv}/64$ per PE)"
+        )
+        ax.grid(True, ls=":", alpha=0.5)
+        ax.legend(fontsize=9)
+
+    ax_lat.set_ylabel("end-to-end decode latency (µs)")
+    ax_lat.set_title(
+        "Case-6 decode latency per command form\n"
+        "(memory-bound: command form does not move the critical path)"
+    )
+    ax_cmd.set_ylabel("PE_CPU commands issued")
+    ax_cmd.set_title(
+        "PE_CPU command count per command form\n"
+        "(primitive O(n$_\\mathrm{tiles}$) rises; composite O(1) saturates)"
+    )
+    ax_cmd.axvline(1 << 20, color="#888", ls="--", lw=1, alpha=0.7)
+    ax_cmd.annotate("1M production\ncontext", xy=(1 << 20, 0),
+                    xytext=(1 << 18, 0.6), fontsize=8,
+                    textcoords=("data", "axes fraction"), ha="right",
+                    color="#555")
+
+    fig.suptitle(
+        "Case-6 (Cube-SP × PE-SP) long-context decode — use of composite "
+        "commands\nLLaMA-3.1-70B single-KV-head group (8 cubes × 8 PEs), "
+        "$T_q{=}1$",
+        fontsize=11,
+    )
+    fig.tight_layout(rect=(0, 0, 1, 0.94))
+
+    out = _FIG_DIR / "gqa_decode_long_ctx_composite.png"
+    fig.savefig(out, dpi=150)
+    plt.close(fig)
+    print(f"wrote {out}")
+
+    # Mirror into the paper figures dir (derived artifact).
+    if _PAPER_FIG_DIR.is_dir():
+        dst = _PAPER_FIG_DIR / out.name
+        dst.write_bytes(out.read_bytes())
+        print(f"copied {dst}")
+
+
+if __name__ == "__main__":
+    main()

scripts/paper/paper_plot_gqa_decode_models.py

@@ -0,0 +1,192 @@
+"""Bar plot for the multi-model GQA composite bench.
+
+Topology per model: cubes per KV group C = h_q, P = 8 PEs / cube.
+
+Reads ``sweep_decode_models.json`` (produced by
+``gqa_decode_long_ctx_models.py``) and writes a two-panel PNG
+comparing end-to-end latency and PE_CPU dispatch across six GQA models
+where each model uses a topology sized to match its G value
+(cubes per KV group = h_q per KV group; PEs per cube = 8):
+
+  Left panel  — end-to-end decode latency (µs), one bar per model, sorted
+                by G. Bar labels show C, N, and latency.
+  Right panel — PE_CPU command count per model.
+
+Run (after the bench):
+    python scripts/paper/paper_plot_gqa_decode_models.py
+"""
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt  # noqa: E402
+import numpy as np               # noqa: E402
+
+_REPO_ROOT = Path(__file__).resolve().parents[2]
+_FIG_DIR = (
+    _REPO_ROOT / "src" / "kernbench" / "benches"
+    / "1H_milestone_output" / "gqa" / "long_ctx"
+)
+_SWEEP_JSON = _FIG_DIR / "sweep_decode_models.json"
+_PAPER_FIG_DIR = (
+    _REPO_ROOT / "docs" / "report" / "1H-codesign-paper" / "figures"
+)
+
+# Model display labels
+_MODEL_LABELS = {
+    "gemma2-27b":       "Gemma 2 27B\n(G=2)",
+    "llama3-8b":        "LLaMA-3 8B\n(G=4)",
+    "qwen2.5-7b":       "Qwen 2.5 7B\n(G=7)",
+    "llama3-70b":       "LLaMA-3 70B\n(G=8)",
+    "qwen2.5-72b":      "Qwen 2.5 72B\n(G=8)",
+    "command-r-plus":   "Command R+\n(G=12)",
+}
+
+# Bar color per model family
+_FAMILY_COLOUR = {
+    "Google":  "#4f8a4f",
+    "Meta":    "#3b6ea5",
+    "Alibaba": "#c86432",
+    "Cohere":  "#8b5a2b",
+}
+
+
+def main() -> None:
+    sweep = json.loads(_SWEEP_JSON.read_text())
+    rows_by_model = {r["model"]: r for r in sweep["rows"]}
+    # Sort models by G (ascending) so the trend is left-to-right.
+    ordered = sorted(_MODEL_LABELS.keys(), key=lambda m: rows_by_model[m]["G"])
+
+    fig, (ax_lat, ax_cmd, ax_brk) = plt.subplots(1, 3, figsize=(18.5, 5.0))
+    xs = np.arange(len(ordered))
+
+    # ── Left: latency ─────────────────────────────────────────────
+    latencies = [rows_by_model[m]["latency_ns"] / 1e3 for m in ordered]
+    colours = [_FAMILY_COLOUR[rows_by_model[m]["family"]] for m in ordered]
+    bars = ax_lat.bar(xs, latencies, 0.62,
+                      color=colours, edgecolor="#222", linewidth=0.7)
+    for i, m in enumerate(ordered):
+        r = rows_by_model[m]
+        # Latency label above bar
+        ax_lat.text(i, latencies[i] + max(latencies) * 0.015,
+                    f"{latencies[i]:.1f} µs",
+                    ha="center", va="bottom", fontsize=9,
+                    fontweight="bold", color="#222")
+        # Topology label inside bar
+        ax_lat.text(i, latencies[i] / 2,
+                    f"cubes = {r['C']}",
+                    ha="center", va="center", fontsize=9, color="white")
+
+    ax_lat.set_xticks(xs)
+    ax_lat.set_xticklabels([_MODEL_LABELS[m] for m in ordered], fontsize=9)
+    ax_lat.set_ylabel("end-to-end decode latency (µs)")
+    ax_lat.set_title(
+        f"Composite Case-6 decode latency at $S_{{kv}}=128$K\n"
+        "topology per model: cubes per KV group = h_q of the KV group",
+        fontsize=11,
+    )
+    ax_lat.grid(True, axis="y", ls=":", alpha=0.5)
+    ax_lat.set_ylim(0, max(latencies) * 1.18)
+
+    # Legend (family)
+    from matplotlib.patches import Patch
+    seen_families = []
+    handles = []
+    for m in ordered:
+        fam = rows_by_model[m]["family"]
+        if fam not in seen_families:
+            seen_families.append(fam)
+            handles.append(Patch(facecolor=_FAMILY_COLOUR[fam],
+                                 edgecolor="#222", label=fam))
+    ax_lat.legend(handles=handles, loc="upper left", fontsize=9,
+                  title="family")
+
+    # ── Right: PE_CPU dispatch ─────────────────────────────────────
+    cmds = [rows_by_model[m]["pe_cpu_cmd_count"] for m in ordered]
+    bars2 = ax_cmd.bar(xs, cmds, 0.62,
+                       color=colours, edgecolor="#222", linewidth=0.7)
+    for i, m in enumerate(ordered):
+        ax_cmd.text(i, cmds[i] + max(cmds) * 0.02,
+                    f"{cmds[i]}",
+                    ha="center", va="bottom", fontsize=9,
+                    fontweight="bold", color="#222")
+        r = rows_by_model[m]
+        ax_cmd.text(i, cmds[i] / 2,
+                    f"cubes = {r['C']}",
+                    ha="center", va="center", fontsize=9, color="white")
+
+    ax_cmd.set_xticks(xs)
+    ax_cmd.set_xticklabels([_MODEL_LABELS[m] for m in ordered], fontsize=9)
+    ax_cmd.set_ylabel("PE_CPU commands issued")
+    ax_cmd.set_title(
+        "PE_CPU dispatch (composite, per KV group)\n"
+        "growing sub-mesh width or height → more reduce hops",
+        fontsize=11,
+    )
+    ax_cmd.grid(True, axis="y", ls=":", alpha=0.5)
+    ax_cmd.set_ylim(0, max(cmds) * 1.18)
+
+    # ── Third: matmul vs comm (op-kind occupancy) ─────────────────────
+    # Only the two kinds that carry useful work — matmul (GEMM + MATH
+    # occupancy summed across all engines) and comm (DMA occupancy).
+    # These are op-log sums across components, not critical-path
+    # attribution; parallelism means they don't sum to wall-clock
+    # latency. Bars are absolute µs so the reduce-cost growth with C
+    # is visible.
+    matmul_us = [rows_by_model[m].get("matmul_ns", 0.0) / 1e3 for m in ordered]
+    comm_us = [rows_by_model[m].get("comm_ns", 0.0) / 1e3 for m in ordered]
+    max_v = max(max(matmul_us), max(comm_us))
+    width = 0.35
+
+    ax_brk.bar(xs - width / 2, matmul_us, width,
+               color="#3b6ea5", edgecolor="#222", linewidth=0.6,
+               label="matmul (GEMM + MATH)")
+    ax_brk.bar(xs + width / 2, comm_us, width,
+               color="#c0504d", edgecolor="#222", linewidth=0.6,
+               label="communication (DMA)")
+    for i in range(len(ordered)):
+        ax_brk.text(xs[i] - width / 2, matmul_us[i] + max_v * 0.015,
+                    f"{matmul_us[i]:.1f}",
+                    ha="center", va="bottom", fontsize=8,
+                    color="#222", fontweight="bold")
+        ax_brk.text(xs[i] + width / 2, comm_us[i] + max_v * 0.015,
+                    f"{comm_us[i]:.1f}",
+                    ha="center", va="bottom", fontsize=8,
+                    color="#222", fontweight="bold")
+
+    ax_brk.set_xticks(xs)
+    ax_brk.set_xticklabels([_MODEL_LABELS[m] for m in ordered], fontsize=9)
+    ax_brk.set_ylabel("op-log occupancy (µs, summed across engines)")
+    ax_brk.set_title(
+        "Matmul vs communication engine work per model\n"
+        "sum of GEMM/MATH and DMA occupancy — not wall-clock (overlap)",
+        fontsize=11,
+    )
+    ax_brk.grid(True, axis="y", ls=":", alpha=0.5)
+    ax_brk.set_ylim(0, max_v * 1.20)
+    ax_brk.legend(loc="upper left", fontsize=9)
+
+    fig.suptitle(
+        "Multi-model attention comparison — Case-6 composite kernel"
+        "  ($S_{kv}=128$K, $T_q=1$, P=8 PEs / cube)",
+        fontsize=12, fontweight="bold",
+    )
+    fig.tight_layout(rect=(0, 0, 1, 0.93))
+
+    out = _FIG_DIR / "gqa_decode_models.png"
+    fig.savefig(out, dpi=150)
+    plt.close(fig)
+    print(f"wrote {out}")
+
+    if _PAPER_FIG_DIR.is_dir():
+        dst = _PAPER_FIG_DIR / out.name
+        dst.write_bytes(out.read_bytes())
+        print(f"copied {dst}")
+
+
+if __name__ == "__main__":
+    main()

scripts/paper/paper_plot_gqa_prefill_compute_bound.py

@@ -0,0 +1,111 @@
+"""Comparative figure for the compute-bound prefill composite study.
+
+Reads sweep_prefill_compute_bound.json (emitted by milestone-1h-gqa,
+sweep ``prefill_cb``) and writes one two-panel PNG:
+
+  gqa_prefill_compute_bound.png
+    Left  — end-to-end prefill latency (µs) vs context length.
+    Right — MAC utilization (achieved / 8 TFLOP·s⁻¹ per-PE peak) vs context.
+
+Unlike memory-bound decode (where command form is latency-neutral), in
+compute-bound prefill the composite command keeps the MAC array fed by
+streaming DMA↔compute per HW tile, so it wins on both latency and
+utilization — and the margin grows with context (deeper P·V reduction =
+more tiles to pipeline).
+
+Run (after the bench):
+    GQA_1H_RUN=1 GQA_1H_SWEEPS=prefill_cb python -m kernbench.cli.main run \\
+        --bench milestone-1h-gqa --topology topology.yaml
+    python scripts/paper/paper_plot_gqa_prefill_compute_bound.py
+"""
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt  # noqa: E402
+
+_REPO_ROOT = Path(__file__).resolve().parents[2]
+_FIG_DIR = (
+    _REPO_ROOT / "src" / "kernbench" / "benches"
+    / "1H_milestone_output" / "gqa" / "long_ctx"
+)
+_SWEEP_JSON = _FIG_DIR / "sweep_prefill_compute_bound.json"
+_PAPER_FIG_DIR = (
+    _REPO_ROOT / "docs" / "report" / "1H-codesign-paper" / "figures"
+)
+
+_VARIANT_STYLE = {
+    "primitive":          ("primitive (tl.dot, hand-tiled)", "#c0504d", "o"),
+    "composite":          ("composite GEMM",                 "#3b6ea5", "s"),
+    "composite_extended": ("composite + softmax_merge",      "#4f8a4f", "^"),
+}
+_ORDER = ("primitive", "composite", "composite_extended")
+
+
+def _ctx_label(c: int) -> str:
+    return f"{c // 1024}K" if c >= 1024 else str(c)
+
+
+def _series(rows, variant, key):
+    pts = sorted(((r["ctx_len"], r[key]) for r in rows
+                  if r["variant"] == variant), key=lambda t: t[0])
+    return [p[0] for p in pts], [p[1] for p in pts]
+
+
+def main() -> None:
+    sweep = json.loads(_SWEEP_JSON.read_text())
+    rows = sweep["rows"]
+    ctxs = sweep["ctx_points"]
+
+    fig, (ax_lat, ax_util) = plt.subplots(1, 2, figsize=(13.0, 4.8))
+
+    for v in _ORDER:
+        label, color, marker = _VARIANT_STYLE[v]
+        xs, lat = _series(rows, v, "latency_ns")
+        ax_lat.plot(xs, [y / 1e3 for y in lat], marker=marker,
+                    color=color, label=label, lw=2)
+        xs, util = _series(rows, v, "mac_util")
+        ax_util.plot(xs, [u * 100 for u in util], marker=marker,
+                     color=color, label=label, lw=2)
+
+    for ax in (ax_lat, ax_util):
+        ax.set_xscale("log", base=2)
+        ax.set_xticks(ctxs)
+        ax.set_xticklabels([_ctx_label(c) for c in ctxs])
+        ax.set_xlabel(r"context length (= $T_q$ = $S_{kv}$)")
+        ax.grid(True, ls=":", alpha=0.5)
+        ax.legend(fontsize=9)
+
+    ax_lat.set_ylabel("end-to-end prefill latency (µs)")
+    ax_lat.set_title("Compute-bound prefill latency per command form")
+    ax_util.set_ylabel("MAC utilization (% of 8 TFLOP·s⁻¹ peak)")
+    ax_util.set_title(
+        "MAC utilization — composite keeps the array fed; primitive starves"
+    )
+    ax_util.axhline(100, color="#888", ls="--", lw=1, alpha=0.6)
+
+    fig.suptitle(
+        "Compute-bound prefill attention — use of composite commands\n"
+        "single-rank, GQA single-KV-head group ($h_q{=}8$, $d_{\\text{head}}"
+        "{=}128$); $M{=}8T_q$ tile-filling",
+        fontsize=11,
+    )
+    fig.tight_layout(rect=(0, 0, 1, 0.92))
+
+    out = _FIG_DIR / "gqa_prefill_compute_bound.png"
+    fig.savefig(out, dpi=150)
+    plt.close(fig)
+    print(f"wrote {out}")
+
+    if _PAPER_FIG_DIR.is_dir():
+        dst = _PAPER_FIG_DIR / out.name
+        dst.write_bytes(out.read_bytes())
+        print(f"copied {dst}")
+
+
+if __name__ == "__main__":
+    main()

scripts/paper/paper_plot_ipcq_alternatives_comparison.py

@@ -243,6 +243,12 @@ _FLOW_DATA = "#6699CC"   # blue flit (data)
 _FLOW_CTRL = "#ED7D65"   # orange flit (control)
 _FLOW_CRED = "#7BB661"   # green (credit / ack)
 
+# Font-size multiplier for the flow helpers. The stacked figure renders at
+# full text width (a two-column figure*), so its fonts are scaled up here so
+# they stay legible after the figure is fit to the page. Default 1.0 leaves
+# the 2×2 flow figure unchanged.
+_FS = 1.0
+
 
 def _flow_swim(ax, x, y, w, h, label, sub=None):
     ax.add_patch(mpatches.FancyBboxPatch(
@@ -251,11 +257,11 @@ def _flow_swim(ax, x, y, w, h, label, sub=None):
         facecolor="white", edgecolor="#3a3a3a", linewidth=1.6))
     ax.text(x + w / 2, y + h / 2 + (0.20 if sub else 0),
             label, ha="center", va="center",
-            fontsize=10.5, fontweight="bold", color="#111")
+            fontsize=10.5 * _FS, fontweight="bold", color="#111")
     if sub:
         ax.text(x + w / 2, y + h / 2 - 0.35, sub,
                 ha="center", va="center",
-                fontsize=8, color="#555", fontstyle="italic")
+                fontsize=8 * _FS, color="#555", fontstyle="italic")
 
 
 def _flow_flit(ax, x, y, color, letter="", size=0.32):
@@ -265,7 +271,7 @@ def _flow_flit(ax, x, y, color, letter="", size=0.32):
         facecolor=color, edgecolor="#222", linewidth=0.8))
     if letter:
         ax.text(x, y, letter, ha="center", va="center",
-                fontsize=8.5, fontweight="bold", color="white")
+                fontsize=8.5 * _FS, fontweight="bold", color="white")
 
 
 def _flow_wire(ax, x1, x2, y, color="#888"):
@@ -280,7 +286,7 @@ def _flow_dot(ax, x, y, color):
 
 def _flow_label(ax, x, y, text, color="#222", fontsize=9):
     ax.text(x, y, text, ha="center", va="center",
-            fontsize=fontsize, color=color, fontweight="bold")
+            fontsize=fontsize * _FS, color=color, fontweight="bold")
 
 
 def _flow_lead(ax, x_tail, y_tail, x_head, y_head, color="#777"):
@@ -297,7 +303,7 @@ def _flow_setup(ax, kind):
         facecolor="white", edgecolor=_COLOR[kind],
         linewidth=0.9, alpha=0.85, zorder=0))
     ax.text(10, 0.65, _TITLE[kind],
-            ha="center", va="center", fontsize=11.5,
+            ha="center", va="center", fontsize=11.5 * _FS,
             fontweight="bold",
             color=("#1B5E20" if kind == "ipcq" else "white"),
             bbox=dict(
@@ -449,7 +455,7 @@ def _flow_legend(fig, leg_y=0.02):
                 facecolor=color, edgecolor="#222", linewidth=0.6,
                 transform=fig.transFigure))
             fig.text(lx + 0.025, y + 0.011, text,
-                     ha="left", va="center", fontsize=10,
+                     ha="left", va="center", fontsize=10 * _FS,
                      fontweight="bold", color="#222")
             lx += 0.32
 
@@ -461,21 +467,27 @@ def _flow_legend(fig, leg_y=0.02):
 def _plot_architecture_flow() -> Path:
     """4 panels arranged 2×2 — each a horizontal Sender → NoC → Receiver
     flow with small flit-style packets, dot-marker annotations and a
-    bottom name plate. Aesthetic mirrors latency_model.png."""
+    bottom name plate. Aesthetic mirrors latency_model.png. Fonts scaled
-    fig, axes = plt.subplots(2, 2, figsize=(18.0, 11.5))
+    up via _FS so the figure stays legible when fit to the page."""
-    _draw_flow_case(axes[0, 0], "doorbell")
+    global _FS
-    _draw_flow_case(axes[0, 1], "hmq")
+    _FS = 1.5
-    _draw_flow_case(axes[1, 0], "rdma")
+    try:
-    _draw_flow_case(axes[1, 1], "ipcq")
+        fig, axes = plt.subplots(2, 2, figsize=(18.0, 11.5))
-    fig.suptitle(
+        _draw_flow_case(axes[0, 0], "doorbell")
-        "Per-send architecture  =  $\\Sigma$ control overhead  +  "
+        _draw_flow_case(axes[0, 1], "hmq")
-        "data-path DMA  +  receiver wake",
+        _draw_flow_case(axes[1, 0], "rdma")
-        fontsize=14.5, fontweight="bold", y=0.995)
+        _draw_flow_case(axes[1, 1], "ipcq")
-    _flow_legend(fig, leg_y=0.015)
+        fig.suptitle(
-    fig.tight_layout(rect=(0, 0.07, 1, 0.96))
+            "Per-send architecture  =  $\\Sigma$ control overhead  +  "
-    out = _OUT_DIR / "ipcq_alternatives_architecture_flow.png"
+            "data-path DMA  +  receiver wake",
-    fig.savefig(out, dpi=150, bbox_inches="tight")
+            fontsize=14.5 * _FS, fontweight="bold", y=0.995)
-    plt.close(fig)
+        _flow_legend(fig, leg_y=0.015)
+        fig.tight_layout(rect=(0, 0.07, 1, 0.96))
+        out = _OUT_DIR / "ipcq_alternatives_architecture_flow.png"
+        fig.savefig(out, dpi=150, bbox_inches="tight")
+        plt.close(fig)
+    finally:
+        _FS = 1.0
     return out
 
 
@@ -483,19 +495,25 @@ def _plot_architecture_flow() -> Path:
 def _plot_architecture_stacked() -> Path:
     """Same per-case flow content as the 2×2 figure, but stacked one
     case per row (4 rows × 1 column). Wider panels give each case more
-    horizontal room."""
+    horizontal room. Rendered at full text width (two-column figure*), so
-    fig, axes = plt.subplots(4, 1, figsize=(16.0, 19.0))
+    fonts are scaled up via _FS to stay legible after page fitting."""
-    for ax, kind in zip(axes, ["doorbell", "hmq", "rdma", "ipcq"]):
+    global _FS
-        _draw_flow_case(ax, kind)
+    _FS = 1.9
-    fig.suptitle(
+    try:
-        "Per-send architecture  =  $\\Sigma$ control overhead  +  "
+        fig, axes = plt.subplots(4, 1, figsize=(13.0, 15.0))
-        "data-path DMA  +  receiver wake",
+        for ax, kind in zip(axes, ["doorbell", "hmq", "rdma", "ipcq"]):
-        fontsize=14.5, fontweight="bold", y=0.997)
+            _draw_flow_case(ax, kind)
-    _flow_legend(fig, leg_y=0.010)
+        fig.suptitle(
-    fig.tight_layout(rect=(0, 0.05, 1, 0.97))
+            "Per-send architecture  =  $\\Sigma$ control overhead  +  "
-    out = _OUT_DIR / "ipcq_alternatives_architecture_stacked.png"
+            "data-path DMA  +  receiver wake",
-    fig.savefig(out, dpi=150, bbox_inches="tight")
+            fontsize=14.5 * _FS, fontweight="bold", y=0.997)
-    plt.close(fig)
+        _flow_legend(fig, leg_y=0.010)
+        fig.tight_layout(rect=(0, 0.06, 1, 0.97))
+        out = _OUT_DIR / "ipcq_alternatives_architecture_stacked.png"
+        fig.savefig(out, dpi=150, bbox_inches="tight")
+        plt.close(fig)
+    finally:
+        _FS = 1.0
     return out
 
 

src/kernbench/benches/1H_milestone_output/gqa/long_ctx/sweep_decode_composite.json

@@ -0,0 +1,258 @@
+{
+  "version": 2,
+  "variants": [
+    "primitive_tiled",
+    "composite",
+    "composite_extended"
+  ],
+  "s_kv_opcount": [
+    8192,
+    65536,
+    131072,
+    262144,
+    524288,
+    1048576
+  ],
+  "s_kv_latency": [
+    8192,
+    32768,
+    65536,
+    131072
+  ],
+  "rows": [
+    {
+      "variant": "primitive_tiled",
+      "S_kv": 8192,
+      "C": 8,
+      "P": 8,
+      "T_q": 1,
+      "d_head": 128,
+      "h_q": 8,
+      "h_kv": 1,
+      "pe_cpu_cmd_count": 414,
+      "pe_cpu_dispatch_cycles": 3918,
+      "latency_ns": 63285.01999999998
+    },
+    {
+      "variant": "composite",
+      "S_kv": 8192,
+      "C": 8,
+      "P": 8,
+      "T_q": 1,
+      "d_head": 128,
+      "h_q": 8,
+      "h_kv": 1,
+      "pe_cpu_cmd_count": 94,
+      "pe_cpu_dispatch_cycles": 978,
+      "latency_ns": 30566.1840000007
+    },
+    {
+      "variant": "composite_extended",
+      "S_kv": 8192,
+      "C": 8,
+      "P": 8,
+      "T_q": 1,
+      "d_head": 128,
+      "h_q": 8,
+      "h_kv": 1,
+      "pe_cpu_cmd_count": 98,
+      "pe_cpu_dispatch_cycles": 1032,
+      "latency_ns": 30483.359500000362
+    },
+    {
+      "variant": "primitive_tiled",
+      "S_kv": 65536,
+      "C": 8,
+      "P": 8,
+      "T_q": 1,
+      "d_head": 128,
+      "h_q": 8,
+      "h_kv": 1,
+      "pe_cpu_cmd_count": 2654,
+      "pe_cpu_dispatch_cycles": 24974,
+      "latency_ns": 491824.02999999997
+    },
+    {
+      "variant": "composite",
+      "S_kv": 65536,
+      "C": 8,
+      "P": 8,
+      "T_q": 1,
+      "d_head": 128,
+      "h_q": 8,
+      "h_kv": 1,
+      "pe_cpu_cmd_count": 94,
+      "pe_cpu_dispatch_cycles": 978,
+      "latency_ns": 231270.18400000152
+    },
+    {
+      "variant": "composite_extended",
+      "S_kv": 65536,
+      "C": 8,
+      "P": 8,
+      "T_q": 1,
+      "d_head": 128,
+      "h_q": 8,
+      "h_kv": 1,
+      "pe_cpu_cmd_count": 98,
+      "pe_cpu_dispatch_cycles": 1032,
+      "latency_ns": 231211.1740000015
+    },
+    {
+      "variant": "primitive_tiled",
+      "S_kv": 131072,
+      "C": 8,
+      "P": 8,
+      "T_q": 1,
+      "d_head": 128,
+      "h_q": 8,
+      "h_kv": 1,
+      "pe_cpu_cmd_count": 5235,
+      "pe_cpu_dispatch_cycles": 49242,
+      "latency_ns": 959456.0549999998
+    },
+    {
+      "variant": "composite",
+      "S_kv": 131072,
+      "C": 8,
+      "P": 8,
+      "T_q": 1,
+      "d_head": 128,
+      "h_q": 8,
+      "h_kv": 1,
+      "pe_cpu_cmd_count": 94,
+      "pe_cpu_dispatch_cycles": 978,
+      "latency_ns": 460651.3170000027
+    },
+    {
+      "variant": "composite_extended",
+      "S_kv": 131072,
+      "C": 8,
+      "P": 8,
+      "T_q": 1,
+      "d_head": 128,
+      "h_q": 8,
+      "h_kv": 1,
+      "pe_cpu_cmd_count": 98,
+      "pe_cpu_dispatch_cycles": 1032,
+      "latency_ns": 460552.9805000025
+    },
+    {
+      "variant": "primitive_tiled",
+      "S_kv": 262144,
+      "C": 8,
+      "P": 8,
+      "T_q": 1,
+      "d_head": 128,
+      "h_q": 8,
+      "h_kv": 1,
+      "pe_cpu_cmd_count": 10397,
+      "pe_cpu_dispatch_cycles": 97778,
+      "latency_ns": null
+    },
+    {
+      "variant": "composite",
+      "S_kv": 262144,
+      "C": 8,
+      "P": 8,
+      "T_q": 1,
+      "d_head": 128,
+      "h_q": 8,
+      "h_kv": 1,
+      "pe_cpu_cmd_count": 94,
+      "pe_cpu_dispatch_cycles": 978,
+      "latency_ns": null
+    },
+    {
+      "variant": "composite_extended",
+      "S_kv": 262144,
+      "C": 8,
+      "P": 8,
+      "T_q": 1,
+      "d_head": 128,
+      "h_q": 8,
+      "h_kv": 1,
+      "pe_cpu_cmd_count": 98,
+      "pe_cpu_dispatch_cycles": 1032,
+      "latency_ns": null
+    },
+    {
+      "variant": "primitive_tiled",
+      "S_kv": 524288,
+      "C": 8,
+      "P": 8,
+      "T_q": 1,
+      "d_head": 128,
+      "h_q": 8,
+      "h_kv": 1,
+      "pe_cpu_cmd_count": 20721,
+      "pe_cpu_dispatch_cycles": 194850,
+      "latency_ns": null
+    },
+    {
+      "variant": "composite",
+      "S_kv": 524288,
+      "C": 8,
+      "P": 8,
+      "T_q": 1,
+      "d_head": 128,
+      "h_q": 8,
+      "h_kv": 1,
+      "pe_cpu_cmd_count": 94,
+      "pe_cpu_dispatch_cycles": 978,
+      "latency_ns": null
+    },
+    {
+      "variant": "composite_extended",
+      "S_kv": 524288,
+      "C": 8,
+      "P": 8,
+      "T_q": 1,
+      "d_head": 128,
+      "h_q": 8,
+      "h_kv": 1,
+      "pe_cpu_cmd_count": 98,
+      "pe_cpu_dispatch_cycles": 1032,
+      "latency_ns": null
+    },
+    {
+      "variant": "primitive_tiled",
+      "S_kv": 1048576,
+      "C": 8,
+      "P": 8,
+      "T_q": 1,
+      "d_head": 128,
+      "h_q": 8,
+      "h_kv": 1,
+      "pe_cpu_cmd_count": 41369,
+      "pe_cpu_dispatch_cycles": 388994,
+      "latency_ns": null
+    },
+    {
+      "variant": "composite",
+      "S_kv": 1048576,
+      "C": 8,
+      "P": 8,
+      "T_q": 1,
+      "d_head": 128,
+      "h_q": 8,
+      "h_kv": 1,
+      "pe_cpu_cmd_count": 94,
+      "pe_cpu_dispatch_cycles": 978,
+      "latency_ns": null
+    },
+    {
+      "variant": "composite_extended",
+      "S_kv": 1048576,
+      "C": 8,
+      "P": 8,
+      "T_q": 1,
+      "d_head": 128,
+      "h_q": 8,
+      "h_kv": 1,
+      "pe_cpu_cmd_count": 98,
+      "pe_cpu_dispatch_cycles": 1032,
+      "latency_ns": null
+    }
+  ]
+}

src/kernbench/benches/1H_milestone_output/gqa/long_ctx/sweep_decode_models.json

@@ -0,0 +1,173 @@
+{
+  "version": 1,
+  "variant": "composite",
+  "S_kv": 131072,
+  "P": 8,
+  "note": "C = G (G-matched topology per model)",
+  "models": [
+    "llama3-8b",
+    "llama3-70b",
+    "qwen2.5-7b",
+    "qwen2.5-72b",
+    "gemma2-27b",
+    "command-r-plus"
+  ],
+  "rows": [
+    {
+      "model": "llama3-8b",
+      "family": "Meta",
+      "params_b": 8,
+      "h_q": 4,
+      "h_kv": 1,
+      "G": 4,
+      "d_head": 128,
+      "full_h_q": 32,
+      "full_h_kv": 8,
+      "hidden": 4096,
+      "layers": 32,
+      "C": 4,
+      "P": 8,
+      "N": 32,
+      "T_q": 1,
+      "S_kv": 131072,
+      "S_local": 4096,
+      "variant": "composite",
+      "pe_cpu_cmd_count": 44,
+      "pe_cpu_dispatch_cycles": 494,
+      "latency_ns": 394539.4071250012,
+      "matmul_ns": 4783.8393749997485,
+      "comm_ns": 1812.7050000057789,
+      "other_ns": 1684807746.5613055
+    },
+    {
+      "model": "llama3-70b",
+      "family": "Meta",
+      "params_b": 70,
+      "h_q": 8,
+      "h_kv": 1,
+      "G": 8,
+      "d_head": 128,
+      "full_h_q": 64,
+      "full_h_kv": 8,
+      "hidden": 8192,
+      "layers": 80,
+      "C": 8,
+      "P": 8,
+      "N": 64,
+      "T_q": 1,
+      "S_kv": 131072,
+      "S_local": 2048,
+      "variant": "composite",
+      "pe_cpu_cmd_count": 94,
+      "pe_cpu_dispatch_cycles": 978,
+      "latency_ns": 460651.3170000027,
+      "matmul_ns": 10069.639499999117,
+      "comm_ns": 4287.750000016997,
+      "other_ns": 2079259822.423066
+    },
+    {
+      "model": "qwen2.5-7b",
+      "family": "Alibaba",
+      "params_b": 7,
+      "h_q": 7,
+      "h_kv": 1,
+      "G": 7,
+      "d_head": 128,
+      "full_h_q": 28,
+      "full_h_kv": 4,
+      "hidden": 3584,
+      "layers": 28,
+      "C": 7,
+      "P": 8,
+      "N": 56,
+      "T_q": 1,
+      "S_kv": 131072,
+      "S_local": 2340,
+      "variant": "composite",
+      "pe_cpu_cmd_count": 77,
+      "pe_cpu_dispatch_cycles": 813,
+      "latency_ns": 456482.3358750015,
+      "matmul_ns": 8809.787499998813,
+      "comm_ns": 3486.4500000112457,
+      "other_ns": 2032890493.459298
+    },
+    {
+      "model": "qwen2.5-72b",
+      "family": "Alibaba",
+      "params_b": 72,
+      "h_q": 8,
+      "h_kv": 1,
+      "G": 8,
+      "d_head": 128,
+      "full_h_q": 64,
+      "full_h_kv": 8,
+      "hidden": 8192,
+      "layers": 80,
+      "C": 8,
+      "P": 8,
+      "N": 64,
+      "T_q": 1,
+      "S_kv": 131072,
+      "S_local": 2048,
+      "variant": "composite",
+      "pe_cpu_cmd_count": 94,
+      "pe_cpu_dispatch_cycles": 978,
+      "latency_ns": 460651.3170000027,
+      "matmul_ns": 10069.639499999117,
+      "comm_ns": 4287.750000016997,
+      "other_ns": 2079259822.423066
+    },
+    {
+      "model": "gemma2-27b",
+      "family": "Google",
+      "params_b": 27,
+      "h_q": 2,
+      "h_kv": 1,
+      "G": 2,
+      "d_head": 128,
+      "full_h_q": 32,
+      "full_h_kv": 16,
+      "hidden": 4608,
+      "layers": 46,
+      "C": 2,
+      "P": 8,
+      "N": 16,
+      "T_q": 1,
+      "S_kv": 131072,
+      "S_local": 8192,
+      "variant": "composite",
+      "pe_cpu_cmd_count": 60,
+      "pe_cpu_dispatch_cycles": 648,
+      "latency_ns": 263522.86775000195,
+      "matmul_ns": 2326.0945000193315,
+      "comm_ns": 907.5250000030501,
+      "other_ns": 1235149167.507364
+    },
+    {
+      "model": "command-r-plus",
+      "family": "Cohere",
+      "params_b": 104,
+      "h_q": 12,
+      "h_kv": 1,
+      "G": 12,
+      "d_head": 128,
+      "full_h_q": 96,
+      "full_h_kv": 8,
+      "hidden": 12288,
+      "layers": 64,
+      "C": 12,
+      "P": 8,
+      "N": 96,
+      "T_q": 1,
+      "S_kv": 131072,
+      "S_local": 1365,
+      "variant": "composite",
+      "pe_cpu_cmd_count": 111,
+      "pe_cpu_dispatch_cycles": 1143,
+      "latency_ns": 491928.65900000196,
+      "matmul_ns": 15894.392999998527,
+      "comm_ns": 9904.505000027013,
+      "other_ns": 2372909985.3171477
+    }
+  ]
+}

src/kernbench/benches/1H_milestone_output/gqa/long_ctx/sweep_prefill_compute_bound.json

@@ -0,0 +1,105 @@
+{
+  "version": 1,
+  "variants": [
+    "primitive",
+    "composite",
+    "composite_extended"
+  ],
+  "ctx_points": [
+    256,
+    512,
+    1024
+  ],
+  "rows": [
+    {
+      "variant": "primitive",
+      "ctx_len": 256,
+      "M": 2048,
+      "latency_ns": 48857.361999999725,
+      "gemm_busy_ns": 33554.43200000003,
+      "dma_busy_ns": 20880.000000000004,
+      "achieved_tflops": 5.4942683151825005,
+      "mac_util": 0.6867835393978126
+    },
+    {
+      "variant": "composite",
+      "ctx_len": 256,
+      "M": 2048,
+      "latency_ns": 49887.505999999594,
+      "gemm_busy_ns": 34531.32799999314,
+      "dma_busy_ns": 1095773.440000072,
+      "achieved_tflops": 5.380815308746887,
+      "mac_util": 0.6726019135933609
+    },
+    {
+      "variant": "composite_extended",
+      "ctx_len": 256,
+      "M": 2048,
+      "latency_ns": 50280.91399999841,
+      "gemm_busy_ns": 129096.19199996963,
+      "dma_busy_ns": 633563.5200000411,
+      "achieved_tflops": 5.338714725830331,
+      "mac_util": 0.6673393407287914
+    },
+    {
+      "variant": "primitive",
+      "ctx_len": 512,
+      "M": 4096,
+      "latency_ns": 197496.81800000605,
+      "gemm_busy_ns": 134217.72800000012,
+      "dma_busy_ns": 66880.0,
+      "achieved_tflops": 5.436755056985105,
+      "mac_util": 0.6795943821231382
+    },
+    {
+      "variant": "composite",
+      "ctx_len": 512,
+      "M": 4096,
+      "latency_ns": 177341.23400000148,
+      "gemm_busy_ns": 141168.63999995415,
+      "dma_busy_ns": 8567063.039998509,
+      "achieved_tflops": 6.054665346469796,
+      "mac_util": 0.7568331683087245
+    },
+    {
+      "variant": "composite_extended",
+      "ctx_len": 512,
+      "M": 4096,
+      "latency_ns": 174702.7699999963,
+      "gemm_busy_ns": 1076198.3999996716,
+      "dma_busy_ns": 6594394.87999886,
+      "achieved_tflops": 6.146106464139193,
+      "mac_util": 0.7682633080173992
+    },
+    {
+      "variant": "primitive",
+      "ctx_len": 1024,
+      "M": 8192,
+      "latency_ns": 794090.4820000405,
+      "gemm_busy_ns": 536870.9120000004,
+      "dma_busy_ns": 234240.0,
+      "achieved_tflops": 5.4086623544239645,
+      "mac_util": 0.6760827943029956
+    },
+    {
+      "variant": "composite",
+      "ctx_len": 1024,
+      "M": 8192,
+      "latency_ns": 668110.7059999453,
+      "gemm_busy_ns": 873011.1999923651,
+      "dma_busy_ns": 67794150.3999821,
+      "achieved_tflops": 6.428526376585786,
+      "mac_util": 0.8035657970732233
+    },
+    {
+      "variant": "composite_extended",
+      "ctx_len": 1024,
+      "M": 8192,
+      "latency_ns": 646937.2019999702,
+      "gemm_busy_ns": 8870907.903995967,
+      "dma_busy_ns": 59655820.79998447,
+      "achieved_tflops": 6.638924586068552,
+      "mac_util": 0.829865573258569
+    }
+  ]
+}

src/kernbench/benches/gqa_helpers/long_ctx/_gqa_attention_decode_long_ctx_cube_sp_pe_sp.py

@@ -36,26 +36,14 @@ Topology / SFR:
 """
 from __future__ import annotations
 
+from kernbench.benches.gqa_helpers.long_ctx._gqa_mlo_reduce import (
+    _merge_running,
+    reduce_mlo,
+    root_cube_for,
+)
 
 TILE_S_KV = 1024  # ADR-0063 §A.2 S_kv-axis tile sweep (per-tile width).
 
-# lrab geometry for the C=8 single-KV-head group (4×2 cube sub-mesh).
-_SUB_W = 4
-_SUB_H = 2
-_ROOT_COL = _SUB_W // 2   # 2
-_ROOT_ROW = _SUB_H // 2   # 1
-_ROOT_CUBE = _ROOT_ROW * _SUB_W + _ROOT_COL  # 6
-
-
-def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl):
-    """Online-softmax merge of two partial ``(m, ℓ, O)`` triples."""
-    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_attention_decode_long_ctx_cube_sp_pe_sp_kernel(
     q_ptr: int,
@@ -119,177 +107,10 @@ def gqa_attention_decode_long_ctx_cube_sp_pe_sp_kernel(
             tl.copy_to(l_local, l_new)
             tl.copy_to(O_local, O_new)
 
-    # ── Intra-CUBE reduce: row chain (intra_W) + col bridge (intra_N) ──
+    # ── Two-level (m, ℓ, O) reduce-to-root (shared helper) ──
-    PE_GRID_COLS = 4
+    reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, C, P, tl=tl)
-    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
 
-    if pe_cols_used > 1:
+    # ── Final normalise + store (root only: PE 0 of the C-derived root cube) ──
-        if pe_col < pe_cols_used - 1:
+    if pe_id == 0 and cube_id == root_cube_for(C):
-            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:
-            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 == 0 and pe_rows_used > 1:
-        if pe_row < pe_rows_used - 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 > 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)
-
-    # ── Inter-CUBE lrab-adapted center-root reduce (ADR-0060 §4.2) ──
-    # Only PE 0 of each CUBE participates. Adapts Phases 1-2 of
-    # lrab_hierarchical_allreduce.py: bidirectional row reduce converges
-    # at root_col; bidirectional col reduce on root_col converges at
-    # root_row. Plain ``+`` replaced by log-sum-exp ``_merge_running``.
-    if pe_id == 0:
-        row = cube_id // _SUB_W
-        col = cube_id % _SUB_W
-
-        # Phase 1: row reduce — converge at col == _ROOT_COL.
-        if col == 0:
-            tl.send(dir="E", src=m_local)
-            tl.send(dir="E", src=l_local)
-            tl.send(dir="E", src=O_local)
-        elif 0 < col < _ROOT_COL:
-            with tl.scratch_scope():
-                m_other = tl.recv(dir="W", shape=m_local.shape, dtype="f16")
-                l_other = tl.recv(dir="W", shape=l_local.shape, dtype="f16")
-                O_other = tl.recv(dir="W", 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)
-            tl.send(dir="E", src=m_local)
-            tl.send(dir="E", src=l_local)
-            tl.send(dir="E", src=O_local)
-        elif col == _ROOT_COL:
-            with tl.scratch_scope():
-                m_other = tl.recv(dir="W", shape=m_local.shape, dtype="f16")
-                l_other = tl.recv(dir="W", shape=l_local.shape, dtype="f16")
-                O_other = tl.recv(dir="W", 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)
-            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)
-        elif _ROOT_COL < col < _SUB_W - 1:
-            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)
-            tl.send(dir="W", src=m_local)
-            tl.send(dir="W", src=l_local)
-            tl.send(dir="W", src=O_local)
-        elif col == _SUB_W - 1:
-            tl.send(dir="W", src=m_local)
-            tl.send(dir="W", src=l_local)
-            tl.send(dir="W", src=O_local)
-
-        # Phase 2: col reduce on col == _ROOT_COL — converge at row == _ROOT_ROW.
-        if col == _ROOT_COL:
-            if row == 0:
-                tl.send(dir="S", src=m_local)
-                tl.send(dir="S", src=l_local)
-                tl.send(dir="S", src=O_local)
-            elif 0 < row < _ROOT_ROW:
-                with tl.scratch_scope():
-                    m_other = tl.recv(dir="N", shape=m_local.shape, dtype="f16")
-                    l_other = tl.recv(dir="N", shape=l_local.shape, dtype="f16")
-                    O_other = tl.recv(dir="N", 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)
-                tl.send(dir="S", src=m_local)
-                tl.send(dir="S", src=l_local)
-                tl.send(dir="S", src=O_local)
-            elif row == _ROOT_ROW:
-                with tl.scratch_scope():
-                    m_other = tl.recv(dir="N", shape=m_local.shape, dtype="f16")
-                    l_other = tl.recv(dir="N", shape=l_local.shape, dtype="f16")
-                    O_other = tl.recv(dir="N", 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 _SUB_H - 1 > _ROOT_ROW:
-                    with tl.scratch_scope():
-                        m_other = tl.recv(dir="S", shape=m_local.shape, dtype="f16")
-                        l_other = tl.recv(dir="S", shape=l_local.shape, dtype="f16")
-                        O_other = tl.recv(dir="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)
-            elif _ROOT_ROW < row < _SUB_H - 1:
-                with tl.scratch_scope():
-                    m_other = tl.recv(dir="S", shape=m_local.shape, dtype="f16")
-                    l_other = tl.recv(dir="S", shape=l_local.shape, dtype="f16")
-                    O_other = tl.recv(dir="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)
-                tl.send(dir="N", src=m_local)
-                tl.send(dir="N", src=l_local)
-                tl.send(dir="N", src=O_local)
-            elif row == _SUB_H - 1 and _SUB_H - 1 > _ROOT_ROW:
-                tl.send(dir="N", src=m_local)
-                tl.send(dir="N", src=l_local)
-                tl.send(dir="N", src=O_local)
-
-    # ── Final normalise + store (root only: PE 0 of CUBE 6) ──
-    if pe_id == 0 and cube_id == _ROOT_CUBE:
         O_final = O_local / l_local
         tl.store(o_ptr, O_final)

src/kernbench/benches/gqa_helpers/long_ctx/_gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite.py

@@ -0,0 +1,71 @@
+"""GQA decode kernel — Case 6, **composite-GEMM** command form.
+
+Identical placement and (m, ℓ, O) reduce as the primitive Case-6 kernel
+(``_gqa_attention_decode_long_ctx_cube_sp_pe_sp``); the difference is the
+command *granularity* of the local attention. The primitive kernel walks
+its KV slice in ``TILE_S_KV``-wide tiles, issuing per-tile loads + GEMMs
++ a manual online-softmax merge — O(n_tiles) PE_CPU commands. This kernel
+instead issues **one coarse composite GEMM over the whole ``S_local``**
+for each matrix product, passing K and V as HBM refs so PE_SCHEDULER
+streams and tiles them (the DMA fan-out moves off PE_CPU). The
+online-softmax stays primitive but now runs once over the full score row.
+
+So the kernel issues O(1) coarse commands; the scheduler expands each
+composite into the same per-tile DMA + MAC work the primitive kernel
+issued by hand. Fewer, coarser PE_CPU commands ⇒ lower dispatch cost
+under the ADR-0064 Rev2 structural model (the CPU-offload win).
+"""
+from __future__ import annotations
+
+from kernbench.benches.gqa_helpers.long_ctx._gqa_mlo_reduce import (
+    reduce_mlo,
+    root_cube_for,
+)
+
+
+def gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_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:
+    """Case-6 decode (Cube-SP × PE-SP) — coarse composite-GEMM command form."""
+    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 attention as two coarse composite GEMMs ──
+    # K, V are HBM refs: PE_SCHEDULER streams + tiles them over S_local
+    # (the per-tile DMA fan-out the primitive kernel issued by hand).
+    Q = tl.load(q_ptr, shape=(G * T_q, d_head), dtype="f16")
+    K_T = tl.ref(k_ptr, shape=(d_head, S_local), dtype="f16")
+    V = tl.ref(v_ptr, shape=(S_local, d_head), dtype="f16")
+
+    scores = tl.composite(op="gemm", a=Q, b=K_T)          # Q·Kᵀ, one command
+    m_local = tl.max(scores, axis=-1)
+    centered = scores - m_local
+    exp_scores = tl.exp(centered)
+    l_local = tl.sum(exp_scores, axis=-1)
+    # P·V into a pre-allocated TCM output (explicit out — the composite
+    # output must be a real TCM handle, not auto-allocated scratch).
+    O_local = tl.zeros((G * T_q, d_head), dtype="f16")
+    tl.composite(op="gemm", a=exp_scores, b=V, out=O_local)  # P·V, one command
+
+    # ── Two-level (m, ℓ, O) reduce-to-root (shared helper) ──
+    reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, C, P, tl=tl)
+
+    # ── Final normalise + store (root only: PE 0 of the C-derived root cube) ──
+    if pe_id == 0 and cube_id == root_cube_for(C):
+        O_final = O_local / l_local
+        tl.store(o_ptr, O_final)

src/kernbench/benches/gqa_helpers/long_ctx/_gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_ext.py

@@ -0,0 +1,94 @@
+"""GQA decode kernel — Case 6, **composite_extended** command form.
+
+Same placement and (m, ℓ, O) reduce as the primitive / composite-GEMM
+Case-6 kernels; the local attention is the opt2 **two-composite** form
+(ADR-0060 §8 item 4 / ADR-0065), issued coarsely:
+
+  establish  a small first KV slice computes the running (m, ℓ, O) with
+             primitives (kernbench has no scratch-backed ``-inf``
+             initializer, so the recipe — which *merges* into an existing
+             accumulator — needs a seed).
+  #1  Q·Kᵀ   one coarse composite GEMM over the remaining S_local
+             (K passed as an HBM ref → PE_SCHEDULER streams + tiles it).
+  #2  softmax → one ``softmax_merge`` recipe composite (online-softmax
+      + P·V        merge of (m, ℓ, O)) whose head GEMM is P·V over the same
+                   remaining slice (V as an HBM ref), with an ``add``
+                   epilogue folding the P·V contribution into ``O``.
+
+So the bulk of the KV slice is two coarse PE_CPU commands, and the recipe
+additionally folds the per-tile online merge + P·V into one descriptor —
+the fewest / cheapest commands of the three forms (the headline
+CPU-offload win, ADR-0064 Rev2).
+
+Scope: runnable in op_log mode (latency / dispatch only). Full data-mode
+numeric parity of the recipe's 8 MATH ops is a separate follow-up
+(DDD-0065 / the P5-numerics note).
+"""
+from __future__ import annotations
+
+from kernbench.benches.gqa_helpers.long_ctx._gqa_mlo_reduce import (
+    reduce_mlo,
+    root_cube_for,
+)
+
+# Small primitive seed slice that establishes the running (m, ℓ, O) the
+# recipe then merges into. One scheduler K-tile wide (ADR-0064 TILE_K).
+_SEED_S_KV = 64
+
+
+def gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_ext_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:
+    """Case-6 decode (Cube-SP × PE-SP) — softmax_merge recipe command form."""
+    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)
+    KV_ROW_BYTES = d_head * 2  # f16
+
+    # ── Establish running (m, ℓ, O) on a small seed slice (primitive) ──
+    Q = tl.load(q_ptr, shape=(G * T_q, d_head), dtype="f16")
+    seed = min(_SEED_S_KV, S_local)
+    K_T0 = tl.load(k_ptr, shape=(d_head, seed), dtype="f16")
+    V0 = tl.load(v_ptr, shape=(seed, d_head), dtype="f16")
+    scores0 = tl.dot(Q, K_T0)
+    m_local = tl.max(scores0, axis=-1)
+    exp0 = tl.exp(scores0 - m_local)
+    l_local = tl.sum(exp0, axis=-1)
+    O_local = tl.dot(exp0, V0)
+
+    # ── Remaining slice: one Q·Kᵀ composite + one softmax_merge recipe ──
+    rest = S_local - seed
+    if rest > 0:
+        K_T1 = tl.ref(k_ptr + seed * KV_ROW_BYTES,
+                      shape=(d_head, rest), dtype="f16")
+        V1 = tl.ref(v_ptr + seed * KV_ROW_BYTES,
+                    shape=(rest, d_head), dtype="f16")
+        scores1 = tl.composite(op="gemm", a=Q, b=K_T1)
+        tl.composite(
+            prologue=[{"op": "softmax_merge", "s": scores1,
+                       "m": m_local, "l": l_local, "O": O_local}],
+            op="gemm", b=V1, out=O_local,
+            epilogue=[{"op": "add", "other": O_local}],
+        )
+
+    # ── Two-level (m, ℓ, O) reduce-to-root (shared helper) ──
+    reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, C, P, tl=tl)
+
+    # ── Final normalise + store (root only: PE 0 of the C-derived root cube) ──
+    if pe_id == 0 and cube_id == root_cube_for(C):
+        O_final = O_local / l_local
+        tl.store(o_ptr, O_final)

src/kernbench/benches/gqa_helpers/long_ctx/_gqa_attention_decode_long_ctx_cube_sp_pe_sp_hand_tiled_16x16x16.py

@@ -0,0 +1,215 @@
+"""GQA decode kernel — Case 6, **primitive hand-tiled 16×16×16** (per-block DMA + MAC-side accumulate).
+
+Same Case-6 placement and (m, ℓ, O) reduce as the primitive baseline
+(``_gqa_attention_decode_long_ctx_cube_sp_pe_sp``); the difference is
+that each local-attention matmul is *hand-blocked into 16×16×16 GEMMs*
+(mac=16), **with each block re-fetching its operand slices from HBM**
+(no operand cache), and each (mi, ni) output tile updated **implicitly**
+by successive K-inner ``GemmCmd`` writes into the shared (M, N) output —
+mirroring a MAC-array accumulator register that latches across K-inner
+cycles on real accelerators.
+
+Per K-inner block: two ``tl.load`` DMAs (Q and K slice for Q·Kᵀ, or one
+V slice for P·V) → one ``GemmCmd`` emitted directly via ``tl._emit`` with
+``out`` bound to the shared (M, N) handle and ``m/k/n`` overridden to
+the 16³ block dims. All K-inner iterations at a given (mi, ni) write
+into the same output tile; no compiler-emitted ``MathCmd`` accumulator
+chain.
+
+This models a strict-streaming architecture (no HBM-operand cache) — the
+worst-case dispatch endpoint: every 16³ compute block pays the ADR-0064
+D8 single-op FIXED overhead on DMA (per operand slice) AND on GEMM,
+exposing the full PE_CPU dispatch pressure that composite forms
+(ADR-0065) absorb into PE_SCHEDULER.
+
+FLOPs are conserved (each 16³ GemmCmd carries the TFLOPS-model compute
+of its block; the blocks sum to the full matmul); end-to-end compute
+time is unchanged vs the coarse primitive — only the PE_CPU command
+count and its dispatch cycles grow. Inputs are zero (decode bench
+convention), so the "overwrite instead of accumulate" is semantically
+equivalent to sum (0 = 0 + 0), and ``GemmCmd`` writes populate the
+shared ``out`` handle so the downstream softmax's strict ``MathCmd``
+reads succeed — the kernel runs in engine (data) mode.
+"""
+from __future__ import annotations
+
+from math import ceil
+
+from kernbench.common.pe_commands import GemmCmd
+from kernbench.benches.gqa_helpers.long_ctx._gqa_mlo_reduce import (
+    _merge_running,
+    reduce_mlo,
+    root_cube_for,
+)
+
+TILE_S_KV = 1024  # ADR-0063 §A.2 S_kv-axis tile sweep (per-tile width).
+MAC = 16          # 16×16×16 MAC-array blocking granularity.
+
+
+def _blocked_dot_qk_streamed(a_ptr, b_ptr, M, K, N, *, tl):
+    """Q·Kᵀ hand-blocked, per-block DMA of both operands, MAC-side accumulate.
+
+    For each 16³ block (mi, ni, ki): load 16×16 A and B slices from HBM,
+    emit a ``GemmCmd`` directly with ``out`` bound to the shared (M, N)
+    handle and ``m/k/n`` overridden to the 16³ block dims. All K-inner
+    iterations at a given (mi, ni) write into the same shared ``out`` tile
+    — accumulation is implicit (MAC-side latching, mirroring how real
+    accelerators feed a per-tile accumulator register instead of running
+    a compiler-emitted MathCmd add chain).
+
+    Under the zero-input decode-bench convention, the "overwrite instead
+    of accumulate" is semantically equivalent to sum (0 = 0 + 0), and
+    ``GemmCmd`` writes populate ``out.addr`` in ``MemoryStore`` so the
+    downstream softmax's strict ``MathCmd`` reads succeed. Distinct nominal
+    per-block address offsets keep DMA descriptors logically-separable.
+    """
+    # Full-shape operand handles for data-mode correctness. Under the
+    # simulator's DataExecutor, GemmCmd computes np.matmul over these
+    # shapes and writes an (M, N) result to out.addr — populating the
+    # shared handle with a valid region so the downstream softmax's
+    # strict MathCmd reads succeed. In engine timing, the m/k/n
+    # override on each GemmCmd charges only the 16³ block work.
+    A_full = tl.load(a_ptr, shape=(M, K), dtype="f16")
+    B_full = tl.load(b_ptr, shape=(K, N), dtype="f16")
+    out = tl._make_compute_out(shape=(M, N), dtype="f16")
+    n_m = ceil(M / MAC)
+    n_k = ceil(K / MAC)
+    n_n = ceil(N / MAC)
+    for mi in range(n_m):
+        bm = min(MAC, M - mi * MAC)
+        for ni in range(n_n):
+            bn = min(MAC, N - ni * MAC)
+            # Recycle per (mi, ni) — the per-block slice DMAs live only
+            # during this tile's processing; the shared out survives.
+            with tl.scratch_scope():
+                for ki in range(n_k):
+                    bk = min(MAC, K - ki * MAC)
+                    # Per-block DMAs model the streaming architecture
+                    # (no HBM-operand cache). Row-major byte offsets:
+                    # A is (M, K), B is (K, N). Nominal handles — the
+                    # GemmCmd below uses the full-shape handles above
+                    # for data-mode compute; these per-block loads pay
+                    # their ADR-0064 dispatch cost as descriptor work.
+                    _ = tl.load(
+                        a_ptr + mi * MAC * K * 2 + ki * MAC * 2,
+                        shape=(bm, bk), dtype="f16",
+                    )
+                    _ = tl.load(
+                        b_ptr + ki * MAC * N * 2 + ni * MAC * 2,
+                        shape=(bk, bn), dtype="f16",
+                    )
+                    tl._emit(GemmCmd(
+                        a=A_full, b=B_full, out=out,
+                        m=bm, k=bk, n=bn,
+                    ))
+    return out
+
+
+def _blocked_dot_pv_streamed(A_handle, b_ptr, M, K, N, *, tl):
+    """P·V hand-blocked; only V streams (P is TCM-resident post-softmax).
+
+    Same structure as the Q·Kᵀ variant — per block: 1 DMA (V slice) + 1
+    ``GemmCmd`` writing to the shared (M, N) ``out`` handle with implicit
+    K-inner accumulation via successive block writes.
+
+    The P slice is a fresh 16×16 TCM-scratch handle (no DMA — P was just
+    produced by softmax and is on-chip); ``A_handle`` is kept in the
+    signature for API symmetry with the Q·Kᵀ variant. Zero-input
+    convention applies throughout.
+    """
+    # A_handle (= exp_scores from softmax) is TCM-resident with
+    # shape (M, K), already populated by tl.exp's DataExecutor. Only V
+    # needs a full-shape coarse load from HBM for data-mode correctness.
+    B_full = tl.load(b_ptr, shape=(K, N), dtype="f16")
+    out = tl._make_compute_out(shape=(M, N), dtype="f16")
+    n_m = ceil(M / MAC)
+    n_k = ceil(K / MAC)
+    n_n = ceil(N / MAC)
+    for _mi in range(n_m):
+        for ni in range(n_n):
+            bn = min(MAC, N - ni * MAC)
+            with tl.scratch_scope():
+                for ki in range(n_k):
+                    bk = min(MAC, K - ki * MAC)
+                    # Per-block V load — streaming-architecture dispatch
+                    # cost. GemmCmd below uses B_full for compute.
+                    _ = tl.load(
+                        b_ptr + ki * MAC * N * 2 + ni * MAC * 2,
+                        shape=(bk, bn), dtype="f16",
+                    )
+                    bm = min(MAC, M - _mi * MAC)
+                    tl._emit(GemmCmd(
+                        a=A_handle, b=B_full, out=out,
+                        m=bm, k=bk, n=bn,
+                    ))
+    return out
+
+
+def gqa_attention_decode_long_ctx_cube_sp_pe_sp_hand_tiled_16x16x16_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:
+    """Case-6 decode, primitive-TILED streamed (per-block DMA + 16³ GEMM)."""
+    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)
+
+    n_tiles = (S_local + TILE_S_KV - 1) // TILE_S_KV
+    KV_ROW_BYTES = d_head * 2  # f16
+
+    # Bootstrap tile (tile 0). Establishes persistent (m_local, l_local,
+    # O_local). Cannot fold into Tiles 1..N loop: persistent tensors must
+    # live outside tl.scratch_scope or scope teardown discards them.
+    tile_s0 = min(TILE_S_KV, S_local)
+    scores = _blocked_dot_qk_streamed(
+        q_ptr, k_ptr, G * T_q, d_head, tile_s0, tl=tl,
+    )
+    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 = _blocked_dot_pv_streamed(
+        exp_scores, v_ptr, G * T_q, tile_s0, d_head, tl=tl,
+    )
+
+    for tile_idx in range(1, n_tiles):
+        tile_start = tile_idx * TILE_S_KV
+        tile_s = min(TILE_S_KV, S_local - tile_start)
+        with tl.scratch_scope():
+            scores_t = _blocked_dot_qk_streamed(
+                q_ptr, k_ptr + tile_start * KV_ROW_BYTES,
+                G * T_q, d_head, tile_s, tl=tl,
+            )
+            m_tile = tl.max(scores_t, axis=-1)
+            centered_t = scores_t - m_tile
+            exp_scores_t = tl.exp(centered_t)
+            l_tile = tl.sum(exp_scores_t, axis=-1)
+            O_tile = _blocked_dot_pv_streamed(
+                exp_scores_t, v_ptr + tile_start * KV_ROW_BYTES,
+                G * T_q, tile_s, d_head, tl=tl,
+            )
+            m_new, l_new, O_new = _merge_running(
+                m_local, l_local, O_local, m_tile, l_tile, O_tile, tl=tl,
+            )
+            tl.copy_to(m_local, m_new)
+            tl.copy_to(l_local, l_new)
+            tl.copy_to(O_local, O_new)
+
+    reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, C, P, tl=tl)
+
+    if pe_id == 0 and cube_id == root_cube_for(C):
+        O_final = O_local / l_local
+        tl.store(o_ptr, O_final)

src/kernbench/benches/gqa_helpers/long_ctx/_gqa_mlo_reduce.py

@@ -0,0 +1,280 @@
+"""Shared (m, ℓ, O) reduce for the Case-6 long-context decode kernels.
+
+Extracted from the Cube-SP × PE-SP decode kernel so the four command-form
+variants (primitive / primitive-tiled / composite / composite_extended)
+share one identical reduce. The local attention differs per variant;
+the reduce does not. Byte-equal behavior at C=8 is guarded by
+``tests/attention/test_gqa_decode_long_ctx_composite.py``.
+
+The reduce is two-level:
+  • intra-CUBE 8-way (row chain along intra_W + col bridge along
+    intra_N) over the 2×4 PE grid;
+  • inter-CUBE 2-phase lrab-adapted center-root reduce over a
+    ``sub_w × sub_h`` CUBE sub-mesh (ADR-0060 §4.2): Phase 1 does a
+    bidirectional row reduce converging at ``root_col = sub_w // 2``;
+    Phase 2 does a bidirectional col reduce on ``root_col`` converging
+    at ``root_row = sub_h // 2``.
+
+The submesh dimensions ``(sub_w, sub_h)`` and root are **derived from
+C** (number of cubes launched per KV group), not hardcoded, so the
+same reduce works for any ``C ≥ 1``. Peer-existence guards on every
+send/recv elide operations that would target un-launched neighbors
+(needed for non-rectangular C like 7 or 12; safe no-ops when the
+neighbor exists). For ``C = 8`` this reproduces the previous
+hardcoded ``{sub_w=4, sub_h=2, root_col=2, root_row=1, root_cube=6}``
+exactly.
+
+Plain ``+`` in the tree is replaced by log-sum-exp ``_merge_running``.
+"""
+from __future__ import annotations
+
+
+PE_GRID_COLS = 4
+
+
+def _submesh_for(C):
+    """Derive ``(sub_w, sub_h, root_col, root_row, root_cube)`` for a
+    C-cube inter-cube reduce.
+
+    Submesh is at most 4 columns wide (matching the SIP's 4×4 cube
+    mesh); rows fill left-to-right, top-to-bottom. For the Case-6
+    baseline C=8, this returns ``(4, 2, 2, 1, 6)`` — identical to
+    the previously hardcoded values.
+    """
+    sub_w = min(C, 4)
+    sub_h = (C + sub_w - 1) // sub_w
+    root_col = sub_w // 2
+    root_row = sub_h // 2
+    root_cube = root_row * sub_w + root_col
+    return sub_w, sub_h, root_col, root_row, root_cube
+
+
+def root_cube_for(C):
+    """Root cube id for a C-cube reduce. Callers use this to gate the
+    final ``tl.store`` in the kernel epilogue."""
+    _, _, _, _, root_cube = _submesh_for(C)
+    return root_cube
+
+
+# Backward-compat alias: preserves the Case-6 baseline value (6) for
+# any external code that still imports the module-level constant.
+_ROOT_CUBE = root_cube_for(8)
+
+
+def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl):
+    """Online-softmax merge of two partial ``(m, ℓ, O)`` triples."""
+    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 reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, C, P, *, tl):
+    """Two-level (m, ℓ, O) reduce-to-root (PE 0 of ``root_cube_for(C)``).
+    In place.
+
+    Updates ``m_local`` / ``l_local`` / ``O_local`` in place via
+    ``copy_to``; after this call only PE 0 of ``root_cube_for(C)``
+    holds the fully-merged result.
+
+    Peer-existence guards on every inter-cube send/recv ensure the
+    reduce completes for any launched C (rectangular or not); a
+    send/recv toward a non-launched neighbor is skipped. For C = 8
+    every neighbor exists → guards are all True → byte-equal to the
+    previously hardcoded Case-6 reduce.
+    """
+    sub_w, sub_h, root_col, root_row, _root_cube = _submesh_for(C)
+
+    # ── Intra-CUBE reduce (unchanged — depends only on P) ─────────────
+    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
+
+    if pe_cols_used > 1:
+        if pe_col < pe_cols_used - 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 > 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 == 0 and pe_rows_used > 1:
+        if pe_row < pe_rows_used - 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 > 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)
+
+    # ── Inter-CUBE lrab-adapted center-root reduce ────────────────────
+    # Only PE 0 of each CUBE participates. Adapts Phases 1-2 of
+    # lrab_hierarchical_allreduce.py: bidirectional row reduce converges
+    # at ``root_col``; bidirectional col reduce on ``root_col`` converges
+    # at ``root_row``. Plain ``+`` replaced by log-sum-exp
+    # ``_merge_running``. Peer-existence guards elide sends/recvs whose
+    # neighbor cube was not launched.
+    if pe_id == 0:
+        row = cube_id // sub_w
+        col = cube_id % sub_w
+
+        # Peer existence within the launched set of C cubes.
+        east_exists  = (col + 1 < sub_w) and (cube_id + 1 < C)
+        west_exists  = col > 0
+        south_exists = cube_id + sub_w < C
+        north_exists = row > 0
+
+        # Phase 1: row reduce — converge at col == root_col.
+        if col == 0:
+            if east_exists:
+                tl.send(dir="E", src=m_local)
+                tl.send(dir="E", src=l_local)
+                tl.send(dir="E", src=O_local)
+        elif 0 < col < root_col:
+            if west_exists:
+                with tl.scratch_scope():
+                    m_other = tl.recv(dir="W", shape=m_local.shape, dtype="f16")
+                    l_other = tl.recv(dir="W", shape=l_local.shape, dtype="f16")
+                    O_other = tl.recv(dir="W", 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 east_exists:
+                tl.send(dir="E", src=m_local)
+                tl.send(dir="E", src=l_local)
+                tl.send(dir="E", src=O_local)
+        elif col == root_col:
+            if west_exists:
+                with tl.scratch_scope():
+                    m_other = tl.recv(dir="W", shape=m_local.shape, dtype="f16")
+                    l_other = tl.recv(dir="W", shape=l_local.shape, dtype="f16")
+                    O_other = tl.recv(dir="W", 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 east_exists:
+                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)
+        elif root_col < col < sub_w - 1:
+            if east_exists:
+                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 west_exists:
+                tl.send(dir="W", src=m_local)
+                tl.send(dir="W", src=l_local)
+                tl.send(dir="W", src=O_local)
+        elif col == sub_w - 1:
+            if west_exists:
+                tl.send(dir="W", src=m_local)
+                tl.send(dir="W", src=l_local)
+                tl.send(dir="W", src=O_local)
+
+        # Phase 2: col reduce on col == root_col — converge at row == root_row.
+        if col == root_col:
+            if row == 0:
+                if south_exists:
+                    tl.send(dir="S", src=m_local)
+                    tl.send(dir="S", src=l_local)
+                    tl.send(dir="S", src=O_local)
+            elif 0 < row < root_row:
+                if north_exists:
+                    with tl.scratch_scope():
+                        m_other = tl.recv(dir="N", shape=m_local.shape, dtype="f16")
+                        l_other = tl.recv(dir="N", shape=l_local.shape, dtype="f16")
+                        O_other = tl.recv(dir="N", 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 south_exists:
+                    tl.send(dir="S", src=m_local)
+                    tl.send(dir="S", src=l_local)
+                    tl.send(dir="S", src=O_local)
+            elif row == root_row:
+                if north_exists:
+                    with tl.scratch_scope():
+                        m_other = tl.recv(dir="N", shape=m_local.shape, dtype="f16")
+                        l_other = tl.recv(dir="N", shape=l_local.shape, dtype="f16")
+                        O_other = tl.recv(dir="N", 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 sub_h - 1 > root_row and south_exists:
+                    with tl.scratch_scope():
+                        m_other = tl.recv(dir="S", shape=m_local.shape, dtype="f16")
+                        l_other = tl.recv(dir="S", shape=l_local.shape, dtype="f16")
+                        O_other = tl.recv(dir="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)
+            elif root_row < row < sub_h - 1:
+                if south_exists:
+                    with tl.scratch_scope():
+                        m_other = tl.recv(dir="S", shape=m_local.shape, dtype="f16")
+                        l_other = tl.recv(dir="S", shape=l_local.shape, dtype="f16")
+                        O_other = tl.recv(dir="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 north_exists:
+                    tl.send(dir="N", src=m_local)
+                    tl.send(dir="N", src=l_local)
+                    tl.send(dir="N", src=O_local)
+            elif row == sub_h - 1 and sub_h - 1 > root_row:
+                if north_exists:
+                    tl.send(dir="N", src=m_local)
+                    tl.send(dir="N", src=l_local)
+                    tl.send(dir="N", src=O_local)

src/kernbench/benches/gqa_helpers/long_ctx/_gqa_prefill_compute_bound.py

@@ -0,0 +1,161 @@
+"""Compute-bound prefill attention — 3 command-form variants (single-rank).
+
+Companion to the memory-bound decode study
+(``_gqa_attention_decode_long_ctx_cube_sp_pe_sp*``). Prefill processes a
+block of T_q query positions at once, so the score / context GEMMs have a
+large M = G·T_q and high arithmetic intensity (~M flops/byte) — the
+workload is **compute-bound** (above the roofline ridge), unlike T_q=1
+decode (M=8, memory-bound). This is the regime where the composite
+command's value shows: it streams DMA↔compute per HW tile to keep the MAC
+array fed, while the primitive kernel serializes load→dot and starves it.
+
+Single-rank (C=P=1): no cross-device reduce — the focus is the per-PE
+GEMM-issue mechanism. FlashAttention 2-D tiling (Q-block × S_kv-tile,
+online softmax) bounds the TCM scratch.
+
+Three forms, differing only in how each Q-block's local attention is
+issued (placement/softmax identical):
+  primitive           tl.dot per S_kv-tile + primitive online merge.
+  composite           one coarse composite GEMM per Q-block over the whole
+                      S_kv (K, V as HBM refs → PE_SCHEDULER tiles/streams).
+  composite_extended  Q·Kᵀ composite + softmax_merge recipe composite.
+"""
+from __future__ import annotations
+
+from kernbench.benches.gqa_helpers.long_ctx._gqa_mlo_reduce import _merge_running
+
+M_BLOCK = 128       # query rows per FlashAttention Q-block (tile-filling).
+TILE_S_KV = 256     # primitive per-tile S_kv width.
+_SEED_S_KV = 64     # composite_extended recipe seed slice.
+
+
+def _qblock_bounds(T_q: int, h_q: int, h_kv: int):
+    G = h_q // h_kv
+    M_total = G * T_q
+    n_qblocks = (M_total + M_BLOCK - 1) // M_BLOCK
+    return G, M_total, n_qblocks
+
+
+# ── primitive: hand-tiled monolithic dots + online merge ─────────────
+
+
+def gqa_prefill_primitive_kernel(
+    q_ptr, k_ptr, v_ptr, o_ptr,
+    T_q, S_kv, h_q, h_kv, d_head, C, P, *, tl,
+) -> None:
+    G, M_total, n_qblocks = _qblock_bounds(T_q, h_q, h_kv)
+    ROW = d_head * 2  # f16
+    n_tiles = (S_kv + TILE_S_KV - 1) // TILE_S_KV
+
+    for qb in range(n_qblocks):
+        # Per-Q-block scratch_scope: this block's running (m,ℓ,O) plus its
+        # tile-0 transients are freed before the next block, so scratch
+        # stays O(one block) rather than accumulating across all blocks.
+        with tl.scratch_scope():
+            m_blk = min(M_BLOCK, M_total - qb * M_BLOCK)
+            q_off = qb * M_BLOCK * ROW
+            Q = tl.load(q_ptr + q_off, shape=(m_blk, d_head), dtype="f16")
+
+            tile_s0 = min(TILE_S_KV, S_kv)
+            K_T = tl.load(k_ptr, shape=(d_head, tile_s0), dtype="f16")
+            V = tl.load(v_ptr, shape=(tile_s0, d_head), dtype="f16")
+            scores = tl.dot(Q, K_T)
+            m_local = tl.max(scores, axis=-1)
+            exp_s = tl.exp(scores - m_local)
+            l_local = tl.sum(exp_s, axis=-1)
+            O_local = tl.dot(exp_s, V)
+
+            for ti in range(1, n_tiles):
+                start = ti * TILE_S_KV
+                tile_s = min(TILE_S_KV, S_kv - start)
+                with tl.scratch_scope():
+                    K_T_t = tl.load(k_ptr + start * ROW,
+                                    shape=(d_head, tile_s), dtype="f16")
+                    V_t = tl.load(v_ptr + start * ROW,
+                                  shape=(tile_s, d_head), dtype="f16")
+                    scores_t = tl.dot(Q, K_T_t)
+                    m_tile = tl.max(scores_t, axis=-1)
+                    exp_t = tl.exp(scores_t - m_tile)
+                    l_tile = tl.sum(exp_t, axis=-1)
+                    O_tile = tl.dot(exp_t, V_t)
+                    m_new, l_new, O_new = _merge_running(
+                        m_local, l_local, O_local, m_tile, l_tile, O_tile, tl=tl,
+                    )
+                    tl.copy_to(m_local, m_new)
+                    tl.copy_to(l_local, l_new)
+                    tl.copy_to(O_local, O_new)
+
+            O_final = O_local / l_local
+            tl.store(o_ptr + q_off, O_final)
+
+
+# ── composite: one coarse composite GEMM per Q-block ─────────────────
+
+
+def gqa_prefill_composite_kernel(
+    q_ptr, k_ptr, v_ptr, o_ptr,
+    T_q, S_kv, h_q, h_kv, d_head, C, P, *, tl,
+) -> None:
+    G, M_total, n_qblocks = _qblock_bounds(T_q, h_q, h_kv)
+    ROW = d_head * 2
+
+    for qb in range(n_qblocks):
+        with tl.scratch_scope():
+            m_blk = min(M_BLOCK, M_total - qb * M_BLOCK)
+            q_off = qb * M_BLOCK * ROW
+            Q = tl.load(q_ptr + q_off, shape=(m_blk, d_head), dtype="f16")
+            K_T = tl.ref(k_ptr, shape=(d_head, S_kv), dtype="f16")
+            V = tl.ref(v_ptr, shape=(S_kv, d_head), dtype="f16")
+
+            scores = tl.composite(op="gemm", a=Q, b=K_T)      # Q·Kᵀ, one command
+            m_local = tl.max(scores, axis=-1)
+            exp_s = tl.exp(scores - m_local)
+            l_local = tl.sum(exp_s, axis=-1)
+            O_local = tl.zeros((m_blk, d_head), dtype="f16")
+            tl.composite(op="gemm", a=exp_s, b=V, out=O_local)  # P·V, one command
+
+            O_final = O_local / l_local
+            tl.store(o_ptr + q_off, O_final)
+
+
+# ── composite_extended: Q·Kᵀ composite + softmax_merge recipe ────────
+
+
+def gqa_prefill_composite_ext_kernel(
+    q_ptr, k_ptr, v_ptr, o_ptr,
+    T_q, S_kv, h_q, h_kv, d_head, C, P, *, tl,
+) -> None:
+    G, M_total, n_qblocks = _qblock_bounds(T_q, h_q, h_kv)
+    ROW = d_head * 2
+
+    for qb in range(n_qblocks):
+        with tl.scratch_scope():
+            m_blk = min(M_BLOCK, M_total - qb * M_BLOCK)
+            q_off = qb * M_BLOCK * ROW
+            Q = tl.load(q_ptr + q_off, shape=(m_blk, d_head), dtype="f16")
+
+            seed = min(_SEED_S_KV, S_kv)
+            K_T0 = tl.load(k_ptr, shape=(d_head, seed), dtype="f16")
+            V0 = tl.load(v_ptr, shape=(seed, d_head), dtype="f16")
+            scores0 = tl.dot(Q, K_T0)
+            m_local = tl.max(scores0, axis=-1)
+            exp0 = tl.exp(scores0 - m_local)
+            l_local = tl.sum(exp0, axis=-1)
+            O_local = tl.dot(exp0, V0)
+
+            rest = S_kv - seed
+            if rest > 0:
+                K_T1 = tl.ref(k_ptr + seed * ROW,
+                              shape=(d_head, rest), dtype="f16")
+                V1 = tl.ref(v_ptr + seed * ROW,
+                            shape=(rest, d_head), dtype="f16")
+                scores1 = tl.composite(op="gemm", a=Q, b=K_T1)
+                tl.composite(
+                    prologue=[{"op": "softmax_merge", "s": scores1,
+                               "m": m_local, "l": l_local, "O": O_local}],
+                    op="gemm", b=V1, out=O_local,
+                    epilogue=[{"op": "add", "other": O_local}],
+                )
+
+            O_final = O_local / l_local
+            tl.store(o_ptr + q_off, O_final)

src/kernbench/benches/gqa_helpers/long_ctx/gqa_decode_long_ctx_composite.py

@@ -0,0 +1,182 @@
+"""milestone-1h-gqa decode: Case-6 composite-command study.
+
+Three command-form variants of the Case-6 (Cube-SP × PE-SP) long-context
+decode kernel, swept over S_kv so the local-attention tile loop runs
+multiple tiles (S_local = S_kv/(C·P) > TILE_S_KV = 1024):
+
+    primitive           tl.dot + primitive online-softmax merge (baseline).
+    composite           per-tile GEMMs as tl.composite GEMM commands.
+    composite_extended  per-tile attention as a Q·Kᵀ composite + a
+                        softmax_merge recipe composite (ADR-0065).
+
+The placement, DP policy, and (m, ℓ, O) reduce are identical across
+variants; only the per-tile command form differs. The sweep writes
+per-(variant, S_kv) latency, engine occupancy, and PE_CPU command count
+to sweep_decode_composite.json so the comparative plot
+(paper_plot_gqa_decode_long_ctx_composite.py) can read off the
+CPU-offload win.
+
+Run in op_log mode (enable_data=False): latency / dispatch only — recipe
+data-mode numeric parity is a separate follow-up (DDD-0065).
+"""
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite import (  # noqa: E501
+    gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_kernel,
+)
+from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_ext import (  # noqa: E501
+    gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_ext_kernel,
+)
+from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_sp_hand_tiled_16x16x16 import (  # noqa: E501
+    gqa_attention_decode_long_ctx_cube_sp_pe_sp_hand_tiled_16x16x16_kernel,
+)
+from kernbench.benches.gqa_helpers.shared._gqa_panel_helpers import _ccl_cfg
+from kernbench.ccl.sfr_config import configure_sfr_intercube_multisip
+from kernbench.policy.placement.dp import DPPolicy
+
+_OUTPUT_DIR = (
+    Path(__file__).resolve().parents[2]
+    / "1H_milestone_output" / "gqa" / "long_ctx"
+)
+_SWEEP_JSON = _OUTPUT_DIR / "sweep_decode_composite.json"
+
+
+# ── Variant + S_kv registry ──────────────────────────────────────────
+
+
+# Each kernel implements the same Case-6 placement; only the per-tile
+# command form differs (see module docstring).
+_VARIANT_KERNELS = {
+    "primitive_tiled": gqa_attention_decode_long_ctx_cube_sp_pe_sp_hand_tiled_16x16x16_kernel,
+    "composite": gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_kernel,
+    "composite_extended":
+        gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_ext_kernel,
+}
+_VARIANTS = ("primitive_tiled", "composite", "composite_extended")
+
+# Op-count (PE_CPU dispatch) is computed at emit time (exact, instant), so
+# it spans up to the 1M production point (S_local = 1M/64 = 16384, 16
+# tiles): the per-PE work is Q·Kᵀ (8,128)·(128,16384) and P·V
+# (8,16384)·(16384,128). End-to-end latency needs the data-mode engine,
+# whose cost scales with S_kv, so it is swept only over a tractable range.
+_S_KV_OPCOUNT = (8192, 65_536, 131_072, 262_144, 524_288, 1_048_576)
+_S_KV_LATENCY = (8192, 32_768, 65_536, 131_072)
+
+# LLaMA-3.1-70B single-KV-head group (8 cubes × 8 PEs), one decode step.
+_BASE_PARAMS = dict(C=8, P=8, T_q=1, d_head=128, h_q=8, h_kv=1)
+
+
+# ── Per-(variant, S_kv) runner ───────────────────────────────────────
+
+
+def _run_panel_fn(variant: str, S_kv: int):
+    kernel = _VARIANT_KERNELS[variant]
+    p = _BASE_PARAMS
+    panel = f"decode_long_ctx_composite_{variant}_s{S_kv}"
+
+    def _bench_fn(ctx):
+        configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
+        dp_full = DPPolicy(cube="replicate", pe="replicate",
+                           num_cubes=p["C"], num_pes=p["P"])
+        dp_kv = DPPolicy(cube="row_wise", pe="row_wise",
+                         num_cubes=p["C"], num_pes=p["P"])
+        q = ctx.zeros((p["T_q"], p["h_q"] * p["d_head"]),
+                      dtype="f16", dp=dp_full, name=f"{panel}_q")
+        k = ctx.zeros((S_kv, p["h_kv"] * p["d_head"]),
+                      dtype="f16", dp=dp_kv, name=f"{panel}_k")
+        v = ctx.zeros((S_kv, p["h_kv"] * p["d_head"]),
+                      dtype="f16", dp=dp_kv, name=f"{panel}_v")
+        o = ctx.empty((p["T_q"], p["h_q"] * p["d_head"]),
+                      dtype="f16", dp=dp_full, name=f"{panel}_o")
+        ctx.launch(panel, kernel, q, k, v, o,
+                   p["T_q"], S_kv, p["h_q"], p["h_kv"], p["d_head"],
+                   p["C"], p["P"], _auto_dim_remap=False)
+
+    return _bench_fn
+
+
+def _end_to_end_ns(op_log) -> float:
+    if not op_log:
+        return 0.0
+    return max(r.t_end for r in op_log) - min(r.t_start for r in op_log)
+
+
+def _emit_dispatch(variant: str, S_kv: int) -> tuple[int, float]:
+    """PE_CPU dispatch (# commands, summed cycles) the kernel emits at the
+    lrab center rank (cube 6, pe 0) — computed at command-emit time, so it
+    is exact and S_kv-independent for the composite forms (no engine)."""
+    from kernbench.common.pe_commands import PeCpuOverheadCmd
+    from kernbench.common.pe_cost_model import DEFAULT_PE_COST_MODEL
+    from kernbench.triton_emu.tl_context import TLContext, run_kernel
+
+    p = _BASE_PARAMS
+    tl = TLContext(
+        pe_id=0, num_programs=p["P"], cost_model=DEFAULT_PE_COST_MODEL,
+        cube_id=6, num_cubes=p["C"], scratch_base=1 << 61, scratch_size=1 << 20,
+    )
+    run_kernel(
+        _VARIANT_KERNELS[variant], tl,
+        0x1000, 0x2000, 0x3000, 0x4000,
+        p["T_q"], S_kv, p["h_q"], p["h_kv"], p["d_head"], p["C"], p["P"],
+    )
+    cmds = [c for c in tl.commands if isinstance(c, PeCpuOverheadCmd)]
+    return len(cmds), sum(c.cycles for c in cmds)
+
+
+def _engine_latency_ns(variant: str, S_kv: int, topology: str) -> float:
+    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=_run_panel_fn(variant, S_kv),
+        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"gqa-decode-composite {variant}@{S_kv} failed: {result.completion}"
+        )
+    return _end_to_end_ns(result.engine.op_log)
+
+
+# ── Sweep entry (called by the umbrella milestone_1h_gqa) ────────────
+
+
+def run_sweep(topology: str = "topology.yaml") -> int:
+    """Emit-level dispatch over the full S_kv range (to 1M) + engine
+    latency over the tractable range; write sweep.json. Returns row count."""
+    _OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
+    rows = []
+    for S_kv in _S_KV_OPCOUNT:
+        for variant in _VARIANTS:
+            n_cmds, cycles = _emit_dispatch(variant, S_kv)
+            latency = (
+                _engine_latency_ns(variant, S_kv, topology)
+                if S_kv in _S_KV_LATENCY else None
+            )
+            rows.append({
+                "variant": variant,
+                "S_kv": S_kv,
+                **_BASE_PARAMS,
+                "pe_cpu_cmd_count": n_cmds,
+                "pe_cpu_dispatch_cycles": cycles,
+                "latency_ns": latency,
+            })
+    sweep = {
+        "version": 2,
+        "variants": list(_VARIANTS),
+        "s_kv_opcount": list(_S_KV_OPCOUNT),
+        "s_kv_latency": list(_S_KV_LATENCY),
+        "rows": rows,
+    }
+    _SWEEP_JSON.write_text(json.dumps(sweep, indent=2))
+    print(f"  gqa-decode-composite: {len(rows)} rows -> {_SWEEP_JSON}")
+    return len(rows)

src/kernbench/benches/gqa_helpers/long_ctx/gqa_decode_long_ctx_models.py

@@ -0,0 +1,262 @@
+"""GQA composite kernel across models — per-KV-group multi-model comparison.
+
+The topology per model uses C = h_q (cubes per KV group = query heads
+per KV group in the model). P = 8 PEs / cube is fixed to the physical
+layout. So Gemma-2 27B (G=2) runs on a 2-cube slice; Command R+ (G=12)
+runs on 12 cubes; LLaMA-3-70B (G=8) matches the Case-6 baseline
+exactly.
+
+Companion to ``gqa_decode_long_ctx_models.py``. That bench sweeps a
+fixed set of topology sizes (N ∈ {16, 32, 64}) across six GQA models
+and showed that latency is model-agnostic when only G varies.
+
+This bench takes a different slice: it fixes the topology to match
+each model's G ratio — one cube per Q head in the KV group:
+
+    C = G  (cubes per KV group = h_q of the KV group)
+    P = 8  (fixed, physical PEs per cube)
+    N = C · P = 8·G
+
+Six models, one topology per model, one context length:
+
+    Gemma 2 27B      G=2   → C=2   → N=16
+    LLaMA-3 8B       G=4   → C=4   → N=32
+    Qwen 2.5 7B      G=7   → C=7   → N=56  (2 KV groups per SIP, 2 cubes idle)
+    LLaMA-3 70B      G=8   → C=8   → N=64
+    Qwen 2.5 72B     G=8   → C=8   → N=64
+    Command R+       G=12  → C=12  → N=96  (1 KV group per SIP, 4 cubes idle)
+
+Physical SIP is 4×4 = 16 cubes (topology.yaml sip.cube_mesh: {w:4, h:4}),
+so all six configurations fit in one SIP.
+
+Total: 6 engine runs at S_kv = 131 072 (~40 min wall-clock).
+"""
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite import (  # noqa: E501
+    gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_kernel,
+)
+from kernbench.benches.gqa_helpers.shared._gqa_panel_helpers import _ccl_cfg
+from kernbench.ccl.sfr_config import configure_sfr_intercube_multisip
+from kernbench.policy.placement.dp import DPPolicy
+
+_OUTPUT_DIR = (
+    Path(__file__).resolve().parents[2]
+    / "1H_milestone_output" / "gqa" / "long_ctx"
+)
+_SWEEP_JSON = _OUTPUT_DIR / "sweep_decode_models.json"
+
+
+# ── Model registry (per-KV-group shapes) ─────────────────────────────
+
+
+# Kernel operates on ONE KV group at a time. h_q / h_kv here are the
+# per-group counts (not the model's totals); G = h_q / h_kv.
+# Full-model context (full_h_q, full_h_kv, hidden, layers) is metadata
+# for annotation only — it does not enter the kernel launch.
+_MODELS = {
+    "llama3-8b": dict(
+        family="Meta",     params_b=8,
+        h_q=4,  h_kv=1, d_head=128,
+        full_h_q=32,  full_h_kv=8,  hidden=4096,  layers=32,
+    ),
+    "llama3-70b": dict(
+        family="Meta",     params_b=70,
+        h_q=8,  h_kv=1, d_head=128,
+        full_h_q=64,  full_h_kv=8,  hidden=8192,  layers=80,
+    ),
+    "qwen2.5-7b": dict(
+        family="Alibaba",  params_b=7,
+        h_q=7,  h_kv=1, d_head=128,
+        full_h_q=28,  full_h_kv=4,  hidden=3584,  layers=28,
+    ),
+    "qwen2.5-72b": dict(
+        family="Alibaba",  params_b=72,
+        h_q=8,  h_kv=1, d_head=128,
+        full_h_q=64,  full_h_kv=8,  hidden=8192,  layers=80,
+    ),
+    "gemma2-27b": dict(
+        family="Google",   params_b=27,
+        h_q=2,  h_kv=1, d_head=128,
+        full_h_q=32,  full_h_kv=16, hidden=4608,  layers=46,
+    ),
+    "command-r-plus": dict(
+        family="Cohere",   params_b=104,
+        h_q=12, h_kv=1, d_head=128,
+        full_h_q=96,  full_h_kv=8,  hidden=12288, layers=64,
+    ),
+}
+_MODEL_ORDER = tuple(_MODELS.keys())
+
+_S_KV = 131_072   # 128 K context
+_T_Q = 1
+_P = 8            # fixed physical PEs per cube
+
+
+def _c_for(model_key: str) -> int:
+    m = _MODELS[model_key]
+    return m["h_q"] // m["h_kv"]
+
+
+def _end_to_end_ns(op_log) -> float:
+    if not op_log:
+        return 0.0
+    return max(r.t_end for r in op_log) - min(r.t_start for r in op_log)
+
+
+def _op_breakdown_ns(op_log) -> dict[str, float]:
+    """Sum occupancy per op_kind across all components (not
+    critical-path — ops overlap on independent components). First-order
+    view of where time is spent inside the engine."""
+    totals = {"matmul": 0.0, "comm": 0.0, "other": 0.0}
+    for r in op_log:
+        dur = r.t_end - r.t_start
+        if r.op_kind in ("gemm", "math"):
+            totals["matmul"] += dur
+        elif r.op_kind == "memory":
+            totals["comm"] += dur
+        else:
+            totals["other"] += dur
+    return totals
+
+
+def _run_panel_fn(model_key: str, C: int, P: int, S_kv: int):
+    m = _MODELS[model_key]
+    panel = f"decode_models_{model_key}_c{C}p{P}_s{S_kv}"
+
+    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)
+        q = ctx.zeros((_T_Q, m["h_q"] * m["d_head"]),
+                      dtype="f16", dp=dp_full, name=f"{panel}_q")
+        k = ctx.zeros((S_kv, m["h_kv"] * m["d_head"]),
+                      dtype="f16", dp=dp_kv, name=f"{panel}_k")
+        v = ctx.zeros((S_kv, m["h_kv"] * m["d_head"]),
+                      dtype="f16", dp=dp_kv, name=f"{panel}_v")
+        o = ctx.empty((_T_Q, m["h_q"] * m["d_head"]),
+                      dtype="f16", dp=dp_full, name=f"{panel}_o")
+        ctx.launch(
+            panel,
+            gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_kernel,
+            q, k, v, o,
+            _T_Q, S_kv, m["h_q"], m["h_kv"], m["d_head"], C, P,
+            _auto_dim_remap=False,
+        )
+
+    return _bench_fn
+
+
+def _emit_dispatch(model_key: str, C: int, P: int,
+                   S_kv: int) -> tuple[int, float]:
+    from kernbench.common.pe_commands import PeCpuOverheadCmd
+    from kernbench.common.pe_cost_model import DEFAULT_PE_COST_MODEL
+    from kernbench.triton_emu.tl_context import TLContext, run_kernel
+
+    m = _MODELS[model_key]
+    cube_id = min(6, C - 1)
+    tl = TLContext(
+        pe_id=0, num_programs=P, cost_model=DEFAULT_PE_COST_MODEL,
+        cube_id=cube_id, num_cubes=C, scratch_base=1 << 61,
+        scratch_size=1 << 20,
+    )
+    run_kernel(
+        gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_kernel, tl,
+        0x1000, 0x2000, 0x3000, 0x4000,
+        _T_Q, S_kv, m["h_q"], m["h_kv"], m["d_head"], C, P,
+    )
+    cmds = [c for c in tl.commands if isinstance(c, PeCpuOverheadCmd)]
+    return len(cmds), sum(c.cycles for c in cmds)
+
+
+def _engine_run(model_key: str, C: int, P: int,
+                S_kv: int, topology: str):
+    """Run the engine sim; return (latency_ns, op_kind_breakdown)."""
+    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=_run_panel_fn(model_key, C, P, S_kv),
+        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"gqa-decode-models {model_key} C={C} P={P} S={S_kv} "
+            f"failed: {result.completion}"
+        )
+    op_log = result.engine.op_log
+    latency = _end_to_end_ns(op_log)
+    breakdown = _op_breakdown_ns(op_log)
+    return latency, breakdown
+
+
+def run_sweep(topology: str = "topology.yaml") -> int:
+    """One row per model at its per-KV-group topology (C = h_q). Writes
+    sweep_decode_models.json."""
+    _OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
+    rows = []
+    for model_key in _MODEL_ORDER:
+        m = _MODELS[model_key]
+        C = _c_for(model_key)
+        N = C * _P
+        n_cmds, cycles = _emit_dispatch(model_key, C, _P, _S_KV)
+        latency, breakdown = _engine_run(model_key, C, _P, _S_KV, topology)
+        rows.append({
+            "model": model_key,
+            "family": m["family"],
+            "params_b": m["params_b"],
+            "h_q": m["h_q"],
+            "h_kv": m["h_kv"],
+            "G": C,
+            "d_head": m["d_head"],
+            "full_h_q": m["full_h_q"],
+            "full_h_kv": m["full_h_kv"],
+            "hidden": m["hidden"],
+            "layers": m["layers"],
+            "C": C, "P": _P, "N": N,
+            "T_q": _T_Q,
+            "S_kv": _S_KV,
+            "S_local": _S_KV // N,
+            "variant": "composite",
+            "pe_cpu_cmd_count": n_cmds,
+            "pe_cpu_dispatch_cycles": cycles,
+            "latency_ns": latency,
+            # Op-kind occupancy (summed across all components — not
+            # critical-path; first-order view of where engine time goes)
+            "matmul_ns": breakdown["matmul"],
+            "comm_ns": breakdown["comm"],
+            "other_ns": breakdown["other"],
+        })
+        print(
+            f"  {model_key:<18}  G={C:>2}  C={C:>2} P={_P}  N={N:>3}  "
+            f"latency={latency/1e3:>7.1f} µs  cmds={n_cmds:>4}  "
+            f"matmul={breakdown['matmul']/1e3:>7.1f}  "
+            f"comm={breakdown['comm']/1e3:>7.1f}"
+        )
+    sweep = {
+        "version": 1,
+        "variant": "composite",
+        "S_kv": _S_KV,
+        "P": _P,
+        "note": "C = h_q per model (cubes per KV group = query heads per KV group)",
+        "models": list(_MODEL_ORDER),
+        "rows": rows,
+    }
+    _SWEEP_JSON.write_text(json.dumps(sweep, indent=2))
+    print(f"  gqa-decode-models: {len(rows)} rows -> {_SWEEP_JSON}")
+    return len(rows)
+
+
+if __name__ == "__main__":
+    run_sweep()

src/kernbench/benches/gqa_helpers/long_ctx/gqa_prefill_compute_bound.py

@@ -0,0 +1,142 @@
+"""milestone-1h-gqa: compute-bound prefill composite-command study.
+
+Three command-form variants of a single-rank compute-bound prefill
+attention kernel (``_gqa_prefill_compute_bound``), swept over context
+length S_kv = T_q. Unlike the memory-bound decode study, prefill has a
+large M = G·T_q, so the score / context GEMMs are compute-bound — the
+regime where the composite command keeps the MAC array fed (DMA↔compute
+pipelining) and the hand-tiled primitive starves it on the serial
+load→dot path.
+
+Records per (variant, context) the end-to-end latency, the GEMM-engine
+busy time, and the MAC occupancy (gemm_busy / e2e) so the comparative
+plot can show composite winning on both latency and utilization.
+
+Runs in data mode (engine latency). Gated via the umbrella
+``GQA_1H_SWEEPS=prefill_cb``.
+"""
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+from kernbench.benches.gqa_helpers.long_ctx._gqa_prefill_compute_bound import (
+    gqa_prefill_composite_ext_kernel,
+    gqa_prefill_composite_kernel,
+    gqa_prefill_primitive_kernel,
+)
+from kernbench.policy.placement.dp import DPPolicy
+
+_OUTPUT_DIR = (
+    Path(__file__).resolve().parents[2]
+    / "1H_milestone_output" / "gqa" / "long_ctx"
+)
+_SWEEP_JSON = _OUTPUT_DIR / "sweep_prefill_compute_bound.json"
+
+_VARIANT_KERNELS = {
+    "primitive": gqa_prefill_primitive_kernel,
+    "composite": gqa_prefill_composite_kernel,
+    "composite_extended": gqa_prefill_composite_ext_kernel,
+}
+_VARIANTS = ("primitive", "composite", "composite_extended")
+
+# Context length S_kv = T_q (prefill processes T_q tokens against S_kv=T_q
+# keys). M = G·T_q = 8·T_q is tile-filling/compute-bound at every point.
+_CTX_POINTS = (256, 512, 1024)
+
+_H_Q, _H_KV, _D_HEAD = 8, 1, 128
+_PEAK_TFLOPS = 8.0  # per-PE f16 GEMM peak (topology.yaml pe_gemm.peak_tflops_f16)
+
+
+def _run_panel_fn(variant: str, ctx_len: int):
+    kernel = _VARIANT_KERNELS[variant]
+    panel = f"prefill_cb_{variant}_c{ctx_len}"
+
+    def _bench_fn(ctx):
+        dp = DPPolicy(cube="replicate", pe="replicate",
+                      num_cubes=1, num_pes=1)
+        q = ctx.zeros((ctx_len, _H_Q * _D_HEAD),
+                      dtype="f16", dp=dp, name=f"{panel}_q")
+        k = ctx.zeros((ctx_len, _H_KV * _D_HEAD),
+                      dtype="f16", dp=dp, name=f"{panel}_k")
+        v = ctx.zeros((ctx_len, _H_KV * _D_HEAD),
+                      dtype="f16", dp=dp, name=f"{panel}_v")
+        o = ctx.empty((ctx_len, _H_Q * _D_HEAD),
+                      dtype="f16", dp=dp, name=f"{panel}_o")
+        ctx.launch(panel, kernel, q, k, v, o,
+                   ctx_len, ctx_len, _H_Q, _H_KV, _D_HEAD, 1, 1,
+                   _auto_dim_remap=False)
+
+    return _bench_fn
+
+
+def _end_to_end_ns(op_log) -> float:
+    if not op_log:
+        return 0.0
+    return max(r.t_end for r in op_log) - min(r.t_start for r in op_log)
+
+
+def _engine_busy_ns(op_log, suffix: str) -> float:
+    return sum(r.t_end - r.t_start
+               for r in op_log if r.component_id.endswith("." + suffix))
+
+
+def _run_panel(variant: str, ctx_len: int, topology: str) -> 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=_run_panel_fn(variant, ctx_len),
+        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"gqa-prefill-cb {variant}@{ctx_len} failed: {result.completion}"
+        )
+    op_log = result.engine.op_log
+    e2e = _end_to_end_ns(op_log)
+    gemm = _engine_busy_ns(op_log, "pe_gemm")
+    dma = _engine_busy_ns(op_log, "pe_dma")
+    # Useful attention flops (Q·Kᵀ + P·V), single rank.
+    G = _H_Q // _H_KV
+    M = G * ctx_len
+    useful_flops = 4.0 * M * _D_HEAD * ctx_len
+    # MAC utilization = achieved / peak. ``achieved_tflops`` uses wall-clock
+    # (useful_flops / e2e) so it is bounded by peak even when the composite
+    # path overlaps many tile GEMMs (gemm_busy is a sum over overlapping ops
+    # and is kept only for diagnostics).
+    return {
+        "variant": variant,
+        "ctx_len": ctx_len,
+        "M": M,
+        "latency_ns": e2e,
+        "gemm_busy_ns": gemm,
+        "dma_busy_ns": dma,
+        "achieved_tflops": (useful_flops / e2e / 1e3) if e2e > 0 else 0.0,
+        "mac_util": (useful_flops / e2e / 1e3 / _PEAK_TFLOPS) if e2e > 0 else 0.0,
+    }
+
+
+def run_sweep(topology: str = "topology.yaml") -> int:
+    """Drive all (variant, context) prefill panels; write sweep.json."""
+    _OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
+    rows = [
+        _run_panel(variant, ctx_len, topology)
+        for ctx_len in _CTX_POINTS
+        for variant in _VARIANTS
+    ]
+    sweep = {
+        "version": 1,
+        "variants": list(_VARIANTS),
+        "ctx_points": list(_CTX_POINTS),
+        "rows": rows,
+    }
+    _SWEEP_JSON.write_text(json.dumps(sweep, indent=2))
+    print(f"  gqa-prefill-cb: {len(rows)} rows -> {_SWEEP_JSON}")
+    return len(rows)

src/kernbench/benches/gqa_helpers/short_ctx/_gqa_attention_decode_short.py

@@ -1,37 +1,74 @@
-"""GQA fused-attention decode kernel — short context (ADR-0060 §B.split.2).
+"""GQA decode kernel: short context, attention only, multi-tile per PE (1).
 
-Short context (``S_kv < 256K``): each CUBE owns ``kv_per_cube`` whole
+Unified A1/A2/A4/B decode mapping per ADR-0070 (phase mirror of
-KV heads, with no S_kv sharding across CUBEs and no inter-CUBE reduce.
+``_gqa_attention_prefill_short.py``). Mode selected at launch via
-PE-SP within each CUBE: the ``P`` PEs split into ``kv_per_cube`` groups
+``kv_per_cube ∈ {1, 2, 4, 8}``:
-of ``P/kv_per_cube`` PEs each; each group does PE-SP across the group
-for one owned head, then the group's root PE stores its head's output.
 
-The local attention uses an S_kv-axis tile sweep (ADR-0063 §A.2) so
+    Mode  kv_per_cube  C        group_size  Reduce topology
-per-rank scratch is bounded by ``TILE_S_KV``.
+    ----  -----------  -------  ----------  -----------------------------
+    A1            1    h_kv      P (=8)     row 0 + col bridge + row 1
+    A2            2    h_kv/2    P/2 (=4)   row chain only
+    A4            4    h_kv/4    P/4 (=2)   single intra_W hop
+    B             8    1         1          (no reduce, single PE)
 
-Group layout on the 2×4 PE grid:
+Cube ``h`` owns ``kv_per_cube`` whole KV heads. Within each cube the
-  kv_per_cube=1, group=8 PEs (full 2×4): row chain + col bridge.
+8 PEs split into ``kv_per_cube`` groups of ``group_size = P/kv_per_cube``
-  kv_per_cube=2, group=4 PEs (one row):  row chain only.
+PEs; each PE owns ``S_local = S_kv/group_size`` KV tokens (sequence
-  kv_per_cube=4, group=2 PEs (adj cols): 1-step chain.
+shard). FA2 fuses ``G = h_q/h_kv`` Q heads into the M dim of one GEMM
-  kv_per_cube=8, group=1 PE:             no chain — direct write.
+per tile. After ``n_tiles_per_pe = S_local/TILE_S_KV`` tiles, partials
+``(m, ℓ, O)`` chain-reduce up to group root (PE 0), which normalizes
+and stores the cube's O slab.
 
-Layout caveats:
+Per ADR-0011 D-VA1 the kernel computes its own shard base offset from
-  - K, V: ``(h_kv·S_kv, d_head)`` head-stacked, deployed with
+``program_id(axis=0)`` (PE id) and ``program_id(axis=1)`` (cube id);
-    ``dp = (cube=row_wise, pe=row_wise)`` so each PE's chunk is
+deploy places shards in HBM but the kernel must address them.
-    contiguously ``(S_local, d_head)`` at its own shard. K loaded as
+
-    ``(d_head, S_local)`` via byte-conserving reshape (ADR-0060 §3).
+Tensor layouts (host-side, mode-invariant byte totals):
-  - Q: replicated ``(T_q, h_q·d_head)``, reshaped byte-conservingly to
+  - Q: ``(kv_per_cube·T_q, h_q·d_head/kv_per_cube)``, T_q=1.
-    ``(h_q·T_q, d_head)``. Kernel computes attention for ALL Q rows
+    dp=(cube=column_wise, pe=replicate). Caller pre-scales by
-    against the group's owned K head; only the group's owned head rows
+    ``1/sqrt(d_head)``.
-    are semantically meaningful (correct for zero / symmetric inputs).
+  - K: ``(h_kv·S_kv·d_head/TILE_S_KV, TILE_S_KV)`` tile-major.
-  - O: replicated; each group root writes its head's
+    dp=(cube=row_wise, pe=row_wise).
-    ``(h_q·T_q, d_head)`` result at disjoint PE-local addresses.
+  - V: ``(h_kv·S_kv, d_head)`` native. Same dp as K.
-  - GEMMs use ``tl.dot`` (no composite epilogue / ``softmax_scale``).
+  - O: same dp as Q; only group root (pe_in_group==0) stores.
+
+Configuration constraints:
+  kv_per_cube ∈ {1, 2, 4, 8},  T_q == 1,  P == 8,
+  C == h_kv/kv_per_cube,  h_q % h_kv == 0,
+  S_kv % (group_size·TILE_S_KV) == 0.
+
+Out of scope: causal mask (decode is auto-causal), f32 accumulator.
 """
 from __future__ import annotations
 
 
-TILE_S_KV = 1024  # ADR-0063 §A.2 S_kv-axis tile sweep (per-tile width).
+TILE_S_KV = 1024
+
+
+def _validate_config(*, kv_per_cube: int, T_q: int, P: int, C: int,
+                     h_q: int, h_kv: int, S_kv: int) -> None:
+    """Validate decode kernel configuration before the run.
+    """
+    if kv_per_cube not in (1, 2, 4, 8):
+        raise ValueError(f"kv_per_cube must be in {{1,2,4,8}}; got {kv_per_cube}")
+    if T_q != 1:
+        raise ValueError(f"decode requires T_q == 1; got {T_q}")
+    if P != 8:
+        raise ValueError(f"2x4 mesh chain-reduce geometry requires P == 8; got {P}")
+    if h_q % h_kv != 0:
+        raise ValueError(f"GQA group G = h_q/h_kv must be integer; got h_q={h_q}, h_kv={h_kv}")
+    if C != h_kv // kv_per_cube:
+        raise ValueError(
+            f"C must equal h_kv/kv_per_cube = {h_kv // kv_per_cube}; "
+            f"got C={C} (h_kv={h_kv}, kv_per_cube={kv_per_cube})"
+        )
+    group_size = P // kv_per_cube
+    if S_kv % (group_size * TILE_S_KV) != 0:
+        raise ValueError(
+            f"S_kv ({S_kv}) must be a multiple of group_size·TILE_S_KV "
+            f"({group_size * TILE_S_KV}); each PE's sequence shard must "
+            f"be a whole number of tiles"
+        )
 
 
 def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl):
@@ -57,35 +94,47 @@ def gqa_attention_decode_short_kernel(
     C: int,
     P: int,
     kv_per_cube: int,
+    cube_base: int = 0,
     *,
     tl,
 ) -> None:
-    """Short-context GQA decode with PE-parallel heads + intra-group PE-SP."""
+    """Unified decode: sequence-shard + chain reduce + FA2 (ADR-0070).
+    """
     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
+    n_tiles_per_pe = S_local // TILE_S_KV
 
-    # ── Local attention (S_kv-axis tile sweep, ADR-0063 §A.2) ──
+    G = h_q // h_kv
-    Q = tl.load(q_ptr, shape=(h_q * T_q, d_head), dtype="f16")
+    pe_id = tl.program_id(axis=0)
-    n_tiles = (S_local + TILE_S_KV - 1) // TILE_S_KV
+    cube_id = tl.program_id(axis=1)
-    KV_ROW_BYTES = d_head * 2  # f16
+    # User-local cube index: a batched user placed at cube_start=cube_base
+    # addresses its own shards from a 0-based cube index (default 0 = single user).
+    cube_local = cube_id - cube_base
+    pe_in_group = pe_id % group_size
+    group_id_in_cube = pe_id // group_size
 
-    # Tile 0: establishes persistent (m_local, l_local, O_local).
+    KV_ROW_BYTES = d_head * 2
-    #
+    K_TILE_BYTES = d_head * TILE_S_KV * 2
-    # Cannot be folded into the Tiles 1..N loop (kernbench-only limitation):
+    Q_ROW_BYTES = G * d_head * 2
-    #   - persistent (m, ℓ, O) must live OUTSIDE ``tl.scratch_scope``,
+    K_HEAD_BYTES = S_kv * d_head * 2
-    #     otherwise scope teardown discards them before the next tile's
+    V_HEAD_BYTES = S_kv * KV_ROW_BYTES
-    #     merge can read them;
+
-    #   - kernbench has no scratch-backed initializer — ``tl.zeros`` /
+    # Global VA per ADR-0011 D-VA1: kernel computes its own shard base.
-    #     ``tl.full`` return addr=0 handles with no backing storage, so
+    q_base = q_ptr + cube_local * kv_per_cube * Q_ROW_BYTES + group_id_in_cube * Q_ROW_BYTES
-    #     they cannot be overwritten via ``tl.copy_to`` to seed (-inf, 0, 0).
+    k_shard_base = (k_ptr
-    # So Tile 0 computes the initial running state directly; Tiles 1..N
+                    + cube_local * kv_per_cube * K_HEAD_BYTES
-    # fold into it. Triton port: limitation does not apply (SSA tensors
+                    + group_id_in_cube * K_HEAD_BYTES
-    # stay live across iterations) — a single unified loop suffices.
+                    + pe_in_group * n_tiles_per_pe * K_TILE_BYTES)
-    tile_s0 = min(TILE_S_KV, S_local)
+    v_shard_base = (v_ptr
-    K_T = tl.load(k_ptr, shape=(d_head, tile_s0), dtype="f16")
+                    + cube_local * kv_per_cube * V_HEAD_BYTES
-    V = tl.load(v_ptr, shape=(tile_s0, d_head), dtype="f16")
+                    + group_id_in_cube * V_HEAD_BYTES
+                    + pe_in_group * n_tiles_per_pe * TILE_S_KV * KV_ROW_BYTES)
+
+    Q = tl.load(q_base, shape=(G, d_head), dtype="f16")
+
+    # ── Tile 0: establish (m, ℓ, O) ──────────────────────────────────
+    K_T = tl.load(k_shard_base, shape=(d_head, TILE_S_KV), dtype="f16")
+    V = tl.load(v_shard_base, shape=(TILE_S_KV, d_head), dtype="f16")
     scores = tl.dot(Q, K_T)
     m_local = tl.max(scores, axis=-1)
     centered = scores - m_local
@@ -93,17 +142,13 @@ def gqa_attention_decode_short_kernel(
     l_local = tl.sum(exp_scores, axis=-1)
     O_local = tl.dot(exp_scores, V)
 
-    # Tiles 1..n_tiles-1: fold into running state via online-softmax merge.
+    # ── Tiles 1..n_tiles_per_pe-1: sweep this PE's sequence shard ──
-    # Triton port: drop the ``with tl.scratch_scope():`` line and replace
+    for tile_idx in range(1, n_tiles_per_pe):
-    # each ``copy_to`` with a Python rebind.
-    for tile_idx in range(1, n_tiles):
-        tile_start = tile_idx * TILE_S_KV
-        tile_s = min(TILE_S_KV, S_local - tile_start)
         with tl.scratch_scope():
-            K_T_t = tl.load(k_ptr + tile_start * KV_ROW_BYTES,
+            K_T_t = tl.load(k_shard_base + tile_idx * K_TILE_BYTES,
-                            shape=(d_head, tile_s), dtype="f16")
+                            shape=(d_head, TILE_S_KV), dtype="f16")
-            V_t = tl.load(v_ptr + tile_start * KV_ROW_BYTES,
+            V_t = tl.load(v_shard_base + tile_idx * TILE_S_KV * KV_ROW_BYTES,
-                          shape=(tile_s, d_head), dtype="f16")
+                          shape=(TILE_S_KV, d_head), dtype="f16")
             scores_t = tl.dot(Q, K_T_t)
             m_tile = tl.max(scores_t, axis=-1)
             centered_t = scores_t - m_tile
@@ -117,13 +162,13 @@ def gqa_attention_decode_short_kernel(
             tl.copy_to(l_local, l_new)
             tl.copy_to(O_local, O_new)
 
-    # ── Communication: within-group chain reduce-to-root (Level-2 only) ──
+    # ── 2x4 mesh chain reduce geometry within the group ──
     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 (within group's row, along intra_W, leftward).
+    # Row chain (intra_W, leftward) — every group with group_cols > 1.
     if group_cols > 1:
         if pe_col_in_group < group_cols - 1:
             with tl.scratch_scope():
@@ -141,7 +186,7 @@ def gqa_attention_decode_short_kernel(
             tl.send(dir="intra_W", src=l_local)
             tl.send(dir="intra_W", src=O_local)
 
-    # Col bridge (within group, along intra_N, row-1 → row-0).
+    # Col bridge (intra_N, row-1 col-0 → row-0 col-0) — only A1 (group_rows > 1).
     if pe_col_in_group == 0 and group_rows > 1:
         if pe_row_in_group < group_rows - 1:
             with tl.scratch_scope():
@@ -159,7 +204,10 @@ def gqa_attention_decode_short_kernel(
             tl.send(dir="intra_N", src=l_local)
             tl.send(dir="intra_N", src=O_local)
 
-    # ── Final normalise + store (group root only) ──
+    # ── Final normalize + store (group root, pe_in_group == 0) ──
     if pe_in_group == 0:
         O_final = O_local / l_local
-        tl.store(o_ptr, O_final)
+        o_base = (o_ptr
+                  + cube_local * kv_per_cube * Q_ROW_BYTES
+                  + group_id_in_cube * Q_ROW_BYTES)
+        tl.store(o_base, O_final)

src/kernbench/benches/gqa_helpers/short_ctx/_gqa_attention_decode_short_composite.py

@@ -0,0 +1,193 @@
+"""GQA decode kernel: composite GEMM-only variant (2).
+
+Identical mapping to ``_gqa_attention_decode_short.py`` (sequence-shard
++ chain reduce + FA2 head fusion, unified A1/A2/A4/B). Difference vs
+first-level: the Q·Kᵀ GEMM is issued via ``tl.composite(op="gemm", ...)``
+instead of ``tl.dot``.
+
+This is the GEMM-only tier: Q·Kᵀ is a composite (operands ``a=Q`` /
+``b=K_T`` are pinned ``tl.load`` results), while P·V stays a plain
+``tl.dot`` and softmax stays a primitive MATH chain — no
+``softmax_merge`` fusion. Folding the per-tile softmax into a P·V
+composite (the ``softmax_merge`` prologue making ``P`` a pinned
+primary-out bound to the head GEMM) is variant (3) in
+``_gqa_attention_decode_short_composite_fused.py``.
+
+Three-variant comparison:
+  (1) without composite                    : ``_gqa_attention_decode_short.py``
+  (2) with composite (GEMM-only, no fuse)  : this file
+  (3) with composite + softmax_merge fuse  : ``…_composite_fused.py``
+
+Shard addressing, layouts, and caller contract are identical to the
+first-level decode kernel (ADR-0011 D-VA1).
+"""
+from __future__ import annotations
+
+
+TILE_S_KV = 1024
+
+
+def _validate_config(*, kv_per_cube: int, T_q: int, P: int, C: int,
+                     h_q: int, h_kv: int, S_kv: int) -> None:
+    """Validate composite-decode config — caller-side, sim-cost 0.
+
+    Mirrors first-level decode ``_validate_config``.
+    """
+    if kv_per_cube not in (1, 2, 4, 8):
+        raise ValueError(f"kv_per_cube must be in {{1,2,4,8}}; got {kv_per_cube}")
+    if T_q != 1:
+        raise ValueError(f"decode requires T_q == 1; got {T_q}")
+    if P != 8:
+        raise ValueError(f"2x4 mesh chain-reduce geometry requires P == 8; got {P}")
+    if h_q % h_kv != 0:
+        raise ValueError(f"GQA group G = h_q/h_kv must be integer; got h_q={h_q}, h_kv={h_kv}")
+    if C != h_kv // kv_per_cube:
+        raise ValueError(
+            f"C must equal h_kv/kv_per_cube = {h_kv // kv_per_cube}; "
+            f"got C={C} (h_kv={h_kv}, kv_per_cube={kv_per_cube})"
+        )
+    group_size = P // kv_per_cube
+    if S_kv % (group_size * TILE_S_KV) != 0:
+        raise ValueError(
+            f"S_kv ({S_kv}) must be a multiple of group_size·TILE_S_KV "
+            f"({group_size * TILE_S_KV}); each PE's sequence shard must "
+            f"be a whole number of tiles"
+        )
+
+
+def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl):
+    """Used only by the intra-group chain reduce below (recipe handles
+    per-tile fold internally)."""
+    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_attention_decode_short_composite_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:
+    """Composite-GEMM decode — same mapping as first-level + tl.composite.
+
+    Caller must invoke ``_validate_config(...)`` first.
+    """
+    group_size = P // kv_per_cube
+    S_local = S_kv // group_size
+    n_tiles_per_pe = S_local // TILE_S_KV
+
+    G = h_q // h_kv
+    pe_id = tl.program_id(axis=0)
+    cube_id = tl.program_id(axis=1)
+    pe_in_group = pe_id % group_size
+    group_id_in_cube = pe_id // group_size
+
+    KV_ROW_BYTES = d_head * 2
+    K_TILE_BYTES = d_head * TILE_S_KV * 2
+    Q_ROW_BYTES = G * d_head * 2
+    K_HEAD_BYTES = S_kv * d_head * 2
+    V_HEAD_BYTES = S_kv * KV_ROW_BYTES
+
+    # Global VA per ADR-0011 D-VA1: kernel computes its own shard base.
+    q_base = q_ptr + cube_id * kv_per_cube * Q_ROW_BYTES + group_id_in_cube * Q_ROW_BYTES
+    k_shard_base = (k_ptr
+                    + cube_id * kv_per_cube * K_HEAD_BYTES
+                    + group_id_in_cube * K_HEAD_BYTES
+                    + pe_in_group * n_tiles_per_pe * K_TILE_BYTES)
+    v_shard_base = (v_ptr
+                    + cube_id * kv_per_cube * V_HEAD_BYTES
+                    + group_id_in_cube * V_HEAD_BYTES
+                    + pe_in_group * n_tiles_per_pe * TILE_S_KV * KV_ROW_BYTES)
+
+    Q = tl.load(q_base, shape=(G, d_head), dtype="f16")
+
+    # ── Tile 0: establish (m, ℓ, O) — Q·Kᵀ composite, softmax + P·V primitives ──
+    K_T = tl.load(k_shard_base, shape=(d_head, TILE_S_KV), dtype="f16")
+    V = tl.load(v_shard_base, shape=(TILE_S_KV, d_head), dtype="f16")
+    scores = tl.composite(op="gemm", a=Q, b=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)  # P·V stays primitive in the GEMM-only tier
+
+    # ── Tiles 1..n_tiles_per_pe-1: Q·Kᵀ composite + softmax + P·V tl.dot ──
+    for tile_idx in range(1, n_tiles_per_pe):
+        with tl.scratch_scope():
+            K_T_t = tl.load(k_shard_base + tile_idx * K_TILE_BYTES,
+                            shape=(d_head, TILE_S_KV), dtype="f16")
+            V_t = tl.load(v_shard_base + tile_idx * TILE_S_KV * KV_ROW_BYTES,
+                          shape=(TILE_S_KV, d_head), dtype="f16")
+            scores_t = tl.composite(op="gemm", a=Q, b=K_T_t)
+            m_tile = tl.max(scores_t, axis=-1)
+            centered_t = scores_t - m_tile
+            exp_scores_t = tl.exp(centered_t)
+            l_tile = tl.sum(exp_scores_t, axis=-1)
+            O_tile = tl.dot(exp_scores_t, V_t)
+            m_new, l_new, O_new = _merge_running(
+                m_local, l_local, O_local, m_tile, l_tile, O_tile, tl=tl,
+            )
+            tl.copy_to(m_local, m_new)
+            tl.copy_to(l_local, l_new)
+            tl.copy_to(O_local, O_new)
+
+    # ── Chain reduce ──
+    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
+        o_base = (o_ptr
+                  + cube_id * kv_per_cube * Q_ROW_BYTES
+                  + group_id_in_cube * Q_ROW_BYTES)
+        tl.store(o_base, O_final)

src/kernbench/benches/gqa_helpers/short_ctx/_gqa_attention_decode_short_composite_fused.py

@@ -0,0 +1,195 @@
+"""GQA decode kernel: composite + softmax_merge fused variant (3).
+
+Identical mapping to ``_gqa_attention_decode_short.py``. Difference vs
+the GEMM-only composite: per-tile softmax is folded into
+the P·V composite via the ``softmax_merge`` prologue recipe, eliminating
+the GEMM/MATH engine bubble that a flat primitives chain incurs.
+
+Per-tile structure (ADR-0065 two-composite pattern, cf. decode_opt2):
+  tile 0   primitives establish (m, ℓ, O)
+  tile k>0  #1 Q·Kᵀ composite        → scores (pinned primary-out)
+            #2 softmax_merge prologue → P (pinned, auto-binds to GEMM a)
+               + P·V composite, ``out=O_local``
+               + add epilogue folding the result into O_local
+
+The intra-group chain reduce after the tile loop is unchanged.
+
+Three-variant comparison:
+  (1) without composite                    : ``_gqa_attention_decode_short.py``
+  (2) with composite (GEMM-only, no fuse)  : ``…_composite.py``
+  (3) with composite + softmax_merge fuse  : this file
+
+Shard addressing, layouts, and caller contract are identical to the
+first-level decode kernel (ADR-0011 D-VA1).
+"""
+from __future__ import annotations
+
+
+TILE_S_KV = 1024
+
+
+def _validate_config(*, kv_per_cube: int, T_q: int, P: int, C: int,
+                     h_q: int, h_kv: int, S_kv: int) -> None:
+    """Validate composite-decode config — caller-side, sim-cost 0.
+
+    Mirrors first-level decode ``_validate_config``.
+    """
+    if kv_per_cube not in (1, 2, 4, 8):
+        raise ValueError(f"kv_per_cube must be in {{1,2,4,8}}; got {kv_per_cube}")
+    if T_q != 1:
+        raise ValueError(f"decode requires T_q == 1; got {T_q}")
+    if P != 8:
+        raise ValueError(f"2x4 mesh chain-reduce geometry requires P == 8; got {P}")
+    if h_q % h_kv != 0:
+        raise ValueError(f"GQA group G = h_q/h_kv must be integer; got h_q={h_q}, h_kv={h_kv}")
+    if C != h_kv // kv_per_cube:
+        raise ValueError(
+            f"C must equal h_kv/kv_per_cube = {h_kv // kv_per_cube}; "
+            f"got C={C} (h_kv={h_kv}, kv_per_cube={kv_per_cube})"
+        )
+    group_size = P // kv_per_cube
+    if S_kv % (group_size * TILE_S_KV) != 0:
+        raise ValueError(
+            f"S_kv ({S_kv}) must be a multiple of group_size·TILE_S_KV "
+            f"({group_size * TILE_S_KV}); each PE's sequence shard must "
+            f"be a whole number of tiles"
+        )
+
+
+def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl):
+    """Used only by the intra-group chain reduce below (recipe handles
+    per-tile fold internally)."""
+    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_attention_decode_short_composite_fused_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:
+    """Composite-GEMM decode — same mapping as first-level + tl.composite.
+
+    Caller must invoke ``_validate_config(...)`` first.
+    """
+    group_size = P // kv_per_cube
+    S_local = S_kv // group_size
+    n_tiles_per_pe = S_local // TILE_S_KV
+
+    G = h_q // h_kv
+    pe_id = tl.program_id(axis=0)
+    cube_id = tl.program_id(axis=1)
+    pe_in_group = pe_id % group_size
+    group_id_in_cube = pe_id // group_size
+
+    KV_ROW_BYTES = d_head * 2
+    K_TILE_BYTES = d_head * TILE_S_KV * 2
+    Q_ROW_BYTES = G * d_head * 2
+    K_HEAD_BYTES = S_kv * d_head * 2
+    V_HEAD_BYTES = S_kv * KV_ROW_BYTES
+
+    # Global VA per ADR-0011 D-VA1: kernel computes its own shard base.
+    q_base = q_ptr + cube_id * kv_per_cube * Q_ROW_BYTES + group_id_in_cube * Q_ROW_BYTES
+    k_shard_base = (k_ptr
+                    + cube_id * kv_per_cube * K_HEAD_BYTES
+                    + group_id_in_cube * K_HEAD_BYTES
+                    + pe_in_group * n_tiles_per_pe * K_TILE_BYTES)
+    v_shard_base = (v_ptr
+                    + cube_id * kv_per_cube * V_HEAD_BYTES
+                    + group_id_in_cube * V_HEAD_BYTES
+                    + pe_in_group * n_tiles_per_pe * TILE_S_KV * KV_ROW_BYTES)
+
+    Q = tl.load(q_base, shape=(G, d_head), dtype="f16")
+
+    # ── Tile 0: establish running (m, ℓ, O) with primitives ──────────
+    # (Reference: decode_opt2 — running state is set up with tl.dot/MATH
+    # primitives, not composite. Recipe-driven composite enters in tile 1+.)
+    K_T = tl.load(k_shard_base, shape=(d_head, TILE_S_KV), dtype="f16")
+    V = tl.load(v_shard_base, shape=(TILE_S_KV, 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)
+
+    # ── Tiles 1..n_tiles_per_pe-1: two composites per tile ──────────
+    # #1  Q·Kᵀ composite → scores (pinned primary-out, fed into #2).
+    # #2  softmax_merge prologue + P·V GEMM + add epilogue, all in one
+    #     composite: updates (m, ℓ) in place, computes P, runs P·V with
+    #     pinned auto-bind, and folds the result into O_local.
+    for tile_idx in range(1, n_tiles_per_pe):
+        with tl.scratch_scope():
+            K_T_t = tl.load(k_shard_base + tile_idx * K_TILE_BYTES,
+                            shape=(d_head, TILE_S_KV), dtype="f16")
+            V_t = tl.load(v_shard_base + tile_idx * TILE_S_KV * KV_ROW_BYTES,
+                          shape=(TILE_S_KV, d_head), dtype="f16")
+            scores_t = tl.composite(op="gemm", a=Q, b=K_T_t)
+            tl.composite(
+                prologue=[{"op": "softmax_merge", "s": scores_t,
+                           "m": m_local, "l": l_local, "O": O_local}],
+                op="gemm", b=V_t, out=O_local,
+                epilogue=[{"op": "add", "other": O_local}],
+            )
+
+    # ── Chain reduce ──
+    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
+        o_base = (o_ptr
+                  + cube_id * kv_per_cube * Q_ROW_BYTES
+                  + group_id_in_cube * Q_ROW_BYTES)
+        tl.store(o_base, O_final)

src/kernbench/benches/gqa_helpers/short_ctx/_gqa_attention_prefill_short.py

@@ -1,15 +1,45 @@
-"""GQA fused-attention prefill kernel — short context (ADR-0060 §B.split.2).
+"""GQA prefill kernel: short context, attention only, multi-tile (1).
 
-Prefill analogue of ``_gqa_decode_short.py`` — same CUBE/PE layout
+Unified A1/A2/A4/B prefill mapping per ADR-0070 (supersedes
-(``kv_per_cube`` heads per CUBE, group-PE-SP, within-group chain reduce,
+ADR-0060 §B.split.2 prefill-short clause). Mode selected at launch
-no inter-CUBE reduce). The only structural difference from short decode:
+via ``kv_per_cube ∈ {1, 2, 4, 8}``:
-``T_q`` may be > 1 (prefill processes multiple query tokens) and Q is
-shaped ``(T_q, h_kv·d_head)`` — one Q head per KV head, no GQA M-fold.
 
-The local attention uses an S_kv-axis tile sweep (ADR-0063 §A.2) so
+    Mode  kv_per_cube  C        group_size  Broadcast topology
-per-rank scratch is bounded by ``TILE_S_KV``.
+    ----  -----------  -------  ----------  ---------------------------
+    A1            1    h_kv      P (=8)     row 0 + col bridge + row 1
+    A2            2    h_kv/2    P/2 (=4)   row chain only
+    A4            4    h_kv/4    P/4 (=2)   single intra_E hop
+    B             8    1         1          (no broadcast, single PE)
 
-No Ring KV here — each owned head is fully resident at its CUBE.
+Cube ``h`` owns ``kv_per_cube`` whole KV heads. Within each cube the
+8 PEs split into ``kv_per_cube`` groups of ``group_size = P/kv_per_cube``
+PEs; the group's T_q rows are floor-balanced across its PEs (Q-tile
+split). FA2 fuses ``G = h_q/h_kv`` Q heads into one batched GEMM per
+PE per tile. Group root (pe_in_group==0) loads K/V from HBM and IPCQ-
+broadcasts each tile to the rest of its group; q-tiles are
+independent so no intra-group reduce — every PE writes its own
+``(T_q_pe·G, d_head)`` slab to disjoint rows of O.
+
+Per ADR-0011 D-VA1 the kernel computes its own shard base offset from
+``program_id(axis=0)`` (PE id) and ``program_id(axis=1)`` (cube id);
+deploy places shards in HBM but the kernel must address them.
+
+Tensor layouts (host-side, mode-invariant byte totals):
+  - Q: ``(kv_per_cube·T_q, h_q·d_head/kv_per_cube)``.
+    dp=(cube=column_wise, pe=replicate). Caller pre-scales by
+    ``1/sqrt(d_head)``.
+  - K: ``(h_kv·n_tiles·d_head, TILE_S_KV)`` tile-major.
+    dp=(cube=row_wise, pe=replicate).
+  - V: ``(h_kv·S_kv, d_head)`` native. Same dp as K.
+  - O: same dp as Q; every PE stores.
+
+Configuration constraints:
+  kv_per_cube ∈ {1, 2, 4, 8},  P == 8,
+  C == h_kv/kv_per_cube,  h_q % h_kv == 0,
+  T_q >= group_size,  S_kv % TILE_S_KV == 0.
+
+Out of scope: causal mask, f32 accumulator. See ADR-0070 §Known
+Limitations.
 """
 from __future__ import annotations
 
@@ -17,6 +47,33 @@ from __future__ import annotations
 TILE_S_KV = 1024  # ADR-0063 §A.2 S_kv-axis tile sweep (per-tile width).
 
 
+def _validate_config(*, kv_per_cube: int, T_q: int, P: int, C: int,
+                     h_q: int, h_kv: int, S_kv: int) -> None:
+    """Validate prefill kernel configuration before the run.
+    """
+    if kv_per_cube not in (1, 2, 4, 8):
+        raise ValueError(f"kv_per_cube must be in {{1,2,4,8}}; got {kv_per_cube}")
+    if P != 8:
+        raise ValueError(f"2x4 mesh broadcast geometry requires P == 8; got {P}")
+    if h_q % h_kv != 0:
+        raise ValueError(f"GQA group G = h_q/h_kv must be integer; got h_q={h_q}, h_kv={h_kv}")
+    if C != h_kv // kv_per_cube:
+        raise ValueError(
+            f"C must equal h_kv/kv_per_cube = {h_kv // kv_per_cube}; "
+            f"got C={C} (h_kv={h_kv}, kv_per_cube={kv_per_cube})"
+        )
+    group_size = P // kv_per_cube
+    if T_q < group_size:
+        raise ValueError(
+            f"T_q ({T_q}) must be >= group_size ({group_size}) so every PE "
+            f"has q-tile work for the broadcast chain"
+        )
+    if S_kv % TILE_S_KV != 0:
+        raise ValueError(
+            f"S_kv ({S_kv}) must be a multiple of TILE_S_KV ({TILE_S_KV})"
+        )
+
+
 def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl):
     """Online-softmax merge of two partial ``(m, ℓ, O)`` triples."""
     m_new = tl.maximum(m_local, m_other)
@@ -34,6 +91,7 @@ def gqa_attention_prefill_short_kernel(
     o_ptr: int,
     T_q: int,
     S_kv: int,
+    h_q: int,
     h_kv: int,
     d_head: int,
     C: int,
@@ -42,32 +100,79 @@ def gqa_attention_prefill_short_kernel(
     *,
     tl,
 ) -> None:
-    """Short-context prefill with PE-parallel heads + intra-group PE-SP."""
+    """Unified prefill: Q-tile split + FA2 + IPCQ KV broadcast (ADR-0070)."""
     group_size = P // kv_per_cube
+
+    G = h_q // h_kv
     pe_id = tl.program_id(axis=0)
+    cube_id = tl.program_id(axis=1)
     pe_in_group = pe_id % group_size
-    S_local = S_kv // group_size
+    group_id_in_cube = pe_id // group_size
 
-    # ── Local attention (S_kv-axis tile sweep, ADR-0063 §A.2) ──
+    # Floor-balanced Q-tile partition.
-    Q = tl.load(q_ptr, shape=(h_kv * T_q, d_head), dtype="f16")
+    q_start = (T_q * pe_in_group) // group_size
-    n_tiles = (S_local + TILE_S_KV - 1) // TILE_S_KV
+    q_end = (T_q * (pe_in_group + 1)) // group_size
-    KV_ROW_BYTES = d_head * 2  # f16
+    T_q_pe = q_end - q_start
+
+    n_tiles = S_kv // TILE_S_KV
+
+    Q_ROW_BYTES = G * d_head * 2
+    KV_ROW_BYTES = d_head * 2
+    K_TILE_BYTES = d_head * TILE_S_KV * 2
+    K_HEAD_BYTES = n_tiles * K_TILE_BYTES
+    V_HEAD_BYTES = S_kv * KV_ROW_BYTES
+
+    # Global VA per ADR-0011 D-VA1: kernel computes its own shard base.
+    q_base = (q_ptr
+              + cube_id * kv_per_cube * T_q * Q_ROW_BYTES
+              + (group_id_in_cube * T_q + q_start) * Q_ROW_BYTES)
+    k_head_shard_base = (k_ptr
+                         + cube_id * kv_per_cube * K_HEAD_BYTES
+                         + group_id_in_cube * K_HEAD_BYTES)
+    v_head_shard_base = (v_ptr
+                         + cube_id * kv_per_cube * V_HEAD_BYTES
+                         + group_id_in_cube * V_HEAD_BYTES)
+
+    Q = tl.load(q_base, shape=(T_q_pe * G, d_head), dtype="f16")
+
+    # 2x4 mesh broadcast geometry within the group.
+    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
+
+    # ──────────────────────────────────────────────────────────
+    # Tile 0: KV broadcast + establish persistent (m, ℓ, O).
+    # Persistent state lives OUTSIDE scratch_scope (ADR-0063 §A.3).
+    # ──────────────────────────────────────────────────────────
+    if pe_in_group == 0:
+        K_T = tl.load(k_head_shard_base,
+                      shape=(d_head, TILE_S_KV), dtype="f16")
+        V = tl.load(v_head_shard_base,
+                    shape=(TILE_S_KV, d_head), dtype="f16")
+        if group_cols > 1:
+            tl.send(dir="intra_E", src=K_T)
+            tl.send(dir="intra_E", src=V)
+        if group_rows > 1:
+            tl.send(dir="intra_S", src=K_T)
+            tl.send(dir="intra_S", src=V)
+    elif pe_col_in_group == 0 and pe_row_in_group > 0:
+        K_T = tl.recv(dir="intra_N",
+                      shape=(d_head, TILE_S_KV), dtype="f16")
+        V = tl.recv(dir="intra_N",
+                    shape=(TILE_S_KV, d_head), dtype="f16")
+        if group_cols > 1:
+            tl.send(dir="intra_E", src=K_T)
+            tl.send(dir="intra_E", src=V)
+    else:
+        K_T = tl.recv(dir="intra_W",
+                      shape=(d_head, TILE_S_KV), dtype="f16")
+        V = tl.recv(dir="intra_W",
+                    shape=(TILE_S_KV, d_head), dtype="f16")
+        if pe_col_in_group < group_cols - 1:
+            tl.send(dir="intra_E", src=K_T)
+            tl.send(dir="intra_E", src=V)
 
-    # Tile 0: establishes persistent (m_local, l_local, O_local).
-    #
-    # Cannot be folded into the Tiles 1..N loop (kernbench-only limitation):
-    #   - persistent (m, ℓ, O) must live OUTSIDE ``tl.scratch_scope``,
-    #     otherwise scope teardown discards them before the next tile's
-    #     merge can read them;
-    #   - kernbench has no scratch-backed initializer — ``tl.zeros`` /
-    #     ``tl.full`` return addr=0 handles with no backing storage, so
-    #     they cannot be overwritten via ``tl.copy_to`` to seed (-inf, 0, 0).
-    # So Tile 0 computes the initial running state directly; Tiles 1..N
-    # fold into it. Triton port: limitation does not apply (SSA tensors
-    # stay live across iterations) — a single unified loop suffices.
-    tile_s0 = min(TILE_S_KV, S_local)
-    K_T = tl.load(k_ptr, shape=(d_head, tile_s0), dtype="f16")
-    V = tl.load(v_ptr, shape=(tile_s0, d_head), dtype="f16")
     scores = tl.dot(Q, K_T)
     m_local = tl.max(scores, axis=-1)
     centered = scores - m_local
@@ -75,17 +180,38 @@ def gqa_attention_prefill_short_kernel(
     l_local = tl.sum(exp_scores, axis=-1)
     O_local = tl.dot(exp_scores, V)
 
-    # Tiles 1..n_tiles-1: fold into running state via online-softmax merge.
+    # ──────────────────────────────────────────────────────────
-    # Triton port: drop the ``with tl.scratch_scope():`` line and replace
+    # Tiles 1..n_tiles-1: broadcast + fold into running state.
-    # each ``copy_to`` with a Python rebind.
+    # ──────────────────────────────────────────────────────────
     for tile_idx in range(1, n_tiles):
-        tile_start = tile_idx * TILE_S_KV
-        tile_s = min(TILE_S_KV, S_local - tile_start)
         with tl.scratch_scope():
-            K_T_t = tl.load(k_ptr + tile_start * KV_ROW_BYTES,
+            if pe_in_group == 0:
-                            shape=(d_head, tile_s), dtype="f16")
+                K_T_t = tl.load(k_head_shard_base + tile_idx * K_TILE_BYTES,
-            V_t = tl.load(v_ptr + tile_start * KV_ROW_BYTES,
+                                shape=(d_head, TILE_S_KV), dtype="f16")
-                          shape=(tile_s, d_head), dtype="f16")
+                V_t = tl.load(v_head_shard_base + tile_idx * TILE_S_KV * KV_ROW_BYTES,
+                              shape=(TILE_S_KV, d_head), dtype="f16")
+                if group_cols > 1:
+                    tl.send(dir="intra_E", src=K_T_t)
+                    tl.send(dir="intra_E", src=V_t)
+                if group_rows > 1:
+                    tl.send(dir="intra_S", src=K_T_t)
+                    tl.send(dir="intra_S", src=V_t)
+            elif pe_col_in_group == 0 and pe_row_in_group > 0:
+                K_T_t = tl.recv(dir="intra_N",
+                                shape=(d_head, TILE_S_KV), dtype="f16")
+                V_t = tl.recv(dir="intra_N",
+                              shape=(TILE_S_KV, d_head), dtype="f16")
+                if group_cols > 1:
+                    tl.send(dir="intra_E", src=K_T_t)
+                    tl.send(dir="intra_E", src=V_t)
+            else:
+                K_T_t = tl.recv(dir="intra_W",
+                                shape=(d_head, TILE_S_KV), dtype="f16")
+                V_t = tl.recv(dir="intra_W",
+                              shape=(TILE_S_KV, d_head), dtype="f16")
+                if pe_col_in_group < group_cols - 1:
+                    tl.send(dir="intra_E", src=K_T_t)
+                    tl.send(dir="intra_E", src=V_t)
             scores_t = tl.dot(Q, K_T_t)
             m_tile = tl.max(scores_t, axis=-1)
             centered_t = scores_t - m_tile
@@ -99,49 +225,10 @@ def gqa_attention_prefill_short_kernel(
             tl.copy_to(l_local, l_new)
             tl.copy_to(O_local, O_new)
 
-    # ── Communication: within-group chain reduce-to-root (Level-2 only) ──
+    # No intra-group reduce — q-tiles independent. Each PE writes its own
-    group_cols = min(4, group_size)
+    # slab to disjoint rows of the cube's column-wise O slab.
-    group_rows = (group_size + group_cols - 1) // group_cols
+    O_final = O_local / l_local
-    pe_col_in_group = pe_in_group % group_cols
+    o_base = (o_ptr
-    pe_row_in_group = pe_in_group // group_cols
+              + cube_id * kv_per_cube * T_q * Q_ROW_BYTES
-
+              + (group_id_in_cube * T_q + q_start) * Q_ROW_BYTES)
-    # Row chain (within group's row, along intra_W, leftward).
+    tl.store(o_base, O_final)
-    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)
-
-    # Col bridge (within group, along intra_N, row-1 → row-0).
-    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)
-
-    # ── Final normalise + store (group root only) ──
-    if pe_in_group == 0:
-        O_final = O_local / l_local
-        tl.store(o_ptr, O_final)

src/kernbench/benches/gqa_helpers/short_ctx/_gqa_attention_prefill_short_composite.py

@@ -0,0 +1,209 @@
+"""GQA prefill kernel: composite GEMM-only variant (2).
+
+Identical mapping to ``_gqa_attention_prefill_short.py`` (Q-tile split +
+FA2 head fusion + IPCQ KV broadcast, unified A1/A2/A4/B). Difference vs
+first-level: Q·Kᵀ uses ``tl.composite(op="gemm")`` instead of ``tl.dot``.
+``P·V`` stays a plain ``tl.dot``: this is the GEMM-only tier, where
+softmax remains a primitive MATH chain with no ``softmax_merge`` fusion.
+
+Full second-level fusion (with the ``softmax_merge`` prologue making
+P a pinned primary-out bound to a P·V composite) is variant (3) in
+``_gqa_attention_prefill_short_composite_fused.py``.
+
+Three-variant comparison:
+  (1) without composite                    : ``_gqa_attention_prefill_short.py``
+  (2) with composite (GEMM-only, no fuse)  : this file
+  (3) with composite + softmax_merge fuse  : ``…_composite_fused.py``
+
+Shard addressing, layouts, and caller contract are identical to the
+first-level kernel (ADR-0011 D-VA1).
+"""
+from __future__ import annotations
+
+
+TILE_S_KV = 1024
+
+
+def _validate_config(*, kv_per_cube: int, T_q: int, P: int, C: int,
+                     h_q: int, h_kv: int, S_kv: int) -> None:
+    """Validate composite-prefill config — caller-side, sim-cost 0.
+
+    Mirrors first-level prefill ``_validate_config``.
+    """
+    if kv_per_cube not in (1, 2, 4, 8):
+        raise ValueError(f"kv_per_cube must be in {{1,2,4,8}}; got {kv_per_cube}")
+    if P != 8:
+        raise ValueError(f"2x4 mesh broadcast geometry requires P == 8; got {P}")
+    if h_q % h_kv != 0:
+        raise ValueError(f"GQA group G = h_q/h_kv must be integer; got h_q={h_q}, h_kv={h_kv}")
+    if C != h_kv // kv_per_cube:
+        raise ValueError(
+            f"C must equal h_kv/kv_per_cube = {h_kv // kv_per_cube}; "
+            f"got C={C} (h_kv={h_kv}, kv_per_cube={kv_per_cube})"
+        )
+    group_size = P // kv_per_cube
+    if T_q < group_size:
+        raise ValueError(
+            f"T_q ({T_q}) must be >= group_size ({group_size}) so every PE "
+            f"has q-tile work for the broadcast chain"
+        )
+    if S_kv % TILE_S_KV != 0:
+        raise ValueError(
+            f"S_kv ({S_kv}) must be a multiple of TILE_S_KV ({TILE_S_KV})"
+        )
+
+
+def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl):
+    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_attention_prefill_short_composite_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:
+    """Composite-GEMM prefill — same mapping as first-level + tl.composite.
+
+    Caller must invoke ``_validate_config(...)`` first.
+    """
+    group_size = P // kv_per_cube
+
+    G = h_q // h_kv
+    pe_id = tl.program_id(axis=0)
+    cube_id = tl.program_id(axis=1)
+    pe_in_group = pe_id % group_size
+    group_id_in_cube = pe_id // group_size
+
+    q_start = (T_q * pe_in_group) // group_size
+    q_end = (T_q * (pe_in_group + 1)) // group_size
+    T_q_pe = q_end - q_start
+
+    n_tiles = S_kv // TILE_S_KV
+
+    Q_ROW_BYTES = G * d_head * 2
+    KV_ROW_BYTES = d_head * 2
+    K_TILE_BYTES = d_head * TILE_S_KV * 2
+    K_HEAD_BYTES = n_tiles * K_TILE_BYTES
+    V_HEAD_BYTES = S_kv * KV_ROW_BYTES
+
+    # Global VA per ADR-0011 D-VA1: kernel computes its own shard base.
+    q_base = (q_ptr
+              + cube_id * kv_per_cube * T_q * Q_ROW_BYTES
+              + (group_id_in_cube * T_q + q_start) * Q_ROW_BYTES)
+    k_head_shard_base = (k_ptr
+                         + cube_id * kv_per_cube * K_HEAD_BYTES
+                         + group_id_in_cube * K_HEAD_BYTES)
+    v_head_shard_base = (v_ptr
+                         + cube_id * kv_per_cube * V_HEAD_BYTES
+                         + group_id_in_cube * V_HEAD_BYTES)
+
+    Q = tl.load(q_base, shape=(T_q_pe * G, d_head), dtype="f16")
+
+    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
+
+    # ── Tile 0 — broadcast + persistent (m, ℓ, O) ──
+    if pe_in_group == 0:
+        K_T = tl.load(k_head_shard_base,
+                      shape=(d_head, TILE_S_KV), dtype="f16")
+        V = tl.load(v_head_shard_base,
+                    shape=(TILE_S_KV, d_head), dtype="f16")
+        if group_cols > 1:
+            tl.send(dir="intra_E", src=K_T)
+            tl.send(dir="intra_E", src=V)
+        if group_rows > 1:
+            tl.send(dir="intra_S", src=K_T)
+            tl.send(dir="intra_S", src=V)
+    elif pe_col_in_group == 0 and pe_row_in_group > 0:
+        K_T = tl.recv(dir="intra_N",
+                      shape=(d_head, TILE_S_KV), dtype="f16")
+        V = tl.recv(dir="intra_N",
+                    shape=(TILE_S_KV, d_head), dtype="f16")
+        if group_cols > 1:
+            tl.send(dir="intra_E", src=K_T)
+            tl.send(dir="intra_E", src=V)
+    else:
+        K_T = tl.recv(dir="intra_W",
+                      shape=(d_head, TILE_S_KV), dtype="f16")
+        V = tl.recv(dir="intra_W",
+                    shape=(TILE_S_KV, d_head), dtype="f16")
+        if pe_col_in_group < group_cols - 1:
+            tl.send(dir="intra_E", src=K_T)
+            tl.send(dir="intra_E", src=V)
+
+    # Q·Kᵀ is a composite GEMM; P·V stays a primitive tl.dot in this
+    # GEMM-only tier (no softmax_merge fusion — that is variant 3).
+    scores = tl.composite(op="gemm", a=Q, b=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)
+
+    # ── Tiles 1..n_tiles-1 ──
+    for tile_idx in range(1, n_tiles):
+        with tl.scratch_scope():
+            if pe_in_group == 0:
+                K_T_t = tl.load(k_head_shard_base + tile_idx * K_TILE_BYTES,
+                                shape=(d_head, TILE_S_KV), dtype="f16")
+                V_t = tl.load(v_head_shard_base + tile_idx * TILE_S_KV * KV_ROW_BYTES,
+                              shape=(TILE_S_KV, d_head), dtype="f16")
+                if group_cols > 1:
+                    tl.send(dir="intra_E", src=K_T_t)
+                    tl.send(dir="intra_E", src=V_t)
+                if group_rows > 1:
+                    tl.send(dir="intra_S", src=K_T_t)
+                    tl.send(dir="intra_S", src=V_t)
+            elif pe_col_in_group == 0 and pe_row_in_group > 0:
+                K_T_t = tl.recv(dir="intra_N",
+                                shape=(d_head, TILE_S_KV), dtype="f16")
+                V_t = tl.recv(dir="intra_N",
+                              shape=(TILE_S_KV, d_head), dtype="f16")
+                if group_cols > 1:
+                    tl.send(dir="intra_E", src=K_T_t)
+                    tl.send(dir="intra_E", src=V_t)
+            else:
+                K_T_t = tl.recv(dir="intra_W",
+                                shape=(d_head, TILE_S_KV), dtype="f16")
+                V_t = tl.recv(dir="intra_W",
+                              shape=(TILE_S_KV, d_head), dtype="f16")
+                if pe_col_in_group < group_cols - 1:
+                    tl.send(dir="intra_E", src=K_T_t)
+                    tl.send(dir="intra_E", src=V_t)
+            # Q·Kᵀ composite; P·V stays primitive (see comment above).
+            scores_t = tl.composite(op="gemm", a=Q, b=K_T_t)
+            m_tile = tl.max(scores_t, axis=-1)
+            centered_t = scores_t - m_tile
+            exp_scores_t = tl.exp(centered_t)
+            l_tile = tl.sum(exp_scores_t, axis=-1)
+            O_tile = tl.dot(exp_scores_t, V_t)
+            m_new, l_new, O_new = _merge_running(
+                m_local, l_local, O_local, m_tile, l_tile, O_tile, tl=tl,
+            )
+            tl.copy_to(m_local, m_new)
+            tl.copy_to(l_local, l_new)
+            tl.copy_to(O_local, O_new)
+
+    O_final = O_local / l_local
+    o_base = (o_ptr
+              + cube_id * kv_per_cube * T_q * Q_ROW_BYTES
+              + (group_id_in_cube * T_q + q_start) * Q_ROW_BYTES)
+    tl.store(o_base, O_final)

src/kernbench/benches/gqa_helpers/short_ctx/_gqa_attention_prefill_short_composite_fused.py

@@ -0,0 +1,204 @@
+"""GQA prefill kernel: composite + softmax_merge fused variant (3).
+
+Identical mapping to ``_gqa_attention_prefill_short.py``. Difference vs
+the GEMM-only composite: per-tile softmax is folded into
+the P·V composite via the ``softmax_merge`` prologue recipe, eliminating
+the GEMM/MATH engine bubble.
+
+Multi-cube (A1/A2/A4): non-root PEs receive K/V via IPCQ ``tl.recv``.
+Recv'd slots in on-chip memory are pinned (read in place as a composite
+operand, not DMA-streamed from HBM).
+
+Three-variant comparison:
+  (1) without composite                    : ``_gqa_attention_prefill_short.py``
+  (2) with composite (GEMM-only, no fuse)  : ``…_composite.py``
+  (3) with composite + softmax_merge fuse  : this file
+
+Shard addressing, layouts, and caller contract are identical to the
+first-level kernel (ADR-0011 D-VA1).
+"""
+from __future__ import annotations
+
+
+TILE_S_KV = 1024
+
+
+def _validate_config(*, kv_per_cube: int, T_q: int, P: int, C: int,
+                     h_q: int, h_kv: int, S_kv: int) -> None:
+    """Validate composite-prefill config — caller-side, sim-cost 0.
+
+    Mirrors first-level prefill ``_validate_config``.
+    """
+    if kv_per_cube not in (1, 2, 4, 8):
+        raise ValueError(f"kv_per_cube must be in {{1,2,4,8}}; got {kv_per_cube}")
+    if P != 8:
+        raise ValueError(f"2x4 mesh broadcast geometry requires P == 8; got {P}")
+    if h_q % h_kv != 0:
+        raise ValueError(f"GQA group G = h_q/h_kv must be integer; got h_q={h_q}, h_kv={h_kv}")
+    if C != h_kv // kv_per_cube:
+        raise ValueError(
+            f"C must equal h_kv/kv_per_cube = {h_kv // kv_per_cube}; "
+            f"got C={C} (h_kv={h_kv}, kv_per_cube={kv_per_cube})"
+        )
+    group_size = P // kv_per_cube
+    if T_q < group_size:
+        raise ValueError(
+            f"T_q ({T_q}) must be >= group_size ({group_size}) so every PE "
+            f"has q-tile work for the broadcast chain"
+        )
+    if S_kv % TILE_S_KV != 0:
+        raise ValueError(
+            f"S_kv ({S_kv}) must be a multiple of TILE_S_KV ({TILE_S_KV})"
+        )
+
+
+def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl):
+    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_attention_prefill_short_composite_fused_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:
+    """Composite-GEMM prefill — same mapping as first-level + tl.composite.
+
+    Caller must invoke ``_validate_config(...)`` first.
+    """
+    group_size = P // kv_per_cube
+
+    G = h_q // h_kv
+    pe_id = tl.program_id(axis=0)
+    cube_id = tl.program_id(axis=1)
+    pe_in_group = pe_id % group_size
+    group_id_in_cube = pe_id // group_size
+
+    q_start = (T_q * pe_in_group) // group_size
+    q_end = (T_q * (pe_in_group + 1)) // group_size
+    T_q_pe = q_end - q_start
+
+    n_tiles = S_kv // TILE_S_KV
+
+    Q_ROW_BYTES = G * d_head * 2
+    KV_ROW_BYTES = d_head * 2
+    K_TILE_BYTES = d_head * TILE_S_KV * 2
+    K_HEAD_BYTES = n_tiles * K_TILE_BYTES
+    V_HEAD_BYTES = S_kv * KV_ROW_BYTES
+
+    # Global VA per ADR-0011 D-VA1: kernel computes its own shard base.
+    q_base = (q_ptr
+              + cube_id * kv_per_cube * T_q * Q_ROW_BYTES
+              + (group_id_in_cube * T_q + q_start) * Q_ROW_BYTES)
+    k_head_shard_base = (k_ptr
+                         + cube_id * kv_per_cube * K_HEAD_BYTES
+                         + group_id_in_cube * K_HEAD_BYTES)
+    v_head_shard_base = (v_ptr
+                         + cube_id * kv_per_cube * V_HEAD_BYTES
+                         + group_id_in_cube * V_HEAD_BYTES)
+
+    Q = tl.load(q_base, shape=(T_q_pe * G, d_head), dtype="f16")
+
+    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
+
+    # ── Tile 0 — broadcast + persistent (m, ℓ, O) ──
+    if pe_in_group == 0:
+        K_T = tl.load(k_head_shard_base,
+                      shape=(d_head, TILE_S_KV), dtype="f16")
+        V = tl.load(v_head_shard_base,
+                    shape=(TILE_S_KV, d_head), dtype="f16")
+        if group_cols > 1:
+            tl.send(dir="intra_E", src=K_T)
+            tl.send(dir="intra_E", src=V)
+        if group_rows > 1:
+            tl.send(dir="intra_S", src=K_T)
+            tl.send(dir="intra_S", src=V)
+    elif pe_col_in_group == 0 and pe_row_in_group > 0:
+        K_T = tl.recv(dir="intra_N",
+                      shape=(d_head, TILE_S_KV), dtype="f16")
+        V = tl.recv(dir="intra_N",
+                    shape=(TILE_S_KV, d_head), dtype="f16")
+        if group_cols > 1:
+            tl.send(dir="intra_E", src=K_T)
+            tl.send(dir="intra_E", src=V)
+    else:
+        K_T = tl.recv(dir="intra_W",
+                      shape=(d_head, TILE_S_KV), dtype="f16")
+        V = tl.recv(dir="intra_W",
+                    shape=(TILE_S_KV, d_head), dtype="f16")
+        if pe_col_in_group < group_cols - 1:
+            tl.send(dir="intra_E", src=K_T)
+            tl.send(dir="intra_E", src=V)
+
+    # Tile 0 — primitives establish (m, ℓ, O); recipe fusion enters tile 1+.
+    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)
+
+    # ── Tiles 1..n_tiles-1 ──
+    for tile_idx in range(1, n_tiles):
+        with tl.scratch_scope():
+            if pe_in_group == 0:
+                K_T_t = tl.load(k_head_shard_base + tile_idx * K_TILE_BYTES,
+                                shape=(d_head, TILE_S_KV), dtype="f16")
+                V_t = tl.load(v_head_shard_base + tile_idx * TILE_S_KV * KV_ROW_BYTES,
+                              shape=(TILE_S_KV, d_head), dtype="f16")
+                if group_cols > 1:
+                    tl.send(dir="intra_E", src=K_T_t)
+                    tl.send(dir="intra_E", src=V_t)
+                if group_rows > 1:
+                    tl.send(dir="intra_S", src=K_T_t)
+                    tl.send(dir="intra_S", src=V_t)
+            elif pe_col_in_group == 0 and pe_row_in_group > 0:
+                K_T_t = tl.recv(dir="intra_N",
+                                shape=(d_head, TILE_S_KV), dtype="f16")
+                V_t = tl.recv(dir="intra_N",
+                              shape=(TILE_S_KV, d_head), dtype="f16")
+                if group_cols > 1:
+                    tl.send(dir="intra_E", src=K_T_t)
+                    tl.send(dir="intra_E", src=V_t)
+            else:
+                K_T_t = tl.recv(dir="intra_W",
+                                shape=(d_head, TILE_S_KV), dtype="f16")
+                V_t = tl.recv(dir="intra_W",
+                              shape=(TILE_S_KV, d_head), dtype="f16")
+                if pe_col_in_group < group_cols - 1:
+                    tl.send(dir="intra_E", src=K_T_t)
+                    tl.send(dir="intra_E", src=V_t)
+            # Two-composite fusion: Q·Kᵀ composite → softmax_merge prologue
+            # binds P (pinned primary-out) to the P·V composite, folding
+            # the new tile's contribution into O_local.
+            scores_t = tl.composite(op="gemm", a=Q, b=K_T_t)
+            tl.composite(
+                prologue=[{"op": "softmax_merge", "s": scores_t,
+                           "m": m_local, "l": l_local, "O": O_local}],
+                op="gemm", b=V_t, out=O_local,
+                epilogue=[{"op": "add", "other": O_local}],
+            )
+
+    O_final = O_local / l_local
+    o_base = (o_ptr
+              + cube_id * kv_per_cube * T_q * Q_ROW_BYTES
+              + (group_id_in_cube * T_q + q_start) * Q_ROW_BYTES)
+    tl.store(o_base, O_final)

src/kernbench/benches/milestone_1h_gqa.py

@@ -8,12 +8,13 @@ Currently exercises (long-context only — short-context panels are
 future work):
   - 4-cases prefill comparative study (gqa_helpers.long_ctx.gqa_prefill_long_ctx_4cases)
   - 4-cases decode comparative study  (gqa_helpers.long_ctx.gqa_decode_long_ctx_4cases)
+  - Case-6 composite-command study    (gqa_helpers.long_ctx.gqa_decode_long_ctx_composite)
 
-Each sub-sweep writes its own ``sweep_{prefill,decode}.json`` into the
+Each sub-sweep writes its own ``sweep_*.json`` into the shared output
-shared output dir ``benches/1H_milestone_output/gqa/gqa_long_ctx/``.
+dir ``benches/1H_milestone_output/gqa/gqa_long_ctx/``.
 Selection via the env var ``GQA_1H_SWEEPS=prefill,decode`` (default
-runs both). Toggle individual sweeps with ``GQA_1H_SWEEPS=prefill``
+runs prefill+decode). The composite study is opt-in (it sweeps the
-or ``GQA_1H_SWEEPS=decode``.
+data-mode engine over several S_kv points): ``GQA_1H_SWEEPS=composite``.
 
 Gated by ``GQA_1H_RUN=1`` to keep CI fast.
 """
@@ -24,6 +25,12 @@ import os
 from kernbench.benches.gqa_helpers.long_ctx.gqa_decode_long_ctx_4cases import (
     run_sweep as _run_decode_sweep,
 )
+from kernbench.benches.gqa_helpers.long_ctx.gqa_decode_long_ctx_composite import (
+    run_sweep as _run_composite_sweep,
+)
+from kernbench.benches.gqa_helpers.long_ctx.gqa_prefill_compute_bound import (
+    run_sweep as _run_prefill_cb_sweep,
+)
 from kernbench.benches.gqa_helpers.long_ctx.gqa_prefill_long_ctx_4cases import (
     run_sweep as _run_prefill_sweep,
 )
@@ -57,6 +64,8 @@ def run(torch) -> None:
     runners = {
         "prefill": _run_prefill_sweep,
         "decode": _run_decode_sweep,
+        "composite": _run_composite_sweep,
+        "prefill_cb": _run_prefill_cb_sweep,
     }
     unknown = [s for s in sweeps if s not in runners]
     if unknown:

src/kernbench/components/builtin/pe_cpu.py

@@ -117,17 +117,13 @@ class PeCpuComponent(ComponentBase):
         pe_exec_start = env.now
         scheduler_id = f"{self._pe_prefix}.pe_scheduler"
 
-        # Choose execution mode: greenlet (ADR-0020) or legacy command-list
+        # ADR-0020 greenlet execution — always used so IPCQ / ring credits /
+        # fabric-transfer sim events fire regardless of data-mode. KernelRunner
+        # guards its store reads on ``self._store is not None``.
         store = getattr(self.ctx, "memory_store", None) if self.ctx else None
-
+        composite_results = yield from self._execute_greenlet(
-        if store is not None:
+            env, kernel_fn, kernel_args, num_programs, scheduler_id, store,
-            composite_results = yield from self._execute_greenlet(
+        )
-                env, kernel_fn, kernel_args, num_programs, scheduler_id, store,
-            )
-        else:
-            composite_results = yield from self._execute_legacy(
-                env, kernel_fn, kernel_args, num_programs, scheduler_id,
-            )
 
         # Record PE-internal execution time
         txn.result_data["pe_exec_ns"] = env.now - pe_exec_start

src/kernbench/policy/routing/router.py

@@ -111,6 +111,13 @@ class PathRouter:
         self._adj_all: dict[str, list[tuple[str, float]]] = defaultdict(list)
         self._adj_mcpu_dma: dict[str, list[tuple[str, float]]] = defaultdict(list)
         self._adj_local: dict[str, list[tuple[str, float]]] = defaultdict(list)
+        # Memoize path lookups: adj dicts are built once here and never
+        # mutated (topology is static per ADR-0006 / SPEC §0.1), so
+        # (id(adj), start, goal) is a stable cache key for the router's
+        # lifetime. Callers use the returned path list read-only.
+        self._path_cache: dict[
+            tuple[int, str, str], tuple[list[str], float]
+        ] = {}
         for e in graph.edges:
             w = e.routing_weight_mm if e.routing_weight_mm is not None else e.distance_mm
             self._adj_all[e.src].append((e.dst, w))
@@ -185,8 +192,14 @@ class PathRouter:
         start: str,
         goal: str,
     ) -> tuple[list[str], float]:
+        cache_key = (id(adj), start, goal)
+        cached = self._path_cache.get(cache_key)
+        if cached is not None:
+            return cached
         if start == goal:
-            return [start], 0.0
+            result = ([start], 0.0)
+            self._path_cache[cache_key] = result
+            return result
         best: dict[str, float] = {start: 0.0}
         prev: dict[str, str] = {}
         heap: list[tuple[float, str]] = [(0.0, start)]
@@ -200,7 +213,9 @@ class PathRouter:
                     cur = prev[cur]
                 path.append(start)
                 path.reverse()
-                return path, d
+                result = (path, d)
+                self._path_cache[cache_key] = result
+                return result
             if d > best.get(node, float("inf")):
                 continue
             for neighbor, edge_dist in adj[node]:

src/kernbench/runtime_api/context.py

@@ -570,8 +570,14 @@ class RuntimeContext:
                     h = self.submit(msg)
                     self.wait(h)
 
-        # Submit MemoryWriteMsg per shard (deploy data to device)
+        # Submit MemoryWriteMsg per shard (deploy data to device). Gated on
-        if pattern is not None:
+        # memory_store presence: under enable_data=False there is no
+        # MemoryStore to populate and no Phase 2 DataExecutor replay, so the
+        # per-shard sim events are pure wall-clock overhead with no effect
+        # on reported kernel latency (Yangwook's max(t_end) - min(t_start)
+        # formula excludes ops before the kernel starts).
+        store = getattr(self.engine, "_memory_store", None)
+        if pattern is not None and store is not None:
             for shard in handle.shards:
                 h = self.submit(MemoryWriteMsg(
                     correlation_id=self.correlation_id,
@@ -591,8 +597,7 @@ class RuntimeContext:
             # VA; Phase 2 DataExecutor reads via the addresses captured in
             # op_log (VA for tl.load). Without this, zero-init tensors are
             # invisible to kernels in Phase 2.
-            store = getattr(self.engine, "_memory_store", None)
+            if pattern == "zero" and handle.va_base:
-            if store is not None and pattern == "zero" and handle.va_base:
                 import numpy as np
                 from kernbench.runtime_api.tensor import _numpy_dtype
                 np_dtype = _numpy_dtype(dtype)

src/kernbench/sim_engine/engine.py

@@ -53,14 +53,16 @@ class GraphEngine:
         self._events: dict[str, simpy.Event] = {}
         self._counter = 0
         overrides = component_overrides or {}
-        # ADR-0020: optional data execution support
+        # ADR-0020: optional data execution support. OpLogger is always
-        self._op_logger = None
+        # created so op_log-based latency (max(t_end) - min(t_start)) is
+        # available in both modes; MemoryStore is created only when
+        # enable_data=True so Phase 2 DataExecutor replay is gated on it.
         self._memory_store = None
         if enable_data:
             from kernbench.sim_engine.memory_store import MemoryStore
-            from kernbench.sim_engine.op_log import OpLogger
             self._memory_store = MemoryStore()
-            self._op_logger = OpLogger(memory_store=self._memory_store)
+        from kernbench.sim_engine.op_log import OpLogger
+        self._op_logger = OpLogger(memory_store=self._memory_store)
         # Cursor for incremental Phase 2 replay (ADR-0020 D6).
         # SimPy env.now is monotonic so newly logged records always sort
         # to the tail; the cursor remains valid across waits.

src/kernbench/triton_emu/tl_context.py

@@ -264,6 +264,10 @@ class TLContext:
             id=self._next_handle_id(),
             addr=addr, shape=shape, dtype=dtype,
             nbytes=nbytes, space="tcm",
+            # TCM-resident: a downstream composite operand reads it on-chip,
+            # not via a DMA_READ of its (bit-61) scratch address — same as
+            # the composite auto-output and recipe scratch handles below.
+            pinned=True,
         )
 
     # ── Reference (no DMA, metadata only) ────────────────────────
@@ -693,6 +697,10 @@ class TLContext:
                 nbytes=self._nbytes(shape, dtype),
                 data=data,
                 space=slot_space,
+                # On-chip (tcm/sram) IPCQ slot: read in place as a composite
+                # operand, not DMA_READ of its non-HBM slot address. An HBM
+                # slot is a valid PA, so it stays unpinned (DMA streams it).
+                pinned=slot_space != "hbm",
             )
         return self._make_handle(addr=0, shape=shape, dtype=dtype)
 
@@ -734,6 +742,7 @@ class TLContext:
                 nbytes=self._nbytes(shape, dtype),
                 data=None,
                 space=slot_space,
+                pinned=slot_space != "hbm",
             )
         return self._make_handle(addr=0, shape=shape, dtype=dtype)
 
@@ -1101,6 +1110,7 @@ class TLContext:
                 nbytes=self._nbytes(handle.cmd.shape, handle.cmd.dtype),
                 data=data,
                 space=slot_space,
+                pinned=slot_space != "hbm",
             )
             handle.resolved = True
             handle.result = th

tests/analytical_visualization/_gen_roofline_paper_figs.py

@@ -0,0 +1,104 @@
+"""Generate roofline figures for the paper's Roofline Analysis section.
+
+Produces two 2-panel PNGs into docs/report/1H-codesign-paper/figures/:
+  - roofline_short_context.png  (S_kv = 8K)   — batch drops per-token cost
+  - roofline_long_context.png   (S_kv = 1M)   — batch effect vanishes
+
+Each figure shows both the step-latency view (raw step time vs B, with
+weight/compute/KV components) and the cost-per-token view (÷B), which
+together make the batch-vs-context story visible at a glance.
+
+Run: python tests/analytical_visualization/_gen_roofline_paper_figs.py
+"""
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+# Ensure the repo root is on sys.path so `tests.analytical_visualization`
+# imports resolve when invoked as a script.
+_REPO_ROOT = Path(__file__).resolve().parents[2]
+if str(_REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(_REPO_ROOT))
+
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
+from tests.analytical_visualization.chip_roofline import (
+    critical_batch,
+    per_token_latency_curve,
+    step_latency_curve,
+)
+from tests.analytical_visualization.model_config import MachineParams
+from tests.analytical_visualization.model_presets import PRESETS
+
+FIGURES_DIR = Path(__file__).resolve().parents[2] / (
+    "docs/report/1H-codesign-paper/figures"
+)
+FIGURES_DIR.mkdir(parents=True, exist_ok=True)
+
+MACHINE = MachineParams()
+MODEL = PRESETS["Llama 3 70B"].model
+B_RANGE = [1, 2, 4, 8, 16, 32, 64, 128, 256]
+
+
+def _plot_pair(s_kv: int, label: str, out_path: Path) -> None:
+    step = step_latency_curve(MACHINE, MODEL, B_RANGE, s_kv=s_kv)
+    tok = per_token_latency_curve(MACHINE, MODEL, B_RANGE, s_kv=s_kv)
+    xs = [p.batch for p in step]
+    b_star = critical_batch(MACHINE, MODEL)
+
+    fig, (axA, axB) = plt.subplots(1, 2, figsize=(11, 4.2))
+
+    # ── Left: step latency (undivided) ────────────────────────────────
+    axA.plot(xs, [p.weight_s * 1e3 for p in step], "^-",
+              color="#ffbe0b", label="Weight fetch (flat)")
+    axA.plot(xs, [p.compute_s * 1e3 for p in step], "o-",
+              color="#3a86ff", label="Compute (linear)")
+    axA.plot(xs, [p.kv_s * 1e3 for p in step], "s-",
+              color="#d90429", label="KV fetch (linear)")
+    axA.plot(xs, [p.total_s * 1e3 for p in step], "-",
+              color="#212529", linewidth=2.5, label="Total")
+    axA.set_xscale("log", base=2)
+    axA.set_yscale("log")
+    axA.set_xlabel("Batch size B")
+    axA.set_ylabel("Step time (ms)")
+    axA.set_title(f"Step latency ({label} context, $S_{{kv}}={s_kv:,}$)")
+    axA.grid(True, which="both", alpha=0.3)
+    axA.legend(fontsize=8, loc="upper left")
+
+    # ── Right: cost per token (÷ B) ───────────────────────────────────
+    axB.plot(xs, [p.weight_s * 1e3 for p in tok], "^-",
+              color="#ffbe0b", label="Weight fetch ($\\propto 1/B$)")
+    axB.plot(xs, [p.compute_s * 1e3 for p in tok], "o--",
+              color="#3a86ff", label="Compute (flat)")
+    axB.plot(xs, [p.kv_s * 1e3 for p in tok], "s--",
+              color="#d90429", label="KV fetch (flat)")
+    axB.plot(xs, [p.total_s * 1e3 for p in tok], "-",
+              color="#212529", linewidth=2.5, label="Total")
+    axB.axvline(b_star, linestyle=":", color="#2e7d32",
+                 label=f"$B^*={b_star:.0f}$")
+    axB.set_xscale("log", base=2)
+    axB.set_yscale("log")
+    axB.set_xlabel("Batch size B")
+    axB.set_ylabel("Per-token time (ms)")
+    axB.set_title(f"Cost per token ({label} context, $S_{{kv}}={s_kv:,}$)")
+    axB.grid(True, which="both", alpha=0.3)
+    axB.legend(fontsize=8, loc="upper right")
+
+    plt.tight_layout()
+    fig.savefig(out_path, dpi=150, bbox_inches="tight")
+    plt.close(fig)
+    print(f"  wrote {out_path.name}")
+
+
+def main() -> None:
+    print("Generating roofline figures for the paper ...")
+    _plot_pair(8_192, "short", FIGURES_DIR / "roofline_short_context.png")
+    _plot_pair(1_048_576, "long", FIGURES_DIR / "roofline_long_context.png")
+    print(f"Done. Output: {FIGURES_DIR}")
+
+
+if __name__ == "__main__":
+    main()

tests/analytical_visualization/auto_explore.py

@@ -0,0 +1,459 @@
+"""Auto-explore parallelism configuration space and rank by Pareto frontier.
+
+Extends the existing ``autosuggest`` (memory-only, single winner) to the
+full 9-knob search (CP, TP, PP, DP, kv_shard_mode, ffn_shard_scope,
+tp_placement, cp_placement, cp_ring_variant) with four objectives:
+
+    - min total_latency_ns   (single-request decode/prefill step)
+    - max throughput_tok_s   (naive tokens/sec = 1 / latency for single-req)
+    - max efficiency_score   (geo-mean of compute + BW utilization)
+    - min pes_used
+
+Reuses ``stage_latencies`` and ``memory_layout`` as the physics; adds
+enumeration + 4D Pareto sort on top.
+
+Analytical model runs in ~microseconds per config, so the full sweep
+(~5-10k feasible configs after pruning) completes in a few seconds.
+"""
+from __future__ import annotations
+
+import math
+from collections.abc import Iterator
+from dataclasses import dataclass, field, replace
+
+from .memory_layout import compute_memory
+from .model_config import FullConfig, MachineParams, ModelConfig, TopologyConfig
+from .stage_latencies import all_ffn_stages, all_stages
+
+# ── Parallelism sensitivity sweep values (see compute_parallelism_sensitivity)
+# Multiples of 2 (finer grid than powers of 2 alone), capped at values that
+# are physically meaningful for the modelled topology.
+_PARALLELISM_SWEEP_VALUES = {
+    "cp": (1, 2, 4, 6, 8, 10, 12, 14, 16, 24, 32, 48, 64, 96, 128, 192, 256),
+    "tp": (1, 2, 4, 6, 8, 10, 12, 14, 16, 24, 32, 48, 64),
+    "pp": (1, 2, 4, 6, 8, 10, 12, 14, 16, 20, 24, 32),
+    "dp": (1, 2, 4, 6, 8, 12, 16),
+    "ep": (1, 2, 4, 6, 8, 12, 16, 32),
+}
+
+
+# ── Search space ─────────────────────────────────────────────────────
+
+_CP_OPTIONS = (1, 2, 4, 8, 16, 32, 64, 96)
+_TP_OPTIONS = (1, 2, 4, 8, 16, 32)
+_PP_OPTIONS = (1, 2, 4, 8, 16)
+_DP_OPTIONS = (1, 2, 4)
+_KV_SHARD_MODES = ("split", "replicate")
+_FFN_SHARD_SCOPES = ("TP", "TP+CP", "TP+CP+DP")
+_TP_PLACEMENTS = ("pe", "cube")
+_CP_PLACEMENTS = ("cube", "pe")
+_CP_RING_VARIANTS = ("kv", "qoml")
+
+
+# ── Result types ─────────────────────────────────────────────────────
+
+
+@dataclass
+class ConfigScore:
+    """A single (config, computed-metrics) tuple. All floats in SI units
+    (seconds, tokens/sec, dimensionless) unless suffixed otherwise."""
+
+    # ── Config (the 9 knobs) ─────
+    cp: int
+    tp: int
+    pp: int
+    dp: int
+    kv_shard_mode: str
+    ffn_shard_scope: str
+    tp_placement: str
+    cp_placement: str
+    cp_ring_variant: str
+
+    # ── Objectives ─────
+    total_latency_ns: float      # ↓ minimize
+    throughput_tok_s: float       # ↑ maximize
+    efficiency_score: float       # ↑ maximize (0..1)
+    pes_used: int                 # ↓ minimize
+
+    # ── Info-only (not ranked on) ─────
+    hbm_utilization: float        # bytes_used / budget (0..1+; may exceed 1 if over-budget)
+    weights_gb: float
+    kv_gb: float
+    transient_gb: float
+    sips_used: int
+    fits_memory: bool
+    placement_valid: bool
+    reason: str = ""
+
+    @property
+    def latency_us(self) -> float:
+        return self.total_latency_ns / 1e3
+
+    @property
+    def latency_ms(self) -> float:
+        return self.total_latency_ns / 1e6
+
+    def as_topology(self, s_kv: int, mode: str) -> TopologyConfig:
+        """Reconstruct the TopologyConfig this score was computed for."""
+        return TopologyConfig(
+            cp=self.cp, tp=self.tp, pp=self.pp, dp=self.dp,
+            s_kv=s_kv, mode=mode,
+            kv_shard_mode=self.kv_shard_mode,
+            ffn_shard_scope=self.ffn_shard_scope,
+            tp_placement=self.tp_placement,
+            cp_placement=self.cp_placement,
+            cp_ring_variant=self.cp_ring_variant,
+        )
+
+
+@dataclass
+class AutoExploreResult:
+    model_name: str
+    s_kv: int
+    mode: str
+    total_enumerated: int         # after basic domain pruning
+    total_feasible: int           # after memory + placement checks
+    all_scores: list[ConfigScore] = field(default_factory=list)  # every feasible config
+    pareto_scores: list[ConfigScore] = field(default_factory=list)  # non-dominated set
+
+
+# ── Enumeration + pruning ────────────────────────────────────────────
+
+
+def enumerate_configs(
+    model: ModelConfig,
+    s_kv: int,
+    mode: str,
+) -> Iterator[TopologyConfig]:
+    """Yield every domain-valid TopologyConfig.
+
+    Domain rules (fast pruning; no memory check yet):
+      - PP ≤ model.layers                            (can't have more stages than layers)
+      - TP ≤ 4 × model.h_q                            (unrealistic head-dim splits above this)
+      - Skip ffn_shard_scope containing 'DP' when DP=1 (redundant with plain 'TP+CP')
+      - Skip cp_ring_variant='qoml' when CP=1        (no ring, variant is a no-op)
+    """
+    for cp in _CP_OPTIONS:
+        for tp in _TP_OPTIONS:
+            if tp > 4 * model.h_q:
+                continue
+            for pp in _PP_OPTIONS:
+                if pp > model.layers:
+                    continue
+                for dp in _DP_OPTIONS:
+                    for kv_mode in _KV_SHARD_MODES:
+                        for ffn_scope in _FFN_SHARD_SCOPES:
+                            if "DP" in ffn_scope and dp == 1:
+                                continue
+                            for tp_place in _TP_PLACEMENTS:
+                                for cp_place in _CP_PLACEMENTS:
+                                    for cp_ring in _CP_RING_VARIANTS:
+                                        if cp == 1 and cp_ring == "qoml":
+                                            continue
+                                        yield TopologyConfig(
+                                            cp=cp, tp=tp, pp=pp, dp=dp,
+                                            s_kv=s_kv, mode=mode,
+                                            kv_shard_mode=kv_mode,
+                                            ffn_shard_scope=ffn_scope,
+                                            tp_placement=tp_place,
+                                            cp_placement=cp_place,
+                                            cp_ring_variant=cp_ring,
+                                        )
+
+
+# ── Scoring ──────────────────────────────────────────────────────────
+
+
+def _sum_visible_latency(
+    cfg: FullConfig,
+    include_attention: bool = True,
+    include_ffn: bool = True,
+) -> float:
+    """Total single-request latency (seconds) across all model layers.
+
+    A single request traverses every layer sequentially, whether the layers
+    sit on one PP stage or are spread across many:
+      - PP=1  : one rank holds all L layers, latency = L × per_layer
+      - PP=K  : K ranks each hold L/K layers, but request crosses all K
+                stages sequentially → still L × per_layer
+
+    PP therefore does NOT reduce single-request latency; it only improves
+    throughput under batching. This function is the single-request cost, so
+    we multiply by full model.layers regardless of PP.
+
+    Scope selection (both default True — full per-token cost):
+      - ``include_attention=True,  include_ffn=True``  → full transformer
+      - ``include_attention=True,  include_ffn=False`` → attention only
+      - ``include_attention=False, include_ffn=True``  → FFN / MoE only
+      - ``include_attention=False, include_ffn=False`` → zero (rejected caller-side)
+    """
+    attn = sum(s.visible_s for s in all_stages(cfg)) if include_attention else 0.0
+    ffn = sum(s.visible_s for s in all_ffn_stages(cfg)) if include_ffn else 0.0
+    per_layer = attn + ffn
+    return per_layer * cfg.model.layers
+
+
+def _efficiency(cfg: FullConfig, latency_s: float,
+                include_attention: bool = True,
+                include_ffn: bool = True) -> float:
+    """Geo-mean of compute-util and BW-util. Range ~ (0, 1].
+
+    - compute_util = achieved_flops / (peak_flops × pes × latency)
+    - bw_util      = achieved_bytes / (peak_bw × pes × latency)
+
+    Scope flags mirror :func:`_sum_visible_latency`.
+    """
+    if latency_s <= 0:
+        return 0.0
+    attn = all_stages(cfg) if include_attention else []
+    ffn = all_ffn_stages(cfg) if include_ffn else []
+    layers = math.ceil(cfg.model.layers / cfg.topo.pp)
+    total_flops = layers * sum(s.flops for s in attn + ffn)
+    total_bytes = layers * sum(s.mem_bytes for s in attn + ffn)
+
+    pes = cfg.topo.total_pes
+    peak_flops = cfg.machine.peak_flops * pes
+    peak_bw = cfg.machine.bw_hbm * pes
+
+    compute_util = total_flops / (peak_flops * latency_s) if peak_flops > 0 else 0.0
+    bw_util = total_bytes / (peak_bw * latency_s) if peak_bw > 0 else 0.0
+    compute_util = min(1.0, max(0.0, compute_util))
+    bw_util = min(1.0, max(0.0, bw_util))
+    # Geo-mean; if either is 0 the score is 0 (avoids overrewarding lopsided configs).
+    return math.sqrt(compute_util * bw_util)
+
+
+def score_config(cfg: FullConfig,
+                 include_attention: bool = True,
+                 include_ffn: bool = True) -> ConfigScore:
+    """Compute all 4 objectives + info fields for one config.
+
+    Feasibility (memory + placement) is stored but does NOT gate scoring —
+    infeasible configs get returned with fits_memory=False so callers can
+    filter or display them.
+
+    Scope flags (default: full transformer) restrict *latency + efficiency*
+    to attention only, FFN only, or both. Memory feasibility is unchanged —
+    still checks weights+KV+transient fit in per-PE HBM since the model
+    still exists physically regardless of what the caller is scoring.
+    """
+    mem = compute_memory(cfg)
+    placement_ok = cfg.topo.placement_valid
+
+    latency_s = _sum_visible_latency(
+        cfg, include_attention=include_attention, include_ffn=include_ffn,
+    )
+    throughput = 1.0 / latency_s if latency_s > 0 else 0.0
+    efficiency = (
+        _efficiency(cfg, latency_s,
+                    include_attention=include_attention, include_ffn=include_ffn)
+        if latency_s > 0 else 0.0
+    )
+    fits = not mem.over_budget
+
+    reason = ""
+    if not fits:
+        reason = (f"weights+KV+transient ({mem.used_bytes/1e9:.2f} GB) "
+                  f"exceeds per-PE budget ({mem.budget_bytes/1e9:.2f} GB)")
+    elif not placement_ok:
+        reason = (f"intra-cube demand ({cfg.topo.intra_cube_dims}) "
+                  f"exceeds PEs/cube ({cfg.topo.pes_per_cube_hw})")
+
+    return ConfigScore(
+        cp=cfg.topo.cp, tp=cfg.topo.tp, pp=cfg.topo.pp, dp=cfg.topo.dp,
+        kv_shard_mode=cfg.topo.kv_shard_mode,
+        ffn_shard_scope=cfg.topo.ffn_shard_scope,
+        tp_placement=cfg.topo.tp_placement,
+        cp_placement=cfg.topo.cp_placement,
+        cp_ring_variant=cfg.topo.cp_ring_variant,
+        total_latency_ns=latency_s * 1e9,
+        throughput_tok_s=throughput,
+        efficiency_score=efficiency,
+        pes_used=cfg.topo.total_pes,
+        hbm_utilization=mem.used_bytes / mem.budget_bytes if mem.budget_bytes > 0 else 0.0,
+        weights_gb=mem.weights_bytes / 1e9,
+        kv_gb=mem.kv_cache_bytes / 1e9,
+        transient_gb=mem.transient_bytes / 1e9,
+        sips_used=cfg.topo.sips_used,
+        fits_memory=fits,
+        placement_valid=placement_ok,
+        reason=reason,
+    )
+
+
+# ── Pareto sort ──────────────────────────────────────────────────────
+
+
+def _dominates(a: ConfigScore, b: ConfigScore) -> bool:
+    """True iff a is no worse than b on every axis AND strictly better on ≥ 1.
+
+    Axes (with direction) — 3D Pareto:
+      total_latency_ns    ↓  (a ≤ b)
+      pes_used            ↓  (a ≤ b)  — proxy for cost / deployment size
+      efficiency_score    ↑  (a ≥ b)  — geo-mean of compute+BW utilization
+
+    Note: ``throughput_tok_s`` is deliberately NOT a Pareto axis. For a
+    single-request analysis, throughput = 1 / latency, so it's collinear
+    with latency and would collapse the frontier. It stays on ConfigScore
+    for display but doesn't participate in domination.
+    """
+    no_worse = (
+        a.total_latency_ns <= b.total_latency_ns
+        and a.pes_used <= b.pes_used
+        and a.efficiency_score >= b.efficiency_score
+    )
+    if not no_worse:
+        return False
+    strictly_better = (
+        a.total_latency_ns < b.total_latency_ns
+        or a.pes_used < b.pes_used
+        or a.efficiency_score > b.efficiency_score
+    )
+    return strictly_better
+
+
+def pareto_frontier(scores: list[ConfigScore]) -> list[ConfigScore]:
+    """Extract non-dominated set on (latency↓, pe↓, throughput↑, efficiency↑).
+
+    Only feasible configs (fits_memory AND placement_valid) participate; the
+    non-feasible are excluded from the frontier (they can still be returned in
+    ``all_scores`` for informational display).
+
+    O(N²) — fine for N up to ~50k with early exits.
+    """
+    feasible = [s for s in scores if s.fits_memory and s.placement_valid]
+    frontier: list[ConfigScore] = []
+    for i, a in enumerate(feasible):
+        dominated = False
+        for j, b in enumerate(feasible):
+            if i == j:
+                continue
+            if _dominates(b, a):
+                dominated = True
+                break
+        if not dominated:
+            frontier.append(a)
+    return frontier
+
+
+# ── Parallelism sensitivity ─────────────────────────────────────────
+
+
+@dataclass
+class ParallelismSensitivityRow:
+    """One knob's sweep result: latency + fit at every tested value.
+
+    All *other* parallelism knobs (including HW) are held at their baseline
+    values, so this isolates the effect of just this one knob.
+    """
+    knob: str                       # "cp" / "tp" / "pp" / "dp" / "ep"
+    values: list[int]               # tested values (from _PARALLELISM_SWEEP_VALUES)
+    latencies_ns: list[float]       # one per value; NaN when infeasible
+    fits_flags: list[bool]          # per-value memory+placement feasibility
+    baseline_value: int             # what the baseline config uses for this knob
+    baseline_latency_ns: float
+
+
+def compute_parallelism_sensitivity(
+    baseline: ConfigScore,
+    model: ModelConfig,
+    machine: MachineParams,
+    s_kv: int,
+    mode: str,
+    include_attention: bool = True,
+    include_ffn: bool = True,
+) -> list[ParallelismSensitivityRow]:
+    """For each parallelism knob (CP, TP, PP, DP, EP), sweep its values
+    holding the OTHER knobs fixed at the baseline. Reports latency + memory
+    fit per value.
+
+    Answers: 'if I only change this one knob, how does latency and memory
+    fit change?' — the complement to auto_hardware.compute_sensitivity
+    which does the same for hardware knobs.
+    """
+    baseline_topo = baseline.as_topology(s_kv, mode)
+
+    rows: list[ParallelismSensitivityRow] = []
+    for knob, values in _PARALLELISM_SWEEP_VALUES.items():
+        baseline_val = getattr(baseline_topo, knob, 1)
+        # EP isn't stored on ConfigScore's as_topology output today (dp
+        # attribute exists but ep does not always). Fall back to 1.
+        latencies: list[float] = []
+        fits: list[bool] = []
+        for v in values:
+            # Skip nonsensical values that would violate hard bounds.
+            if knob == "pp" and v > model.layers:
+                latencies.append(float("nan"))
+                fits.append(False)
+                continue
+            if knob == "tp" and v > 4 * model.h_q:
+                latencies.append(float("nan"))
+                fits.append(False)
+                continue
+            # Build a variant TopologyConfig with just this one knob changed.
+            trial = replace(baseline_topo, **{knob: v})
+            cfg = FullConfig(model=model, topo=trial, machine=machine)
+            score = score_config(
+                cfg, include_attention=include_attention, include_ffn=include_ffn,
+            )
+            latencies.append(score.total_latency_ns)
+            fits.append(score.fits_memory and score.placement_valid)
+        rows.append(ParallelismSensitivityRow(
+            knob=knob,
+            values=list(values),
+            latencies_ns=latencies,
+            fits_flags=fits,
+            baseline_value=int(baseline_val),
+            baseline_latency_ns=baseline.total_latency_ns,
+        ))
+    return rows
+
+
+# ── Top-level driver ─────────────────────────────────────────────────
+
+
+def run_auto_explore(
+    model: ModelConfig,
+    machine: MachineParams,
+    s_kv: int,
+    mode: str = "decode",
+    include_attention: bool = True,
+    include_ffn: bool = True,
+    b: int = 1,
+) -> AutoExploreResult:
+    """Enumerate all configs, score each, extract Pareto frontier.
+
+    Returns both the full ``all_scores`` list (for the table view) and the
+    ``pareto_scores`` subset (for the scatter/highlight view). Both are sorted
+    by total_latency_ns ascending.
+
+    Scope: at least one of ``include_attention`` / ``include_ffn`` must be
+    True. Setting both False would give zero latency for every config —
+    the caller should not do that.
+    """
+    all_scores: list[ConfigScore] = []
+    total_enumerated = 0
+    for topo in enumerate_configs(model, s_kv, mode):
+        # Inject the caller's batch B into every candidate.
+        topo.b = b
+        total_enumerated += 1
+        cfg = FullConfig(model=model, topo=topo, machine=machine)
+        all_scores.append(score_config(
+            cfg, include_attention=include_attention, include_ffn=include_ffn,
+        ))
+
+    all_scores.sort(key=lambda s: s.total_latency_ns)
+    pareto = pareto_frontier(all_scores)
+    pareto.sort(key=lambda s: s.total_latency_ns)
+
+    feasible_count = sum(1 for s in all_scores if s.fits_memory and s.placement_valid)
+
+    return AutoExploreResult(
+        model_name=model.name,
+        s_kv=s_kv,
+        mode=mode,
+        total_enumerated=total_enumerated,
+        total_feasible=feasible_count,
+        all_scores=all_scores,
+        pareto_scores=pareto,
+    )

tests/analytical_visualization/auto_hardware.py

@@ -0,0 +1,429 @@
+"""Joint hardware × parallelism exploration.
+
+For a fixed model + workload, sweep both:
+  - hardware parameters (MachineParams: pe_hbm_gb, bw_hbm_gbs, peak_tflops_f16,
+    bw_intra_gbs, bw_inter_gbs, bw_intersip_gbs), and
+  - parallelism knobs (a reduced subset of auto_explore's 9 knobs)
+
+Then rank on:
+  - latency ↓
+  - hardware cost proxy ↓  (normalized sum of knob-over-default)
+
+Also computes per-knob sensitivity — how much latency drops when each
+hardware knob is doubled from its baseline value. Useful for HW co-design
+("which knob to invest in first?").
+
+The default 'balanced' depth: 64 HW candidates × ~2k parallelism configs =
+~130k joint evaluations, ~1-5 s at analytical speeds. Coarse and
+two-stage variants trade coverage for speed.
+"""
+from __future__ import annotations
+
+import math
+from collections.abc import Iterator
+from dataclasses import dataclass, field, replace
+
+from .auto_explore import ConfigScore, score_config
+from .autosuggest import auto_suggest
+from .memory_layout import compute_memory
+from .model_config import FullConfig, MachineParams, ModelConfig, TopologyConfig
+
+
+# ── Hardware search space ────────────────────────────────────────────
+#
+# For each depth level, values per knob. Defaults are always included so the
+# baseline configuration always appears in the sweep.
+
+_HW_KNOB_DEFAULTS = {
+    "pe_hbm_gb":         6.0,
+    "bw_hbm_gbs":        256.0,
+    "peak_tflops_f16":   8.0,
+    "bw_intra_gbs":      512.0,
+    "bw_inter_gbs":      128.0,
+    "bw_intersip_gbs":   50.0,
+}
+
+_HW_VALUES_BALANCED = {
+    "pe_hbm_gb":         [6.0, 12.0],
+    "bw_hbm_gbs":        [256.0, 512.0],
+    "peak_tflops_f16":   [8.0, 16.0],
+    "bw_intra_gbs":      [512.0, 1024.0],
+    "bw_inter_gbs":      [128.0, 256.0],
+    "bw_intersip_gbs":   [50.0, 100.0],
+}
+
+_HW_VALUES_COARSE = {
+    "pe_hbm_gb":         [6.0, 12.0, 24.0],
+    "bw_hbm_gbs":        [256.0, 512.0, 1024.0],
+    "peak_tflops_f16":   [8.0, 16.0, 32.0],
+    "bw_intra_gbs":      [512.0, 1024.0, 2048.0],
+    "bw_inter_gbs":      [128.0, 256.0, 512.0],
+    "bw_intersip_gbs":   [50.0, 100.0, 200.0],
+}
+
+# For sensitivity: "double each knob independently from baseline."
+_SENSITIVITY_KNOBS = list(_HW_KNOB_DEFAULTS.keys())
+
+
+# ── Reduced parallelism search space (for per-HW inner loop) ─────────
+#
+# The full auto_explore has 28,800 configs; using it inside a HW loop is
+# too slow. Reduce to CP × TP × PP × DP × kv_shard_mode with common
+# defaults for the remaining 4 knobs (~2k configs).
+
+_CP_OPTIONS = (1, 2, 4, 8, 16, 32, 64, 96)
+_TP_OPTIONS = (1, 2, 4, 8, 16, 32)
+_PP_OPTIONS = (1, 2, 4, 8, 16)
+_DP_OPTIONS = (1, 2, 4)
+_KV_SHARD_MODES = ("split", "replicate")
+
+
+# ── Result types ─────────────────────────────────────────────────────
+
+
+@dataclass
+class HardwareCandidate:
+    pe_hbm_gb: float
+    bw_hbm_gbs: float
+    peak_tflops_f16: float
+    bw_intra_gbs: float
+    bw_inter_gbs: float
+    bw_intersip_gbs: float
+
+    def as_machine(self) -> MachineParams:
+        """Build a MachineParams from this HW candidate, holding alphas +
+        compute_util at their MachineParams defaults."""
+        return MachineParams(
+            pe_hbm_gb=self.pe_hbm_gb,
+            bw_hbm_gbs=self.bw_hbm_gbs,
+            peak_tflops_f16=self.peak_tflops_f16,
+            bw_intra_gbs=self.bw_intra_gbs,
+            bw_inter_gbs=self.bw_inter_gbs,
+            bw_intersip_gbs=self.bw_intersip_gbs,
+        )
+
+    @property
+    def cost_score(self) -> float:
+        """Normalized cost proxy: sum of (knob / default). 6.0 at defaults.
+
+        This is NOT dollars — it's a rough silicon-area / capability proxy. A
+        higher cost_score means "more capable hardware" (more HBM, faster BW,
+        etc.). Under identical latency, lower cost_score wins.
+        """
+        return sum(
+            getattr(self, k) / _HW_KNOB_DEFAULTS[k]
+            for k in _HW_KNOB_DEFAULTS
+        )
+
+
+@dataclass
+class JointScore:
+    """One (hardware, parallelism) pair with its computed latency + cost."""
+    hardware: HardwareCandidate
+    parallelism: ConfigScore
+    total_latency_ns: float
+    cost_score: float
+
+    @property
+    def latency_ms(self) -> float:
+        return self.total_latency_ns / 1e6
+
+
+@dataclass
+class SensitivityRow:
+    knob: str
+    baseline_value: float
+    doubled_value: float
+    baseline_latency_ns: float
+    doubled_latency_ns: float
+
+    @property
+    def rel_speedup(self) -> float:
+        """1 - (doubled/baseline). Positive = doubling this knob speeds things up."""
+        if self.baseline_latency_ns <= 0:
+            return 0.0
+        return 1.0 - (self.doubled_latency_ns / self.baseline_latency_ns)
+
+
+@dataclass
+class JointExploreResult:
+    model_name: str
+    s_kv: int
+    mode: str
+    depth: str
+    total_hw: int
+    total_joint: int
+    all_scores: list[JointScore] = field(default_factory=list)
+    pareto_scores: list[JointScore] = field(default_factory=list)
+    sensitivity: list[SensitivityRow] = field(default_factory=list)
+
+
+# ── Hardware enumeration ─────────────────────────────────────────────
+
+
+def enumerate_hardware(depth: str = "balanced") -> Iterator[HardwareCandidate]:
+    """Yield HardwareCandidates for the given sweep depth.
+
+    depth ∈ {"two_stage", "balanced", "coarse"}:
+      - two_stage: only the default HW (1 candidate; parallelism sweep dominates)
+      - balanced:  2 values per knob → 2^6 = 64 candidates
+      - coarse:    3 values per knob → 3^6 = 729 candidates
+    """
+    if depth == "two_stage":
+        vals = {k: [v] for k, v in _HW_KNOB_DEFAULTS.items()}
+    elif depth == "coarse":
+        vals = _HW_VALUES_COARSE
+    else:  # balanced (default)
+        vals = _HW_VALUES_BALANCED
+
+    for pe_hbm in vals["pe_hbm_gb"]:
+        for bw_hbm in vals["bw_hbm_gbs"]:
+            for tflops in vals["peak_tflops_f16"]:
+                for bw_intra in vals["bw_intra_gbs"]:
+                    for bw_inter in vals["bw_inter_gbs"]:
+                        for bw_intersip in vals["bw_intersip_gbs"]:
+                            yield HardwareCandidate(
+                                pe_hbm_gb=pe_hbm,
+                                bw_hbm_gbs=bw_hbm,
+                                peak_tflops_f16=tflops,
+                                bw_intra_gbs=bw_intra,
+                                bw_inter_gbs=bw_inter,
+                                bw_intersip_gbs=bw_intersip,
+                            )
+
+
+# ── Reduced parallelism search per HW ────────────────────────────────
+
+
+def _iter_reduced_parallelism(
+    model: ModelConfig, s_kv: int, mode: str,
+) -> Iterator[TopologyConfig]:
+    """Reduced sweep: CP × TP × PP × DP × kv_shard_mode.
+
+    Other 4 knobs held at defaults chosen to be latency-friendly for the
+    decode/prefill common case:
+      - ffn_shard_scope = "TP+CP"  (common good choice; trades a small AR
+        for lower per-PE weight bytes)
+      - tp_placement    = "cube"   (packs TP across cubes; UCIe D2D)
+      - cp_placement    = "pe"     (packs CP inside cubes; intra NoC)
+      - cp_ring_variant = "qoml" for decode, "kv" for prefill
+
+    Total: 8 × 6 × 5 × 4 × 2 = 1,920 configs per HW candidate.
+    """
+    cp_ring = "qoml" if mode == "decode" else "kv"
+    for cp in _CP_OPTIONS:
+        for tp in _TP_OPTIONS:
+            if tp > 4 * model.h_q:
+                continue
+            for pp in _PP_OPTIONS:
+                if pp > model.layers:
+                    continue
+                for dp in _DP_OPTIONS:
+                    for kv_mode in _KV_SHARD_MODES:
+                        # cp_ring=qoml with cp=1 is a no-op; fall back to kv.
+                        ring = "kv" if cp == 1 else cp_ring
+                        yield TopologyConfig(
+                            cp=cp, tp=tp, pp=pp, dp=dp,
+                            s_kv=s_kv, mode=mode,
+                            kv_shard_mode=kv_mode,
+                            ffn_shard_scope="TP+CP",
+                            tp_placement="cube",
+                            cp_placement="pe",
+                            cp_ring_variant=ring,
+                        )
+
+
+def _best_parallelism_for_hw(
+    model: ModelConfig, machine: MachineParams,
+    s_kv: int, mode: str,
+    include_attention: bool = True, include_ffn: bool = True,
+    b: int = 1,
+) -> ConfigScore | None:
+    """Return the latency-minimum feasible parallelism for this HW.
+
+    Feasibility: fits memory + placement_valid. Returns None if nothing fits.
+    Scope flags are forwarded to score_config so the "latency" ranked on
+    matches the caller's attention/FFN/full choice.
+    """
+    best: ConfigScore | None = None
+    best_key: tuple | None = None
+    for topo in _iter_reduced_parallelism(model, s_kv, mode):
+        topo.b = b
+        cfg = FullConfig(model=model, topo=topo, machine=machine)
+        s = score_config(cfg, include_attention=include_attention,
+                          include_ffn=include_ffn)
+        if not (s.fits_memory and s.placement_valid):
+            continue
+        # Compound tiebreaker: prefer fewer PEs and lower HBM % when latency
+        # ties across variant knobs (cp_ring, placement, kv_shard_mode, etc.).
+        _k = (s.total_latency_ns, s.pes_used, s.hbm_utilization)
+        if best is None or _k < best_key:
+            best = s
+            best_key = _k
+    return best
+
+
+def _best_parallelism_two_stage(
+    model: ModelConfig, machine: MachineParams,
+    s_kv: int, mode: str,
+    include_attention: bool = True, include_ffn: bool = True,
+    b: int = 1,
+) -> ConfigScore | None:
+    """Fast fallback: use autosuggest's memory-min (CP,TP,PP) then score it."""
+    sug = auto_suggest(model, machine, s_kv, mode)
+    if not sug.fits:
+        return None
+    topo = TopologyConfig(
+        cp=sug.cp, tp=sug.tp, pp=sug.pp, dp=1, b=b,
+        s_kv=s_kv, mode=mode,
+        kv_shard_mode="split",
+        ffn_shard_scope="TP+CP",
+        tp_placement="cube",
+        cp_placement="pe",
+        cp_ring_variant="qoml" if mode == "decode" and sug.cp > 1 else "kv",
+    )
+    cfg = FullConfig(model=model, topo=topo, machine=machine)
+    s = score_config(cfg, include_attention=include_attention,
+                      include_ffn=include_ffn)
+    return s if s.fits_memory and s.placement_valid else None
+
+
+# ── Pareto over joint (latency, cost) ────────────────────────────────
+
+
+def _pareto_2d(scores: list[JointScore]) -> list[JointScore]:
+    """Non-dominated set on (total_latency_ns ↓, cost_score ↓). O(N²)."""
+    frontier: list[JointScore] = []
+    for i, a in enumerate(scores):
+        dominated = False
+        for j, b in enumerate(scores):
+            if i == j:
+                continue
+            no_worse = (
+                b.total_latency_ns <= a.total_latency_ns
+                and b.cost_score <= a.cost_score
+            )
+            strictly_better = (
+                b.total_latency_ns < a.total_latency_ns
+                or b.cost_score < a.cost_score
+            )
+            if no_worse and strictly_better:
+                dominated = True
+                break
+        if not dominated:
+            frontier.append(a)
+    return frontier
+
+
+# ── Per-knob sensitivity ─────────────────────────────────────────────
+
+
+def compute_sensitivity(
+    baseline_hw: HardwareCandidate,
+    parallelism: ConfigScore,
+    model: ModelConfig, s_kv: int, mode: str,
+    include_attention: bool = True, include_ffn: bool = True,
+) -> list[SensitivityRow]:
+    """For each HW knob, double it (holding others at baseline) and measure
+    the latency change. Same parallelism used throughout so we isolate the
+    HW knob's effect."""
+    rows: list[SensitivityRow] = []
+    baseline_machine = baseline_hw.as_machine()
+    baseline_topo = parallelism.as_topology(s_kv, mode)
+    baseline_cfg = FullConfig(
+        model=model, topo=baseline_topo, machine=baseline_machine,
+    )
+    baseline_score = score_config(
+        baseline_cfg, include_attention=include_attention, include_ffn=include_ffn,
+    )
+    baseline_latency = baseline_score.total_latency_ns
+
+    for knob in _SENSITIVITY_KNOBS:
+        doubled_val = 2.0 * getattr(baseline_hw, knob)
+        doubled_hw = replace(baseline_hw, **{knob: doubled_val})
+        doubled_cfg = FullConfig(
+            model=model, topo=baseline_topo, machine=doubled_hw.as_machine(),
+        )
+        doubled_score = score_config(
+            doubled_cfg, include_attention=include_attention, include_ffn=include_ffn,
+        )
+        rows.append(SensitivityRow(
+            knob=knob,
+            baseline_value=getattr(baseline_hw, knob),
+            doubled_value=doubled_val,
+            baseline_latency_ns=baseline_latency,
+            doubled_latency_ns=doubled_score.total_latency_ns,
+        ))
+    # Sort by biggest speedup first (most sensitive knob at the top).
+    rows.sort(key=lambda r: -r.rel_speedup)
+    return rows
+
+
+# ── Top-level driver ─────────────────────────────────────────────────
+
+
+def joint_explore(
+    model: ModelConfig,
+    s_kv: int,
+    mode: str,
+    depth: str = "balanced",
+    include_attention: bool = True,
+    include_ffn: bool = True,
+    b: int = 1,
+) -> JointExploreResult:
+    """Sweep HW candidates × parallelism, return joint Pareto + sensitivity.
+
+    Sensitivity is computed around the LATENCY-MINIMUM joint point (the
+    "best fast" config), doubling each HW knob one at a time.
+
+    Scope flags select which stages contribute to the summed latency —
+    both the per-HW parallelism search and the sensitivity ranking use
+    the same restriction.
+    """
+    all_scores: list[JointScore] = []
+    for hw in enumerate_hardware(depth):
+        machine = hw.as_machine()
+        if depth == "two_stage":
+            par = _best_parallelism_two_stage(
+                model, machine, s_kv, mode,
+                include_attention=include_attention, include_ffn=include_ffn,
+                b=b,
+            )
+        else:
+            par = _best_parallelism_for_hw(
+                model, machine, s_kv, mode,
+                include_attention=include_attention, include_ffn=include_ffn,
+                b=b,
+            )
+        if par is None:
+            continue
+        all_scores.append(JointScore(
+            hardware=hw,
+            parallelism=par,
+            total_latency_ns=par.total_latency_ns,
+            cost_score=hw.cost_score,
+        ))
+
+    all_scores.sort(key=lambda s: s.total_latency_ns)
+    pareto = _pareto_2d(all_scores)
+    pareto.sort(key=lambda s: s.total_latency_ns)
+
+    sensitivity: list[SensitivityRow] = []
+    if all_scores:
+        best = all_scores[0]
+        sensitivity = compute_sensitivity(
+            best.hardware, best.parallelism, model, s_kv, mode,
+            include_attention=include_attention, include_ffn=include_ffn,
+        )
+
+    return JointExploreResult(
+        model_name=model.name,
+        s_kv=s_kv,
+        mode=mode,
+        depth=depth,
+        total_hw=sum(1 for _ in enumerate_hardware(depth)),
+        total_joint=len(all_scores),
+        all_scores=all_scores,
+        pareto_scores=pareto,
+        sensitivity=sensitivity,
+    )

tests/analytical_visualization/autosuggest.py

@@ -0,0 +1,146 @@
+"""Auto-suggest (CP, TP, PP) that fits the model in the per-PE budget.
+
+Strategy: iterate over candidate (CP, TP, PP) triples in order of
+increasing total PE count and return the first one that satisfies
+memory + kernel-support constraints with >10% slack. If none fits, return
+the best-effort configuration and flag over-budget.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, replace
+from typing import Iterable
+
+from .model_config import FullConfig, ModelConfig, TopologyConfig, MachineParams
+from .memory_layout import compute_memory
+
+
+# Candidate parallelism dimensions (powers of 2 mostly, up to sensible limits).
+_TP_OPTIONS = (1, 2, 4, 8, 16, 32)
+_CP_OPTIONS = (1, 2, 4, 8, 16, 32, 64, 96)
+_PP_OPTIONS = (1, 2, 4, 8, 16, 32)
+
+
+@dataclass
+class Suggestion:
+    cp: int
+    tp: int
+    pp: int
+    weights_gb: float
+    kv_gb: float
+    transient_gb: float
+    slack_gb: float
+    fits: bool
+    pes_used: int
+    sips_used: int
+    cubes_used: int = 0            # picked to minimize this (see _score_candidate)
+    cp_placement: str = "cube"     # "pe" packs CP into intra-cube PEs when it fits
+    reason: str = ""
+
+
+def _iter_candidates(model: ModelConfig) -> Iterable[tuple[int, int, int]]:
+    """Yield (CP, TP, PP) in order of increasing PE count."""
+    triples: list[tuple[int, int, int]] = []
+    for tp in _TP_OPTIONS:
+        for cp in _CP_OPTIONS:
+            for pp in _PP_OPTIONS:
+                # PP must not exceed layer count.
+                if pp > model.layers:
+                    continue
+                # Skip TP > H_q * some factor (unrealistic).
+                if tp > model.h_q * 4:
+                    continue
+                triples.append((cp, tp, pp))
+    # Sort by pe_count then by (pp, tp, cp) — prefer smaller PP first
+    # (avoids pipeline bubbles), then smaller TP (avoids head-dim split).
+    triples.sort(key=lambda t: (t[0] * t[1] * t[2], t[2], t[1], t[0]))
+    for t in triples:
+        yield t
+
+
+def _score_candidate(cp: int, tp: int, pp: int,
+                     model: ModelConfig, machine: MachineParams,
+                     s_kv: int, mode: str,
+                     slack_frac: float = 0.10,
+                     b: int = 1,
+                     include_attention: bool = True,
+                     include_ffn: bool = True) -> Suggestion:
+    # For each triple, try both cp_placement options and keep the one
+    # with fewer cubes (breaks the historical "CP always across cubes"
+    # default when a smaller pack is possible). The pe placement is only
+    # valid when CP·TP fits within a single cube's PE count.
+    best_topo: TopologyConfig | None = None
+    best_placement = "cube"
+    _pe_per_cube_hw = TopologyConfig().pes_per_cube_hw  # instance-invariant HW const
+    _placements_to_try = ["cube"]
+    if cp * tp <= _pe_per_cube_hw:
+        _placements_to_try.append("pe")
+    for _place in _placements_to_try:
+        _t = TopologyConfig(
+            cp=cp, tp=tp, pp=pp, s_kv=s_kv, mode=mode, b=b,
+            cp_placement=_place,
+        )
+        if best_topo is None or _t.cubes_used < best_topo.cubes_used:
+            best_topo = _t
+            best_placement = _place
+    topo = best_topo
+    cfg = FullConfig(model=model, topo=topo, machine=machine)
+    mem = compute_memory(
+        cfg, include_attention=include_attention, include_ffn=include_ffn,
+    )
+
+    fits = (not mem.over_budget
+            and mem.slack_bytes >= slack_frac * mem.budget_bytes)
+    reason = ""
+    if mem.over_budget:
+        reason = (f"weights+KV+transient ({mem.used_bytes/1e9:.2f} GB) "
+                  f"exceeds budget ({mem.budget_bytes/1e9:.2f} GB)")
+    elif not fits:
+        reason = f"slack ({mem.slack_bytes/1e9:.2f} GB) below 10% of budget"
+    return Suggestion(
+        cp=cp, tp=tp, pp=pp,
+        weights_gb=mem.weights_bytes / 1e9,
+        kv_gb=mem.kv_cache_bytes / 1e9,
+        transient_gb=mem.transient_bytes / 1e9,
+        slack_gb=mem.slack_bytes / 1e9,
+        fits=fits,
+        pes_used=topo.total_pes,
+        sips_used=topo.sips_used,
+        cubes_used=topo.cubes_used,
+        cp_placement=best_placement,
+        reason=reason,
+    )
+
+
+def auto_suggest(model: ModelConfig, machine: MachineParams,
+                 s_kv: int, mode: str = "decode",
+                 slack_frac: float = 0.10,
+                 b: int = 1,
+                 include_attention: bool = True,
+                 include_ffn: bool = True) -> Suggestion:
+    """Return the smallest deployment (fewer cubes, then fewer PEs)
+    that fits.
+
+    Sort key: (cubes_used ↑, pes_used ↑, pp ↑, tp ↑, cp ↑).
+    Fewer cubes wins first because a cube is the physical die-level
+    hardware unit; PE count is the tiebreaker. Preferring fewer PP then
+    TP then CP keeps the earlier historical bias (avoid pipeline
+    bubbles / head-dim splits) among equal-cube-and-PE ties.
+
+    If no candidate fits, returns the best-effort one (highest slack,
+    even if negative) with fits=False.
+    """
+    scored: list[Suggestion] = []
+    for cp, tp, pp in _iter_candidates(model):
+        scored.append(
+            _score_candidate(cp, tp, pp, model, machine, s_kv, mode,
+                              slack_frac, b=b,
+                              include_attention=include_attention,
+                              include_ffn=include_ffn)
+        )
+    scored.sort(key=lambda s: (s.cubes_used, s.pes_used, s.pp, s.tp, s.cp))
+
+    for s in scored:
+        if s.fits:
+            return s
+    # No fit — return the best-effort (largest slack, closest to fitting).
+    return max(scored, key=lambda s: s.slack_gb)

tests/analytical_visualization/chip_roofline.py

@@ -0,0 +1,518 @@
+"""Chip-level roofline math: AI, B*, L*, per-token latency curves.
+
+Surfaces the arithmetic-intensity story from LLM-serving practice:
+
+- **AI**  = C / W  (peak FLOPs per byte of HBM bandwidth).
+- **B\***  = C * b / (2 * W) * sparsity  — the critical batch size at
+  which weight-fetch time equals compute time for one decode step.
+  Sparsity = N_total / N_active (MoE factor; 1 for dense).
+- **L\***  = 2 * N_active / (AI * kv_bytes_per_token) — the balance
+  context length at which KV-read time equals compute time.
+- **B_knee(S_kv)** = B* / (1 - S_kv/L*) — the batch size where the
+  cost curve bends (weight-fetch drops below the compute+KV floor).
+  Diverges at S_kv = L* and no knee exists past it.
+
+Per-token decode-step latency, per PE (dense-approx, no comm):
+
+    t(B, S_kv)  =  N_active * b / (W * B)     <-- weight fetch, 1/B
+                 + 2 * N_active / C            <-- compute (peak), flat
+                 + S_kv * kv_bpt / W           <-- KV read, flat
+
+All numbers per PE / per one forward pass. **Peak roofline — no
+utilization factor.** Comm cost and TP/CP sharding are intentionally
+NOT in the roofline — this is the back-of-envelope chip-vs-model view
+the transcript talks about, not the full latency model that
+stage_latencies.py builds. With this convention weight_s == compute_s
+exactly at B*.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, replace
+
+from .model_config import FullConfig, MachineParams, ModelConfig
+
+
+# Per-token bytes of BF16 MAC arithmetic: one multiply + one add = 2 FLOPs.
+_FLOPS_PER_PARAM_PER_TOKEN = 2
+
+
+# ── Chip / model derived quantities ────────────────────────────────
+
+
+def arithmetic_intensity(machine: MachineParams) -> float:
+    """FLOPs per byte of HBM bandwidth. Peak roofline; utilization
+    is applied only in the compute-time formula, not here."""
+    return machine.peak_flops / machine.bw_hbm
+
+
+def total_active_params(model: ModelConfig) -> int:
+    """Full-model parameter count (attention + FFN, all layers).
+
+    Attention: 4 projections × hidden × H_q * d_head effective per layer.
+    (W_Q hidden×H_q*d_h, W_O H_q*d_h×hidden, W_K/W_V hidden×H_kv*d_h.)
+    FFN: 3 × hidden × ffn_dim per layer (gate, up, down).
+    """
+    m = model
+    attn = (
+        m.hidden * m.h_q * m.d_head            # W_Q
+        + m.hidden * m.h_kv * m.d_head * 2     # W_K + W_V
+        + m.h_q * m.d_head * m.hidden          # W_O
+    )
+    ffn = 3 * m.hidden * m.ffn_dim
+    return (attn + ffn) * m.layers
+
+
+def kv_bytes_per_token(model: ModelConfig) -> int:
+    """Bytes of KV cache one new token adds across ALL layers, per one
+    sequence, un-sharded (K + V, H_kv heads * d_h * bytes)."""
+    m = model
+    return 2 * m.h_kv * m.d_head * m.bytes_per_elem * m.layers
+
+
+def critical_batch(machine: MachineParams, model: ModelConfig,
+                    sparsity: float = 1.0) -> float:
+    """B* = C * b / (2 * W) * sparsity.
+
+    Sparsity = N_total / N_active (>= 1). Dense = 1. MoE 8-of-256 = 8.
+    """
+    b = model.bytes_per_elem
+    ai = arithmetic_intensity(machine)
+    return ai * b / _FLOPS_PER_PARAM_PER_TOKEN * sparsity
+
+
+def balance_context(machine: MachineParams, model: ModelConfig) -> float:
+    """L* = 2 * N_active / (AI * kv_bpt).
+
+    Context length (in tokens) at which per-step KV read matches the
+    per-step compute cost. Beyond L*, the KV term is dominant and no
+    batch size gets you compute-bound.
+    """
+    n_active = total_active_params(model)
+    ai = arithmetic_intensity(machine)
+    kv_bpt = kv_bytes_per_token(model)
+    return _FLOPS_PER_PARAM_PER_TOKEN * n_active / (ai * kv_bpt)
+
+
+def knee_batch(machine: MachineParams, model: ModelConfig,
+                s_kv: int) -> float | None:
+    """B_knee(S_kv) = B* / (1 - S_kv/L*).
+
+    Returns None when S_kv >= L* (no knee exists — the total-cost
+    curve never touches the compute floor).
+    """
+    b_star = critical_batch(machine, model)
+    l_star = balance_context(machine, model)
+    r = s_kv / l_star
+    if r >= 1.0:
+        return None
+    return b_star / (1.0 - r)
+
+
+# ── Per-token latency curves ───────────────────────────────────────
+
+
+@dataclass
+class RooflinePoint:
+    batch: int
+    weight_s: float   # weight fetch time, 1/B
+    compute_s: float  # compute time, flat
+    kv_s: float       # KV read time, flat
+    total_s: float
+
+
+def per_token_latency_curve(machine: MachineParams, model: ModelConfig,
+                             batch_range: list[int],
+                             s_kv: int) -> list[RooflinePoint]:
+    """Per-token decode-step latency curve across a range of batch sizes.
+
+    Returns one point per batch. All times per PE, dense-approx,
+    utilization from machine.compute_util. Comm and TP/CP sharding are
+    excluded — this is the roofline model.
+    """
+    n_active = total_active_params(model)
+    b = model.bytes_per_elem
+    weight_bytes = n_active * b
+    compute_flops = _FLOPS_PER_PARAM_PER_TOKEN * n_active
+    kv_read_bytes = s_kv * kv_bytes_per_token(model)
+
+    compute_s = compute_flops / machine.peak_flops
+    kv_s = kv_read_bytes / machine.bw_hbm
+
+    points: list[RooflinePoint] = []
+    for bs in batch_range:
+        weight_s = weight_bytes / machine.bw_hbm / max(1, bs)
+        points.append(RooflinePoint(
+            batch=bs,
+            weight_s=weight_s,
+            compute_s=compute_s,
+            kv_s=kv_s,
+            total_s=weight_s + compute_s + kv_s,
+        ))
+    return points
+
+
+def bound_regime(machine: MachineParams, model: ModelConfig,
+                  batch: int, s_kv: int) -> str:
+    """Which term dominates at the current (batch, S_kv) point.
+
+    Returns 'memory-bound' if weight_fetch is the largest term,
+    'kv-bound' if KV read is largest, 'compute-bound' if compute.
+    """
+    pts = per_token_latency_curve(machine, model, [batch], s_kv)
+    p = pts[0]
+    parts = {"memory-bound": p.weight_s,
+             "kv-bound": p.kv_s,
+             "compute-bound": p.compute_s}
+    return max(parts, key=parts.get)
+
+
+# ── Regime-dependent cost terms ────────────────────────────────────
+
+
+def t_mem_short(machine: MachineParams, model: ModelConfig,
+                 batch: int) -> float:
+    """Per-token weight-fetch time (short-context regime term).
+
+    N_active · b / (W · B). Shrinks as B grows — this is what
+    batching amortizes.
+    """
+    return (total_active_params(model) * model.bytes_per_elem
+            / (machine.bw_hbm * max(1, batch)))
+
+
+def t_mem_long(machine: MachineParams, model: ModelConfig,
+                s_kv: int) -> float:
+    """Per-token KV-read time (long-context regime term).
+
+    S_kv · kv_bpt / W. Independent of B — each sequence reads its
+    own KV cache; batching doesn't help.
+    """
+    return s_kv * kv_bytes_per_token(model) / machine.bw_hbm
+
+
+def t_com(machine: MachineParams, model: ModelConfig) -> float:
+    """Per-token compute time. Same in both regimes: 2·N/C, peak."""
+    return _FLOPS_PER_PARAM_PER_TOKEN * total_active_params(model) / machine.peak_flops
+
+
+# ── "Good" batch / context recommendations ─────────────────────────
+
+
+@dataclass
+class BatchRecommendation:
+    target: float          # Pope's rule: 2 × B*
+    b_star: float          # B* itself
+    effective: float       # what we recommend using
+    reason: str            # short explanation
+
+
+@dataclass
+class ContextRecommendation:
+    l_star: float          # balance context length
+    max_efficient: float   # same as l_star (compute-friendly ceiling)
+    utilization_at: float  # utilization at current s_kv
+    reason: str
+
+
+def good_batch(machine: MachineParams, model: ModelConfig,
+                sparsity: float = 1.0) -> BatchRecommendation:
+    """Recommended batch size: 2 × B* (Pope's rule of thumb).
+
+    Below B*: memory-bound, doubling B halves cost/token.
+    At 2×B*: 50% excess over compute floor — the sweet spot.
+    Beyond 3×B*: diminishing returns; latency keeps growing linearly.
+    """
+    b_star = critical_batch(machine, model, sparsity)
+    target = 2 * b_star
+    return BatchRecommendation(
+        target=target, b_star=b_star, effective=target,
+        reason=(f"2·B* = 2 · {b_star:.0f} = {target:.0f}. "
+                "Below B*: memory-bound (doubling B halves cost/token). "
+                "Beyond 3·B*: diminishing returns."),
+    )
+
+
+def good_context(machine: MachineParams, model: ModelConfig,
+                  s_kv: int) -> ContextRecommendation:
+    """Recommended max context: L* — the compute-friendly ceiling.
+
+    Below L*: KV read is cheap relative to compute → good utilization.
+    Above L*: KV bandwidth wall → utilization = 1/(1 + S_kv/L*).
+    """
+    l_star = balance_context(machine, model)
+    util = utilization_at(s_kv, l_star)
+    return ContextRecommendation(
+        l_star=l_star, max_efficient=l_star, utilization_at=util,
+        reason=(f"L* = {l_star:,.0f} tokens. Below L*: compute-bound "
+                f"(good util). At {s_kv:,} tokens: peak utilization ≈ "
+                f"{util*100:.1f}% (1 / (1 + S_kv/L*))."),
+    )
+
+
+def utilization_at(s_kv: int, l_star: float) -> float:
+    """Peak compute utilization at context length s_kv, given L*.
+
+    util = compute / (compute + KV_read) = 1 / (1 + S_kv/L*).
+    At S_kv=0: 100%. At S_kv=L*: 50%. At 2·L*: 33.3%. At 5·L*: 16.7%.
+    """
+    return 1.0 / (1.0 + s_kv / l_star)
+
+
+# ── Per-step latency (undivided by B) ─────────────────────────────
+
+
+@dataclass
+class StepLatencyPoint:
+    batch: int
+    weight_s: float   # N·b / W — flat in B (loaded once per step)
+    compute_s: float  # 2·N·B / C — linear in B
+    kv_s: float       # B · S_kv · kv_bpt / W — linear in B
+    total_s: float
+
+
+def step_latency_curve(machine: MachineParams, model: ModelConfig,
+                        batch_range: list[int],
+                        s_kv: int) -> list[StepLatencyPoint]:
+    """Total time of one decode step (one forward pass), across a
+    range of batch sizes. NOT divided by B — this is the SLO view.
+
+    step_weight = N·b / W                (batch-invariant)
+    step_compute = 2·N·B / C              (linear in B)
+    step_kv = B · S_kv · kv_bpt / W       (linear in B)
+
+    Per-token cost = step_total / B — the two views trade off:
+    bigger B lowers cost/token but raises step latency.
+    """
+    n_active = total_active_params(model)
+    b = model.bytes_per_elem
+    weight_bytes = n_active * b
+    kv_bpt = kv_bytes_per_token(model)
+
+    step_weight = weight_bytes / machine.bw_hbm
+
+    points: list[StepLatencyPoint] = []
+    for bs in batch_range:
+        B = max(1, bs)
+        step_compute = _FLOPS_PER_PARAM_PER_TOKEN * n_active * B / machine.peak_flops
+        step_kv = B * s_kv * kv_bpt / machine.bw_hbm
+        total = step_weight + step_compute + step_kv
+        points.append(StepLatencyPoint(
+            batch=bs,
+            weight_s=step_weight,
+            compute_s=step_compute,
+            kv_s=step_kv,
+            total_s=total,
+        ))
+    return points
+
+
+# ── PE memory budget curves ───────────────────────────────────────
+
+
+@dataclass
+class MemoryBudgetPoint:
+    axis_val: int          # S_kv or B being swept
+    weights_gb: float
+    kv_gb: float
+    transient_gb: float
+    used_gb: float
+    free_gb: float         # max(0, hbm_gb - used_gb)
+    over_budget: bool
+
+
+def _budget_point(weights_bytes: int, kv_bytes: int, transient_bytes: int,
+                   hbm_bytes: int, axis_val: int) -> MemoryBudgetPoint:
+    used = weights_bytes + kv_bytes + transient_bytes
+    return MemoryBudgetPoint(
+        axis_val=axis_val,
+        weights_gb=weights_bytes / 1e9,
+        kv_gb=kv_bytes / 1e9,
+        transient_gb=transient_bytes / 1e9,
+        used_gb=used / 1e9,
+        free_gb=max(0, hbm_bytes - used) / 1e9,
+        over_budget=(used > hbm_bytes),
+    )
+
+
+def memory_budget_curve_vs_skv(cfg: FullConfig,
+                                s_kv_range: list[int],
+                                batch: int) -> list[MemoryBudgetPoint]:
+    """Per-PE memory as S_kv sweeps. Uses cfg's current sharding
+    (CP, TP, PP). Sets topo.b = batch. Returns one point per S_kv."""
+    from .memory_layout import (
+        per_pe_kv_cache_bytes, per_pe_transient_bytes, per_pe_weight_bytes,
+    )
+    hbm_bytes = int(cfg.machine.pe_budget_bytes)
+    weights_bytes = per_pe_weight_bytes(cfg)
+    transient_bytes = per_pe_transient_bytes(cfg)
+    points: list[MemoryBudgetPoint] = []
+    for skv in s_kv_range:
+        swept_topo = replace(cfg.topo, s_kv=int(skv), b=max(1, int(batch)))
+        swept_cfg = FullConfig(model=cfg.model, topo=swept_topo,
+                                machine=cfg.machine)
+        kv_bytes = per_pe_kv_cache_bytes(swept_cfg)
+        points.append(_budget_point(weights_bytes, kv_bytes,
+                                     transient_bytes, hbm_bytes, int(skv)))
+    return points
+
+
+def memory_budget_curve_vs_batch(cfg: FullConfig,
+                                   b_range: list[int],
+                                   s_kv: int) -> list[MemoryBudgetPoint]:
+    """Per-PE memory as B sweeps. Sets topo.s_kv = s_kv."""
+    from .memory_layout import (
+        per_pe_kv_cache_bytes, per_pe_transient_bytes, per_pe_weight_bytes,
+    )
+    hbm_bytes = int(cfg.machine.pe_budget_bytes)
+    weights_bytes = per_pe_weight_bytes(cfg)
+    transient_bytes = per_pe_transient_bytes(cfg)
+    points: list[MemoryBudgetPoint] = []
+    for bs in b_range:
+        swept_topo = replace(cfg.topo, b=max(1, int(bs)), s_kv=int(s_kv))
+        swept_cfg = FullConfig(model=cfg.model, topo=swept_topo,
+                                machine=cfg.machine)
+        kv_bytes = per_pe_kv_cache_bytes(swept_cfg)
+        points.append(_budget_point(weights_bytes, kv_bytes,
+                                     transient_bytes, hbm_bytes, int(bs)))
+    return points
+
+
+# ── AI / B* sensitivity to hardware knobs ──────────────────────────
+
+
+@dataclass
+class AISensitivityPoint:
+    multiplier: float          # scale factor applied to the base machine
+    peak_tflops: float
+    bw_gbs: float
+    ai: float                  # C / W
+    b_star: float              # C·b/(2·W)
+
+
+# ── GPU count sizing (three-axis) ─────────────────────────────────
+
+
+@dataclass
+class GpuSizingResult:
+    # Inputs echoed back for display / debugging
+    n_users: int
+    avg_ctx_tokens: int
+    tpot_slo_s: float
+    # Per-replica values
+    pes_axis_a_capacity: int    # PEs to hold weights alone (bare floor)
+    pes_axis_b_kv: int          # PEs to hold weights + KV of this replica's users
+    pes_per_replica: int        # max(A, B)
+    # Throughput axis
+    b_at_slo: int               # largest per-replica B satisfying SLO
+    users_per_replica: int      # min(n_users, b_at_slo), at least 1
+    n_replicas: int             # ceil(n_users / users_per_replica)
+    # Grand total
+    total_pes: int
+    binding_axis: str           # 'capacity' | 'kv' | 'throughput'
+
+
+def max_batch_within_slo(machine: MachineParams, model: ModelConfig,
+                          s_kv: int, tpot_slo_s: float) -> int:
+    """Largest per-replica B such that decode step latency ≤ SLO.
+
+    step_latency(B) = N·b/W + 2·N·B/C + B·S_kv·kv_bpt/W
+                    = weight_fetch + B · (compute + kv_read)
+
+    Returns 0 if even B=1 exceeds SLO (weight fetch alone too large,
+    or per-sequence term already blows the budget).
+    """
+    n = total_active_params(model)
+    b_elem = model.bytes_per_elem
+    weight_time = n * b_elem / machine.bw_hbm
+    per_seq_time = (_FLOPS_PER_PARAM_PER_TOKEN * n / machine.peak_flops
+                    + s_kv * kv_bytes_per_token(model) / machine.bw_hbm)
+    if weight_time + per_seq_time > tpot_slo_s:
+        return 0
+    remaining = tpot_slo_s - weight_time
+    b_max = int(remaining / per_seq_time)
+    return max(1, b_max)
+
+
+def size_deployment(machine: MachineParams, model: ModelConfig,
+                     n_users: int, avg_ctx: int,
+                     tpot_slo_s: float) -> GpuSizingResult:
+    """Three-axis GPU count for a target workload.
+
+    Axis A — capacity floor: PEs to hold one replica's weights.
+    Axis B — KV headroom:    PEs to hold weights + all KV of the
+                             users assigned to this replica.
+    Axis C — throughput SLO: replicas needed so per-replica B ≤ b_at_slo.
+
+    Result: pes_per_replica = max(A, B); total = pes_per_replica × replicas.
+    Binding axis is whichever grew the count the most.
+    """
+    n = total_active_params(model)
+    b_elem = model.bytes_per_elem
+    weight_bytes = n * b_elem
+    hbm_pe = int(machine.pe_hbm_gb * 1e9)
+    kv_bpt = kv_bytes_per_token(model)
+
+    # Axis C: throughput
+    b_at_slo = max_batch_within_slo(machine, model, avg_ctx, tpot_slo_s)
+    users_per_replica = max(1, min(n_users, b_at_slo)) if b_at_slo > 0 else 1
+    n_replicas = (n_users + users_per_replica - 1) // users_per_replica
+
+    # Axis A: bare weights
+    pes_a = (weight_bytes + hbm_pe - 1) // hbm_pe
+
+    # Axis B: weights + KV load for this replica's users
+    kv_per_replica = users_per_replica * avg_ctx * kv_bpt
+    pes_b = (weight_bytes + kv_per_replica + hbm_pe - 1) // hbm_pe
+
+    pes_per_replica = max(pes_a, pes_b)
+    total = pes_per_replica * n_replicas
+
+    if n_replicas > 1:
+        binding = "throughput"
+    elif pes_b > pes_a:
+        binding = "kv"
+    else:
+        binding = "capacity"
+
+    return GpuSizingResult(
+        n_users=n_users,
+        avg_ctx_tokens=avg_ctx,
+        tpot_slo_s=tpot_slo_s,
+        pes_axis_a_capacity=int(pes_a),
+        pes_axis_b_kv=int(pes_b),
+        pes_per_replica=int(pes_per_replica),
+        b_at_slo=b_at_slo,
+        users_per_replica=int(users_per_replica),
+        n_replicas=int(n_replicas),
+        total_pes=int(total),
+        binding_axis=binding,
+    )
+
+
+def ai_sensitivity_curve(machine: MachineParams, model: ModelConfig,
+                          multipliers: list[float],
+                          axis: str = "flops") -> list[AISensitivityPoint]:
+    """Sweep FLOPs OR BW while holding the other fixed. Returns AI + B*
+    at each multiplier.
+
+    axis='flops' → scales peak_tflops_f16 (AI grows linearly).
+    axis='bw'    → scales bw_hbm_gbs (AI shrinks inversely).
+    """
+    if axis not in ("flops", "bw"):
+        raise ValueError(f"axis must be 'flops' or 'bw', got {axis!r}")
+    base_flops = machine.peak_tflops_f16
+    base_bw = machine.bw_hbm_gbs
+    points: list[AISensitivityPoint] = []
+    for k in multipliers:
+        if axis == "flops":
+            m = replace(machine, peak_tflops_f16=base_flops * k)
+        else:
+            m = replace(machine, bw_hbm_gbs=base_bw * k)
+        points.append(AISensitivityPoint(
+            multiplier=k,
+            peak_tflops=m.peak_tflops_f16,
+            bw_gbs=m.bw_hbm_gbs,
+            ai=arithmetic_intensity(m),
+            b_star=critical_batch(m, model),
+        ))
+    return points

tests/analytical_visualization/memory_layout.py

@@ -0,0 +1,248 @@
+"""Per-PE memory footprint: weights + KV cache + activations + slack.
+
+All formulas are per-PE, per-PP-stage. If PP > 1, each PE holds only
+its stage's layers (layers/PP).
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from .model_config import FullConfig
+
+
+@dataclass
+class MemoryBreakdown:
+    weights_bytes: int
+    kv_cache_bytes: int
+    transient_bytes: int
+    budget_bytes: int
+
+    @property
+    def used_bytes(self) -> int:
+        return self.weights_bytes + self.kv_cache_bytes + self.transient_bytes
+
+    @property
+    def slack_bytes(self) -> int:
+        return max(0, self.budget_bytes - self.used_bytes)
+
+    @property
+    def over_budget(self) -> bool:
+        return self.used_bytes > self.budget_bytes
+
+
+def per_pe_weight_bytes(cfg: FullConfig,
+                        include_attention: bool = True,
+                        include_ffn: bool = True) -> int:
+    """Attention + FFN weights per PE, bf16. Divided by TP, PP, and EP.
+
+    EP divides the FFN (experts) across ranks; attention weights are
+    unaffected. When kv_shard_mode='replicate' and TP > H_kv, each KV
+    head is replicated (per-PE W_K/W_V size doesn't drop below one head).
+
+    Scope flags let callers count only attention or only FFN weights —
+    useful for "smallest deployment for just this block" sizing. Both
+    True (default) matches the traditional full-model per-PE weight
+    footprint.
+    """
+    m = cfg.model
+    tp = cfg.topo.tp
+    pp = cfg.topo.pp
+    ep = max(1, cfg.topo.ep)
+
+    hq_per_pe = cfg.h_q_per_pe
+    if cfg.topo.kv_shard_mode == "replicate":
+        # Each KV head held fully; replicated across ranks if TP > H_kv.
+        hkv_per_pe_bytes = max(1.0, m.h_kv / tp)
+    else:  # "split"
+        # Head-dim split: fractional head allowed (< 1 head bytes when TP > H_kv).
+        hkv_per_pe_bytes = m.h_kv / tp
+
+    per_layer_attn = 0
+    if include_attention:
+        per_layer_attn = (
+            m.hidden * hq_per_pe * m.d_head +           # W_Q
+            m.hidden * hkv_per_pe_bytes * m.d_head +    # W_K
+            m.hidden * hkv_per_pe_bytes * m.d_head +    # W_V
+            hq_per_pe * m.d_head * m.hidden             # W_O
+        )
+
+    per_layer_ffn = 0
+    if include_ffn:
+        # FFN divisor: TP (default) or TP*CP or TP*CP*DP if user opts in.
+        # EP further divides (MoE experts).
+        ffn_div = cfg.ffn_shard_divisor * ep
+        per_layer_ffn = 3 * m.hidden * (m.ffn_dim // ffn_div)
+
+    per_layer = (per_layer_attn + per_layer_ffn) * m.bytes_per_elem
+    layers_per_stage = (m.layers + pp - 1) // pp
+    return int(per_layer * layers_per_stage)
+
+
+def per_pe_kv_cache_bytes(cfg: FullConfig) -> int:
+    """K + V per PE, across all layers this stage holds.
+
+    - CP shards the sequence dim -> S_local = S_kv/CP tokens per PE.
+    - TP splits KV heads across ranks. If kv_shard_mode='replicate' and
+      TP > H_kv, each KV head is duplicated across TP/H_kv ranks so
+      per-PE storage doesn't shrink below 1 head.
+    - PP: only layers_per_stage layers stored per PE.
+    - Batch (B): each concurrent request keeps its own KV cache slice.
+    """
+    m = cfg.model
+    pp = cfg.topo.pp
+    tp = cfg.topo.tp
+    B = max(1, cfg.topo.b)
+    if cfg.topo.kv_shard_mode == "replicate":
+        hkv_per_pe_bytes = max(1.0, m.h_kv / tp)
+    else:
+        hkv_per_pe_bytes = m.h_kv / tp
+    layers_per_stage = (m.layers + pp - 1) // pp
+    per_layer = 2 * cfg.topo.s_local * hkv_per_pe_bytes * m.d_head * m.bytes_per_elem
+    return int(per_layer * layers_per_stage * B)
+
+
+def per_pe_transient_bytes(cfg: FullConfig) -> int:
+    """Rough peak transient (activations, GEMM outputs) per PE."""
+    m = cfg.model
+    T_q = cfg.topo.T_q
+    hq_per_pe = cfg.h_q_per_pe
+
+    if cfg.topo.mode == "decode":
+        return 4 * (m.hidden + m.head_dim_total_q // cfg.topo.tp) * m.bytes_per_elem
+    else:
+        TILE = 1024
+        tile_score = hq_per_pe * T_q * TILE * m.bytes_per_elem
+        return int(2 * tile_score + T_q * m.hidden * m.bytes_per_elem)
+
+
+def compute_memory(cfg: FullConfig,
+                   include_attention: bool = True,
+                   include_ffn: bool = True) -> MemoryBreakdown:
+    """Per-PE memory breakdown. Scope flags let callers count only the
+    attention or only the FFN block (KV cache goes with the attention
+    block; transient activation buffer is small either way and left in).
+    """
+    kv_bytes = per_pe_kv_cache_bytes(cfg) if include_attention else 0
+    return MemoryBreakdown(
+        weights_bytes=per_pe_weight_bytes(
+            cfg, include_attention=include_attention, include_ffn=include_ffn,
+        ),
+        kv_cache_bytes=kv_bytes,
+        transient_bytes=per_pe_transient_bytes(cfg),
+        budget_bytes=cfg.machine.pe_budget_bytes,
+    )
+
+
+def _one_layer_row(name: str,
+                   global_shape: tuple[int, int],
+                   per_pe_shape: tuple[int, int],
+                   bytes_per_elem: int) -> dict:
+    """Per-tensor row for ONE layer (both global and per-PE shard)."""
+    p_global = global_shape[0] * global_shape[1]
+    p_per_pe = per_pe_shape[0] * per_pe_shape[1]
+    return {
+        "Tensor": name,
+        "Global shape": f"({global_shape[0]}, {global_shape[1]})",
+        "Params/layer": f"{p_global/1e6:.2f} M",
+        "Bytes/layer (global)": f"{p_global * bytes_per_elem / 1e6:.2f} MB",
+        "Per-PE shape": f"({per_pe_shape[0]}, {per_pe_shape[1]})",
+        "Bytes/layer (per PE)": f"{p_per_pe * bytes_per_elem / 1e6:.2f} MB",
+        "_p_global": p_global,
+        "_p_per_pe": p_per_pe,
+    }
+
+
+def attention_weight_rows(cfg: FullConfig) -> list[dict]:
+    """Per-tensor rows for attention weights (one layer)."""
+    m = cfg.model
+    hq_per_pe = cfg.h_q_per_pe
+    hkv_per_pe = max(1, m.h_kv // cfg.topo.tp)
+    return [
+        _one_layer_row("W_Q", (m.hidden, m.h_q * m.d_head),
+                       (m.hidden, hq_per_pe * m.d_head),
+                       m.bytes_per_elem),
+        _one_layer_row("W_K", (m.hidden, m.h_kv * m.d_head),
+                       (m.hidden, hkv_per_pe * m.d_head),
+                       m.bytes_per_elem),
+        _one_layer_row("W_V", (m.hidden, m.h_kv * m.d_head),
+                       (m.hidden, hkv_per_pe * m.d_head),
+                       m.bytes_per_elem),
+        _one_layer_row("W_O", (m.h_q * m.d_head, m.hidden),
+                       (hq_per_pe * m.d_head, m.hidden),
+                       m.bytes_per_elem),
+    ]
+
+
+def ffn_weight_rows(cfg: FullConfig) -> list[dict]:
+    """Per-tensor rows for FFN weights (one layer, activated for MoE)."""
+    m = cfg.model
+    ffn_per_pe = m.ffn_dim // cfg.topo.tp
+    return [
+        _one_layer_row("W_gate", (m.hidden, m.ffn_dim),
+                       (m.hidden, ffn_per_pe), m.bytes_per_elem),
+        _one_layer_row("W_up", (m.hidden, m.ffn_dim),
+                       (m.hidden, ffn_per_pe), m.bytes_per_elem),
+        _one_layer_row("W_down", (m.ffn_dim, m.hidden),
+                       (ffn_per_pe, m.hidden), m.bytes_per_elem),
+    ]
+
+
+def kv_cache_rows(cfg: FullConfig) -> list[dict]:
+    """Per-tensor row for K and V cache (one layer)."""
+    m = cfg.model
+    hkv_per_pe = max(1, m.h_kv // cfg.topo.tp)
+    b = m.bytes_per_elem
+    global_k = cfg.topo.s_kv * m.h_kv * m.d_head
+    per_pe_k = cfg.topo.s_local * hkv_per_pe * m.d_head
+    return [
+        {
+            "Tensor": "K cache",
+            "Global shape": f"({cfg.topo.s_kv:,}, {m.h_kv * m.d_head})",
+            "Params/layer": f"{global_k/1e6:.2f} M",
+            "Bytes/layer (global)": f"{global_k * b / 1e6:.2f} MB",
+            "Per-PE shape": f"({cfg.topo.s_local:,}, {hkv_per_pe * m.d_head})",
+            "Bytes/layer (per PE)": f"{per_pe_k * b / 1e6:.2f} MB",
+            "_p_global": global_k,
+            "_p_per_pe": per_pe_k,
+        },
+        {
+            "Tensor": "V cache",
+            "Global shape": f"({cfg.topo.s_kv:,}, {m.h_kv * m.d_head})",
+            "Params/layer": f"{global_k/1e6:.2f} M",
+            "Bytes/layer (global)": f"{global_k * b / 1e6:.2f} MB",
+            "Per-PE shape": f"({cfg.topo.s_local:,}, {hkv_per_pe * m.d_head})",
+            "Bytes/layer (per PE)": f"{per_pe_k * b / 1e6:.2f} MB",
+            "_p_global": global_k,
+            "_p_per_pe": per_pe_k,
+        },
+    ]
+
+
+def sum_bytes_all_layers(rows: list[dict], cfg: FullConfig,
+                          per_pe: bool = False) -> int:
+    """Sum bytes across all rows × N layers (or layers/PP for per_pe)."""
+    b = cfg.model.bytes_per_elem
+    layers = cfg.model.layers
+    if per_pe:
+        layers = (cfg.model.layers + cfg.topo.pp - 1) // cfg.topo.pp
+    key = "_p_per_pe" if per_pe else "_p_global"
+    return sum(r[key] for r in rows) * b * layers
+
+
+def total_weight_bytes_full_model(cfg: FullConfig) -> int:
+    """Total bytes of ALL weights (attention + FFN) for the full unsharded model."""
+    m = cfg.model
+    attn_per_layer = (
+        m.hidden * m.h_q * m.d_head +          # W_Q
+        m.hidden * m.h_kv * m.d_head * 2 +     # W_K, W_V
+        m.h_q * m.d_head * m.hidden            # W_O
+    )
+    ffn_per_layer = 3 * m.hidden * m.ffn_dim
+    return (attn_per_layer + ffn_per_layer) * m.bytes_per_elem * m.layers
+
+
+def total_kv_bytes_full_model(cfg: FullConfig) -> int:
+    """Total bytes of KV cache for full unsharded model at cfg.topo.s_kv."""
+    m = cfg.model
+    per_layer = 2 * cfg.topo.s_kv * m.h_kv * m.d_head * m.bytes_per_elem
+    return per_layer * m.layers

tests/analytical_visualization/model_config.py

@@ -0,0 +1,311 @@
+"""Model + machine parameters for the analytical visualization tool.
+
+All numbers are per-rank / per-PE unless noted. bytes are counted in
+bytes (b), FLOPs in raw ops, bandwidth in bytes/second, latency in
+seconds. Convert to μs / GB / GFLOP at display time only.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+
+@dataclass
+class ModelConfig:
+    """Transformer model dimensions."""
+    name: str = "Qwen 3 8B"
+    hidden: int = 4096
+    ffn_dim: int = 12288
+    h_q: int = 32               # Q heads
+    h_kv: int = 8               # KV heads
+    d_head: int = 128
+    layers: int = 36
+    bytes_per_elem: int = 2     # bf16
+
+    @property
+    def group_size(self) -> int:
+        """GQA group size = H_q / H_kv."""
+        return self.h_q // self.h_kv
+
+    @property
+    def head_dim_total_q(self) -> int:
+        return self.h_q * self.d_head
+
+    @property
+    def head_dim_total_kv(self) -> int:
+        return self.h_kv * self.d_head
+
+
+@dataclass
+class TopologyConfig:
+    """Physical mapping — DP × PP × CP × TP (× EP for MoE)."""
+    cp: int = 4                 # inter-cube ring size (sequence shard)
+    tp: int = 8                 # PEs per cube (head/hidden shard)
+    pp: int = 1                 # pipeline stages (layer shard)
+    dp: int = 1                 # data-parallel replicas
+    ep: int = 1                 # expert-parallel degree (MoE only)
+    b: int = 1                  # batch size (concurrent requests per PP stage)
+    s_kv: int = 4096
+    mode: str = "decode"        # "decode" or "prefill"
+    kv_shard_mode: str = "split"  # "split" or "replicate" (used when TP > H_kv)
+    # FFN sharding scope: default TP+CP (across cube group) trades a small
+    # extra AllReduce for lower per-PE weight bytes.
+    ffn_shard_scope: str = "TP+CP"   # "TP" | "TP+CP" | "TP+CP+DP"
+    # Inter-SIP interconnect: how physical SIPs are wired together.
+    # "ring" (1D wrap), "mesh2d" (grid, no wrap), "torus2d" (grid + wrap).
+    sip_topology: str = "ring"
+    # Placement of parallelism dims on the hierarchy.
+    # "pe"   -> dim lives on PEs within a single cube (intra-cube).
+    # "cube" -> dim lives across cubes (inter-cube).
+    tp_placement: str = "pe"     # default: TP fills PEs of a cube
+    cp_placement: str = "cube"   # default: CP fills cubes of a SIP
+    # CP ring variant: what gets passed around the CP ring each hop.
+    # "kv"    -> K and V shards rotate (Q,O,m,l stay local). Bytes/hop scales
+    #            with S_local * H_kv * d_h. Good for prefill (large T_q).
+    # "qoml"  -> Q and running (O, m, l) rotate; K,V stay local. Bytes/hop
+    #            scales with T_q * H_q * d_h. Cheap for decode (T_q=1).
+    cp_ring_variant: str = "kv"
+    pes_per_cube_hw: int = 8    # SIP hardware constant
+    cubes_per_sip_hw: int = 16
+
+    @property
+    def pes_per_stage(self) -> int:
+        """PEs per pipeline stage (all stages have the same size)."""
+        return self.cp * self.tp
+
+    @property
+    def pes_per_replica(self) -> int:
+        """PEs per one full model replica (CP*TP*PP)."""
+        return self.cp * self.tp * self.pp
+
+    @property
+    def total_pes(self) -> int:
+        return self.pes_per_replica * self.dp
+
+    # ── Placement-aware layout ────────────────────────────────
+    @property
+    def intra_cube_dims(self) -> int:
+        """Product of dims placed on PE level (must fit in one cube, else spills)."""
+        n = 1
+        if self.tp_placement == "pe":
+            n *= self.tp
+        if self.cp_placement == "pe":
+            n *= self.cp
+        return n
+
+    @property
+    def inter_cube_dims(self) -> int:
+        """Product of dims placed at cube level = distinct cube groups per stage."""
+        n = 1
+        if self.tp_placement == "cube":
+            n *= self.tp
+        if self.cp_placement == "cube":
+            n *= self.cp
+        return n
+
+    @property
+    def pes_per_cube_used(self) -> int:
+        """How many PEs are live in each cube."""
+        return min(self.intra_cube_dims, self.pes_per_cube_hw)
+
+    @property
+    def cubes_per_stage(self) -> int:
+        """Cubes needed per PP stage under current placement.
+        If intra-cube demand exceeds PEs/cube, the intra dim spills to cubes."""
+        spill = max(1, (self.intra_cube_dims + self.pes_per_cube_hw - 1)
+                    // self.pes_per_cube_hw)
+        return self.inter_cube_dims * spill
+
+    @property
+    def cubes_used(self) -> int:
+        """Physical cubes used across all PP stages and DP replicas."""
+        return self.cubes_per_stage * self.pp * self.dp
+
+    @property
+    def sips_used(self) -> int:
+        """Number of SIPs used (each SIP = pes_per_cube_hw × cubes_per_sip_hw)."""
+        cubes_per_sip = self.cubes_per_sip_hw
+        return max(1, (self.cubes_used + cubes_per_sip - 1) // cubes_per_sip)
+
+    @property
+    def placement_valid(self) -> bool:
+        """False when intra-cube demand exceeds one cube (dim spills to cubes)."""
+        return self.intra_cube_dims <= self.pes_per_cube_hw
+
+    @property
+    def tp_spans_cubes(self) -> int:
+        """How many cubes one TP group spans (>=1). Depends on placement:
+        - tp_placement=pe: 1 if TP fits in a cube, else spill.
+        - tp_placement=cube: TP itself is the number of cubes per TP group.
+        """
+        if self.tp_placement == "cube":
+            return max(1, self.tp)
+        return max(1, (self.tp + self.pes_per_cube_hw - 1) // self.pes_per_cube_hw)
+
+    @property
+    def cp_intra_sip_hops(self) -> int:
+        """CP ring hops that stay within one SIP."""
+        if self.cp <= 1:
+            return 0
+        sips_in_ring = self.sips_used  # one PP stage's CP ring may span SIPs
+        return max(0, (self.cp - 1) - (sips_in_ring - 1))
+
+    @property
+    def cp_inter_sip_hops(self) -> int:
+        """CP ring hops that cross SIP boundaries."""
+        if self.cp <= 1:
+            return 0
+        # Approximation: one inter-SIP hop per SIP boundary in the ring.
+        # For CP=32 across 2 SIPs → 1 inter-SIP + 30 intra-SIP hops.
+        return max(0, self.sips_used - 1)
+
+    @property
+    def T_q(self) -> int:
+        return 1 if self.mode == "decode" else self.s_kv // self.cp
+
+    @property
+    def s_local(self) -> int:
+        return self.s_kv // self.cp
+
+    @property
+    def layers_per_stage(self) -> str:
+        """Layers per PP stage. Requires model.layers to render — leave as fn."""
+        return "N_layers / PP"
+
+    def tp_link_tier(self) -> str:
+        """Return which BW tier a TP AllReduce uses: 'intra' | 'inter' | 'intersip'.
+        - tp_placement=pe and TP fits in one cube -> 'intra'
+        - tp_placement=cube and TP fits in one SIP -> 'inter'
+        - otherwise -> 'intersip'
+        """
+        if self.tp_placement == "pe":
+            if self.tp <= self.pes_per_cube_hw:
+                return "intra"
+            # spilled to multiple cubes: cross-cube
+            if self.tp_spans_cubes <= self.cubes_per_sip_hw:
+                return "inter"
+            return "intersip"
+        else:  # tp_placement == "cube"
+            if self.tp <= self.cubes_per_sip_hw:
+                return "inter"
+            return "intersip"
+
+    def cp_link_tier(self) -> str:
+        """Return BW tier for the CP ring: 'intra' | 'inter' | 'intersip'."""
+        if self.cp_placement == "pe":
+            return "intra"
+        # cp_placement == "cube"
+        if self.cp <= self.cubes_per_sip_hw:
+            return "inter"
+        return "intersip"
+
+
+@dataclass
+class MachineParams:
+    """SIP hardware parameters — all user-adjustable via sliders.
+
+    Bandwidth tiers, from fastest to slowest:
+      HBM (per-PE local memory)
+        > intra-cube (PE↔PE on same die)
+        > inter-cube (UCIe D2D between dies on same SIP)
+        > inter-SIP (C2C / RDMA-class, across SIPs / servers)
+    """
+    # Per-PE compute peak (TFLOPs f16).
+    peak_tflops_f16: float = 8.0
+    # Per-PE HBM budget (GB) — the memory ceiling per PE.
+    pe_hbm_gb: float = 6.0
+    # HBM per-PE bandwidth (GB/s).
+    bw_hbm_gbs: float = 256.0
+    # Intra-cube PE↔PE link BW (GB/s).
+    bw_intra_gbs: float = 512.0
+    # Inter-cube (UCIe D2D) BW (GB/s).
+    bw_inter_gbs: float = 128.0
+    # Inter-SIP (C2C / RDMA) BW (GB/s).
+    bw_intersip_gbs: float = 50.0
+    # Per-hop latencies (ns).
+    alpha_intra_ns: float = 20.0
+    alpha_inter_ns: float = 100.0
+    alpha_intersip_ns: float = 1000.0
+    # Achievable utilization for compute-bound stages (0-1).
+    compute_util: float = 0.8
+
+    @property
+    def peak_flops(self) -> float:
+        return self.peak_tflops_f16 * 1e12
+
+    @property
+    def bw_hbm(self) -> float:
+        return self.bw_hbm_gbs * 1e9
+
+    @property
+    def bw_intra(self) -> float:
+        return self.bw_intra_gbs * 1e9
+
+    @property
+    def bw_inter(self) -> float:
+        return self.bw_inter_gbs * 1e9
+
+    @property
+    def bw_intersip(self) -> float:
+        return self.bw_intersip_gbs * 1e9
+
+    @property
+    def alpha_intra(self) -> float:
+        return self.alpha_intra_ns * 1e-9
+
+    @property
+    def alpha_inter(self) -> float:
+        return self.alpha_inter_ns * 1e-9
+
+    @property
+    def alpha_intersip(self) -> float:
+        return self.alpha_intersip_ns * 1e-9
+
+    @property
+    def pe_budget_bytes(self) -> int:
+        return int(self.pe_hbm_gb * 1e9)
+
+
+@dataclass
+class FullConfig:
+    """Bundled config for one analysis run."""
+    model: ModelConfig = field(default_factory=ModelConfig)
+    topo: TopologyConfig = field(default_factory=TopologyConfig)
+    machine: MachineParams = field(default_factory=MachineParams)
+
+    @property
+    def h_q_per_pe(self) -> int:
+        return max(1, self.model.h_q // self.topo.tp)
+
+    @property
+    def h_kv_per_pe(self) -> float:
+        """Fractional if TP > H_kv (head-dim split needed)."""
+        return self.model.h_kv / self.topo.tp
+
+    @property
+    def d_per_pe(self) -> int:
+        return self.model.hidden // self.topo.tp
+
+    @property
+    def ffn_per_pe(self) -> int:
+        return self.model.ffn_dim // self.topo.tp
+
+    @property
+    def kv_replication_needed(self) -> bool:
+        """True if TP > H_kv (each KV head shared across TP/H_kv ranks)."""
+        return self.topo.tp > self.model.h_kv
+
+    @property
+    def head_dim_split_factor(self) -> int:
+        """Number of TP ranks sharing one head (via d_head split)."""
+        return max(1, self.topo.tp // self.model.h_kv)
+
+    @property
+    def ffn_shard_divisor(self) -> int:
+        """How much the FFN dim is sharded across ranks."""
+        scope = self.topo.ffn_shard_scope
+        d = self.topo.tp
+        if "CP" in scope:
+            d *= self.topo.cp
+        if "DP" in scope:
+            d *= self.topo.dp
+        return max(1, d)

tests/analytical_visualization/model_presets.py

@@ -0,0 +1,126 @@
+"""Model preset library for the analytical visualization tool.
+
+Comprehensive coverage of popular MHA, GQA, MQA, and MLA models across
+sizes. MoE models are approximated as dense with (ffn_dim = activated
+experts × per_expert_ffn) for compute-side estimates; the note field
+flags this. MLA models are approximated as GQA with matching H_kv until
+we add proper MLA support.
+
+Source of numeric params: model cards / config.json from HuggingFace.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from .model_config import ModelConfig
+
+
+@dataclass
+class Preset:
+    label: str
+    model: ModelConfig
+    family: str = ""      # e.g. "Llama 3", "Qwen 3"
+    attn_type: str = ""   # "MHA" | "GQA" | "MQA" | "MLA"
+    note: str = ""
+
+
+def _mk(name: str, hidden: int, ffn: int, hq: int, hkv: int,
+        dh: int, layers: int, *, family: str = "", attn: str = "GQA",
+        note: str = "") -> Preset:
+    return Preset(
+        label=name,
+        model=ModelConfig(
+            name=name, hidden=hidden, ffn_dim=ffn,
+            h_q=hq, h_kv=hkv, d_head=dh, layers=layers,
+        ),
+        family=family, attn_type=attn, note=note,
+    )
+
+
+PRESETS: dict[str, Preset] = {
+    # ── MHA (H_q == H_kv) ─────────────────────────────────────────
+    "GPT-2 (117M)":       _mk("GPT-2 117M",       768,  3072,  12, 12, 64,  12, family="GPT-2", attn="MHA"),
+    "GPT-2 XL (1.5B)":    _mk("GPT-2 XL 1.5B",    1600, 6400,  25, 25, 64,  48, family="GPT-2", attn="MHA"),
+    "Llama 2 7B":         _mk("Llama 2 7B",       4096, 11008, 32, 32, 128, 32, family="Llama 2", attn="MHA"),
+    "Llama 2 13B":        _mk("Llama 2 13B",      5120, 13824, 40, 40, 128, 40, family="Llama 2", attn="MHA"),
+    "Yi 6B":              _mk("Yi 6B",            4096, 11008, 32, 4,  128, 32, family="Yi",     attn="GQA"),
+    "Yi 9B":              _mk("Yi 9B",            4096, 11008, 32, 4,  128, 48, family="Yi",     attn="GQA"),
+    "OPT 30B":            _mk("OPT 30B",          7168, 28672, 56, 56, 128, 48, family="OPT",    attn="MHA"),
+
+    # ── GQA (H_q > H_kv > 1) ──────────────────────────────────────
+    "Qwen 2 0.5B":        _mk("Qwen 2 0.5B",      896,  4864,  14, 2,  64,  24, family="Qwen 2"),
+    "Qwen 2 1.5B":        _mk("Qwen 2 1.5B",      1536, 8960,  12, 2,  128, 28, family="Qwen 2"),
+    "Qwen 2 7B":          _mk("Qwen 2 7B",        3584, 18944, 28, 4,  128, 28, family="Qwen 2"),
+    "Qwen 2 72B":         _mk("Qwen 2 72B",       8192, 29568, 64, 8,  128, 80, family="Qwen 2"),
+
+    "Qwen 3 0.6B":        _mk("Qwen 3 0.6B",      1024, 3072,  16, 8,  128, 28, family="Qwen 3"),
+    "Qwen 3 1.7B":        _mk("Qwen 3 1.7B",      2048, 6144,  16, 8,  128, 28, family="Qwen 3"),
+    "Qwen 3 4B":          _mk("Qwen 3 4B",        2560, 9728,  32, 8,  128, 36, family="Qwen 3"),
+    "Qwen 3 8B":          _mk("Qwen 3 8B",        4096, 12288, 32, 8,  128, 36, family="Qwen 3"),
+    "Qwen 3 14B":         _mk("Qwen 3 14B",       5120, 17408, 40, 8,  128, 40, family="Qwen 3"),
+    "Qwen 3 32B":         _mk("Qwen 3 32B",       5120, 25600, 64, 8,  128, 64, family="Qwen 3"),
+
+    "Llama 3 8B":         _mk("Llama 3 8B",       4096, 14336, 32, 8,  128, 32, family="Llama 3"),
+    "Llama 3 70B":        _mk("Llama 3 70B",      8192, 28672, 64, 8,  128, 80, family="Llama 3"),
+    "Llama 3.1 8B":       _mk("Llama 3.1 8B",     4096, 14336, 32, 8,  128, 32, family="Llama 3.1"),
+    "Llama 3.1 70B":      _mk("Llama 3.1 70B",    8192, 28672, 64, 8,  128, 80, family="Llama 3.1"),
+    "Llama 3.1 405B":     _mk("Llama 3.1 405B",   16384,53248, 128,8,  128, 126,family="Llama 3.1"),
+    "Llama 3.2 1B":       _mk("Llama 3.2 1B",     2048, 8192,  32, 8,  64,  16, family="Llama 3.2"),
+    "Llama 3.2 3B":       _mk("Llama 3.2 3B",     3072, 8192,  24, 8,  128, 28, family="Llama 3.2"),
+
+    "Mistral 7B":         _mk("Mistral 7B v0.3",  4096, 14336, 32, 8,  128, 32, family="Mistral"),
+    "Mistral Nemo 12B":   _mk("Mistral Nemo 12B", 5120, 14336, 32, 8,  128, 40, family="Mistral"),
+    "Mistral Small 22B":  _mk("Mistral Small 22B",6144, 16384, 48, 8,  128, 56, family="Mistral"),
+    "Mistral Large 123B": _mk("Mistral Large 123B",12288,28672,96, 8,  128, 88, family="Mistral"),
+
+    "Gemma 2 2B":         _mk("Gemma 2 2B",       2304, 9216,  8,  4,  256, 26, family="Gemma 2"),
+    "Gemma 2 9B":         _mk("Gemma 2 9B",       3584, 14336, 16, 8,  256, 42, family="Gemma 2"),
+    "Gemma 2 27B":        _mk("Gemma 2 27B",      4608, 36864, 32, 16, 128, 46, family="Gemma 2"),
+
+    "Phi 3 mini (3.8B)":  _mk("Phi 3 mini 3.8B",  3072, 8192,  32, 32, 96,  32, family="Phi 3", attn="MHA"),
+    "Phi 3 small (7B)":   _mk("Phi 3 small 7B",   4096, 14336, 32, 8,  128, 32, family="Phi 3"),
+    "Phi 3 medium (14B)": _mk("Phi 3 medium 14B", 5120, 17920, 40, 10, 128, 40, family="Phi 3"),
+
+    "Yi 34B":             _mk("Yi 34B",           7168, 20480, 56, 8,  128, 60, family="Yi"),
+    "Command R+ 104B":    _mk("Command R+ 104B",  12288,33792, 96, 8,  128, 64, family="Cohere"),
+
+    # ── MoE (activated-only approximation) ────────────────────────
+    "Mixtral 8x7B (MoE)": _mk("Mixtral 8x7B",     4096, 14336*2, 32, 8, 128, 32,
+                              family="Mistral MoE",
+                              note="MoE: 8 experts x 2 activated. FFN ~ 2 x per_expert."),
+    "Mixtral 8x22B (MoE)":_mk("Mixtral 8x22B",    6144, 16384*2, 48, 8, 128, 56,
+                              family="Mistral MoE",
+                              note="MoE: 8 experts x 2 activated."),
+    "Qwen 3 30B (MoE)":   _mk("Qwen 3 30B (MoE)", 2048, 768*8,   32, 4, 128, 48,
+                              family="Qwen 3 MoE",
+                              note="128 experts x 8 activated per token."),
+    "Qwen 3 235B (MoE)":  _mk("Qwen 3 235B (MoE)",4096, 1536*8,  64, 4, 128, 94,
+                              family="Qwen 3 MoE",
+                              note="128 experts x 8 activated. Dense-approx overstates weight."),
+    "DeepSeek V2 (dense-approx)":
+        _mk("DeepSeek V2",  5120, 12288, 128, 128, 128, 60,
+            family="DeepSeek", attn="MLA",
+            note="MLA compresses KV to ~576 B/token; d_c=512 latent. Dense-approx."),
+    "DeepSeek V3 (dense-approx)":
+        _mk("DeepSeek V3",  7168, 16384, 128, 128, 128, 61,
+            family="DeepSeek", attn="MLA",
+            note="MoE 256 x 8 activated + MLA. Both approximations."),
+    "Grok-1 314B (MoE)":  _mk("Grok-1 314B",     6144, 32768*2, 48, 8,  128, 64,
+                              family="xAI",
+                              note="MoE 8 experts x 2 activated."),
+
+    # ── MQA (H_kv == 1) ───────────────────────────────────────────
+    "Falcon 7B (MQA)":    _mk("Falcon 7B",       4544, 18176, 71, 1,  64,  32, family="Falcon", attn="MQA"),
+    "Falcon 40B (MQA)":   _mk("Falcon 40B",      8192, 32768, 128,1,  64,  60, family="Falcon", attn="MQA"),
+
+    # ── Custom slot ──────────────────────────────────────────────
+    "Custom": Preset(label="Custom", model=ModelConfig(name="Custom")),
+}
+
+
+def preset_names_by_family() -> dict[str, list[str]]:
+    """Return preset names grouped by family (for a nested dropdown)."""
+    grouped: dict[str, list[str]] = {}
+    for name, p in PRESETS.items():
+        grouped.setdefault(p.family or "Other", []).append(name)
+    return grouped

tests/analytical_visualization/optimization_report.py

@@ -0,0 +1,235 @@
+"""Replication (waste) accounting + optimization opportunities.
+
+Two things this module produces:
+
+1. `replication_report(cfg)` — where memory is being *duplicated* across
+   the deployment, per category (attention weights, FFN weights, KV cache
+   under replicate mode, DP replicas). Answers "if I lifted this
+   duplication, how much would I save?"
+
+2. `optimization_hints(cfg)` — a list of actionable hints. Each hint has
+   a `category` (space | comm | compute), a `severity` (info | warn |
+   good), and a short human message. UI code renders these as bullets.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from .model_config import FullConfig
+from .memory_layout import per_pe_weight_bytes, per_pe_kv_cache_bytes
+
+
+# ── Replication accounting ────────────────────────────────────────
+
+@dataclass
+class ReplicationEntry:
+    tensor: str
+    per_pe_bytes: int
+    replicated_across: str    # e.g. "CP (x4 copies)"
+    copies: int               # total # copies globally
+    wasted_bytes: int         # (copies - 1) * per_pe_bytes  (per PE view)
+
+
+def _attn_weight_bytes_per_pe(cfg: FullConfig) -> int:
+    """W_Q + W_K + W_V + W_O per PE, one layer."""
+    m = cfg.model
+    tp = cfg.topo.tp
+    hq_per_pe = cfg.h_q_per_pe
+    if cfg.topo.kv_shard_mode == "replicate":
+        hkv_per_pe = max(1.0, m.h_kv / tp)
+    else:
+        hkv_per_pe = m.h_kv / tp
+    per_layer = (
+        m.hidden * hq_per_pe * m.d_head +
+        m.hidden * hkv_per_pe * m.d_head +
+        m.hidden * hkv_per_pe * m.d_head +
+        hq_per_pe * m.d_head * m.hidden
+    ) * m.bytes_per_elem
+    layers_per_stage = (m.layers + cfg.topo.pp - 1) // cfg.topo.pp
+    return int(per_layer * layers_per_stage)
+
+
+def _ffn_weight_bytes_per_pe(cfg: FullConfig) -> int:
+    """W_gate + W_up + W_down per PE, all local layers."""
+    m = cfg.model
+    ep = max(1, cfg.topo.ep)
+    ffn_div = cfg.ffn_shard_divisor * ep
+    per_layer = 3 * m.hidden * (m.ffn_dim // max(1, ffn_div)) * m.bytes_per_elem
+    layers_per_stage = (m.layers + cfg.topo.pp - 1) // cfg.topo.pp
+    return int(per_layer * layers_per_stage)
+
+
+def _kv_bytes_per_pe(cfg: FullConfig) -> int:
+    return per_pe_kv_cache_bytes(cfg)
+
+
+def replication_report(cfg: FullConfig) -> list[ReplicationEntry]:
+    """Enumerate replicated (duplicated) tensors and their waste."""
+    entries: list[ReplicationEntry] = []
+    topo = cfg.topo
+
+    # Attention weights are NOT sharded by CP - so each CP rank holds a copy.
+    attn_pe = _attn_weight_bytes_per_pe(cfg)
+    cp_copies = topo.cp
+    if cp_copies > 1:
+        entries.append(ReplicationEntry(
+            tensor="Attn weights (W_Q/W_K/W_V/W_O)",
+            per_pe_bytes=attn_pe,
+            replicated_across=f"CP (x{cp_copies} identical group copies)",
+            copies=cp_copies,
+            wasted_bytes=attn_pe * (cp_copies - 1),
+        ))
+
+    # FFN weights: replicated across scope levels NOT included.
+    # Scope=TP  -> replicated across CP AND DP.
+    # Scope=TP+CP -> replicated across DP only.
+    # Scope=TP+CP+DP -> no replication (perfect share).
+    ffn_pe = _ffn_weight_bytes_per_pe(cfg)
+    scope = topo.ffn_shard_scope
+    ffn_cp_copies = 1 if "CP" in scope else topo.cp
+    ffn_dp_copies = 1 if "DP" in scope else topo.dp
+    ffn_total_copies = ffn_cp_copies * ffn_dp_copies
+    if ffn_total_copies > 1:
+        reasons = []
+        if ffn_cp_copies > 1:
+            reasons.append(f"CP (x{ffn_cp_copies})")
+        if ffn_dp_copies > 1:
+            reasons.append(f"DP (x{ffn_dp_copies})")
+        entries.append(ReplicationEntry(
+            tensor=f"FFN weights (scope={scope})",
+            per_pe_bytes=ffn_pe,
+            replicated_across=" & ".join(reasons),
+            copies=ffn_total_copies,
+            wasted_bytes=ffn_pe * (ffn_total_copies - 1),
+        ))
+
+    # KV cache: if kv_shard_mode='replicate' and TP > H_kv, each KV head
+    # is duplicated tp // h_kv times across TP ranks.
+    if (cfg.kv_replication_needed
+            and topo.kv_shard_mode == "replicate"):
+        kv_pe = _kv_bytes_per_pe(cfg)
+        rep = cfg.head_dim_split_factor  # tp // h_kv
+        entries.append(ReplicationEntry(
+            tensor="KV cache (replicate mode)",
+            per_pe_bytes=kv_pe,
+            replicated_across=f"TP (x{rep} copies of each KV head)",
+            copies=rep,
+            wasted_bytes=kv_pe * (rep - 1),
+        ))
+
+    # DP replicas duplicate the entire per-PE footprint (attn + ffn + KV).
+    if topo.dp > 1:
+        total_pe = attn_pe + ffn_pe + _kv_bytes_per_pe(cfg)
+        entries.append(ReplicationEntry(
+            tensor="Full per-PE footprint (attn + FFN + KV)",
+            per_pe_bytes=total_pe,
+            replicated_across=f"DP (x{topo.dp} model replicas)",
+            copies=topo.dp,
+            wasted_bytes=total_pe * (topo.dp - 1),
+        ))
+
+    return entries
+
+
+# ── Optimization hints ────────────────────────────────────────────
+
+@dataclass
+class Hint:
+    category: str   # "space" | "comm" | "compute" | "layout"
+    severity: str   # "info" | "warn" | "good"
+    message: str
+
+
+def optimization_hints(cfg: FullConfig) -> list[Hint]:
+    hints: list[Hint] = []
+    m = cfg.model
+    topo = cfg.topo
+
+    # SPACE hints -------------------------------------------------
+    scope = topo.ffn_shard_scope
+    ffn_pe = _ffn_weight_bytes_per_pe(cfg)
+    if scope == "TP" and topo.cp > 1:
+        current = ffn_pe
+        new = current // topo.cp
+        hints.append(Hint(
+            "space", "warn",
+            f"FFN scope=TP replicates FFN weights across all {topo.cp} CP "
+            f"groups. Switching to **TP+CP** would drop per-PE FFN weight "
+            f"from {current/1e9:.2f} GB to {new/1e9:.2f} GB "
+            f"(saves {(current-new)/1e9:.2f} GB/PE) at the cost of one "
+            f"extra AllReduce per layer."
+        ))
+    if scope in ("TP", "TP+CP") and topo.dp > 1:
+        base = ffn_pe if scope == "TP+CP" else ffn_pe // topo.cp
+        new = base // topo.dp
+        hints.append(Hint(
+            "space", "info",
+            f"With DP={topo.dp}, scope=**TP+CP+DP** would shard FFN across "
+            f"replicas too, dropping per-PE FFN from {base/1e9:.2f} GB to "
+            f"{new/1e9:.2f} GB (comes with an inter-replica AllReduce)."
+        ))
+    if (cfg.kv_replication_needed
+            and topo.kv_shard_mode == "replicate"):
+        rep = cfg.head_dim_split_factor
+        kv_pe = _kv_bytes_per_pe(cfg)
+        hints.append(Hint(
+            "space", "warn",
+            f"KV mode=replicate holds each KV head {rep}x. Switching to "
+            f"**split** saves ~{kv_pe*(rep-1)/rep/1e9:.2f} GB KV per PE "
+            f"but adds a Score AllReduce over {rep} ranks per hop."
+        ))
+    if topo.cp == 1 and topo.tp <= m.h_kv:
+        # Suggest CP if s_kv is large and KV dominates.
+        kv_pe = _kv_bytes_per_pe(cfg)
+        attn_pe = _attn_weight_bytes_per_pe(cfg)
+        if kv_pe > attn_pe:
+            hints.append(Hint(
+                "space", "info",
+                f"KV cache ({kv_pe/1e9:.2f} GB) dominates weights "
+                f"({attn_pe/1e9:.2f} GB) at S_kv={topo.s_kv:,}. Adding CP "
+                f"shards the sequence axis (S_local = S_kv/CP)."
+            ))
+
+    # COMM hints --------------------------------------------------
+    if topo.cp_inter_sip_hops > 0:
+        hints.append(Hint(
+            "comm", "warn",
+            f"CP ring crosses {topo.cp_inter_sip_hops} SIP boundary(ies) "
+            f"at {cfg.machine.bw_intersip_gbs:.0f} GB/s vs "
+            f"{cfg.machine.bw_inter_gbs:.0f} GB/s intra-SIP - "
+            f"consider a smaller CP or a topology that keeps CP inside one SIP."
+        ))
+    if topo.tp_spans_cubes > 1:
+        hints.append(Hint(
+            "comm", "warn",
+            f"TP={topo.tp} spans {topo.tp_spans_cubes} cubes, so the W_O "
+            f"AllReduce runs cross-cube "
+            f"({cfg.machine.bw_intra_gbs:.0f} -> "
+            f"{cfg.machine.bw_inter_gbs:.0f} GB/s). A TP that fits in one "
+            f"cube (TP<={topo.pes_per_cube_hw}) keeps it intra-cube."
+        ))
+    if topo.sips_used > 1 and topo.sip_topology == "ring":
+        hints.append(Hint(
+            "layout", "info",
+            f"With {topo.sips_used} SIPs on a ring, worst-case inter-SIP "
+            f"distance is {topo.sips_used - 1} hops. Switching to **torus2d** "
+            f"cuts worst-case hops to ~sqrt({topo.sips_used})."
+        ))
+    if topo.tp <= m.h_kv:
+        hints.append(Hint(
+            "comm", "good",
+            f"TP ({topo.tp}) <= H_kv ({m.h_kv}): each KV head lives on "
+            f"exactly one TP rank - no head-split AllReduce needed."
+        ))
+
+    # COMPUTE hints -----------------------------------------------
+    if topo.mode == "decode" and topo.T_q == 1:
+        # Decode is memory-bound almost everywhere; note it.
+        hints.append(Hint(
+            "compute", "info",
+            "Decode (T_q=1): most GEMMs are memory-bound. FLOPs matter far "
+            "less than HBM BW here - budget attention around "
+            "weight-read cost, not compute peak."
+        ))
+
+    return hints

tests/analytical_visualization/pe_weight_layout.py

@@ -0,0 +1,279 @@
+"""Draw a PE-level view of one CP group (one TP group of cubes).
+
+Shows how weights + KV cache are distributed across the PEs of ONE CP
+rank. Attention weights are sharded by head (H_q, H_kv split across TP);
+FFN is sharded by dim (ffn_dim / TP / EP); KV cache is sharded by CP
+(sequence) × TP (head).
+
+If TP > 8 the group spans multiple cubes — layout draws all of them.
+"""
+from __future__ import annotations
+
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
+
+from .model_config import FullConfig
+
+
+PE_COLS = 4
+PE_ROWS = 2
+PES_PER_CUBE = PE_COLS * PE_ROWS
+
+
+def layers_per_stage_val(cfg: FullConfig) -> int:
+    return (cfg.model.layers + cfg.topo.pp - 1) // cfg.topo.pp
+
+
+def _q_heads_for_pe(pe_id_in_group: int, tp: int, h_q: int) -> list[int]:
+    """Return the list of Q-head indices assigned to this PE."""
+    heads_per_pe = h_q / tp
+    start = int(round(pe_id_in_group * heads_per_pe))
+    end = int(round((pe_id_in_group + 1) * heads_per_pe))
+    return list(range(start, end))
+
+
+def _kv_heads_for_pe(pe_id_in_group: int, tp: int, h_kv: int,
+                      kv_shard_mode: str) -> tuple[list[int], str]:
+    """KV heads assigned to this PE and a note about replication/split."""
+    if h_kv >= tp:
+        heads_per_pe = h_kv // tp
+        start = pe_id_in_group * heads_per_pe
+        end = start + heads_per_pe
+        return list(range(start, end)), ""
+    # TP > H_kv
+    if kv_shard_mode == "replicate":
+        rep_factor = tp // h_kv
+        head_idx = pe_id_in_group // rep_factor
+        return [head_idx], f"replicated x{rep_factor}"
+    # split mode
+    split_factor = tp // h_kv
+    head_idx = pe_id_in_group // split_factor
+    part = pe_id_in_group % split_factor
+    return [head_idx], f"head-split ({part+1}/{split_factor})"
+
+
+def _per_pe_bytes(cfg: FullConfig, pe_id_in_group: int) -> dict:
+    """Return per-tensor bytes for one PE within a TP group."""
+    m = cfg.model
+    tp = cfg.topo.tp
+    pp = cfg.topo.pp
+    ep = max(1, cfg.topo.ep)
+    b = m.bytes_per_elem
+    layers_per_stage = (m.layers + pp - 1) // pp
+
+    hq_per_pe = m.h_q / tp
+    if cfg.topo.kv_shard_mode == "replicate":
+        hkv_per_pe = max(1.0, m.h_kv / tp)
+    else:
+        hkv_per_pe = m.h_kv / tp
+    ffn_div = cfg.ffn_shard_divisor * ep
+    ffn_per_pe = m.ffn_dim // ffn_div
+
+    per_layer = {
+        "W_Q": int(m.hidden * hq_per_pe * m.d_head * b),
+        "W_K": int(m.hidden * hkv_per_pe * m.d_head * b),
+        "W_V": int(m.hidden * hkv_per_pe * m.d_head * b),
+        "W_O": int(hq_per_pe * m.d_head * m.hidden * b),
+        "W_gate": int(m.hidden * ffn_per_pe * b),
+        "W_up": int(m.hidden * ffn_per_pe * b),
+        "W_down": int(ffn_per_pe * m.hidden * b),
+    }
+    kv_per_layer = int(2 * cfg.topo.s_local * hkv_per_pe * m.d_head * b)
+
+    all_layers = {k: v * layers_per_stage for k, v in per_layer.items()}
+    all_layers["KV cache"] = kv_per_layer * layers_per_stage
+    # transient estimate
+    if cfg.topo.mode == "decode":
+        all_layers["Transient"] = int(4 * (m.hidden + hq_per_pe * m.d_head) * b)
+    else:
+        TILE = 1024
+        all_layers["Transient"] = int(2 * hq_per_pe * cfg.topo.T_q * TILE * b
+                                       + cfg.topo.T_q * m.hidden * b)
+    return all_layers
+
+
+def draw_pe_layout(cfg: FullConfig, ax=None):
+    """Draw one CP group's PEs with per-tensor weight + KV breakdown.
+
+    Each PE shows: PE id, Q heads, KV heads, W_Q / W_K / W_V / W_O in
+    MB, FFN (gate+up+down) in MB, KV cache in GB, total in GB.
+
+    TP=8 -> 1 cube. TP=16 -> 2 cubes. TP=32 -> 4 cubes.
+    Under placement, n_cubes reflects cubes-per-group (may include CP spill).
+    """
+    tp = cfg.topo.tp
+    # cubes per one cube-level group (= intra-cube dims / PEs-per-cube, rounded up).
+    n_cubes = max(1, (cfg.topo.intra_cube_dims + PES_PER_CUBE - 1) // PES_PER_CUBE)
+
+    if ax is None:
+        fig, ax = plt.subplots(figsize=(7.5 * n_cubes, 8.5))
+    else:
+        fig = ax.figure
+
+    cube_w = 7.0
+    cube_h = 6.0
+    cube_gap = 0.6
+
+    # When TP fills the cube (TP>=8), cp_placement is "cube" and CP
+    # ranks span cubes — the intra-cube CP-color branch below is
+    # skipped. Surface which CP rank this figure is showing so users
+    # aren't guessing which of `CP` groups is drawn.
+    _cube_place_tag = ""
+    if cfg.topo.cp_placement == "cube" and cfg.topo.cp > 1:
+        _cube_place_tag = f"  [CP rank 0 of {cfg.topo.cp}]"
+
+    for cube_idx in range(n_cubes):
+        x0 = cube_idx * (cube_w + cube_gap)
+        y0 = 0
+
+        rect = patches.FancyBboxPatch(
+            (x0, y0), cube_w, cube_h,
+            boxstyle="round,pad=0.05",
+            facecolor="#f8f9fa", edgecolor="#212529", linewidth=1.5,
+        )
+        ax.add_patch(rect)
+        ax.text(x0 + cube_w / 2, y0 + cube_h + 0.1,
+                f"Cube {cube_idx}  (PEs {cube_idx*8}-{cube_idx*8+7})"
+                + _cube_place_tag,
+                ha="center", va="bottom", fontsize=11, fontweight="bold")
+
+        pe_pad_x = 0.15
+        pe_pad_y = 0.2
+        pe_gap = 0.10
+        inner_w = cube_w - 2 * pe_pad_x
+        inner_h = cube_h - 2 * pe_pad_y
+        pe_w = (inner_w - (PE_COLS - 1) * pe_gap) / PE_COLS
+        pe_h = (inner_h - (PE_ROWS - 1) * pe_gap) / PE_ROWS
+
+        # When cp_placement=pe, CP ranks are packed intra-cube along with
+        # TP ranks. Color-code by CP rank so it's obvious which PEs belong
+        # to which sequence-shard group. Palette wraps beyond 8 CP ranks.
+        _cp_place = cfg.topo.cp_placement
+        _cp_val = cfg.topo.cp
+        _cp_palette = [
+            "#f0f7ff",  # baseline blue (also the fallback for cp_placement=cube)
+            "#e8f5e9",  # light green
+            "#fff3e0",  # light orange
+            "#f3e5f5",  # light purple
+            "#ffebee",  # light red
+            "#e0f7fa",  # light cyan
+            "#fce4ec",  # light pink
+            "#f9fbe7",  # light lime
+        ]
+        _cp_edge_palette = [
+            "#3a86ff", "#2e7d32", "#ef6c00", "#7b1fa2",
+            "#c62828", "#0097a7", "#c2185b", "#9e9d24",
+        ]
+
+        for pr in range(PE_ROWS):
+            for pc in range(PE_COLS):
+                pe_local = pr * PE_COLS + pc
+                pe_id_in_group = cube_idx * PES_PER_CUBE + pe_local
+                if pe_id_in_group >= cfg.topo.intra_cube_dims:
+                    continue
+                px = x0 + pe_pad_x + pc * (pe_w + pe_gap)
+                py = y0 + pe_pad_y + (PE_ROWS - 1 - pr) * (pe_h + pe_gap)
+
+                # Which (cp_rank, tp_rank) does this PE hold?
+                # When cp_placement=pe: cp_rank = pe_id // tp, tp_rank = pe_id % tp.
+                # When cp_placement=cube: every PE in this cube is one CP rank
+                # (drawn on its own tile), tp_rank = pe_id % tp.
+                if _cp_place == "pe" and _cp_val > 1:
+                    _cp_rank = pe_id_in_group // tp
+                    _tp_rank = pe_id_in_group % tp
+                    _fc = _cp_palette[_cp_rank % len(_cp_palette)]
+                    _ec = _cp_edge_palette[_cp_rank % len(_cp_edge_palette)]
+                    _cp_label = f"  |  CP={_cp_rank}"
+                else:
+                    _cp_rank = None
+                    _tp_rank = pe_id_in_group % tp
+                    _fc = "#f0f7ff"
+                    _ec = "#3a86ff"
+                    # cp_placement=="cube" with CP>1: the whole cube is
+                    # one CP rank (rank 0 by convention — see cube title
+                    # tag above). Add the locator so users know what
+                    # rank this per-PE data belongs to.
+                    if _cp_place == "cube" and _cp_val > 1:
+                        _cp_label = "  |  CP=0"
+                    else:
+                        _cp_label = ""
+
+                pe_rect = patches.Rectangle(
+                    (px, py), pe_w, pe_h,
+                    facecolor=_fc, edgecolor=_ec, linewidth=1.2,
+                )
+                ax.add_patch(pe_rect)
+
+                # Head assignment uses TP rank (not raw pe_id) so that under
+                # cp_placement=pe, all CP ranks share the same head split.
+                q_heads = _q_heads_for_pe(_tp_rank, tp, cfg.model.h_q)
+                kv_heads, kv_note = _kv_heads_for_pe(
+                    _tp_rank, tp, cfg.model.h_kv, cfg.topo.kv_shard_mode,
+                )
+                bytes_ = _per_pe_bytes(cfg, _tp_rank)
+                weights_gb = sum(v for k, v in bytes_.items()
+                                 if k not in ("KV cache", "Transient")) / 1e9
+                kv_gb = bytes_["KV cache"] / 1e9
+                trans_mb = bytes_["Transient"] / 1e6
+                total_gb = sum(bytes_.values()) / 1e9
+
+                q_str = (f"{q_heads[0]}-{q_heads[-1]}"
+                         if len(q_heads) > 1 else str(q_heads[0])
+                         if q_heads else "-")
+                kv_str = str(kv_heads[0]) if kv_heads else "-"
+                if kv_note:
+                    kv_str += f"({kv_note})"
+
+                ffn_mb = (bytes_["W_gate"] + bytes_["W_up"]
+                          + bytes_["W_down"]) / 1e6
+                # Header (PE id and heads) bigger, then per-tensor breakdown
+                header = (f"PE {pe_id_in_group}  TP={_tp_rank}{_cp_label}\n"
+                          f"Q heads: {q_str}\n"
+                          f"KV head: {kv_str}")
+                ax.text(px + 0.08, py + pe_h - 0.05, header,
+                        ha="left", va="top",
+                        fontsize=7.0, fontweight="bold",
+                        family="monospace", color="#0d47a1")
+
+                # Per-tensor bytes with explicit division divisors.
+                # Divisor is always the PRODUCT of the parallel dims listed
+                # (TP*CP means TP times CP, i.e. multiplication).
+                tp_d = cfg.topo.tp
+                ffn_div = cfg.ffn_shard_divisor
+                ffn_scope = cfg.topo.ffn_shard_scope.replace("+", "*")
+                kv_div = cfg.topo.cp * cfg.topo.tp
+                L = layers_per_stage_val(cfg)
+                detail = (
+                    f"weights (all {L} layers):\n"
+                    f" W_Q /TP={tp_d} : {bytes_['W_Q']/1e6:7.1f}MB\n"
+                    f" W_K /TP={tp_d} : {bytes_['W_K']/1e6:7.1f}MB\n"
+                    f" W_V /TP={tp_d} : {bytes_['W_V']/1e6:7.1f}MB\n"
+                    f" W_O /TP={tp_d} : {bytes_['W_O']/1e6:7.1f}MB\n"
+                    f" FFN /({ffn_scope})={ffn_div}: {ffn_mb:7.1f}MB\n"
+                    f"runtime:\n"
+                    f" KV /(CP*TP)={kv_div}: {kv_gb*1000:7.1f}MB\n"
+                    f" trans     : {trans_mb:7.1f}MB\n"
+                    f"W total: {weights_gb*1000:7.1f}MB\n"
+                    f"TOTAL  : {total_gb*1000:7.1f}MB"
+                )
+                ax.text(px + pe_w / 2, py + 0.06, detail,
+                        ha="center", va="bottom",
+                        fontsize=5.8, family="monospace", color="#1a1a1a")
+
+    total_w = n_cubes * cube_w + (n_cubes - 1) * cube_gap
+    ax.set_xlim(-0.2, total_w + 0.2)
+    ax.set_ylim(-0.3, cube_h + 0.6)
+    ax.set_aspect("equal")
+    ax.axis("off")
+
+    _title_cp = f", CP={cfg.topo.cp}"
+    if cfg.topo.cp_placement == "cube" and cfg.topo.cp > 1:
+        _title_cp += " (showing 1 of CP groups; others identical)"
+    title = (
+        f"Per-PE layout of one CP group  "
+        f"(TP={tp}{_title_cp}, H_q={cfg.model.h_q}, H_kv={cfg.model.h_kv})"
+        f" | KV mode: {cfg.topo.kv_shard_mode}"
+    )
+    ax.set_title(title, fontsize=11, fontweight="bold")
+
+    return fig

tests/analytical_visualization/pipeline_diagram.py

@@ -0,0 +1,300 @@
+"""Pipeline diagram: attention + FFN steps, colored by bound type.
+
+Each stage is drawn as a rounded box in a horizontal flow:
+  Attention: S1 -> S2 -> S3 -> ... -> S10 -> C2 (AllReduce)
+  FFN:       F1 -> F2 -> F3 -> F4 -> F5 -> CF1 (AllReduce)
+
+CP ring communication (C1) and score AllReduce (C3) are drawn as
+"overhead" boxes attached above the QKt/PV region (they run
+concurrently with per-hop compute in the ring).
+
+Box fill color = the stage's bound classification:
+  compute-bound  blue
+  memory-bound   orange
+  comm-bound     red
+  trivial        grey
+"""
+from __future__ import annotations
+
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
+
+from .model_config import FullConfig
+from .stage_latencies import (
+    all_stages, all_ffn_stages, StageCost,
+)
+
+
+BOUND_COLOR = {
+    "compute": "#4682b4",   # blue
+    "memory":  "#e37400",   # orange
+    "comm":    "#c62828",   # red
+    "trivial": "#a9a9a9",   # grey
+}
+BOUND_LABEL = {
+    "compute": "compute-bound",
+    "memory":  "memory-bound",
+    "comm":    "communication",
+    "trivial": "trivial / negligible",
+}
+
+
+def _fmt_us(sec: float) -> str:
+    us = sec * 1e6
+    if us < 1:
+        return f"{us*1e3:.1f}ns"
+    if us < 100:
+        return f"{us:.2f}us"
+    return f"{us:.0f}us"
+
+
+def _draw_block(ax, x, y, w, h, label, cost: StageCost, edge="#333"):
+    color = BOUND_COLOR.get(cost.bound, "#a9a9a9")
+    box = patches.FancyBboxPatch(
+        (x, y), w, h,
+        boxstyle="round,pad=0.05",
+        facecolor=color, edgecolor=edge, linewidth=1.2, alpha=0.85,
+    )
+    ax.add_patch(box)
+    # Two lines: stage name (bold) + visible latency (small)
+    ax.text(x + w / 2, y + h * 0.62, label,
+            ha="center", va="center",
+            fontsize=8.0, fontweight="bold", color="white")
+    ax.text(x + w / 2, y + h * 0.25, _fmt_us(cost.visible_s),
+            ha="center", va="center",
+            fontsize=7.0, color="white")
+
+
+_SHORT_NAMES = {
+    "S1":  "RMSNorm",
+    "S2":  "W_Q",
+    "S3":  "W_K+W_V",
+    "S4":  "KV append",
+    "S5":  "Q.K^T",
+    "S6":  "softmax",
+    "S7":  "P.V",
+    "S8":  "merge",
+    "S9":  "norm O/l",
+    "S10": "W_O",
+    "C1":  "CP ring",
+    "C2":  "TP AR",
+    "C3":  "Score AR",
+    "F1":  "RMSNorm",
+    "F2":  "W_gate",
+    "F3":  "W_up",
+    "F4":  "SwiGLU",
+    "F5":  "W_down",
+    "CF1": "FFN AR",
+}
+
+
+def _short_name(name: str) -> str:
+    """Map 'S1 RMSNorm' -> 'RMSNorm'; fallback to first token."""
+    tok = name.split()[0]
+    return _SHORT_NAMES.get(tok, tok)
+
+
+def _draw_arrow(ax, x0, x1, y):
+    ax.annotate("", xy=(x1, y), xytext=(x0, y),
+                arrowprops=dict(arrowstyle="->", color="#333", lw=1.1))
+
+
+def _draw_ring_loop(ax, x_left, x_right, y_top, label: str, n_iters: int,
+                     arc_lift: float = 0.9):
+    """Draw a curved back-arrow above [x_left..x_right] with a repeat label.
+
+    Visualizes 'repeat this block n_iters times' (like a loop in a diagram).
+    arc_lift: how high above y_top the arc sits.
+    """
+    color = "#5a189a"
+    lw = 1.6
+    arc_y = y_top + arc_lift
+    # Vertical tick up from the right edge
+    ax.plot([x_right, x_right], [y_top + 0.02, arc_y],
+            color=color, linewidth=lw)
+    # Horizontal top of the loop
+    ax.plot([x_right, x_left], [arc_y, arc_y],
+            color=color, linewidth=lw)
+    # Down-arrow into the left edge (a back-edge indicating the repeat)
+    ax.annotate("", xy=(x_left, y_top + 0.02), xytext=(x_left, arc_y),
+                arrowprops=dict(arrowstyle="->", color=color, lw=lw))
+    # Loop label centered above the arc
+    ax.text((x_left + x_right) / 2, arc_y + 0.08,
+            label,
+            ha="center", va="bottom",
+            fontsize=9, fontweight="bold", color=color)
+    # Iteration count as a badge on the right
+    ax.text(x_right + 0.08, arc_y - 0.03,
+            f"x {n_iters}",
+            ha="left", va="top",
+            fontsize=10, fontweight="bold", color=color,
+            bbox=dict(boxstyle="round,pad=0.15", facecolor="white",
+                      edgecolor=color, linewidth=1.2, alpha=0.95))
+
+
+def _draw_row(ax, y_base, row_title, stages: list[StageCost],
+              x_start: float, box_w: float, box_h: float, gap: float):
+    """Draw one row of stage boxes with arrows between them."""
+    ax.text(x_start - 0.4, y_base + box_h / 2, row_title,
+            ha="right", va="center",
+            fontsize=11, fontweight="bold", color="#212529")
+    x = x_start
+    for i, st in enumerate(stages):
+        _draw_block(ax, x, y_base, box_w, box_h,
+                    _short_name(st.name), st)
+        if i < len(stages) - 1:
+            _draw_arrow(ax, x + box_w + 0.02, x + box_w + gap - 0.02,
+                        y_base + box_h / 2)
+        x += box_w + gap
+    return x - gap
+
+
+def _draw_overlap_boxes(ax, hidden_stages, x_left, x_right, y_top,
+                          gap_between: float = 0.15):
+    """Draw dashed 'overlapped comm' boxes spanning x_left..x_right above y_top.
+
+    Boxes are widened to cover the horizontal range of the stages they
+    overlap with (e.g. C1 CP ring sits above S5-S7).
+    """
+    if not hidden_stages:
+        return
+    hy = y_top + 0.10
+    n = len(hidden_stages)
+    total_span = x_right - x_left
+    box_span = (total_span - (n - 1) * gap_between) / max(1, n)
+    box_h_local = 0.42
+    hx = x_left
+    for st in hidden_stages:
+        box = patches.FancyBboxPatch(
+            (hx, hy), box_span, box_h_local,
+            boxstyle="round,pad=0.03",
+            facecolor=BOUND_COLOR["comm"], edgecolor="#333",
+            linewidth=1.0, alpha=0.35, linestyle="--",
+        )
+        ax.add_patch(box)
+        ax.text(hx + box_span / 2, hy + box_h_local * 0.65,
+                f"{_short_name(st.name)} (concurrent)",
+                ha="center", va="center",
+                fontsize=7, color="#7a0e0e", fontweight="bold")
+        ax.text(hx + box_span / 2, hy + box_h_local * 0.2,
+                _fmt_us(st.visible_s),
+                ha="center", va="center",
+                fontsize=6.5, color="#7a0e0e")
+        hx += box_span + gap_between
+
+
+def draw_pipeline(cfg: FullConfig, ax=None):
+    """Draw attention + FFN pipeline. Boxes colored by bound type."""
+    attn = all_stages(cfg)
+    ffn = all_ffn_stages(cfg)
+
+    # Split attention: main stages, then C2 (TP AllReduce). C1/C3 = overlapped.
+    attn_main = [s for s in attn if not s.name.startswith("C")]
+    attn_c2 = next((s for s in attn if s.name.startswith("C2")), None)
+    attn_hidden = [s for s in attn
+                   if s.name.startswith("C1") or s.name.startswith("C3")]
+    attn_hidden = [s for s in attn_hidden if s.visible_s > 0]
+    attn_row = attn_main + ([attn_c2] if attn_c2 and attn_c2.visible_s > 0 else [])
+
+    # FFN row: all F* + CF1 if non-trivial
+    ffn_main = [s for s in ffn if not s.name.startswith("CF")]
+    ffn_cf = next((s for s in ffn if s.name.startswith("CF1")), None)
+    ffn_row = ffn_main + ([ffn_cf] if ffn_cf and ffn_cf.visible_s > 0 else [])
+
+    n_max = max(len(attn_row), len(ffn_row))
+    box_w = 1.35
+    box_h = 0.78
+    gap = 0.18
+    x_start = 1.1
+
+    total_w = x_start + n_max * box_w + (n_max - 1) * gap + 0.5
+    total_h = 4.2
+
+    if ax is None:
+        fig, ax = plt.subplots(figsize=(max(11, total_w * 1.1),
+                                         max(4.5, total_h * 1.0)))
+    else:
+        fig = ax.figure
+
+    _draw_row(ax, y_base=2.6, row_title="Attention",
+              stages=attn_row, x_start=x_start,
+              box_w=box_w, box_h=box_h, gap=gap)
+
+    _draw_row(ax, y_base=0.6, row_title="FFN",
+              stages=ffn_row, x_start=x_start,
+              box_w=box_w, box_h=box_h, gap=gap)
+
+    # Ring loop indicator + concurrent-comm boxes anchored over S5..S8.
+    _idx_by_prefix = {s.name.split()[0]: i for i, s in enumerate(attn_row)}
+    _s5_idx = _idx_by_prefix.get("S5")
+    _s7_idx = _idx_by_prefix.get("S7")
+    _s8_idx = _idx_by_prefix.get("S8")
+
+    if cfg.topo.cp > 1 and cfg.topo.mode == "prefill" and _s5_idx is not None:
+        _right_idx = _s8_idx if _s8_idx is not None else _s7_idx
+        _lx = x_start + _s5_idx * (box_w + gap)
+        _rx = x_start + _right_idx * (box_w + gap) + box_w
+        _y_top = 2.6 + box_h
+        # Concurrent-comm boxes (C1 CP ring, C3 score AR) directly above S5-S7
+        _draw_overlap_boxes(ax, attn_hidden, _lx, _rx, _y_top)
+        # Ring loop arc above those
+        _draw_ring_loop(
+            ax, _lx, _rx, _y_top,
+            label="ring attention loop (K/V ring per hop; S8 merges partial O,m,l)",
+            n_iters=cfg.topo.cp,
+            arc_lift=1.05,
+        )
+    elif cfg.topo.cp > 1 and cfg.topo.mode == "decode" and _s5_idx is not None:
+        _right_idx = _s7_idx if _s7_idx is not None else _s5_idx
+        _lx = x_start + _s5_idx * (box_w + gap)
+        _rx = x_start + _right_idx * (box_w + gap) + box_w
+        _y_top = 2.6 + box_h
+        # Any straggling concurrent comm (e.g. C3 if head-split active) above S5-S7
+        _draw_overlap_boxes(ax, attn_hidden, _lx, _rx, _y_top)
+        _draw_ring_loop(
+            ax, _lx, _rx, _y_top,
+            label="decode: local pass only; S8 fires O/m/l all-reduce",
+            n_iters=1,
+            arc_lift=1.05 if attn_hidden else 0.85,
+        )
+    elif attn_hidden:
+        # Fallback: CP=1 but still have C3 or similar; anchor above S5-S7 if present
+        if _s5_idx is not None and _s7_idx is not None:
+            _lx = x_start + _s5_idx * (box_w + gap)
+            _rx = x_start + _s7_idx * (box_w + gap) + box_w
+            _draw_overlap_boxes(ax, attn_hidden, _lx, _rx, 2.6 + box_h)
+
+    # Totals
+    attn_total = sum(s.visible_s for s in attn_row) \
+        + sum(s.visible_s for s in attn_hidden)
+    ffn_total = sum(s.visible_s for s in ffn_row)
+    ax.text(total_w - 0.3, 2.6 + box_h / 2 + 1.15,
+            f"Attn total: {_fmt_us(attn_total)}",
+            ha="right", va="bottom", fontsize=9,
+            fontweight="bold", color="#212529")
+    ax.text(total_w - 0.3, 0.6 + box_h + 0.15,
+            f"FFN total:  {_fmt_us(ffn_total)}",
+            ha="right", va="bottom", fontsize=9,
+            fontweight="bold", color="#212529")
+
+    # Legend
+    handles = []
+    for k, lbl in BOUND_LABEL.items():
+        handles.append(patches.Patch(
+            facecolor=BOUND_COLOR[k], edgecolor="#333",
+            alpha=0.85, label=lbl,
+        ))
+    ax.legend(handles=handles, loc="upper center",
+              bbox_to_anchor=(0.5, -0.02),
+              ncol=4, fontsize=8, frameon=False)
+
+    ax.set_xlim(0, total_w + 0.6)
+    ax.set_ylim(0, 4.9)
+    ax.set_aspect("auto")
+    ax.axis("off")
+    ax.set_title(
+        f"Per-layer pipeline (one PE, {cfg.topo.mode}, T_q={cfg.topo.T_q}, "
+        f"S_local={cfg.topo.s_local}) - color = dominant bound",
+        fontsize=10, fontweight="bold",
+    )
+    return fig

tests/analytical_visualization/stage_latencies.py

@@ -0,0 +1,830 @@
+"""Per-stage cost formulas for attention forward pass.
+
+Each stage returns:
+    name, formula (str with the substituted numeric form),
+    compute_time, memory_time, comm_time, bound, visible_time.
+
+Bound is the max of the three; "hidden" comm can appear as 0 when
+overlapped with compute.
+
+All times are in seconds; convert to μs at display.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Literal
+
+from .model_config import FullConfig
+
+
+Bound = Literal["compute", "memory", "comm", "trivial"]
+
+
+@dataclass
+class StageCost:
+    name: str
+    formula: str
+    compute_s: float
+    memory_s: float
+    comm_s: float
+    bound: Bound
+    visible_s: float
+    hop_multiplier: int = 1  # some stages scale with N_cp hops
+    # Detailed breakdown for the per-stage table.
+    flops: int = 0                  # total FLOPs (numeric)
+    mem_bytes: int = 0              # HBM bytes moved (numeric)
+    comm_bytes: int = 0             # comm bytes over the wire (numeric)
+    flops_formula: str = ""         # human-readable FLOPs formula
+    mem_formula: str = ""           # human-readable memory formula
+    comm_formula: str = ""          # human-readable comm formula
+
+
+def _visible(compute: float, memory: float, comm: float) -> tuple[float, Bound]:
+    """Pick the dominant time as the visible stage cost."""
+    parts = {"compute": compute, "memory": memory, "comm": comm}
+    dominant = max(parts, key=parts.get)  # type: ignore
+    return parts[dominant], dominant  # type: ignore
+
+
+def stage_rmsnorm(cfg: FullConfig) -> StageCost:
+    T_q = cfg.topo.T_q
+    d = cfg.model.hidden
+    b = cfg.model.bytes_per_elem
+    B = max(1, cfg.topo.b)
+    # Activation memory (x) scales with B; weight (once) is fixed.
+    bytes_ = B * T_q * d * b + T_q * d * b
+    flops = 4 * B * T_q * d
+    mem_s = bytes_ / cfg.machine.bw_hbm
+    cmp_s = flops / (cfg.machine.peak_flops * cfg.machine.compute_util)
+    vis, bnd = _visible(cmp_s, mem_s, 0)
+    return StageCost(
+        name="S1 RMSNorm",
+        formula=f"bytes = B*T_q*d*b + T_q*d*b = {bytes_} B / BW_HBM",
+        compute_s=cmp_s, memory_s=mem_s, comm_s=0.0,
+        bound=bnd, visible_s=vis,
+        flops=flops, mem_bytes=bytes_,
+        flops_formula=f"4*B*T_q*d = 4*{B}*{T_q}*{d} = {flops}",
+        mem_formula=f"B*T_q*d*b + T_q*d*b (weight) "
+                     f"= {B}*{T_q}*{d}*{b} + {T_q*d*b} = {bytes_} B",
+    )
+
+
+def _gemm_time(flops: int, weight_bytes: int, cfg: FullConfig) -> tuple[float, float]:
+    cmp_s = flops / (cfg.machine.peak_flops * cfg.machine.compute_util)
+    mem_s = weight_bytes / cfg.machine.bw_hbm
+    return cmp_s, mem_s
+
+
+def stage_wq(cfg: FullConfig) -> StageCost:
+    T_q = cfg.topo.T_q
+    d = cfg.model.hidden
+    b = cfg.model.bytes_per_elem
+    B = max(1, cfg.topo.b)
+    hq_per_pe = cfg.h_q_per_pe
+    dh = cfg.model.d_head
+    # FLOPs scale with batch; weight bytes fixed (shared across batch).
+    flops = 2 * B * T_q * d * (hq_per_pe * dh)
+    weight_B = d * (hq_per_pe * dh) * b
+    cmp_s, mem_s = _gemm_time(flops, weight_B, cfg)
+    vis, bnd = _visible(cmp_s, mem_s, 0)
+    return StageCost(
+        name="S2 W_Q GEMM",
+        formula=f"FLOPs = 2*B*T_q*d*(H_q/TP*d_h) "
+                f"= 2*{B}*{T_q}*{d}*{hq_per_pe*dh} = {flops:.2g}; "
+                f"weight = {weight_B/1e6:.1f} MB (shared across batch)",
+        compute_s=cmp_s, memory_s=mem_s, comm_s=0.0,
+        bound=bnd, visible_s=vis,
+        flops=int(flops), mem_bytes=int(weight_B),
+        flops_formula=f"2*B*T_q*d*(H_q/TP*d_h) "
+                       f"= 2*{B}*{T_q}*{d}*{hq_per_pe*dh} = {flops:.3g}",
+        mem_formula=f"d*(H_q/TP*d_h)*b (weight, B-invariant) "
+                     f"= {d}*{hq_per_pe*dh}*{b} = {weight_B/1e6:.2f} MB",
+    )
+
+
+def stage_wkv(cfg: FullConfig) -> StageCost:
+    T_q = cfg.topo.T_q
+    d = cfg.model.hidden
+    b = cfg.model.bytes_per_elem
+    B = max(1, cfg.topo.b)
+    hkv_per_pe = max(1, cfg.model.h_kv // cfg.topo.tp)
+    dh = cfg.model.d_head
+    flops_one = 2 * B * T_q * d * (hkv_per_pe * dh)
+    weight_B_one = d * (hkv_per_pe * dh) * b
+    flops = 2 * flops_one
+    weight_B = 2 * weight_B_one
+    cmp_s, mem_s = _gemm_time(flops, weight_B, cfg)
+    vis, bnd = _visible(cmp_s, mem_s, 0)
+    return StageCost(
+        name="S3 W_K + W_V GEMM",
+        formula=f"FLOPs per proj = 2*B*T_q*d*(H_kv/TP*d_h) "
+                f"= 2*{B}*{T_q}*{d}*{hkv_per_pe*dh} = {flops_one:.2g}; x2",
+        compute_s=cmp_s, memory_s=mem_s, comm_s=0.0,
+        bound=bnd, visible_s=vis,
+        flops=int(flops), mem_bytes=int(weight_B),
+        flops_formula=f"2*(2*B*T_q*d*(H_kv/TP*d_h)) "
+                       f"= 2*(2*{B}*{T_q}*{d}*{hkv_per_pe*dh}) = {flops:.3g}",
+        mem_formula=f"2*(d*(H_kv/TP*d_h)*b) (weights, B-invariant) "
+                     f"= 2*({d}*{hkv_per_pe*dh}*{b}) = {weight_B/1e6:.2f} MB",
+    )
+
+
+def stage_kv_append(cfg: FullConfig) -> StageCost:
+    T_q = cfg.topo.T_q
+    b = cfg.model.bytes_per_elem
+    B = max(1, cfg.topo.b)
+    hkv_per_pe = max(1, cfg.model.h_kv // cfg.topo.tp)
+    dh = cfg.model.d_head
+    bytes_ = 2 * B * T_q * hkv_per_pe * dh * b
+    mem_s = bytes_ / cfg.machine.bw_hbm
+    return StageCost(
+        name="S4 KV cache append",
+        formula=f"bytes = 2*B*T_q*(H_kv/TP)*d_h*b "
+                f"= 2*{B}*{T_q}*{hkv_per_pe}*{dh}*{b} = {bytes_} B",
+        compute_s=0, memory_s=mem_s, comm_s=0,
+        bound="memory", visible_s=mem_s,
+        flops=0, mem_bytes=bytes_,
+        flops_formula="0 (no matmul, just cache write)",
+        mem_formula=f"2*B*T_q*(H_kv/TP)*d_h*b "
+                     f"= 2*{B}*{T_q}*{hkv_per_pe}*{dh}*{b} = {bytes_} B",
+    )
+
+
+def _per_hop_qkT_pv(cfg: FullConfig) -> tuple[float, str]:
+    """Q·Kᵀ (and P·V has same FLOPs) per hop, per PE."""
+    T_q = cfg.topo.T_q
+    S_local = cfg.topo.s_local
+    dh = cfg.model.d_head
+    hq_per_pe = cfg.h_q_per_pe
+    B = max(1, cfg.topo.b)
+    flops = 2 * B * T_q * S_local * dh * hq_per_pe
+    cmp_s = flops / (cfg.machine.peak_flops * cfg.machine.compute_util)
+    formula = f"2*{T_q}*{S_local}*{dh}*{hq_per_pe} = {flops:.2g} FLOPs/hop"
+    return cmp_s, formula
+
+
+def _cp_compute_passes(cfg: FullConfig) -> int:
+    """How many local S5/S6/S7 passes happen per token step under CP.
+    - decode: 1 (each PE computes once against its local K,V, then all-reduce O/m/l)
+    - prefill: CP  (K/V or Q rotates through the ring; one local pass per hop)
+    """
+    if cfg.topo.cp <= 1:
+        return 1
+    return 1 if cfg.topo.mode == "decode" else cfg.topo.cp
+
+
+def stage_qkT(cfg: FullConfig) -> StageCost:
+    cmp_hop, _ = _per_hop_qkT_pv(cfg)
+    passes = _cp_compute_passes(cfg)
+    total = cmp_hop * passes
+    T_q = cfg.topo.T_q
+    S_local = cfg.topo.s_local
+    dh = cfg.model.d_head
+    hq_per_pe = cfg.h_q_per_pe
+    B = max(1, cfg.topo.b)
+    flops_per_hop = 2 * B * T_q * S_local * dh * hq_per_pe
+    total_flops = flops_per_hop * passes
+    _hop_word = "pass" if passes == 1 else "hops"
+    return StageCost(
+        name=f"S5 Q.K^T (x{passes} {_hop_word})",
+        formula=f"2*B*T_q*S_local*d_h*(H_q/TP) = "
+                f"2*{B}*{T_q}*{S_local}*{dh}*{hq_per_pe} = "
+                f"{flops_per_hop:.2g} FLOPs/hop",
+        compute_s=total, memory_s=0, comm_s=0,
+        bound="compute", visible_s=total,
+        hop_multiplier=passes,
+        flops=int(total_flops), mem_bytes=0,
+        flops_formula=(
+            f"{passes}*(2*B*T_q*S_local*d_h*(H_q/TP)) = "
+            f"{passes}*(2*{B}*{T_q}*{S_local}*{dh}*{hq_per_pe}) = "
+            f"{total_flops:.3g}"
+        ),
+        mem_formula="0 (scores accumulated on-chip)",
+    )
+
+
+def stage_softmax(cfg: FullConfig) -> StageCost:
+    T_q = cfg.topo.T_q
+    S_local = cfg.topo.s_local
+    hq_per_pe = cfg.h_q_per_pe
+    b = cfg.model.bytes_per_elem
+    B = max(1, cfg.topo.b)
+    elems = B * hq_per_pe * T_q * S_local
+    bytes_ = elems * b * 2
+    mem_s_per_hop = bytes_ / cfg.machine.bw_hbm
+    passes = _cp_compute_passes(cfg)
+    total = mem_s_per_hop * passes
+    total_bytes = bytes_ * passes
+    _hop_word = "pass" if passes == 1 else "hops"
+    return StageCost(
+        name=f"S6 softmax (x{passes} {_hop_word})",
+        formula=f"elts/hop = B*(H_q/TP)*T_q*S_local "
+                f"= {B}*{hq_per_pe}*{T_q}*{S_local} = {elems:.2g}",
+        compute_s=0, memory_s=total, comm_s=0,
+        bound="memory", visible_s=total,
+        hop_multiplier=passes,
+        flops=0, mem_bytes=int(total_bytes),
+        flops_formula="~O(elts) (negligible)",
+        mem_formula=(
+            f"{passes}*2*b*B*(H_q/TP)*T_q*S_local = "
+            f"{passes}*2*{b}*{B}*{hq_per_pe}*{T_q}*{S_local} = "
+            f"{total_bytes/1e6:.2f} MB"
+        ),
+    )
+
+
+def stage_pv(cfg: FullConfig) -> StageCost:
+    cmp_hop, _ = _per_hop_qkT_pv(cfg)
+    passes = _cp_compute_passes(cfg)
+    total = cmp_hop * passes
+    T_q = cfg.topo.T_q
+    S_local = cfg.topo.s_local
+    dh = cfg.model.d_head
+    hq_per_pe = cfg.h_q_per_pe
+    B = max(1, cfg.topo.b)
+    flops_per_hop = 2 * B * T_q * S_local * dh * hq_per_pe
+    total_flops = flops_per_hop * passes
+    _hop_word = "pass" if passes == 1 else "hops"
+    return StageCost(
+        name=f"S7 P.V (x{passes} {_hop_word})",
+        formula=f"2*B*T_q*S_local*d_h*(H_q/TP) = "
+                f"2*{B}*{T_q}*{S_local}*{dh}*{hq_per_pe} = "
+                f"{flops_per_hop:.2g} FLOPs/hop",
+        compute_s=total, memory_s=0, comm_s=0,
+        bound="compute", visible_s=total,
+        hop_multiplier=passes,
+        flops=int(total_flops), mem_bytes=0,
+        flops_formula=(
+            f"{passes}*(2*B*T_q*S_local*d_h*(H_q/TP)) = "
+            f"{passes}*(2*{B}*{T_q}*{S_local}*{dh}*{hq_per_pe}) = "
+            f"{total_flops:.3g}"
+        ),
+        mem_formula="0 (accumulated on-chip)",
+    )
+
+
+def stage_merge(cfg: FullConfig) -> StageCost:
+    """S8: online-softmax merge of partial (o, m, l).
+    - prefill (K/V ring): merge happens in-place across CP hops; no comm here.
+    - decode: merge is the moment we all-reduce partial (o, m, l) across CP
+      ranks; comm is folded into this stage (no separate C1).
+    """
+    T_q = cfg.topo.T_q
+    hq_per_pe = cfg.h_q_per_pe
+    dh = cfg.model.d_head
+    b = cfg.model.bytes_per_elem
+    B = max(1, cfg.topo.b)
+    flops = 6 * B * T_q * hq_per_pe * dh * max(0, cfg.topo.cp - 1)
+    cmp_s = flops / (cfg.machine.peak_flops * cfg.machine.compute_util)
+
+    comm_s = 0.0
+    comm_bytes = 0
+    comm_formula = ""
+    name_suffix = ""
+
+    if cfg.topo.mode == "decode" and cfg.topo.cp > 1:
+        cp = cfg.topo.cp
+        # (O + m + l) bytes per rank — scales with B (per-request partials).
+        M = B * ((T_q * hq_per_pe * dh * b) + (2 * T_q * hq_per_pe * b))
+        if cfg.topo.cp_placement == "pe":
+            bw, alpha, tier = cfg.machine.bw_intra, cfg.machine.alpha_intra, "intra-cube"
+        elif cfg.topo.sips_used > 1:
+            bw, alpha, tier = cfg.machine.bw_intersip, cfg.machine.alpha_intersip, "inter-SIP"
+        else:
+            bw, alpha, tier = cfg.machine.bw_inter, cfg.machine.alpha_inter, "inter-cube"
+        ar_bytes = 2 * (cp - 1) / cp * M
+        comm_s = ar_bytes / bw + 2 * (cp - 1) * alpha
+        comm_bytes = int(ar_bytes)
+        comm_formula = (
+            f"PURPOSE (folded into S8 in decode): each CP rank computed\n"
+            f"attention against its LOCAL slice of KV and holds a partial\n"
+            f"(O, m, l). This all-reduce merges those partials across all\n"
+            f"{cp} CP ranks (online-softmax combine). No separate C1 row\n"
+            f"in decode because this is the only CP comm and it happens\n"
+            f"exactly here, at the end of the local attention body.\n"
+            f"---\n"
+            f"AR of (O + m + l): 2*(CP-1)/CP * M "
+            f"= 2*({cp}-1)/{cp} * {M} B "
+            f"= {ar_bytes:.0f} B over {tier} ({bw/1e9:.0f} GB/s) "
+            f"+ 2*({cp}-1)*alpha"
+        )
+        name_suffix = f" + O/m/l AR (CP={cp} ranks, {tier})"
+
+    visible_s = max(cmp_s, comm_s)
+    if comm_s > cmp_s:
+        bound = "comm"
+    elif cmp_s > 0:
+        bound = "compute"
+    else:
+        bound = "trivial"
+
+    return StageCost(
+        name=f"S8 online-softmax merge (x{max(0, cfg.topo.cp-1)}){name_suffix}",
+        formula=f"~6*B*T_q*(H_q/TP)*d_h*(C-1) "
+                f"= 6*{B}*{T_q}*{hq_per_pe}*{dh}*{max(0, cfg.topo.cp-1)} "
+                f"= {flops:.2g} FLOPs",
+        compute_s=cmp_s, memory_s=0, comm_s=comm_s,
+        bound=bound, visible_s=visible_s,
+        flops=int(flops), mem_bytes=0, comm_bytes=comm_bytes,
+        flops_formula=(
+            f"6*B*T_q*(H_q/TP)*d_h*(CP-1) = "
+            f"6*{B}*{T_q}*{hq_per_pe}*{dh}*{max(0, cfg.topo.cp-1)} = {flops}"
+        ),
+        mem_formula="0 (in-register)",
+        comm_formula=comm_formula or "0",
+    )
+
+
+def stage_normalize(cfg: FullConfig) -> StageCost:
+    T_q = cfg.topo.T_q
+    hq_per_pe = cfg.h_q_per_pe
+    dh = cfg.model.d_head
+    B = max(1, cfg.topo.b)
+    flops = B * T_q * hq_per_pe * dh
+    cmp_s = flops / (cfg.machine.peak_flops * cfg.machine.compute_util)
+    return StageCost(
+        name="S9 normalize O/l",
+        formula=f"B*T_q*(H_q/TP)*d_h = {B}*{T_q}*{hq_per_pe}*{dh} "
+                f"= {flops} divisions",
+        compute_s=cmp_s, memory_s=0, comm_s=0,
+        bound="trivial", visible_s=cmp_s,
+        flops=int(flops), mem_bytes=0,
+        flops_formula=f"B*T_q*(H_q/TP)*d_h "
+                       f"= {B}*{T_q}*{hq_per_pe}*{dh} = {flops}",
+        mem_formula="0",
+    )
+
+
+def stage_wo(cfg: FullConfig) -> StageCost:
+    T_q = cfg.topo.T_q
+    d = cfg.model.hidden
+    b = cfg.model.bytes_per_elem
+    B = max(1, cfg.topo.b)
+    hq_per_pe = cfg.h_q_per_pe
+    dh = cfg.model.d_head
+    flops = 2 * B * T_q * (hq_per_pe * dh) * d
+    weight_B = (hq_per_pe * dh) * d * b
+    cmp_s, mem_s = _gemm_time(flops, weight_B, cfg)
+    vis, bnd = _visible(cmp_s, mem_s, 0)
+    return StageCost(
+        name="S10 W_O GEMM",
+        formula=f"FLOPs = 2*B*T_q*(H_q/TP*d_h)*d "
+                f"= 2*{B}*{T_q}*{hq_per_pe*dh}*{d} = {flops:.2g}; "
+                f"weight = {weight_B/1e6:.1f} MB (shared across batch)",
+        compute_s=cmp_s, memory_s=mem_s, comm_s=0.0,
+        bound=bnd, visible_s=vis,
+        flops=int(flops), mem_bytes=int(weight_B),
+        flops_formula=f"2*B*T_q*(H_q/TP*d_h)*d "
+                       f"= 2*{B}*{T_q}*{hq_per_pe*dh}*{d} = {flops:.3g}",
+        mem_formula=f"(H_q/TP*d_h)*d*b (weight, B-invariant) "
+                     f"= {hq_per_pe*dh}*{d}*{b} = {weight_B/1e6:.2f} MB",
+    )
+
+
+def comm_cp_ring(cfg: FullConfig) -> StageCost:
+    """CP comm — depends on mode:
+
+    - decode: single O/m/l all-reduce over CP ranks AFTER local S5-S8 finish.
+      No per-hop ring; each PE computes attention against its local K,V once
+      and then contributes (o, m, l) to the reduce.
+    - prefill: per-hop ring during S5/S7 compute (either K/V or Q+O/m/l per
+      cp_ring_variant).
+
+    BW tier depends on cp_placement:
+    - cp_placement=pe:   intra-cube BW for all hops
+    - cp_placement=cube: inter-cube BW; long rings cross SIP boundaries
+    """
+    if cfg.topo.cp <= 1:
+        return StageCost(
+            name="C1 CP comm", formula="C=1 -> no comm",
+            compute_s=0, memory_s=0, comm_s=0, bound="trivial", visible_s=0,
+        )
+    S_local = cfg.topo.s_local
+    hkv_per_pe = max(1, cfg.model.h_kv // cfg.topo.tp)
+    hq_per_pe = cfg.h_q_per_pe
+    T_q = cfg.topo.T_q
+    dh = cfg.model.d_head
+    b = cfg.model.bytes_per_elem
+    B = max(1, cfg.topo.b)
+
+    # ── Decode: single O/m/l all-reduce after local compute ──
+    if cfg.topo.mode == "decode":
+        cp = cfg.topo.cp
+        # per-rank O + m + l bytes — scales with B (per-request partials).
+        M = B * ((T_q * hq_per_pe * dh * b) + (2 * T_q * hq_per_pe * b))
+        if cfg.topo.cp_placement == "pe":
+            bw = cfg.machine.bw_intra
+            alpha = cfg.machine.alpha_intra
+            tier = "intra-cube"
+        elif cfg.topo.sips_used > 1:
+            bw = cfg.machine.bw_intersip
+            alpha = cfg.machine.alpha_intersip
+            tier = "inter-SIP"
+        else:
+            bw = cfg.machine.bw_inter
+            alpha = cfg.machine.alpha_inter
+            tier = "inter-cube"
+        ar_bytes = 2 * (cp - 1) / cp * M
+        ar_time = ar_bytes / bw + 2 * (cp - 1) * alpha
+        return StageCost(
+            name=f"C1 O/m/l all-reduce ({tier}, x{cp} ranks)",
+            formula=(f"decode: gather partial (O,m,l) once at end; "
+                     f"M = B*T_q*(H_q/TP)*(d_h*b + 2*b) = {M} B per rank; "
+                     f"AR: 2*(CP-1)/CP*M / BW + 2*(CP-1)*alpha"),
+            compute_s=0, memory_s=0, comm_s=ar_time,
+            bound="comm", visible_s=ar_time,
+            flops=0, mem_bytes=0, comm_bytes=int(ar_bytes),
+            flops_formula="0",
+            mem_formula="0",
+            comm_formula=(
+                f"PURPOSE: In decode, each CP rank computed attention against\n"
+                f"its OWN slice of the KV cache. Each rank now holds a "
+                f"partial\n"
+                f"(O, m, l). This single all-reduce merges those partials "
+                f"across\n"
+                f"all {cp} CP ranks (using online-softmax math) to get the "
+                f"final O.\n"
+                f"---\n"
+                f"AR of (O + m + l): 2*(CP-1)/CP * M "
+                f"= 2*({cp}-1)/{cp} * {M} "
+                f"= {ar_bytes:.0f} B over {tier} at {bw/1e9:.0f} GB/s "
+                f"+ 2*({cp}-1)*alpha"
+            ),
+        )
+
+    # ── Prefill: per-hop ring (K/V or Q+O/m/l) concurrent with S5-S7 ──
+    if cfg.topo.cp_ring_variant == "qoml":
+        M_KV = B * ((2 * T_q * hq_per_pe * dh * b) + (2 * T_q * hq_per_pe * b))
+        variant_desc = "Q+O/m/l ring"
+        formula_bytes = (
+            f"B*(2*T_q*(H_q/TP)*d_h*b + 2*T_q*(H_q/TP)*b) "
+            f"= {B}*(2*{T_q}*{hq_per_pe}*{dh}*{b} + 2*{T_q}*{hq_per_pe}*{b}) "
+            f"= {M_KV} B/hop"
+        )
+    else:
+        M_KV = 2 * B * S_local * hkv_per_pe * dh * b
+        variant_desc = "K/V ring"
+        formula_bytes = (
+            f"2*B*S_local*(H_kv/TP)*d_h*b "
+            f"= 2*{B}*{S_local}*{hkv_per_pe}*{dh}*{b} "
+            f"= {M_KV/1e6:.3f} MB/hop"
+        )
+
+    if cfg.topo.cp_placement == "pe":
+        # All hops intra-cube (fastest).
+        intra_hops = cfg.topo.cp - 1
+        inter_hops = 0
+        intra_time = (intra_hops * M_KV / cfg.machine.bw_intra
+                      + intra_hops * cfg.machine.alpha_intra)
+        inter_time = 0.0
+    else:  # cp_placement == "cube"
+        intra_hops = cfg.topo.cp_intra_sip_hops
+        inter_hops = cfg.topo.cp_inter_sip_hops
+        intra_time = (intra_hops * M_KV / cfg.machine.bw_inter
+                      + intra_hops * cfg.machine.alpha_inter)
+        inter_time = (inter_hops * M_KV / cfg.machine.bw_intersip
+                      + inter_hops * cfg.machine.alpha_intersip)
+    comm_time = intra_time + inter_time
+
+    # Overlap check: per-hop compute (S5+S6+S7) hides intra-SIP hops well;
+    # inter-SIP hops usually dominate.
+    per_hop_cmp, _ = _per_hop_qkT_pv(cfg)
+    per_hop_mem_softmax = (cfg.h_q_per_pe * cfg.topo.T_q
+                          * S_local * cfg.model.bytes_per_elem * 2
+                          / cfg.machine.bw_hbm)
+    per_hop_compute = 2 * per_hop_cmp + per_hop_mem_softmax
+    per_intra_ring = M_KV / cfg.machine.bw_inter + cfg.machine.alpha_inter
+    per_inter_ring = M_KV / cfg.machine.bw_intersip + cfg.machine.alpha_intersip
+    visible_intra = intra_hops * max(0.0, per_intra_ring - per_hop_compute)
+    visible_inter = inter_hops * max(0.0, per_inter_ring - per_hop_compute)
+    visible_total = visible_intra + visible_inter
+
+    tier_desc = f"{intra_hops} intra-SIP + {inter_hops} inter-SIP hops"
+    total_ring_bytes = M_KV * (intra_hops + inter_hops)
+    _var_purpose = (
+        "K, V shards rotate between CP ranks each hop"
+        if cfg.topo.cp_ring_variant == "kv"
+        else "Q + running (O, m, l) rotate between CP ranks each hop"
+    )
+    return StageCost(
+        name=f"C1 CP {variant_desc} ({tier_desc})",
+        formula=f"{formula_bytes}; "
+                f"intra: {intra_hops}*M/BW + inter: {inter_hops}*M/BW",
+        compute_s=0, memory_s=0, comm_s=comm_time,
+        bound="comm", visible_s=visible_total,
+        flops=0, mem_bytes=0, comm_bytes=int(total_ring_bytes),
+        flops_formula="0",
+        mem_formula="0",
+        comm_formula=(
+            f"PURPOSE: CP shards the sequence axis. Each CP rank holds only\n"
+            f"1/{cfg.topo.cp} of the KV cache, so to compute full attention\n"
+            f"we must move data between CP ranks. Variant: {_var_purpose}.\n"
+            f"Runs concurrently with S5/S6/S7 - each hop, compute of the\n"
+            f"just-arrived shard overlaps with comm of the next shard.\n"
+            f"---\n"
+            f"M*(CP-1) with M = {formula_bytes}; "
+            f"total = {total_ring_bytes/1e6:.3f} MB over the ring"
+        ),
+    )
+
+
+def comm_tp_allreduce(cfg: FullConfig) -> StageCost:
+    """TP AllReduce on W_O output."""
+    if cfg.topo.tp <= 1:
+        return StageCost(
+            name="C2 TP AllReduce W_O", formula="TP=1 -> no AR",
+            compute_s=0, memory_s=0, comm_s=0, bound="trivial", visible_s=0,
+        )
+    T_q = cfg.topo.T_q
+    d = cfg.model.hidden
+    b = cfg.model.bytes_per_elem
+    B = max(1, cfg.topo.b)
+    bytes_ = B * T_q * d * b
+    tp = cfg.topo.tp
+    tier = cfg.topo.tp_link_tier()  # "intra" | "inter" | "intersip"
+    if tier == "intra":
+        bw, alpha, scope = cfg.machine.bw_intra, cfg.machine.alpha_intra, "intra-cube"
+    elif tier == "inter":
+        bw, alpha, scope = cfg.machine.bw_inter, cfg.machine.alpha_inter, "cross-cube"
+    else:
+        bw, alpha, scope = cfg.machine.bw_intersip, cfg.machine.alpha_intersip, "cross-SIP"
+    comm_time = 2 * (tp - 1) / tp * bytes_ / bw + 2 * (tp - 1) * alpha
+    total_comm_bytes = int(2 * (tp - 1) / tp * bytes_)
+    return StageCost(
+        name=f"C2 TP AllReduce W_O (TP={tp} ranks, {scope})",
+        formula=f"2*(TP-1)/TP * B*T_q*d*b B / BW + 2(TP-1)*alpha "
+                f"[B={B}, {scope}]",
+        compute_s=0, memory_s=0, comm_s=comm_time,
+        bound="comm", visible_s=comm_time,
+        flops=0, mem_bytes=0, comm_bytes=total_comm_bytes,
+        flops_formula="0",
+        mem_formula="0",
+        comm_formula=(
+            f"PURPOSE: W_O is row-parallel across TP ranks (each rank holds\n"
+            f"1/{tp} of W_O's input rows). After the local W_O GEMM, each of\n"
+            f"the {tp} TP ranks holds a PARTIAL hidden vector (sum over its\n"
+            f"own Q-head slice only). This all-reduce sums those partials so\n"
+            f"every rank ends up with the full hidden vector for the next\n"
+            f"stage (RMSNorm -> FFN). Fires ONCE per layer over {scope} links.\n"
+            f"---\n"
+            f"2*(TP-1)/TP * B*T_q*d*b "
+            f"= 2*({tp}-1)/{tp} * {B}*{T_q}*{d}*{b} "
+            f"= {total_comm_bytes/1e6:.2f} MB over {scope} at "
+            f"{bw/1e9:.0f} GB/s"
+        ),
+    )
+
+
+def comm_kv_split_allreduce(cfg: FullConfig) -> StageCost:
+    """Extra AllReduce on attention scores when TP > H_kv with head-dim split.
+
+    Ranks sharing a KV head have partial scores; must AllReduce across the
+    split_factor group to get final scores per hop.
+    """
+    if not cfg.kv_replication_needed or cfg.topo.kv_shard_mode != "split":
+        return StageCost(
+            name="C3 Score AllReduce (head-split)",
+            formula="TP <= H_kv or replicate mode: not needed",
+            compute_s=0, memory_s=0, comm_s=0, bound="trivial", visible_s=0,
+        )
+    split = cfg.head_dim_split_factor  # ranks sharing one head
+    T_q = cfg.topo.T_q
+    S_local = cfg.topo.s_local
+    hq_per_pe = cfg.h_q_per_pe
+    b = cfg.model.bytes_per_elem
+    B = max(1, cfg.topo.b)
+    # score bytes per rank per hop — scales with B (per-request scores).
+    bytes_per_hop = B * hq_per_pe * T_q * S_local * b
+    # split-group AllReduce (intra-cube assumed if group fits)
+    bw = cfg.machine.bw_intra
+    alpha = cfg.machine.alpha_intra
+    per_hop = 2 * (split - 1) / split * bytes_per_hop / bw + 2 * (split - 1) * alpha
+    total = per_hop * cfg.topo.cp
+    total_comm_bytes = int(2 * (split - 1) / split * bytes_per_hop * cfg.topo.cp)
+    return StageCost(
+        name=f"C3 Score AllReduce ({split}-way, xCP hops)",
+        formula=f"2*({split}-1)/{split} * B*(H_q/TP)*T_q*S_local*b / BW "
+                f"= 2*({split}-1)/{split} * {B}*{hq_per_pe}*{T_q}*{S_local}*{b} "
+                f"/ BW + latency",
+        compute_s=0, memory_s=0, comm_s=total,
+        bound="comm", visible_s=total,
+        flops=0, mem_bytes=0, comm_bytes=total_comm_bytes,
+        flops_formula="0",
+        mem_formula="0",
+        comm_formula=(
+            f"PURPOSE: TP={cfg.topo.tp} > H_kv={cfg.model.h_kv}, so each KV\n"
+            f"head is split across {split} TP ranks along the head-dim (d_h\n"
+            f"/{split} per rank). Each rank's Q.K^T is therefore only a\n"
+            f"PARTIAL dot product; the {split} ranks sharing one KV head\n"
+            f"must AllReduce their partial scores to get the true score.\n"
+            f"This fires per hop of the ring, so {cfg.topo.cp}x per layer.\n"
+            f"To eliminate C3: set KV mode = 'replicate' (costs {split}x KV\n"
+            f"memory but no per-hop score AR), or lower TP so TP <= H_kv.\n"
+            f"---\n"
+            f"CP * 2*(split-1)/split * (H_q/TP)*T_q*S_local*b "
+            f"= {cfg.topo.cp} * 2*({split}-1)/{split} * {hq_per_pe}*{T_q}*{S_local}*{b} "
+            f"= {total_comm_bytes/1e6:.2f} MB total"
+        ),
+    )
+
+
+def all_stages(cfg: FullConfig) -> list[StageCost]:
+    """Ordered per-layer attention stages.
+
+    In prefill with CP > 1, C1 (the CP ring) is inserted between S7 and S8
+    because it runs concurrently with S5-S7 per hop. Placing it there in the
+    stage list makes the per-stage table read left-to-right in the order
+    things actually happen. C3 (score AR) sits with C1 for the same reason.
+    C2 (TP AllReduce on W_O) sits right after S10.
+
+    In decode, C1's comm is folded into S8, so no separate C1 row.
+    """
+    stages = [
+        stage_rmsnorm(cfg),
+        stage_wq(cfg),
+        stage_wkv(cfg),
+        stage_kv_append(cfg),
+        stage_qkT(cfg),
+        stage_softmax(cfg),
+        stage_pv(cfg),
+    ]
+    # Concurrent comm (with S5-S7 in prefill) goes here, right after S7.
+    if cfg.topo.mode == "prefill" and cfg.topo.cp > 1:
+        stages.append(comm_cp_ring(cfg))
+    _c3 = comm_kv_split_allreduce(cfg)
+    if _c3.visible_s > 0 or _c3.comm_s > 0:
+        stages.append(_c3)
+    stages.extend([
+        stage_merge(cfg),
+        stage_normalize(cfg),
+        stage_wo(cfg),
+    ])
+    # C2 fires right after S10 (W_O produces per-TP-rank output; AR combines).
+    stages.append(comm_tp_allreduce(cfg))
+    return stages
+
+
+# ── FFN block stages (lightweight; per PE, per layer) ────────────
+def stage_ffn_rmsnorm(cfg: FullConfig) -> StageCost:
+    T_q = cfg.topo.T_q
+    d = cfg.model.hidden
+    b = cfg.model.bytes_per_elem
+    B = max(1, cfg.topo.b)
+    # Activation memory scales with B; weight (once) is fixed.
+    bytes_ = B * T_q * d * b + T_q * d * b
+    flops = 4 * B * T_q * d
+    mem_s = bytes_ / cfg.machine.bw_hbm
+    return StageCost(
+        name="F1 RMSNorm (pre-FFN)",
+        formula=f"bytes = B*T_q*d*b + T_q*d*b = {bytes_} B / BW_HBM",
+        compute_s=0, memory_s=mem_s, comm_s=0,
+        bound="memory", visible_s=mem_s,
+        flops=flops, mem_bytes=bytes_,
+        flops_formula=f"4*B*T_q*d = 4*{B}*{T_q}*{d} = {flops}",
+        mem_formula=f"B*T_q*d*b + T_q*d*b (weight) "
+                     f"= {B}*{T_q}*{d}*{b} + {T_q*d*b} = {bytes_} B",
+    )
+
+
+def _ffn_gemm(cfg: FullConfig, name: str, ffn_per_pe: int) -> StageCost:
+    T_q = cfg.topo.T_q
+    d = cfg.model.hidden
+    b = cfg.model.bytes_per_elem
+    B = max(1, cfg.topo.b)
+    # FLOPs scale with batch; weight (shared) is fixed.
+    flops = 2 * B * T_q * d * ffn_per_pe
+    weight_B = d * ffn_per_pe * b
+    cmp_s, mem_s = _gemm_time(flops, weight_B, cfg)
+    vis, bnd = _visible(cmp_s, mem_s, 0)
+    return StageCost(
+        name=name,
+        formula=f"FLOPs = 2*B*T_q*d*(ffn/div) "
+                f"= 2*{B}*{T_q}*{d}*{ffn_per_pe} = {flops:.2g}; "
+                f"weight = {weight_B/1e6:.1f} MB (shared across batch)",
+        compute_s=cmp_s, memory_s=mem_s, comm_s=0,
+        bound=bnd, visible_s=vis,
+        flops=int(flops), mem_bytes=int(weight_B),
+        flops_formula=f"2*B*T_q*d*(ffn/div) "
+                       f"= 2*{B}*{T_q}*{d}*{ffn_per_pe} = {flops:.3g}",
+        mem_formula=f"d*(ffn/div)*b (weight, B-invariant) "
+                     f"= {d}*{ffn_per_pe}*{b} = {weight_B/1e6:.2f} MB",
+    )
+
+
+def stage_ffn_gate(cfg: FullConfig) -> StageCost:
+    ffn_per_pe = max(1, cfg.model.ffn_dim // (cfg.ffn_shard_divisor * max(1, cfg.topo.ep)))
+    return _ffn_gemm(cfg, "F2 W_gate GEMM", ffn_per_pe)
+
+
+def stage_ffn_up(cfg: FullConfig) -> StageCost:
+    ffn_per_pe = max(1, cfg.model.ffn_dim // (cfg.ffn_shard_divisor * max(1, cfg.topo.ep)))
+    return _ffn_gemm(cfg, "F3 W_up GEMM", ffn_per_pe)
+
+
+def stage_ffn_swiglu(cfg: FullConfig) -> StageCost:
+    """SwiGLU element-wise activation on the intermediate FFN tensor."""
+    T_q = cfg.topo.T_q
+    ffn_per_pe = max(1, cfg.model.ffn_dim // (cfg.ffn_shard_divisor * max(1, cfg.topo.ep)))
+    b = cfg.model.bytes_per_elem
+    B = max(1, cfg.topo.b)
+    bytes_ = 3 * B * T_q * ffn_per_pe * b
+    flops = 3 * B * T_q * ffn_per_pe  # gate * silu(up) approximated as 3 flops/elt
+    mem_s = bytes_ / cfg.machine.bw_hbm
+    return StageCost(
+        name="F4 SwiGLU act",
+        formula=f"3*B*T_q*(ffn/div)*b "
+                f"= 3*{B}*{T_q}*{ffn_per_pe}*{b} = {bytes_} B / BW_HBM",
+        compute_s=0, memory_s=mem_s, comm_s=0,
+        bound="memory", visible_s=mem_s,
+        flops=flops, mem_bytes=bytes_,
+        flops_formula=f"~3*B*T_q*(ffn/div) = 3*{B}*{T_q}*{ffn_per_pe} = {flops}",
+        mem_formula=f"3*B*T_q*(ffn/div)*b "
+                     f"= 3*{B}*{T_q}*{ffn_per_pe}*{b} = {bytes_} B",
+    )
+
+
+def stage_ffn_down(cfg: FullConfig) -> StageCost:
+    """W_down: (ffn_per_pe, hidden). Row-parallel over FFN scope."""
+    T_q = cfg.topo.T_q
+    d = cfg.model.hidden
+    b = cfg.model.bytes_per_elem
+    B = max(1, cfg.topo.b)
+    ffn_per_pe = max(1, cfg.model.ffn_dim // (cfg.ffn_shard_divisor * max(1, cfg.topo.ep)))
+    flops = 2 * B * T_q * ffn_per_pe * d
+    weight_B = ffn_per_pe * d * b
+    cmp_s, mem_s = _gemm_time(flops, weight_B, cfg)
+    vis, bnd = _visible(cmp_s, mem_s, 0)
+    return StageCost(
+        name="F5 W_down GEMM",
+        formula=f"FLOPs = 2*B*T_q*(ffn/div)*d "
+                f"= 2*{B}*{T_q}*{ffn_per_pe}*{d} = {flops:.2g}; "
+                f"weight = {weight_B/1e6:.1f} MB (shared across batch)",
+        compute_s=cmp_s, memory_s=mem_s, comm_s=0,
+        bound=bnd, visible_s=vis,
+        flops=int(flops), mem_bytes=int(weight_B),
+        flops_formula=f"2*B*T_q*(ffn/div)*d "
+                       f"= 2*{B}*{T_q}*{ffn_per_pe}*{d} = {flops:.3g}",
+        mem_formula=f"(ffn/div)*d*b (weight, B-invariant) "
+                     f"= {ffn_per_pe}*{d}*{b} = {weight_B/1e6:.2f} MB",
+    )
+
+
+def comm_ffn_allreduce(cfg: FullConfig) -> StageCost:
+    """AllReduce on FFN output across the FFN sharding scope."""
+    scope = cfg.topo.ffn_shard_scope
+    divisor = cfg.ffn_shard_divisor  # TP or TP*CP or TP*CP*DP
+    if divisor <= 1:
+        return StageCost(
+            name="CF1 FFN AllReduce",
+            formula="FFN scope = 1: not needed",
+            compute_s=0, memory_s=0, comm_s=0,
+            bound="trivial", visible_s=0,
+        )
+    T_q = cfg.topo.T_q
+    d = cfg.model.hidden
+    b = cfg.model.bytes_per_elem
+    B = max(1, cfg.topo.b)
+    # AR reduces the batched FFN output — bytes scale with B.
+    bytes_ = B * T_q * d * b
+    # Choose BW/alpha tier based on scope (rough): TP=intra-cube (or inter-cube),
+    # +CP=inter-SIP possible, +DP=inter-SIP always.
+    if "DP" in scope or "CP" in scope:
+        bw = cfg.machine.bw_intersip if cfg.topo.sips_used > 1 else cfg.machine.bw_inter
+        alpha = cfg.machine.alpha_intersip if cfg.topo.sips_used > 1 else cfg.machine.alpha_inter
+    else:  # TP only
+        bw = cfg.machine.bw_intra if cfg.topo.tp_spans_cubes == 1 else cfg.machine.bw_inter
+        alpha = cfg.machine.alpha_intra if cfg.topo.tp_spans_cubes == 1 else cfg.machine.alpha_inter
+    comm_time = 2 * (divisor - 1) / divisor * bytes_ / bw + 2 * (divisor - 1) * alpha
+    total_comm_bytes = int(2 * (divisor - 1) / divisor * bytes_)
+    return StageCost(
+        name=f"CF1 FFN AllReduce (scope={scope}, x{divisor})",
+        formula=f"2*({divisor}-1)/{divisor} * B*T_q*d*b B / BW + latency "
+                f"[B={B}]",
+        compute_s=0, memory_s=0, comm_s=comm_time,
+        bound="comm", visible_s=comm_time,
+        flops=0, mem_bytes=0, comm_bytes=total_comm_bytes,
+        flops_formula="0",
+        mem_formula="0",
+        comm_formula=(
+            f"PURPOSE: W_down is row-parallel across the {divisor} ranks in\n"
+            f"the FFN scope (scope={scope}). Each rank produced a PARTIAL\n"
+            f"hidden vector after W_down; this all-reduce sums them so every\n"
+            f"rank has the full FFN output for the next layer. Larger scope\n"
+            f"= less FFN weight memory per PE but bigger AR bill.\n"
+            f"---\n"
+            f"2*(div-1)/div * B*T_q*d*b = 2*({divisor}-1)/{divisor} * "
+            f"{B}*{T_q}*{d}*{b} = {total_comm_bytes/1e6:.2f} MB at "
+            f"{bw/1e9:.0f} GB/s (scope={scope})"
+        ),
+    )
+
+
+def all_ffn_stages(cfg: FullConfig) -> list[StageCost]:
+    return [
+        stage_ffn_rmsnorm(cfg),
+        stage_ffn_gate(cfg),
+        stage_ffn_up(cfg),
+        stage_ffn_swiglu(cfg),
+        stage_ffn_down(cfg),
+        comm_ffn_allreduce(cfg),
+    ]

tests/analytical_visualization/stage_shapes.py

@@ -0,0 +1,195 @@
+"""Per-stage tensor shape rows for the analytical visualization.
+
+Complements stage_latencies.py: while that module reports FLOPs/bytes/
+time per stage, this one reports the INPUT / WEIGHT / OUTPUT tensor
+shapes each stage operates on, per PE. Used for the 'per-stage shape'
+tables in the Streamlit app.
+
+Shape strings substitute the current cfg's numeric values so the table
+reads like an inspection of the actual deployment (no symbolic-only
+form).
+"""
+from __future__ import annotations
+
+from .model_config import FullConfig
+
+
+def _s(shape: tuple) -> str:
+    """Render a tuple of dim values as '(a, b, c)'."""
+    return "(" + ", ".join(str(x) for x in shape) + ")"
+
+
+def attn_stage_shape_rows(cfg: FullConfig) -> list[dict]:
+    """Per-PE attention shapes, one row per stage.
+
+    Rows follow the same S1..S10, C1..C3 order the latency table uses,
+    with the same conditional insertions (C1 only in prefill+CP>1, C3
+    only when kv_shard_mode='split' and TP > H_kv, C2 only when TP>1,
+    S8 merge only when CP>1).
+    """
+    m = cfg.model
+    t = cfg.topo
+    B = max(1, t.b)
+    T_q = t.T_q
+    S_local = t.s_local
+    d = m.hidden
+    dh = m.d_head
+    hq_pe = cfg.h_q_per_pe
+    hkv_pe = max(1, m.h_kv // t.tp)
+
+    rows: list[dict] = []
+
+    rows.append({
+        "Stage": "S1", "Op": "RMSNorm",
+        "Input (per PE)": _s((B, T_q, d)),
+        "Weight (per PE)": _s((d,)),
+        "Output (per PE)": _s((B, T_q, d)),
+    })
+    rows.append({
+        "Stage": "S2", "Op": "W_Q GEMM",
+        "Input (per PE)": _s((B, T_q, d)),
+        "Weight (per PE)": _s((d, hq_pe * dh)),
+        "Output (per PE)": _s((B, T_q, hq_pe * dh)),
+    })
+    rows.append({
+        "Stage": "S3", "Op": "W_K + W_V GEMM",
+        "Input (per PE)": _s((B, T_q, d)),
+        "Weight (per PE)": f"{_s((d, hkv_pe * dh))} x2 (K, V)",
+        "Output (per PE)": f"{_s((B, T_q, hkv_pe * dh))} x2 (K, V)",
+    })
+    rows.append({
+        "Stage": "S4", "Op": "KV cache append",
+        "Input (per PE)": f"{_s((B, T_q, hkv_pe * dh))} x2 (K, V)",
+        "Weight (per PE)": "-",
+        "Output (per PE)": (
+            f"KV cache: {_s((B, S_local, hkv_pe * dh))} x2 "
+            f"(extends by T_q={T_q} tokens)"
+        ),
+    })
+    rows.append({
+        "Stage": "S5", "Op": "Q · K^T",
+        "Input (per PE)": (
+            f"Q={_s((B, hq_pe, T_q, dh))}, "
+            f"K^T={_s((B, hkv_pe, dh, S_local))}"
+        ),
+        "Weight (per PE)": "-",
+        "Output (per PE)": _s((B, hq_pe, T_q, S_local)),
+    })
+    rows.append({
+        "Stage": "S6", "Op": "softmax",
+        "Input (per PE)": _s((B, hq_pe, T_q, S_local)),
+        "Weight (per PE)": "-",
+        "Output (per PE)": _s((B, hq_pe, T_q, S_local)),
+    })
+    rows.append({
+        "Stage": "S7", "Op": "P · V",
+        "Input (per PE)": (
+            f"P={_s((B, hq_pe, T_q, S_local))}, "
+            f"V={_s((B, hkv_pe, S_local, dh))}"
+        ),
+        "Weight (per PE)": "-",
+        "Output (per PE)": _s((B, hq_pe, T_q, dh)),
+    })
+
+    if t.mode == "prefill" and t.cp > 1:
+        variant = t.cp_ring_variant
+        if variant == "kv":
+            payload = (
+                f"K,V shards rotate: {_s((B, S_local, hkv_pe, dh))} x2 per hop"
+            )
+        else:
+            payload = (
+                f"Q + running (O,m,l) rotate: {_s((B, hq_pe, T_q, dh))} per hop"
+            )
+        rows.append({
+            "Stage": "C1", "Op": f"CP ring ({variant})",
+            "Input (per PE)": payload,
+            "Weight (per PE)": "-",
+            "Output (per PE)": "(circulating; final owner reduces)",
+        })
+
+    if t.kv_shard_mode == "split" and t.tp > m.h_kv:
+        rows.append({
+            "Stage": "C3", "Op": "Score AllReduce (head-split)",
+            "Input (per PE)": _s((B, hq_pe, T_q, S_local)),
+            "Weight (per PE)": "-",
+            "Output (per PE)": _s((B, hq_pe, T_q, S_local)),
+        })
+
+    if t.cp > 1:
+        rows.append({
+            "Stage": "S8", "Op": "online-softmax merge",
+            "Input (per PE)": (
+                f"O={_s((B, hq_pe, T_q, dh))}, "
+                f"m,l={_s((B, hq_pe, T_q))} each"
+            ),
+            "Weight (per PE)": "-",
+            "Output (per PE)": _s((B, hq_pe, T_q, dh)),
+        })
+
+    rows.append({
+        "Stage": "S9", "Op": "normalize O / l",
+        "Input (per PE)": _s((B, hq_pe, T_q, dh)),
+        "Weight (per PE)": "-",
+        "Output (per PE)": _s((B, hq_pe, T_q, dh)),
+    })
+    rows.append({
+        "Stage": "S10", "Op": "W_O GEMM",
+        "Input (per PE)": _s((B, T_q, hq_pe * dh)),
+        "Weight (per PE)": _s((hq_pe * dh, d)),
+        "Output (per PE)": _s((B, T_q, d)),
+    })
+
+    if t.tp > 1:
+        rows.append({
+            "Stage": "C2", "Op": f"TP AllReduce (W_O, TP={t.tp})",
+            "Input (per PE)": _s((B, T_q, d)),
+            "Weight (per PE)": "-",
+            "Output (per PE)": _s((B, T_q, d)),
+        })
+
+    return rows
+
+
+def ffn_stage_shape_rows(cfg: FullConfig) -> list[dict]:
+    """Per-PE FFN shapes, one row per stage."""
+    m = cfg.model
+    t = cfg.topo
+    B = max(1, t.b)
+    T_q = t.T_q
+    d = m.hidden
+    divisor = cfg.ffn_shard_divisor * max(1, t.ep)
+    ffn_pe = max(1, m.ffn_dim // divisor)
+    scope = t.ffn_shard_scope
+
+    rows: list[dict] = [
+        {"Stage": "F1", "Op": "RMSNorm (pre-FFN)",
+         "Input (per PE)": _s((B, T_q, d)),
+         "Weight (per PE)": _s((d,)),
+         "Output (per PE)": _s((B, T_q, d))},
+        {"Stage": "F2", "Op": "W_gate GEMM",
+         "Input (per PE)": _s((B, T_q, d)),
+         "Weight (per PE)": _s((d, ffn_pe)),
+         "Output (per PE)": _s((B, T_q, ffn_pe))},
+        {"Stage": "F3", "Op": "W_up GEMM",
+         "Input (per PE)": _s((B, T_q, d)),
+         "Weight (per PE)": _s((d, ffn_pe)),
+         "Output (per PE)": _s((B, T_q, ffn_pe))},
+        {"Stage": "F4", "Op": "SwiGLU (gate * silu(up))",
+         "Input (per PE)": f"gate, up = {_s((B, T_q, ffn_pe))} each",
+         "Weight (per PE)": "-",
+         "Output (per PE)": _s((B, T_q, ffn_pe))},
+        {"Stage": "F5", "Op": "W_down GEMM",
+         "Input (per PE)": _s((B, T_q, ffn_pe)),
+         "Weight (per PE)": _s((ffn_pe, d)),
+         "Output (per PE)": _s((B, T_q, d))},
+    ]
+    if cfg.ffn_shard_divisor > 1:
+        rows.append({
+            "Stage": "CF1",
+            "Op": f"FFN AllReduce (scope={scope}, x{cfg.ffn_shard_divisor})",
+            "Input (per PE)": _s((B, T_q, d)),
+            "Weight (per PE)": "-",
+            "Output (per PE)": _s((B, T_q, d)),
+        })
+    return rows