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/06-discussion}
-\input{sections/07-conclusion}
-\input{sections/08-future-work}
+\input{sections/05-gqa}                  % section header + intro
+\input{sections/05a-roofline}            % 5.1 Roofline Analysis
+\input{sections/05b-capacity-planning}   % 5.2 Capacity Planning
+\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/03-gemm.tex

@@ -172,16 +172,16 @@ kernel of \S\ref{sec:gqa} reaches for next — keeping the right
 working set on-chip so the composite pipeline lands in the
 compute-rich regime rather than the BW-bound one.
 
-\subsection{Why composite, and not user-orchestrated async loading?}
+\subsection{Why composite, and not kernel-orchestrated async loading?}
 \label{sec:gemm-vs-async}
 
 A reader familiar with double-buffered GEMM kernels on conventional
 hardware may ask: why a hardware-side composite command at all? Why
-isn't the obvious user-level pattern --- async-load each operand,
+isn't the obvious kernel-level pattern --- async-load each operand,
 overlap with compute, accumulate --- sufficient?
 
 To answer this concretely we contrast composite against two
-user-orchestrated baselines that have access to the same single-op
+kernel-orchestrated baselines that have access to the same single-op
 primitives the platform exposes (\textsf{tl.load} for async DMA into
 TCM, \textsf{tl.dot} for a single-op GEMM command on TCM-resident
 operands, \textsf{tl.store} for a DMA write-back). Both baselines
@@ -199,7 +199,7 @@ way at the runtime API surface to express ``start computing on
 tile 0 of $B$ while tile 1 is still in flight.'' Load-of-$B$ and
 GEMM therefore serialize.
 
-\paragraph{Async-tiled (chunked prefetch).} The user-level workaround is to
+\paragraph{Async-tiled (chunked prefetch).} The kernel-level workaround is to
 split $B$ along $K$ into \textsf{TILE\_K}-sized chunks, issue async
 \textsf{tl.load}s for those chunks, issue one \textsf{tl.dot} per
 chunk (each blocking only on its own $b_i$), and accumulate via
@@ -227,7 +227,7 @@ $K_{\text{KV}}=4096, d_{\text{head}}=128$ already needs
 --- past the cap. async-full and queue-all async-tiled are therefore
 not just slower than composite but \emph{architecturally infeasible}
 at LLM context length. The depth-2 async-tiled kernel is the only
-user-level
+kernel-level
 variant whose peak TCM footprint stays
 $O(2 \cdot \textsf{TILE\_K} \cdot N)$ regardless of $K$, the same
 order as composite's per-tile streaming buffer. It is the apples-to-apples
@@ -249,7 +249,7 @@ onto $K / \textsf{TILE\_K} = 48$ hardware tiles, so per-tile costs
 amplify into the largest measurable gap. The work content is
 identical for all four kernels: $\sim$6.3 M f16 MACs and
 $\sim$386 KiB of $B$ traffic from HBM, which together require
-$\sim$786 ns of GEMM-engine compute and $\sim$750 ns of DMA on a
+$\sim$786 ns of GEMM-engine compute and $\sim$781 ns of DMA on a
 saturated per-PE link. What differs is the number of host commands
 the same work is decomposed into --- 2 for composite (one
 \textsf{tl.load(A)} plus one composite), 4 for async-full, and 192
@@ -257,7 +257,7 @@ for either async-tiled variant (48 $A$-loads + 48 $B$-loads + 48
 \textsf{tl.dot}s + 47 elementwise adds + 1 store). The
 engine-pipeline-window throughput tracks that decomposition closely:
 composite reaches \SI{7.18}{\tera\flop\per\second} (post-overlap
-limit, only \SI{12}{\percent} below the \SI{8}{\tera\flop\per\second}
+limit, only \SI{10}{\percent} below the \SI{8}{\tera\flop\per\second}
 per-PE GEMM peak), async-full \SI{3.91}{\tera\flop\per\second}
 (DMA and compute serialize on a single big dot), and both async-tiled
 variants $\sim$\SI{2.53}{\tera\flop\per\second} (192 commands' worth of
@@ -272,7 +272,7 @@ specific simulator mechanisms.
 four issuance patterns: composite (one command, scheduler streams
 per-tile internally), async-full (one \textsf{tl.dot} on
 fully-loaded $B$), async-tiled with depth-2 double-buffer
-(TCM-bounded; the only user-level variant that scales to LLM
+(TCM-bounded; the only kernel-level variant that scales to LLM
 context length), and async-tiled with depth-$\infty$
 (all B-tiles queued up front; included as a sanity check that the
 prefetch depth is \emph{not} what separates the async-tiled kernel from
@@ -293,7 +293,7 @@ and all three async kernels beat it.}
 \end{figure*}
 
 \paragraph{Decomposing the gap.} Three structural mechanisms separate
-composite from the user-level baselines, and they layer.
+composite from the kernel-level baselines, and they layer.
 
 \emph{1. Inter-engine token routing happens below the host-side
 dispatch path.} The composite encodes the full
@@ -309,7 +309,7 @@ hand-offs total --- behind a single command from the host's point of
 view.
 
 \emph{2. \textsf{tl.dot} cannot replicate that per-tile pipeline at the
-user level.} A single-op GEMM command is handled on the GEMM engine as
+kernel level.} A single-op GEMM command is handled on the GEMM engine as
 a single monolithic compute timeout for the supplied $M{\times}K{\times}N$;
 there is no internal token loop that would let a streaming DMA of
 $B[i{+}1]$ overlap with the GEMM of $B[i]$ inside one
@@ -388,7 +388,7 @@ the four kernels to combine (a) macro-command dispatch at the host
 boundary (amortizing the structural CPU cost across all the work a
 single GEMM does), (b) scheduler-internal per-HW-tile streaming of
 DMA$\rightleftarrows$compute, and (c) TCM-bounded streaming buffer.
-User-orchestrated async kernels can have any two of those, not all
+Kernel-orchestrated async kernels can have any two of those, not all
 three: async-full pays one host dispatch (a) but forfeits per-tile
 overlap (b) and pins all of $B$ in TCM (c); depth-$\infty$ async-tiled
 achieves inter-chunk overlap but at $N_K$ host dispatches and

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

@@ -127,6 +127,63 @@ across whatever inter-device topology the configuration specifies,
 with the IPCQ ring buffer placed in on-PE TCM, PE-local HBM, or
 cube-shared SRAM---the third knob the results section sweeps.
 
+\subsection{Design alternatives}
+\label{sec:ipcq-alternatives}
+
+PE\_IPCQ is one point in a small space of hardware mechanisms for
+moving a short message from one PE to a neighbor and signalling its
+arrival. Three established alternatives anchor the space, each the
+HW realization of a familiar host-networking idea: a \emph{doorbell +
+polling} scheme (the classic MMIO doorbell---write the payload by DMA,
+write a doorbell, let the peer poll or take an interrupt); a
+\emph{hardware message queue} (HMQ, the NVLink-style descriptor engine
+that pushes a queue entry to the peer, with large payloads still
+riding a second DMA); and a \emph{completion-queue} design (RDMA-CQ,
+the InfiniBand/RoCE pattern where a DMA write auto-posts a completion
+entry the peer's CQ polls). PE\_IPCQ is the fourth: a hardware ring
+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]
+\centering
+\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
+doorbell / completion) and leave the peer polling or taking an
+interrupt; HMQ adds a dedicated descriptor engine but still moves large
+payloads on a second DMA; PE\_IPCQ folds head-pointer signalling into
+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*}
+
+\begin{figure}[t]
+\centering
+\includegraphics[width=\linewidth]{ipcq_alternatives_decision_matrix.png}
+\caption{Why the ring+credit design was chosen, across five criteria:
+single-send latency, whether the host CPU sits on the critical path,
+whether the receiver must poll or take a wake-up interrupt, whether the
+control and data datapaths are duplicated, and whether the mechanism is
+right-sized for single-owner PE-to-PE traffic (rather than a
+multi-tenant fabric). PE\_IPCQ is the only design that clears every
+criterion. The accompanying per-send step-count tally
+($\sim$28 control events for IPCQ versus $\sim$38 for HMQ, $\sim$53 for
+RDMA-CQ, and $\sim$56 for doorbell+polling) is an \emph{illustrative}
+order-of-magnitude comparator over hand-counted pipeline steps---not a
+simulator measurement. The measured, simulator-grounded results follow
+in the next subsection.}
+\label{fig:ipcq-decision}
+\end{figure}
+
+The qualitative comparison motivates the design but is not a
+quantitative claim: the cycle-step tallies above are hand-counted
+control events, deliberately separated from the measured latencies that
+follow. Everything in the results subsection runs on the PE\_IPCQ
+substrate and is simulator-grounded.
+
 \subsection{Results}
 
 All measurements in this section run on the PE\_IPCQ substrate
@@ -138,8 +195,21 @@ the collective sweep builds its own six-device (six-SIP, $2\times3$)
 configurations---distinct from the two-SIP default of
 Table~\ref{tab:hw}---and measures all-reduce latency as a function of
 payload size for three inter-device topologies: a 1D ring, a 2D mesh
-(no wrap), and a 2D torus. Table~\ref{tab:allreduce} and
-Figure~\ref{fig:allreduce-cmp} report the result.
+(no wrap), and a 2D torus (Figure~\ref{fig:allreduce-topo}).
+Table~\ref{tab:allreduce} and Figure~\ref{fig:allreduce-cmp} report the
+result.
+
+\begin{figure}[t]
+\centering
+\includegraphics[width=\linewidth]{allreduce_topology.png}
+\caption{The three six-device ($2\times3$) inter-device topologies the
+collective sweep runs over, and the hierarchical local-reduce /
+global all-reduce-broadcast schedule mapped onto each: a 1D ring, a 2D
+mesh (no wrap-around), and a 2D torus (wrap-around links on both axes).
+The torus's wrap links shorten the worst-case reduction path, which is
+what the latency sweep below rewards.}
+\label{fig:allreduce-topo}
+\end{figure}
 
 \begin{table}[t]
 \centering

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

@@ -24,151 +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}
-
-% TODO: compare data-placement options that apply across both the
-% short- and long-context regimes. Candidate axes:
-%   - KV cache:  per-CUBE shard vs. replicate; per-PE shard vs. replicate
-%   - Q / W_qkv / W_o weights:  static partition across CUBEs and PEs
-%   - Workspace (m, l, O softmax state):  scratch arena placement
-% The 4-case taxonomy used for long-context decode in
-% \S\ref{sec:gqa-long} (Cube-{SP,Repl} x PE-{TP,SP}) is one instantiation
-% of this framework; the short-context mapping in \S\ref{sec:gqa-short}
-% is another.
-
-\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.
-
-The four headline panels above stress the kernel at moderate context
-lengths. Long-context decode---the regime where KV cache size, not
-attention compute, sets serving cost---turns the choice of how to
-parallelize across cubes and PEs into a first-order design knob. We
-compare four strategies on the LLaMA-3.1-70B single-KV-head-group
-target (8 CUBEs $\times$ 8 PEs, one KV-head group):
-
-\begin{itemize}\setlength\itemsep{1pt}
-  \item \textbf{Case 1} (Cube-SP $\times$ PE-TP): KV split by $S_{kv}$
-        across CUBEs; PEs tensor-parallel on the batch dimension
-        (wastes PE-TP work at $B{=}1$).
-  \item \textbf{Case 2} (Cube-Repl $\times$ PE-TP): full KV
-        replicated to every CUBE; PEs tensor-parallel on batch.
-  \item \textbf{Case 3} (Cube-Repl $\times$ PE-SP): full KV
-        replicated; PEs sequence-parallel on $S_{kv}$ with an
-        intra-CUBE all-reduce.
-  \item \textbf{Case 4} ($\star$, Cube-SP $\times$ PE-SP): KV split
-        64-way (across both CUBEs and PEs) with a two-phase all-reduce
-        on the running softmax state $(m, \ell, O)$.
-\end{itemize}
-
-\begin{figure}[t]
-\centering
-\includegraphics[width=\linewidth]{gqa_decode_long_ctx_4cases_latency.png}
-\caption{End-to-end decode latency per parallelism strategy
-(LLaMA-3.1-70B single-KV-head group, 8 CUBEs $\times$ 8 PEs).
-Replication into CUBEs (Cases 2/3) wins the latency race
-(\SI{20.2}{\micro\second} for Case 3), but Case~4 ($\star$, KV split
-64-way) finishes within \SI{14}{\micro\second} of the leader while
-paying a different cost---visible in Figure~\ref{fig:gqa-4cases-mem}.}
-\label{fig:gqa-4cases-lat}
-\end{figure}
-
-\begin{figure}[t]
-\centering
-\includegraphics[width=\linewidth]{gqa_decode_long_ctx_4cases_memory.png}
-\caption{Per-CUBE KV memory footprint for the four cases. Cases 1
-and 4---both with the KV cache split across cubes (Cube-SP)---hold
-only \SI{0.5}{\mebi\byte} of KV state per CUBE; Cases 2 and 3, which
-replicate the full KV, hold \SI{4}{\mebi\byte} per CUBE, an
-\textbf{8$\times$} blowup at this configuration that scales linearly
-with context length.}
-\label{fig:gqa-4cases-mem}
-\end{figure}
-
-\begin{figure}[t]
-\centering
-\includegraphics[width=\linewidth]{gqa_decode_long_ctx_4cases_traffic.png}
-\caption{Per-case op-count breakdown. The replicated-KV PE-TP design
-(Case 2) avoids almost all on-device communication
-(\textasciitilde0 IPCQ copies), but at the cost of KV memory.
-Case~4's two-phase reduce charges \textasciitilde190 IPCQ copies and
-\textasciitilde190 DMA reads---this is the traffic that PE\_IPCQ
-(\S\ref{sec:allreduce}) is built to absorb at on-device speed.}
-\label{fig:gqa-4cases-traffic}
-\end{figure}
-
-Three things stand out. First, the fastest case in pure latency
-(Case~3, \SI{20.2}{\micro\second}) is also the most memory-hungry,
-requiring the full KV state on every CUBE---an option that fails to
-scale once context length blows past the per-CUBE budget. Second,
-Case~4's KV-split design gives back roughly \SI{14}{\micro\second}
-versus Case~3 in exchange for an \textbf{8$\times$} KV-memory
-reduction; for practical long-context serving where KV capacity is
-the binding constraint, this is the trade the design chooses
-(marked $\star$). Third, Case~4 pays its way in
-\emph{communication}: the op-count panel shows \textasciitilde190
-IPCQ copies and \textasciitilde190 DMA reads, precisely the
-on-device collective traffic that PE\_IPCQ and the torus links of
-\S\ref{sec:allreduce} are provisioned to move quickly---so the
-``slower'' strategy is in fact the one that fully cashes in the
-communication-side codesign work of this report.
-
-\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
-short, bursty decode requests with tool use and long shared contexts,
-which stresses the system differently from a single long generation:
-context reuse across requests, dynamic batching, and uneven expert load all
-change how compute and data should be dispersed. Characterizing how total
-compute and data movement distribute across the SIP/CUBE/PE hierarchy under
-such workloads---and which of the 1H hardware levers still dominate when the
-workload is this irregular---is the broader 2H agenda.
+\paragraph{Agentic workloads: from design to implementation.}
+Section~\ref{sec:agentic} established \emph{how} the fused GQA design
+extends to agentic fan-out/fan-in and \emph{which} responsibilities fall to
+the framework, the runtime, and the kernel. The 2H step is no longer to
+analyse the workload but to \emph{build} that support: the detailed design
+and implementation of all three levels. This is necessary work rather than
+optional, because today's agentic frameworks are effectively all
+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.
 
-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
     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/measure_gqa_decode_placement_comm.py

@@ -0,0 +1,362 @@
+"""Measured comm overlay for all 6 GQA decode KV placements.
+
+Runs the simulator for Cases 1-6 (the same 6 placements the chart in
+paper_plot_gqa_4cases_summary.py covers analytically) at S_kv = 64 K,
+sums actual IPCQ-copy bytes from the engine op_log, projects to
+per-token (x80 layers), scales the partial-score-AR component of
+Cases 4/5 from S_kv = 64 K -> S_kv = 1 M (linear in S_kv; other
+cases are S_kv-independent), adds the constant Wo + FFN AR
+(1.25 MB / token), and writes the result to JSON for
+paper_plot_gqa_4cases_summary.py to overlay on the analytical bars.
+
+Single layer of decode attention only — the projection × 80 takes
+the per-layer measurement to a per-token total.
+
+Usage:
+  python scripts/paper/measure_gqa_decode_placement_comm.py
+"""
+from __future__ import annotations
+
+import json
+import os
+import sys
+from pathlib import Path
+
+from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_repl_pe_sp import (
+    gqa_attention_decode_long_ctx_cube_repl_pe_sp_kernel as _case3_kernel,
+)
+from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_repl_pe_tp import (
+    gqa_attention_decode_long_ctx_cube_repl_pe_tp_kernel as _case1_kernel,
+)
+from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_sp import (
+    gqa_attention_decode_long_ctx_cube_sp_pe_sp_kernel as _case6_kernel,
+)
+from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_tp import (
+    gqa_attention_decode_long_ctx_cube_sp_pe_tp_kernel as _case2_kernel,
+)
+from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_tp_dhead import (
+    gqa_attention_decode_long_ctx_cube_sp_pe_tp_dhead_kernel as _case4_kernel,
+)
+from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_tp_dhead_pe_sp import (
+    gqa_attention_decode_long_ctx_cube_tp_dhead_pe_sp_kernel as _case5_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
+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
+
+_C = 8
+_P = 8
+_N_LAYERS = 80
+_S_KV_MEAS = 64 * 1024       # per-run simulator S_kv (1/16th of headline)
+_S_KV_HEADLINE = 1 << 20     # 1 Mi tokens, the chart's headline S_kv
+
+# Per-cube S_kv share for the d_head-TP partial-score AR cost.
+#   Case 4 (Cube-SP × PE-TP_dhead): per-cube = S_kv / C
+#   Case 5 (Cube-TP_dhead × PE-SP): per-cube = S_kv (KV replicated across
+#       cubes for the cube-axis d_head-TP), so partial-score AR scales
+#       with the full S_kv.
+# Headline / measured per-cube ratios give the partial-score-AR scale-up
+# from S_kv = 64 K to S_kv = 1 M. (m,ℓ,O) AR is S_kv-independent.
+_PARTIAL_SCORE_SCALE = _S_KV_HEADLINE / _S_KV_MEAS   # = 16
+
+# Per-token Wo + FFN AR (constant across all cases, comes from the
+# attn-output and FFN-down all-reduces NOT measured by the attention-
+# only kernel run here).
+_WO_PER_LAYER_BYTES = 8 * 1024
+_FFN_PER_LAYER_BYTES = 8 * 1024
+_WO_FFN_PER_TOKEN_BYTES = (
+    (_WO_PER_LAYER_BYTES + _FFN_PER_LAYER_BYTES) * _N_LAYERS
+)   # 1.25 MB
+
+# Total PE count in one KV-head group — average per-PE comm = total / N.
+_NUM_PES = _C * _P
+
+# Total partial-score-AR slice produced by the attention compute when
+# d_head is sharded. Used to split measured IPCQ traffic into the
+# S_kv-scaling component (partial scores) vs the S_kv-independent
+# component ((m,ℓ,O) merge). Same as the analytical formula in
+# paper_plot_gqa_4cases_summary.py: h_q · S_q · per_cube_S_kv · 2 bytes.
+_H_Q = 8
+_S_Q = 1
+_BYTES_PER_ELEM = 2
+
+_PARAMS = dict(C=_C, P=_P, T_q=_S_Q, S_kv=_S_KV_MEAS,
+               d_head=128, h_q=_H_Q, h_kv=1)
+
+
+def _bench_fn_case1(ctx):
+    """Case 1: Cube-Repl x PE-repl (PE-TP doesn't shard KV)."""
+    p = _PARAMS
+    configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
+    dp = DPPolicy(cube="replicate", pe="replicate",
+                  num_cubes=p["C"], num_pes=p["P"])
+    q = ctx.zeros((p["T_q"], p["h_q"] * p["d_head"]),
+                  dtype="f16", dp=dp, name="q_c1")
+    k = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
+                  dtype="f16", dp=dp, name="k_c1")
+    v = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
+                  dtype="f16", dp=dp, name="v_c1")
+    o = ctx.empty((p["T_q"], p["h_q"] * p["d_head"]),
+                  dtype="f16", dp=dp, name="o_c1")
+    ctx.launch("case1_repl_repl", _case1_kernel,
+               q, k, v, o,
+               p["T_q"], p["S_kv"], p["h_q"], p["h_kv"],
+               p["d_head"], p["C"], p["P"],
+               _auto_dim_remap=False)
+
+
+def _bench_fn_case2(ctx):
+    """Case 2: Cube-SP x PE-repl (PE-TP doesn't shard KV)."""
+    p = _PARAMS
+    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="replicate",
+                     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="q_c2")
+    k = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
+                  dtype="f16", dp=dp_kv, name="k_c2")
+    v = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
+                  dtype="f16", dp=dp_kv, name="v_c2")
+    o = ctx.empty((p["T_q"], p["h_q"] * p["d_head"]),
+                  dtype="f16", dp=dp_full, name="o_c2")
+    ctx.launch("case2_sp_repl", _case2_kernel,
+               q, k, v, o,
+               p["T_q"], p["S_kv"], p["h_q"], p["h_kv"],
+               p["d_head"], p["C"], p["P"],
+               _auto_dim_remap=False)
+
+
+def _bench_fn_case3(ctx):
+    """Case 3: Cube-Repl x PE-SP."""
+    p = _PARAMS
+    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="replicate", 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="q_c3")
+    k = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
+                  dtype="f16", dp=dp_kv, name="k_c3")
+    v = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
+                  dtype="f16", dp=dp_kv, name="v_c3")
+    o = ctx.empty((p["T_q"], p["h_q"] * p["d_head"]),
+                  dtype="f16", dp=dp_full, name="o_c3")
+    ctx.launch("case3_repl_sp", _case3_kernel,
+               q, k, v, o,
+               p["T_q"], p["S_kv"], p["h_q"], p["h_kv"],
+               p["d_head"], p["C"], p["P"],
+               _auto_dim_remap=False)
+
+
+def _bench_fn_case4(ctx):
+    p = _PARAMS
+    configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
+    dp_full = DPPolicy(cube="replicate", pe="column_wise",
+                       num_cubes=p["C"], num_pes=p["P"])
+    dp_kv = DPPolicy(cube="row_wise", pe="column_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="q_c4")
+    k = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
+                  dtype="f16", dp=dp_kv, name="k_c4")
+    v = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
+                  dtype="f16", dp=dp_kv, name="v_c4")
+    o = ctx.empty((p["T_q"], p["h_q"] * p["d_head"]),
+                  dtype="f16", dp=dp_full, name="o_c4")
+    ctx.launch("case4_dhead_tp", _case4_kernel,
+               q, k, v, o,
+               p["T_q"], p["S_kv"], p["h_q"], p["h_kv"],
+               p["d_head"], p["C"], p["P"],
+               _auto_dim_remap=False)
+
+
+def _bench_fn_case5(ctx):
+    p = _PARAMS
+    configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
+    dp_q = DPPolicy(cube="column_wise", pe="replicate",
+                    num_cubes=p["C"], num_pes=p["P"])
+    dp_kv = DPPolicy(cube="column_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_q, name="q_c5")
+    k = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
+                  dtype="f16", dp=dp_kv, name="k_c5")
+    v = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
+                  dtype="f16", dp=dp_kv, name="v_c5")
+    o = ctx.empty((p["T_q"], p["h_q"] * p["d_head"]),
+                  dtype="f16", dp=dp_q, name="o_c5")
+    ctx.launch("case5_dhead_tp_inter", _case5_kernel,
+               q, k, v, o,
+               p["T_q"], p["S_kv"], p["h_q"], p["h_kv"],
+               p["d_head"], p["C"], p["P"],
+               _auto_dim_remap=False)
+
+
+def _bench_fn_case6(ctx):
+    p = _PARAMS
+    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="q_c6")
+    k = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
+                  dtype="f16", dp=dp_kv, name="k_c6")
+    v = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
+                  dtype="f16", dp=dp_kv, name="v_c6")
+    o = ctx.empty((p["T_q"], p["h_q"] * p["d_head"]),
+                  dtype="f16", dp=dp_full, name="o_c6")
+    ctx.launch("case6_sp_sp", _case6_kernel,
+               q, k, v, o,
+               p["T_q"], p["S_kv"], p["h_q"], p["h_kv"],
+               p["d_head"], p["C"], p["P"],
+               _auto_dim_remap=False)
+
+
+def _sum_ipcq_bytes(op_log) -> int:
+    """Sum nbytes across all ipcq_copy records."""
+    return sum(
+        r.params.get("nbytes", 0)
+        for r in op_log
+        if r.op_kind == "memory" and r.op_name == "ipcq_copy"
+    )
+
+
+def _partial_score_slices(case: int) -> int:
+    """Divisor that splits S_kv into partial-score tiles, per the
+    analytical model in paper_plot_gqa_4cases_summary.py:
+
+      partial_score_per_PE = h_q * S_q * (s_kv / slices) * bytes
+
+    Case 4 (Cube-SP x PE-TP-dhead): intra-cube AR over d_head shards
+        on PE axis -> partial tile per PE has per-cube S_kv = s_kv/C.
+    Case 5 (Cube-TP-dhead x PE-SP): inter-cube AR over d_head shards
+        on cube axis -> partial tile per PE has per-PE S_kv = s_kv/P.
+    Cases 1, 2, 3, 6: no partial-score AR (only (m,l,O) merge).
+    """
+    if case == 4:
+        return _C
+    if case == 5:
+        return _P
+    return 0
+
+
+def _split_attn_layer_bytes(case: int, total_ipcq_bytes: int,
+                            s_kv: int) -> tuple[int, int]:
+    """Split per-layer attention-time IPCQ bytes into:
+      (partial_score_component, mlo_component).
+
+    The partial-score component scales with s_kv (so it must be scaled
+    when projecting from the measure-time s_kv to the headline s_kv);
+    the (m,l,O) component is constant in s_kv.
+
+    Partial-score size is analytically known per-case (formula in
+    _partial_score_slices); the remainder is treated as (m,l,O) + any
+    other S_kv-independent overhead. Per-PE = total / NUM_PES.
+    """
+    per_pe_total = total_ipcq_bytes // _NUM_PES
+    slices = _partial_score_slices(case)
+    if slices == 0:
+        # No partial-score AR for this case.
+        return 0, per_pe_total
+    partial_score_per_pe = (
+        _H_Q * _S_Q * (s_kv // slices) * _BYTES_PER_ELEM
+    )
+    partial_score_per_pe = min(partial_score_per_pe, per_pe_total)
+    mlo_per_pe = per_pe_total - partial_score_per_pe
+    return partial_score_per_pe, mlo_per_pe
+
+
+_KERNELS = (
+    (1, "Case 1 (Cube-Repl x PE-repl)",      _bench_fn_case1),
+    (2, "Case 2 (Cube-SP   x PE-repl)",      _bench_fn_case2),
+    (3, "Case 3 (Cube-Repl x PE-SP)",        _bench_fn_case3),
+    (4, "Case 4 (Cube-SP   x PE-TP d_head)", _bench_fn_case4),
+    (5, "Case 5 (Cube-TP d_head x PE-SP)",   _bench_fn_case5),
+    (6, "Case 6 (Cube-SP   x PE-SP)  [*]",   _bench_fn_case6),
+)
+
+
+def main() -> int:
+    topology = os.environ.get("GQA_1H_TOPOLOGY", "topology.yaml")
+    topo = resolve_topology(topology)
+
+    out: dict = {
+        "S_kv_measured": _S_KV_MEAS,
+        "S_kv_headline": _S_KV_HEADLINE,
+        "n_layers": _N_LAYERS,
+        "num_pes": _NUM_PES,
+        "wo_ffn_per_token_bytes": _WO_FFN_PER_TOKEN_BYTES,
+        "cases": {},
+    }
+
+    print(f"Measuring at S_kv={_S_KV_MEAS:,} ; scaling partial-score AR "
+          f"to S_kv={_S_KV_HEADLINE:,} (×{int(_PARTIAL_SCORE_SCALE)})")
+    print()
+
+    for case_id, label, bench_fn in _KERNELS:
+        try:
+            res = run_bench(
+                topology=topo, bench_fn=bench_fn,
+                device=resolve_device(None),
+                engine_factory=lambda t, d: GraphEngine(
+                    getattr(t, "topology_obj", t), enable_data=True,
+                ),
+            )
+        except Exception as e:
+            print(f"  {label:<42} FAIL: {type(e).__name__}: {e}")
+            return 1
+        if not res.completion.ok:
+            print(f"  {label:<42} ENGINE FAIL: {res.completion}")
+            return 1
+
+        total_ipcq = _sum_ipcq_bytes(res.engine.op_log)
+        partial_pe, mlo_pe = _split_attn_layer_bytes(
+            case_id, total_ipcq, _S_KV_MEAS,
+        )
+        # Per-token attention-time comm at S_kv = 1 M:
+        #   (partial_score_per_layer × scale  +  mlo_per_layer) × 80 layers
+        scaled_partial_per_token = (
+            partial_pe * int(_PARTIAL_SCORE_SCALE) * _N_LAYERS
+        )
+        mlo_per_token = mlo_pe * _N_LAYERS
+        attn_per_token = scaled_partial_per_token + mlo_per_token
+        total_per_token = attn_per_token + _WO_FFN_PER_TOKEN_BYTES
+
+        out["cases"][str(case_id)] = {
+            "label": label,
+            "total_ipcq_bytes_one_layer": total_ipcq,
+            "per_pe_partial_score_bytes_one_layer": partial_pe,
+            "per_pe_mlo_bytes_one_layer": mlo_pe,
+            "per_pe_attn_bytes_per_token_at_1M": attn_per_token,
+            "per_pe_total_bytes_per_token_at_1M": total_per_token,
+        }
+
+        print(f"  {label:<42} "
+              f"ipcq_total={total_ipcq:>10,}  "
+              f"per_pe_attn(1L)={(partial_pe + mlo_pe):>9,}  "
+              f"per_pe_total/tok@1M={total_per_token / (1<<20):>7.2f} MB")
+
+    out_path = (
+        Path(__file__).resolve().parents[2]
+        / "src" / "kernbench" / "benches"
+        / "1H_milestone_output" / "gqa" / "long_ctx"
+        / "gqa_long_ctx_6cases_measured_comm.json"
+    )
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    out_path.write_text(json.dumps(out, indent=2))
+    print()
+    print(f"wrote {out_path}")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

scripts/paper/paper_plot_gqa_4cases_summary.py

@@ -114,6 +114,8 @@ Output PNG:
 """
 from __future__ import annotations
 
+import json
+import textwrap
 from pathlib import Path
 
 import matplotlib.patches as mpatches
@@ -186,9 +188,36 @@ _WV_GB = _WV_MB_PER_LAYER * _N_LAYERS / 1000    # 0.16 GB
 _WO_GB = _WO_MB_PER_LAYER * _N_LAYERS / 1000    # 0.16 GB
 _WEIGHTS_GB = _WQ_GB + _WK_GB + _WV_GB + _WO_GB  # 1.76 GB
 
-# (m, ℓ, O) AR payload — used by Cases 1, 3, 4 (different axes per case).
-_MLO_INTRA_BYTES_PER_LAYER = 32 * 1024   # ~32 KB intra-cube AR
-_MLO_INTER_BYTES_PER_LAYER = 32 * 1024   # ~32 KB inter-cube AR
+# (m, ℓ, O) AR cost per layer, derived from the kernel topology.
+#
+# Per-PE payload  T  for one (m, ℓ, O) merge step:
+#     O  — h_q · S_q · d_head · 2 bytes  (FP16)
+#     m  — h_q · S_q · 4 bytes           (FP32)
+#     ℓ  — h_q · S_q · 4 bytes           (FP32)
+#     T  = h_q · S_q · (d_head · 2 + 8)  ≈ 2.1 KB
+#
+# Reduce algorithm: the decode kernels use hierarchical reduce-only
+# (chain / tree) rather than ring all-reduce, because the merged result
+# only needs to land on the cube that runs the downstream Wo gemm —
+# not on every PE.  For a chain of N participants the total traffic is
+# (N-1)·T and the per-PE average is (N-1)/N · T.
+#
+# The previous version used a hard-coded 32 KB / layer placeholder which
+# overestimated the per-PE cost by ~6× (and ~3× even under a hypothetical
+# ring-AR assumption).  See `gqa_long_ctx_6cases_measured_comm.json` for
+# the measured numbers this matches against.
+_MLO_PAYLOAD_BYTES = _H_Q * _S_Q * (_D_HEAD * _BYTES_PER_ELEM + 8)
+
+
+def _reduce_chain_per_pe_bytes(n_participants: int) -> int:
+    """Per-PE average bytes for a single-stage chain/tree reduce of T."""
+    if n_participants <= 1:
+        return 0
+    return _MLO_PAYLOAD_BYTES * (n_participants - 1) // n_participants
+
+
+_MLO_INTRA_BYTES_PER_LAYER = _reduce_chain_per_pe_bytes(_P)  # PE-axis
+_MLO_INTER_BYTES_PER_LAYER = _reduce_chain_per_pe_bytes(_C)  # cube-axis
 
 # Cases — renumbered in MEMORY-DESCENDING ORDER (left to right):
 #   Case 1   (40 GB) : no sharding
@@ -224,11 +253,11 @@ _CASE_COLOR = {
 }
 _ATTN_DESC = {
     1: "none",
-    2: "online-softmax\nmerge of\n(m,ℓ,O)\ninter-cube",
-    3: "online-softmax\nmerge of\n(m,ℓ,O)\nintra-cube",
-    4: "partial\nscores\n+ (m,ℓ,O)\nmerge\n(d_head-TP)",
-    5: "partial\nscores\n+ (m,ℓ,O)\nmerge\n(d_head-TP)",
-    6: "online-softmax\nmerge of\n(m,ℓ,O)\nintra + inter",
+    2: "online-softmax (m,ℓ,O) — inter-cube",
+    3: "online-softmax (m,ℓ,O) — intra-cube",
+    4: "partial scores + (m,ℓ,O) merge  (d_head-TP)",
+    5: "partial scores + (m,ℓ,O) merge  (d_head-TP)",
+    6: "online-softmax (m,ℓ,O) — intra + inter",
 }
 
 _WO_COLOR   = "#9EC5E8"
@@ -240,6 +269,22 @@ _OUT_DIR = (
     / "src" / "kernbench" / "benches"
     / "1H_milestone_output" / "gqa" / "long_ctx"
 )
+_MEASURED_JSON = _OUT_DIR / "gqa_long_ctx_6cases_measured_comm.json"
+
+
+def _load_measured() -> dict[int, float] | None:
+    """Load measured per-PE comm bytes (already scaled to S_kv=1M).
+
+    Returns {case_id: per_token_total_bytes} or None if JSON missing.
+    Produced by scripts/paper/measure_gqa_decode_placement_comm.py.
+    """
+    if not _MEASURED_JSON.exists():
+        return None
+    data = json.loads(_MEASURED_JSON.read_text())
+    return {
+        int(cid): info["per_pe_total_bytes_per_token_at_1M"]
+        for cid, info in data["cases"].items()
+    }
 
 
 # ── Formulae ────────────────────────────────────────────────────────
@@ -388,51 +433,118 @@ def _plot_memory(ax) -> None:
     ax.legend(loc="upper right", fontsize=9)
 
 
-def _plot_comm(ax) -> None:
-    wo_mb = []
-    ffn_mb = []
-    attn_mb = []
+def _plot_comm(ax, *, mode: str = "analytical") -> None:
+    """Per-PE comm panel.
+
+    mode = "analytical": single solid bars from per_token_bytes formula.
+    mode = "paired"    : analytical (solid) + simulator-measured
+                         (hatched) side-by-side per case, when the
+                         measurement JSON is available.
+    """
+    labels = [_CASE_LABEL[c] for c in _CASES]
+
+    wo_mb_list: list[float] = []
+    ffn_mb_list: list[float] = []
+    attn_mb: list[float] = []
     for c in _CASES:
         wo, ffn, attn = per_token_bytes(c, _HEADLINE_S_KV)
-        wo_mb.append(wo / (1 << 20))
-        ffn_mb.append(ffn / (1 << 20))
+        wo_mb_list.append(wo / (1 << 20))
+        ffn_mb_list.append(ffn / (1 << 20))
         attn_mb.append(attn / (1 << 20))
 
-    labels = [_CASE_LABEL[c] for c in _CASES]
-    x = list(range(len(_CASES)))
+    measured = _load_measured() if mode == "paired" else None
+    paired = measured is not None
+    source_tag = "analytical (solid) vs simulator-measured (hatched)" \
+                 if paired else "analytical"
 
-    ax.bar(x, wo_mb, color=_WO_COLOR, edgecolor="black")
-    ax.bar(x, ffn_mb, bottom=wo_mb, color=_FFN_COLOR, edgecolor="black")
-    bottoms_attn = [w + f for w, f in zip(wo_mb, ffn_mb)]
-    ax.bar(x, attn_mb, bottom=bottoms_attn,
+    n_cases = len(_CASES)
+    x = list(range(n_cases))
+    bar_w = 0.36 if paired else 0.65
+    x_ana  = [xi - bar_w / 2 for xi in x] if paired else x
+    x_meas = [xi + bar_w / 2 for xi in x] if paired else None
+
+    # Analytical bars (solid).
+    ax.bar(x_ana, wo_mb_list,  width=bar_w,
+           color=_WO_COLOR,  edgecolor="black")
+    ax.bar(x_ana, ffn_mb_list, width=bar_w, bottom=wo_mb_list,
+           color=_FFN_COLOR, edgecolor="black")
+    bottoms_attn = [w + f for w, f in zip(wo_mb_list, ffn_mb_list)]
+    ax.bar(x_ana, attn_mb, width=bar_w, bottom=bottoms_attn,
            color=_ATTN_COLOR, edgecolor="black")
 
+    # Measured bars (hatched) — same Wo+FFN base, attn from op_log.
+    meas_attn_mb: list[float] = []
+    if paired:
+        for i, c in enumerate(_CASES):
+            meas_total = measured.get(c, 0) / (1 << 20)
+            meas_attn_mb.append(
+                max(meas_total - wo_mb_list[i] - ffn_mb_list[i], 0.0))
+        ax.bar(x_meas, wo_mb_list,  width=bar_w,
+               color=_WO_COLOR,  edgecolor="black",
+               hatch="///", alpha=0.85)
+        ax.bar(x_meas, ffn_mb_list, width=bar_w, bottom=wo_mb_list,
+               color=_FFN_COLOR, edgecolor="black",
+               hatch="///", alpha=0.85)
+        ax.bar(x_meas, meas_attn_mb, width=bar_w, bottom=bottoms_attn,
+               color=_ATTN_COLOR, edgecolor="black",
+               hatch="///", alpha=0.85)
+
     ax.set_xticks(x)
     ax.set_xticklabels(labels, fontsize=9)
     ax.set_ylabel("Comm per token per PE (MB, log)")
     ax.set_yscale("log")
-    ax.set_title(
+    title = (
         f"Per-PE communication at S_kv = {_HEADLINE_S_KV:,} tokens "
-        f"(decode S_q=1, B=1; {_N_LAYERS} layers)",
-        fontsize=11,
+        f"(decode S_q=1, B=1; {_N_LAYERS} layers)  —  {source_tag}"
     )
+    if paired:
+        title += "\n(simulator measured at S_kv = 8K; " \
+                 "partial-score AR scaled ×128 to S_kv = 1M)"
+    ax.set_title(title, fontsize=10)
     ax.grid(axis="y", ls=":", alpha=0.5, which="both")
 
-    totals = [wo_mb[i] + ffn_mb[i] + attn_mb[i] for i in range(len(_CASES))]
-    ax.set_ylim(top=max(totals) * 10)
+    totals_ana = [wo_mb_list[i] + ffn_mb_list[i] + attn_mb[i]
+                  for i in range(n_cases)]
+    ymax = max(totals_ana)
+    if paired:
+        ymax = max(ymax, max(
+            (measured.get(c, 0) / (1 << 20)) for c in _CASES))
+    ax.set_ylim(top=ymax * 22)
 
     for i, c in enumerate(_CASES):
-        total_bytes = (wo_mb[i] + ffn_mb[i] + attn_mb[i]) * (1 << 20)
-        ax.text(x[i], totals[i] * 1.5, _fmt_bytes(total_bytes),
-                ha="center", fontsize=9, weight="bold")
+        ana_bytes = totals_ana[i] * (1 << 20)
+        if paired:
+            meas_bytes = measured.get(c, 0)
+            top_y = max(totals_ana[i], meas_bytes / (1 << 20))
+            label = (f"ana: {_fmt_bytes(ana_bytes)}\n"
+                     f"sim: {_fmt_bytes(meas_bytes)}")
+        else:
+            top_y = totals_ana[i]
+            label = _fmt_bytes(ana_bytes)
+        ax.text(x[i], top_y * 1.5, label,
+                ha="center", va="bottom",
+                fontsize=8 if paired else 9,
+                weight="bold", linespacing=1.05)
+
+        # Attention-time AR descriptor — placed at the attention-segment
+        # midpoint, centered between the analytical and measured bars so
+        # it visually spans both. Coloured the same as the attention bar
+        # (no background box) so it lives within the case's bar zone.
+        # Wrapped narrow so each line fits inside the bar-pair width.
+        wrapped = textwrap.fill(_ATTN_DESC[c], width=14)
         if attn_mb[i] > 0:
             mid = bottoms_attn[i] + attn_mb[i] / 2
-            ax.text(x[i], mid, _ATTN_DESC[c],
-                    ha="center", va="center", fontsize=7,
-                    color="black", weight="bold")
+            ax.text(x[i], mid, wrapped,
+                    ha="center", va="center",
+                    fontsize=7 if paired else 8,
+                    color="black", weight="bold",
+                    linespacing=1.0)
         else:
-            ax.text(x[i], totals[i] * 0.4, _ATTN_DESC[c],
-                    ha="center", fontsize=7.5, color="grey", style="italic")
+            ax.text(x[i], totals_ana[i] * 0.4, wrapped,
+                    ha="center",
+                    fontsize=7 if paired else 8,
+                    color="grey", style="italic",
+                    linespacing=1.0)
 
     legend_handles = [
         mpatches.Patch(facecolor=_WO_COLOR, edgecolor="black",
@@ -442,40 +554,67 @@ def _plot_comm(ax) -> None:
         mpatches.Patch(facecolor=_ATTN_COLOR, edgecolor="black",
                        label="Attn-time collective"),
     ]
+    if paired:
+        legend_handles.append(
+            mpatches.Patch(facecolor="white", edgecolor="black",
+                           hatch="///", label="simulator-measured"))
     ax.legend(handles=legend_handles, loc="upper right", fontsize=8,
               framealpha=0.92)
 
 
 def main() -> Path:
     _OUT_DIR.mkdir(parents=True, exist_ok=True)
-    fig = plt.figure(figsize=(21.0, 6.0))
+
+    # (a) Per-PE HBM budget — standalone PNG.
+    fig_b, ax_b = plt.subplots(figsize=(4.0, 6.0))
+    _plot_budget(ax_b)
+    fig_b.tight_layout()
+    out_b = _OUT_DIR / "gqa_long_ctx_6cases_hbm_budget.png"
+    fig_b.savefig(out_b, dpi=150)
+    plt.close(fig_b)
+    print(f"wrote {out_b}")
+
+    # (b) Combined 3-panel summary — HBM budget + KV memory + comm.
+    fig = plt.figure(figsize=(22.0, 6.5))
     gs = fig.add_gridspec(1, 3, width_ratios=[0.7, 1.6, 1.6], wspace=0.22)
-    ax_b = fig.add_subplot(gs[0, 0])
+    ax_b2 = fig.add_subplot(gs[0, 0])
     ax_m  = fig.add_subplot(gs[0, 1])
     ax_c  = fig.add_subplot(gs[0, 2])
-    _plot_budget(ax_b)
+    _plot_budget(ax_b2)
     _plot_memory(ax_m)
     _plot_comm(ax_c)
-
-    fig.suptitle(
-        f"GQA per-PE memory + communication  —  LLaMA-3.1-70B "
-        f"single-KV-head group  (C={_C}, P={_P}, {_N_LAYERS} layers, "
-        f"S_kv = 1 M, FP16)",
-        fontsize=12, y=0.985,
-    )
-    fig.text(
-        0.5, 0.94,
-        "Cases ordered by memory  ·  Case 1 (40 GB, no sharding)  →  "
-        "Cases 2-3 (5 GB, 1-axis sharded)  →  Cases 4-6 (640 MB, 2-axis sharded)  "
-        "·  Case 6 ★ = lowest comm among memory-feasible cases",
-        ha="center", fontsize=9.5, color="#444",
-    )
-    fig.tight_layout(rect=(0, 0, 1, 0.92))
-    out = _OUT_DIR / "gqa_4cases_summary.png"
+    fig.tight_layout()
+    out = _OUT_DIR / "gqa_long_ctx_6cases_summary.png"
     fig.savefig(out, dpi=150)
     plt.close(fig)
     print(f"wrote {out}")
 
+    # (c) 2-panel companion (analytical only).
+    fig2 = plt.figure(figsize=(18.0, 6.5))
+    gs2 = fig2.add_gridspec(1, 2, width_ratios=[1.0, 1.0], wspace=0.18)
+    ax_m2 = fig2.add_subplot(gs2[0, 0])
+    ax_c2 = fig2.add_subplot(gs2[0, 1])
+    _plot_memory(ax_m2)
+    _plot_comm(ax_c2, mode="analytical")
+    fig2.tight_layout()
+    out2 = _OUT_DIR / "gqa_long_ctx_6cases_memory_comm_analytical.png"
+    fig2.savefig(out2, dpi=150)
+    plt.close(fig2)
+    print(f"wrote {out2}")
+
+    # (d) 2-panel companion — analytical vs simulator-measured paired.
+    fig3 = plt.figure(figsize=(19.0, 6.5))
+    gs3 = fig3.add_gridspec(1, 2, width_ratios=[1.0, 1.0], wspace=0.18)
+    ax_m3 = fig3.add_subplot(gs3[0, 0])
+    ax_c3 = fig3.add_subplot(gs3[0, 1])
+    _plot_memory(ax_m3)
+    _plot_comm(ax_c3, mode="paired")
+    fig3.tight_layout()
+    out3 = _OUT_DIR / "gqa_long_ctx_6cases_memory_comm_paired.png"
+    fig3.savefig(out3, dpi=150)
+    plt.close(fig3)
+    print(f"wrote {out3}")
+
     # Paper-ready table to stdout.
     print()
     print(f"  {'Case':<27} {'KV/tok·PE':>12} {'KV @ 1M':>11} "

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_4cases.py

@@ -1,13 +1,25 @@
 """Comparative figures for milestone-gqa-decode-long-ctx-4cases.
 
-Reads sweep.json (emitted by ``kernbench run --bench
-milestone-gqa-decode-long-ctx-4cases``) and writes four PNGs into
-``docs/report/1H-codesign-paper/figures/``:
+Reads sweep_decode.json (emitted by the milestone-1h-gqa bench) and
+writes four PNGs into the same bench-output dir
+(src/kernbench/benches/1H_milestone_output/gqa/long_ctx/):
 
-  gqa_decode_long_ctx_4cases_latency.png       end-to-end latency per case
-  gqa_decode_long_ctx_4cases_traffic.png       ipcq/dma op-count breakdown
-  gqa_decode_long_ctx_4cases_memory.png        per-PE KV bytes per case
-  gqa_decode_long_ctx_4cases_parallelism.png   per-PE S_local (compute work)
+  gqa_decode_long_ctx_6cases_latency.png       end-to-end latency per case
+  gqa_decode_long_ctx_6cases_traffic.png       ipcq/dma op-count breakdown
+  gqa_decode_long_ctx_6cases_memory.png        per-PE KV bytes per case
+  gqa_decode_long_ctx_6cases_parallelism.png   per-PE S_local (compute work)
+
+Filename still says "4cases" for backwards compat, but the script now
+covers all SIX kv-sharding strategies from the analytical chart
+(`gqa_4cases_summary.png`) — the original 4 plus the two new
+d_head-TP variants:
+
+  Case 1   Cube-Repl × PE-repl      (PE-TP doesn't shard KV)
+  Case 2   Cube-SP   × PE-repl
+  Case 3   Cube-Repl × PE-SP
+  Case 4   Cube-SP   × PE-TP-d_head      ← NEW
+  Case 5   Cube-TP-d_head × PE-SP        ← NEW
+  Case 6 ★ Cube-SP   × PE-SP             (Pareto-best)
 
 Run (after the bench):
     GQA_DECODE_LONG_CTX_4CASES_RUN=1 python -m kernbench.cli.main run \\
@@ -32,16 +44,27 @@ _FIG_DIR = (
 )
 _SWEEP_JSON = _FIG_DIR / "sweep_decode.json"
 
-# Panel name → (short label, case ordinal for left-to-right plot order).
+# Panel name → (short label, case ordinal, accent flag) using the
+# analytical chart's memory-descending ordering. PE-TP doesn't shard
+# KV memory, so the cube_repl_pe_tp panel maps to Case 1 (no
+# sharding, KV-wise) and cube_sp_pe_tp panel maps to Case 2.
+_NORMAL, _OVERFLOW, _PARETO = "normal", "overflow", "pareto"
+
 _CASE_INFO = {
-    "single_kv_group_decode_long_ctx_gqa_cube_sp_pe_tp": (
-        "Case 1\nCube-SP × PE-TP", 1),
-    "single_kv_group_decode_long_ctx_gqa_cube_repl_pe_tp": (
-        "Case 2\nCube-Repl × PE-TP", 2),
-    "single_kv_group_decode_long_ctx_gqa_cube_repl_pe_sp": (
-        "Case 3\nCube-Repl × PE-SP", 3),
-    "single_kv_group_decode_long_ctx_gqa_cube_sp_pe_sp": (
-        "Case 4 ★\nCube-SP × PE-SP", 4),
+    # panel name                                                       label                            ord  flag
+    "single_kv_group_decode_long_ctx_gqa_cube_repl_pe_tp":            ("Case 1\nCube-Repl × PE-repl",   1, _OVERFLOW),
+    "single_kv_group_decode_long_ctx_gqa_cube_sp_pe_tp":              ("Case 2\nCube-SP × PE-repl",     2, _OVERFLOW),
+    "single_kv_group_decode_long_ctx_gqa_cube_repl_pe_sp":            ("Case 3\nCube-Repl × PE-SP",     3, _OVERFLOW),
+    "single_kv_group_decode_long_ctx_gqa_cube_sp_pe_tp_dhead":        ("Case 4\nCube-SP × PE-TP-d_head", 4, _NORMAL),
+    "single_kv_group_decode_long_ctx_gqa_cube_tp_dhead_pe_sp":        ("Case 5\nCube-TP-d_head × PE-SP", 5, _NORMAL),
+    "single_kv_group_decode_long_ctx_gqa_cube_sp_pe_sp":              ("Case 6 ★\nCube-SP × PE-SP",     6, _PARETO),
+}
+
+# Bar fill colour per flag (used by every panel).
+_FLAG_COLOR = {
+    _NORMAL:   "#888888",     # neutral grey
+    _OVERFLOW: "#c0504d",     # red — fails the per-PE HBM budget
+    _PARETO:   "#3b6ea5",     # blue — Pareto-best
 }
 
 
@@ -53,23 +76,26 @@ def _sorted_by_case(rows: list[dict]) -> list[dict]:
     return sorted(rows, key=lambda r: _CASE_INFO[r["panel"]][1])
 
 
+def _bar_colors(rows: list[dict]) -> list[str]:
+    return [_FLAG_COLOR[_CASE_INFO[r["panel"]][2]] for r in rows]
+
+
 def _plot_latency(rows: list[dict]) -> Path:
     rows = _sorted_by_case(rows)
     labels = [_CASE_INFO[r["panel"]][0] for r in rows]
     lat_us = [r["latency_ns"] / 1e3 for r in rows]
-    colors = ["#888", "#888", "#888", "#3b6ea5"]  # Case 4 highlighted
-    fig, ax = plt.subplots(figsize=(8.0, 4.5))
-    bars = ax.bar(labels, lat_us, color=colors, width=0.6)
+    fig, ax = plt.subplots(figsize=(12.0, 4.8))
+    bars = ax.bar(labels, lat_us, color=_bar_colors(rows), width=0.6)
     ax.set_ylabel("end-to-end latency (µs)")
     ax.set_title(
-        "Long-context decode 4-cases — end-to-end latency per case\n"
+        "Long-context decode 6-cases — end-to-end latency per case\n"
         "LLaMA-3.1-70B single-KV-head group (8 cubes × 8 PEs)"
     )
     ax.bar_label(bars, fmt="%.1f", padding=3, fontsize=9)
     ax.grid(axis="y", ls=":", alpha=0.5)
     ax.set_ylim(0, max(lat_us) * 1.15)
     fig.tight_layout()
-    out = _FIG_DIR / "gqa_decode_long_ctx_4cases_latency.png"
+    out = _FIG_DIR / "gqa_decode_long_ctx_6cases_latency.png"
     fig.savefig(out, dpi=150)
     plt.close(fig)
     return out
@@ -83,18 +109,18 @@ def _plot_traffic(rows: list[dict]) -> Path:
     disp = ["IPCQ copy", "DMA read", "DMA write"]
     colors = ["#c0504d", "#9bbb59", "#8064a2"]
     w = 0.25
-    fig, ax = plt.subplots(figsize=(9.0, 4.5))
+    fig, ax = plt.subplots(figsize=(11.0, 4.5))
     for i, (k, d, c) in enumerate(zip(keys, disp, colors)):
         vals = [r["op_log_summary"][k] for r in rows]
         ax.bar([xi + (i - 1) * w for xi in x], vals, width=w, label=d, color=c)
     ax.set_xticks(list(x))
     ax.set_xticklabels(labels, fontsize=9)
     ax.set_ylabel("op count")
-    ax.set_title("Long-context decode 4-cases — op-count breakdown per case")
+    ax.set_title("Long-context decode 6-cases — op-count breakdown per case")
     ax.legend(fontsize=9)
     ax.grid(axis="y", ls=":", alpha=0.5)
     fig.tight_layout()
-    out = _FIG_DIR / "gqa_decode_long_ctx_4cases_traffic.png"
+    out = _FIG_DIR / "gqa_decode_long_ctx_6cases_traffic.png"
     fig.savefig(out, dpi=150)
     plt.close(fig)
     return out
@@ -103,27 +129,41 @@ def _plot_traffic(rows: list[dict]) -> Path:
 def _s_local_per_pe(panel: str, *, S_kv: int, C: int, P: int) -> int:
     """S_local (token count) each PE attends over locally.
 
-    Encodes the cube/pe sharding axes from the panel name:
-      cube_sp_pe_tp (Case 1): S_kv / C       (pe=replicate within cube)
-      cube_repl_pe_tp (Case 2): S_kv         (pe=replicate; only 1 PE works)
-      cube_repl_pe_sp (Case 3): S_kv / P     (pe=row_wise within cube)
-      cube_sp_pe_sp (Case 4): S_kv / (C·P)   (★ 64-way split)
+      cube_repl_pe_tp  (Case 1): S_kv        (no sharding, KV-wise)
+      cube_sp_pe_tp    (Case 2): S_kv / C    (cube splits S_kv, PEs replicate)
+      cube_repl_pe_sp  (Case 3): S_kv / P
+      cube_sp_pe_tp_dhead (Case 4): S_kv / C (cube splits S_kv, PE splits d_head)
+      cube_tp_dhead_pe_sp (Case 5): S_kv / P (cube splits d_head, PE splits S_kv)
+      cube_sp_pe_sp    (Case 6 ★): S_kv / (C·P)
     """
-    S_per_cube = S_kv if "cube_repl" in panel else S_kv // C
-    return S_per_cube // P if "pe_sp" in panel else S_per_cube
+    cube_splits_s = "cube_sp" in panel
+    pe_splits_s   = "pe_sp" in panel
+    S_per_cube = S_kv // C if cube_splits_s else S_kv
+    return S_per_cube // P if pe_splits_s else S_per_cube
+
+
+def _d_head_per_pe(panel: str, *, d_head: int, C: int, P: int) -> int:
+    """d_head dims each PE owns (Cases 4 and 5 shard d_head)."""
+    if "cube_tp_dhead" in panel:    # Case 5: cube shards d_head
+        return d_head // C
+    if "pe_tp_dhead" in panel:      # Case 4: PE shards d_head
+        return d_head // P
+    return d_head                   # Cases 1, 2, 3, 6: full d_head per PE
 
 
 def _active_pe_count(panel: str, *, C: int, P: int) -> int:
     """Number of PEs doing non-idle attention work.
 
-      cube_sp_pe_tp (Case 1): C       (PE 0 of each cube; 7 PEs idle per cube)
-      cube_repl_pe_tp (Case 2): 1     (only PE 0 of CUBE 0)
-      cube_repl_pe_sp (Case 3): C·P   (all PEs busy, but cubes are redundant)
-      cube_sp_pe_sp (Case 4): C·P     (all 64 PEs doing unique work)
+      cube_repl_pe_tp  (Case 1): 1     (PE-TP idle for B=1; only one PE works)
+      cube_sp_pe_tp    (Case 2): C     (PE 0 of each cube; 7 PEs idle per cube)
+      cube_repl_pe_sp  (Case 3): C·P   (all PEs busy, cube-side redundant)
+      cube_sp_pe_tp_dhead (Case 4): C·P (PE shards d_head — all 64 active)
+      cube_tp_dhead_pe_sp (Case 5): C·P (PE shards S_kv — all active)
+      cube_sp_pe_sp    (Case 6 ★): C·P (all 64 PEs doing unique work)
     """
-    if "cube_repl" in panel and "pe_tp" in panel:
+    if "cube_repl" in panel and "pe_tp" in panel and "dhead" not in panel:
         return 1
-    if "cube_sp" in panel and "pe_tp" in panel:
+    if "cube_sp" in panel and "pe_tp" in panel and "dhead" not in panel:
         return C
     return C * P
 
@@ -132,11 +172,12 @@ def _kv_bytes_per_pe(panel: str, *, S_kv: int, h_kv: int,
                      d_head: int, C: int, P: int) -> int:
     """KV bytes a single PE references (K + V, f16, 2 B/elem)."""
     s_local = _s_local_per_pe(panel, S_kv=S_kv, C=C, P=P)
-    return 2 * s_local * h_kv * d_head * 2
+    d_local = _d_head_per_pe(panel, d_head=d_head, C=C, P=P)
+    return 2 * s_local * h_kv * d_local * 2
 
 
 def _plot_memory(rows: list[dict]) -> Path:
-    """Per-PE KV bytes — Case 4 wins (64-way split)."""
+    """Per-PE KV bytes — Case 6 ★ wins (64-way split)."""
     rows = _sorted_by_case(rows)
     labels = [_CASE_INFO[r["panel"]][0] for r in rows]
     mib_per_pe = [
@@ -146,26 +187,25 @@ def _plot_memory(rows: list[dict]) -> Path:
         ) / (1024 * 1024)
         for r in rows
     ]
-    colors = ["#888", "#c0504d", "#888", "#3b6ea5"]  # 4 highlighted, 2 marked red
-    fig, ax = plt.subplots(figsize=(8.0, 4.5))
-    bars = ax.bar(labels, mib_per_pe, color=colors, width=0.6)
+    fig, ax = plt.subplots(figsize=(12.0, 4.8))
+    bars = ax.bar(labels, mib_per_pe, color=_bar_colors(rows), width=0.6)
     ax.set_ylabel("KV bytes per PE (MiB, K + V, f16)")
     ax.set_title(
-        "Long-context decode 4-cases — KV memory per PE\n"
+        "Long-context decode 6-cases — KV memory per PE\n"
         "(one KV-head group; per-layer, per-token state)"
     )
     ax.bar_label(bars, fmt="%.3f", padding=3, fontsize=9)
     ax.grid(axis="y", ls=":", alpha=0.5)
     ax.set_ylim(0, max(mib_per_pe) * 1.15)
     fig.tight_layout()
-    out = _FIG_DIR / "gqa_decode_long_ctx_4cases_memory.png"
+    out = _FIG_DIR / "gqa_decode_long_ctx_6cases_memory.png"
     fig.savefig(out, dpi=150)
     plt.close(fig)
     return out
 
 
 def _plot_parallelism(rows: list[dict]) -> Path:
-    """Total active PE-token compute load — exposes Case 3's redundancy."""
+    """Total active PE-token compute load — exposes redundant-work cases."""
     rows = _sorted_by_case(rows)
     labels = [_CASE_INFO[r["panel"]][0] for r in rows]
     total_work = [
@@ -173,19 +213,19 @@ def _plot_parallelism(rows: list[dict]) -> Path:
         * _s_local_per_pe(r["panel"], S_kv=r["S_kv"], C=r["C"], P=r["P"])
         for r in rows
     ]
-    colors = ["#888", "#888", "#c0504d", "#3b6ea5"]  # 4 highlighted, 3 marked red
-    fig, ax = plt.subplots(figsize=(8.0, 4.5))
-    bars = ax.bar(labels, total_work, color=colors, width=0.6)
+    fig, ax = plt.subplots(figsize=(12.0, 4.8))
+    bars = ax.bar(labels, total_work, color=_bar_colors(rows), width=0.6)
     ax.set_ylabel("active-PE × S_local (PE-tokens; lower ⇒ less wasted work)")
     ax.set_title(
-        "Long-context decode 4-cases — total compute load across active PEs\n"
-        "(Case 3 replicates the full K/V across 8 cubes ⇒ 8× wasted PE-tokens)"
+        "Long-context decode 6-cases — total compute load across active PEs\n"
+        "(Case 3 replicates KV across 8 cubes → 8× wasted PE-tokens; "
+        "Case 6 ★ is fully parallel without replication)"
     )
     ax.bar_label(bars, fmt="%d", padding=3, fontsize=9)
     ax.grid(axis="y", ls=":", alpha=0.5)
     ax.set_ylim(0, max(total_work) * 1.15)
     fig.tight_layout()
-    out = _FIG_DIR / "gqa_decode_long_ctx_4cases_parallelism.png"
+    out = _FIG_DIR / "gqa_decode_long_ctx_6cases_parallelism.png"
     fig.savefig(out, dpi=150)
     plt.close(fig)
     return out

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_kv_sharding_diagram.py

@@ -0,0 +1,324 @@
+"""6-case KV-sharding tensor diagram (the slide-13 PNG export).
+
+Flat 2-D rectangles, one per sharding case, with:
+    Y axis = S_kv (vertical)        — Cube-SP / PE-SP slice it
+    X axis = d_head (horizontal)    — Cube-TP-d_head / PE-TP-d_head slice it
+
+Drops the batch axis entirely (decode: B = 1, T_q = 1). Same case set
+and visual encoding as slide 13 of GQA_full_deck.pptx; matplotlib
+renders it cleanly so the PNG sits next to the other GQA summary
+artifacts in 1H_milestone_output/gqa/long_ctx/.
+"""
+from __future__ import annotations
+
+from pathlib import Path
+
+import matplotlib.patches as mpatches
+import matplotlib.pyplot as plt
+
+_C = 8
+_P = 8
+
+_GROUP_FILLS = [
+    "#A5D8FF", "#B2F2BB", "#FFD8A8", "#FFC9C9",
+    "#D0BFFF", "#99E9F2", "#FCC2D7", "#FFEC99",
+]
+
+_ACC = {
+    "red":    "#E03131",
+    "orange": "#FD7E14",
+    "blue":   "#1C7ED6",
+    "green":  "#37B24D",
+}
+
+# (label, accent, kv, comm, overflow, encoding-flags, axis-spec)
+#   y_split   = 8 horizontal Y bands  (Cube-SP on S_kv)
+#   x_split   = 8 vertical X bands    (Cube-TP-d_head)
+#   pe_y      = 7 fine horizontal dividers within each Y band
+#   pe_x      = 7 fine vertical dividers within each X band
+#   axes      = small annotation under the chip naming the axes
+#               that the cube/PE actually shard, so the reader can
+#               parse Case 5 (where cube colour fills run X instead
+#               of Y, breaking the visual symmetry of the rest).
+_CASES = [
+    dict(label="Case 1\nCube-Repl / PE-repl", accent=_ACC["red"],
+         kv="40 GB", comm="1.2 MB", overflow=True,
+         y_split=False, x_split=False, pe_y=False, pe_x=False,
+         axes="Cube: replicated     PE: replicated"),
+    dict(label="Case 2\nCube-SP / PE-repl",   accent=_ACC["orange"],
+         kv="5 GB",  comm="3.8 MB", overflow=True,
+         y_split=True,  x_split=False, pe_y=False, pe_x=False,
+         axes="Cube → Y (S_kv)     PE: replicated"),
+    dict(label="Case 3\nCube-Repl / PE-SP",   accent=_ACC["orange"],
+         kv="5 GB",  comm="3.8 MB", overflow=True,
+         y_split=False, x_split=False, pe_y=True,  pe_x=False,
+         axes="Cube: replicated     PE → Y (S_kv)"),
+    dict(label="Case 4\nCube-SP / PE-TP-d_head", accent=_ACC["blue"],
+         kv="640 MB", comm="166 MB", overflow=False,
+         y_split=True,  x_split=False, pe_y=False, pe_x=True,
+         axes="Cube → Y (S_kv)     PE → X (d_head)"),
+    dict(label="Case 5\nCube-TP-d_head / PE-SP", accent=_ACC["blue"],
+         kv="640 MB", comm="166 MB", overflow=False,
+         y_split=False, x_split=True,  pe_y=True,  pe_x=False,
+         axes="Cube → X (d_head)     PE → Y (S_kv)"),
+    dict(label="Case 6 ★\nCube-SP / PE-SP",   accent=_ACC["green"],
+         kv="640 MB", comm="6.2 MB", overflow=False,
+         y_split=True,  x_split=False, pe_y=True,  pe_x=False,
+         axes="Cube → Y (S_kv)     PE → Y (S_kv)"),
+]
+
+_OUT_DIR = (
+    Path(__file__).resolve().parents[2]
+    / "src" / "kernbench" / "benches"
+    / "1H_milestone_output" / "gqa" / "long_ctx"
+)
+
+
+def _draw_panel(ax, cfg):
+    """Draw one case's 2-D KV-tensor rectangle into a panel ax."""
+    ax.set_xlim(0, 1)
+    ax.set_ylim(1, 0)        # Y points down (S_kv ↓)
+    ax.set_aspect("auto")
+    ax.set_xticks([])
+    ax.set_yticks([])
+
+    cube_repl = not cfg["y_split"] and not cfg["x_split"]
+    pe_repl   = not cfg["pe_y"]   and not cfg["pe_x"]
+
+    # Cube-level colour fill.
+    if cfg["y_split"] and not cfg["x_split"]:
+        # 8 horizontal Y bands.
+        for c in range(_C):
+            ax.add_patch(mpatches.Rectangle(
+                (0, c / _C), 1, 1 / _C,
+                facecolor=_GROUP_FILLS[c], edgecolor="black", linewidth=0.6))
+            ax.text(0.04, c / _C + 0.5 / _C, f"C{c}",
+                    ha="left", va="center", fontsize=8,
+                    fontweight="bold", color="#333")
+    elif cfg["x_split"] and not cfg["y_split"]:
+        # 8 vertical X bands.
+        for c in range(_C):
+            ax.add_patch(mpatches.Rectangle(
+                (c / _C, 0), 1 / _C, 1,
+                facecolor=_GROUP_FILLS[c], edgecolor="black", linewidth=0.6))
+            ax.text(c / _C + 0.5 / _C, 0.04, f"C{c}",
+                    ha="center", va="top", fontsize=8,
+                    fontweight="bold", color="#333")
+    else:
+        ax.add_patch(mpatches.Rectangle(
+            (0, 0), 1, 1,
+            facecolor="#F5F5F5", edgecolor="black", linewidth=0.8))
+        ax.text(0.5, 0.5, "× 8 cubes\nfull KV",
+                ha="center", va="center",
+                fontsize=10, fontweight="bold",
+                fontstyle="italic", color="#666")
+
+    # PE-level fine dividers — distinguished from cube boundaries by
+    # using a dashed style + slightly stronger contrast. This is what
+    # makes Case 5's PE-SP (horizontal lines across vertical cube
+    # bands) read as "different axis from the cubes" at a glance.
+    if cfg["pe_y"]:
+        outer = _C if cfg["y_split"] else 1
+        band = 1 / outer
+        for o in range(outer):
+            for p in range(1, _P):
+                y = o * band + band * p / _P
+                ax.axhline(y, color="#222", linewidth=0.8,
+                           linestyle=(0, (3, 2)), alpha=0.75)
+    if cfg["pe_x"]:
+        outer = _C if cfg["x_split"] else 1
+        band = 1 / outer
+        for o in range(outer):
+            for p in range(1, _P):
+                x = o * band + band * p / _P
+                ax.axvline(x, color="#222", linewidth=0.8,
+                           linestyle=(0, (3, 2)), alpha=0.75)
+
+    # Heavy outline on top.
+    ax.add_patch(mpatches.Rectangle(
+        (0, 0), 1, 1, facecolor="none",
+        edgecolor="black", linewidth=1.2))
+
+    # Replication badges — small text-only badges in the corners of
+    # the rectangle, no ghost-card stacking (which mis-reads as a
+    # larger enclosing tensor).
+    badges: list[str] = []
+    if cube_repl:
+        badges.append("× 8 cube copies")
+    if pe_repl and (cfg["y_split"] or cfg["x_split"]):
+        # Cube is sharded but PEs in each cube replicate that shard.
+        badges.append("× 8 PEs / cube replicate")
+    elif pe_repl and cube_repl:
+        # Both replicated — PE replication adds to the cube one.
+        badges.append("× 8 PEs / cube replicate")
+    if badges:
+        ax.text(0.98, 0.02, "\n".join(badges),
+                ha="right", va="top", fontsize=7,
+                fontweight="bold", color="#444",
+                fontstyle="italic",
+                bbox=dict(facecolor="white", edgecolor="#888",
+                          boxstyle="round,pad=0.20", linewidth=0.5))
+
+
+def _make_table_png() -> Path:
+    """Slide-14 companion table: per-PE memory + comm for all 6 cases."""
+    headers = ["Case", "Sharding", "KV / PE", "Fit",
+               "Comm/tok\n(analytical)", "Notes"]
+    rows = [
+        ("Case 1",   "Cube-Repl  ·  PE-repl",      "40 GB",  "✗",
+         "1.2 MB",
+         "no sharding —\nfull KV on every PE"),
+        ("Case 2",   "Cube-SP    ·  PE-repl",      "5 GB",   "✗",
+         "3.8 MB",
+         "cube-axis\nsharded only"),
+        ("Case 3",   "Cube-Repl  ·  PE-SP",        "5 GB",   "✗",
+         "3.8 MB",
+         "PE-axis\nsharded only"),
+        ("Case 4",   "Cube-SP    ·  PE-TP-d_head", "640 MB", "✓",
+         "166 MB",
+         "d_head split intra-cube\npartial-score AR ∝ S_kv"),
+        ("Case 5",   "Cube-TP-d_head  ·  PE-SP",   "640 MB", "✓",
+         "166 MB",
+         "d_head split inter-cube\npartial-score AR on UCIe"),
+        ("Case 6 ★", "Cube-SP    ·  PE-SP",        "640 MB", "✓",
+         "6.2 MB",
+         "S_kv split both axes\n(m,ℓ,O) AR only"),
+    ]
+    accents = [_ACC["red"], _ACC["orange"], _ACC["orange"],
+               _ACC["blue"], _ACC["blue"], _ACC["green"]]
+
+    fig, ax = plt.subplots(figsize=(15.0, 5.0))
+    ax.set_axis_off()
+
+    cell_data = [headers] + [list(r) for r in rows]
+    tbl = ax.table(cellText=cell_data,
+                   colWidths=[0.07, 0.20, 0.09, 0.05, 0.14, 0.28],
+                   cellLoc="center", loc="center")
+    tbl.auto_set_font_size(False)
+    tbl.set_fontsize(10.5)
+    tbl.scale(1.0, 2.4)
+
+    n_cols = len(headers)
+    n_rows = len(rows) + 1   # +1 header
+    # Header styling.
+    for ci in range(n_cols):
+        cell = tbl[(0, ci)]
+        cell.set_facecolor("#1F4E79")
+        cell.set_text_props(color="white", weight="bold")
+        cell.set_edgecolor("#1F4E79")
+    # Body styling.
+    for ri, row in enumerate(rows, start=1):
+        is_pareto = row[0].endswith("★")
+        row_fill = "#E8F5E9" if is_pareto else (
+            "white" if ri % 2 == 1 else "#F5F5F7")
+        # Case-name cell uses accent.
+        case_cell = tbl[(ri, 0)]
+        case_cell.set_facecolor(accents[ri - 1])
+        case_cell.set_text_props(color="white", weight="bold")
+        # Remaining cells.
+        for ci in range(1, n_cols):
+            cell = tbl[(ri, ci)]
+            cell.set_facecolor(row_fill)
+            txt_kwargs = {"weight": "bold" if is_pareto else "normal",
+                          "color": "#333"}
+            if ci == 2:                              # KV / PE
+                txt_kwargs["color"] = (
+                    "#C62828" if row[3] == "✗" else "#2E7D32")
+                txt_kwargs["weight"] = "bold"
+            if ci == 3:                              # Fit
+                txt_kwargs["color"] = (
+                    "#C62828" if row[3] == "✗" else "#2E7D32")
+                txt_kwargs["weight"] = "bold"
+            cell.set_text_props(**txt_kwargs)
+        # Last-column (Notes) cells left-aligned for readability.
+        tbl[(ri, n_cols - 1)].get_text().set_ha("left")
+
+    # Force left-align on the Notes header too.
+    tbl[(0, n_cols - 1)].get_text().set_ha("left")
+
+    fig.suptitle(
+        "GQA decode KV-sharding — per-PE memory & communication\n"
+        "(LLaMA 70B GQA single KV-head group  ·  S_kv = 1 M, FP16, "
+        "80 layers)",
+        fontsize=11.5, y=0.94,
+    )
+    out = _OUT_DIR / "gqa_long_ctx_6cases_kv_sharding_table.png"
+    fig.savefig(out, dpi=150, bbox_inches="tight")
+    plt.close(fig)
+    print(f"wrote {out}")
+    return out
+
+
+def main() -> Path:
+    _OUT_DIR.mkdir(parents=True, exist_ok=True)
+    n = len(_CASES)
+    fig = plt.figure(figsize=(20.0, 7.0))
+    # Three rows per column: case chip · axis-spec annotation · rectangle.
+    gs = fig.add_gridspec(3, n,
+                          height_ratios=[0.55, 0.32, 8.5],
+                          hspace=0.05, wspace=0.20,
+                          left=0.04, right=0.99,
+                          top=0.93, bottom=0.06)
+
+    for i, cfg in enumerate(_CASES):
+        # Top: case chip header.
+        ax_chip = fig.add_subplot(gs[0, i])
+        ax_chip.set_xticks([])
+        ax_chip.set_yticks([])
+        for spine in ax_chip.spines.values():
+            spine.set_visible(False)
+        ax_chip.add_patch(mpatches.Rectangle(
+            (0, 0), 1, 1, transform=ax_chip.transAxes,
+            facecolor=cfg["accent"], edgecolor=cfg["accent"]))
+        ax_chip.text(0.5, 0.5, cfg["label"],
+                     ha="center", va="center",
+                     fontsize=10, fontweight="bold",
+                     color="white")
+
+        # Middle: axis-spec annotation — names which axis the cube
+        # shards on and which axis the PE shards on (essential for
+        # parsing Case 5 where the cube colour fills run X instead
+        # of Y, breaking the visual symmetry of the rest).
+        ax_axes = fig.add_subplot(gs[1, i])
+        ax_axes.set_xticks([])
+        ax_axes.set_yticks([])
+        for spine in ax_axes.spines.values():
+            spine.set_visible(False)
+        ax_axes.add_patch(mpatches.Rectangle(
+            (0, 0), 1, 1, transform=ax_axes.transAxes,
+            facecolor="#F5F5F7", edgecolor="#CCCCCC",
+            linewidth=0.6))
+        ax_axes.text(0.5, 0.5, cfg["axes"],
+                     ha="center", va="center",
+                     fontsize=8.5, fontweight="bold",
+                     color="#1F4E79")
+
+        # Bottom: the tensor rectangle.
+        ax = fig.add_subplot(gs[2, i])
+        _draw_panel(ax, cfg)
+        ax.set_xlabel("X : d_head = 128  →",
+                      fontsize=9, fontweight="bold",
+                      fontstyle="italic", color="#1F4E79")
+        ax.set_ylabel("Y : S_kv = 1 M  ↓",
+                      fontsize=9, fontweight="bold",
+                      fontstyle="italic", color="#1F4E79")
+
+    fig.suptitle(
+        "GQA decode KV-tensor sharding — 6 cases  ·  "
+        "LLaMA 70B GQA single KV-head group  ·  "
+        "C = 8 cubes × P = 8 PEs  ·  S_kv = 1 M, FP16, 80 layers",
+        fontsize=12, y=0.99,
+    )
+
+    out = _OUT_DIR / "gqa_long_ctx_6cases_kv_sharding_diagram.png"
+    fig.savefig(out, dpi=150, bbox_inches="tight")
+    plt.close(fig)
+    print(f"wrote {out}")
+
+    # Companion table PNG (slide-14 export).
+    _make_table_png()
+    return out
+
+
+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/verify_case4_dhead_tp.py

@@ -0,0 +1,148 @@
+"""Phase 1a smoke test for the Case 4 d_head-TP decode kernel.
+
+Runs `gqa_attention_decode_long_ctx_cube_sp_pe_tp_dhead_kernel` at small
+S_kv (2K) to verify:
+  1. it imports cleanly,
+  2. it runs without error under the same harness as Case 6,
+  3. captures op_log_summary (gemm/ipcq/dma counts),
+  4. compares vs analytical predictions and vs Case 6 (same memory tier).
+
+Usage:
+  python scripts/verify_case4_dhead_tp.py
+"""
+from __future__ import annotations
+
+import os
+import sys
+from pathlib import Path
+
+from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_sp import (
+    gqa_attention_decode_long_ctx_cube_sp_pe_sp_kernel as _case6_kernel,
+)
+from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_tp_dhead import (
+    gqa_attention_decode_long_ctx_cube_sp_pe_tp_dhead_kernel as _case4_kernel,
+)
+from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_tp_dhead_pe_sp import (
+    gqa_attention_decode_long_ctx_cube_tp_dhead_pe_sp_kernel as _case5_kernel,
+)
+from kernbench.benches.gqa_helpers.shared._gqa_panel_helpers import (
+    _ccl_cfg, _summarize_op_log,
+)
+from kernbench.ccl.sfr_config import configure_sfr_intercube_multisip
+from kernbench.policy.placement.dp import DPPolicy
+from kernbench.runtime_api.bench_runner import run_bench
+from kernbench.runtime_api.types import resolve_device
+from kernbench.sim_engine.engine import GraphEngine
+from kernbench.topology.builder import resolve_topology
+
+# Small smoke-test params (S_kv=2K is enough to exercise tile-loop + AR).
+_PARAMS = dict(C=8, P=8, T_q=1, S_kv=2_048,
+               d_head=128, h_q=8, h_kv=1)
+
+
+def _bench_fn_case4(ctx):
+    p = _PARAMS
+    configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
+    dp_full = DPPolicy(cube="replicate", pe="column_wise",
+                       num_cubes=p["C"], num_pes=p["P"])
+    dp_kv = DPPolicy(cube="row_wise", pe="column_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="q_c4")
+    k = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
+                  dtype="f16", dp=dp_kv, name="k_c4")
+    v = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
+                  dtype="f16", dp=dp_kv, name="v_c4")
+    o = ctx.empty((p["T_q"], p["h_q"] * p["d_head"]),
+                  dtype="f16", dp=dp_full, name="o_c4")
+    ctx.launch("case4_dhead_tp", _case4_kernel,
+               q, k, v, o,
+               p["T_q"], p["S_kv"], p["h_q"], p["h_kv"],
+               p["d_head"], p["C"], p["P"],
+               _auto_dim_remap=False)
+
+
+def _bench_fn_case5(ctx):
+    p = _PARAMS
+    configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
+    dp_q = DPPolicy(cube="column_wise", pe="replicate",
+                    num_cubes=p["C"], num_pes=p["P"])
+    dp_kv = DPPolicy(cube="column_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_q, name="q_c5")
+    k = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
+                  dtype="f16", dp=dp_kv, name="k_c5")
+    v = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
+                  dtype="f16", dp=dp_kv, name="v_c5")
+    o = ctx.empty((p["T_q"], p["h_q"] * p["d_head"]),
+                  dtype="f16", dp=dp_q, name="o_c5")
+    ctx.launch("case5_dhead_tp_inter", _case5_kernel,
+               q, k, v, o,
+               p["T_q"], p["S_kv"], p["h_q"], p["h_kv"],
+               p["d_head"], p["C"], p["P"],
+               _auto_dim_remap=False)
+
+
+def _bench_fn_case6(ctx):
+    p = _PARAMS
+    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="q_c6")
+    k = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
+                  dtype="f16", dp=dp_kv, name="k_c6")
+    v = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
+                  dtype="f16", dp=dp_kv, name="v_c6")
+    o = ctx.empty((p["T_q"], p["h_q"] * p["d_head"]),
+                  dtype="f16", dp=dp_full, name="o_c6")
+    ctx.launch("case6_sp_sp", _case6_kernel,
+               q, k, v, o,
+               p["T_q"], p["S_kv"], p["h_q"], p["h_kv"],
+               p["d_head"], p["C"], p["P"],
+               _auto_dim_remap=False)
+
+
+def main() -> int:
+    topology = os.environ.get("GQA_1H_TOPOLOGY", "topology.yaml")
+    topo = resolve_topology(topology)
+
+    print(f"Smoke params: {_PARAMS}")
+    print()
+
+    for label, bench_fn in (
+        ("Case 4 (Cube-SP × PE-TP d_head)",   _bench_fn_case4),
+        ("Case 5 (Cube-TP d_head × PE-SP)",   _bench_fn_case5),
+        ("Case 6 (Cube-SP × PE-SP S_kv)",     _bench_fn_case6),
+    ):
+        try:
+            res = run_bench(
+                topology=topo, bench_fn=bench_fn,
+                device=resolve_device(None),
+                engine_factory=lambda t, d: GraphEngine(
+                    getattr(t, "topology_obj", t), enable_data=True,
+                ),
+            )
+        except Exception as e:
+            print(f"  {label:<42} FAIL: {type(e).__name__}: {e}")
+            continue
+        if not res.completion.ok:
+            print(f"  {label:<42} ENGINE FAIL: {res.completion}")
+            continue
+        s = _summarize_op_log(res.engine.op_log)
+        lat = (res.engine.op_log[-1].t_end if res.engine.op_log else 0.0)
+        print(f"  {label:<42} "
+              f"gemm={s['gemm_count']:>4}  "
+              f"ipcq={s['ipcq_copy_count']:>4}  "
+              f"dma_r={s['dma_read_count']:>4}  "
+              f"dma_w={s['dma_write_count']:>3}  "
+              f"latency={lat:.1f} ns  "
+              f"(n_ops={len(res.engine.op_log)})")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

src/kernbench/benches/1H_milestone_output/gemm/gemm_sweep.json

@@ -33,7 +33,7 @@
       "bytes_hbm": 6144,
       "arith_intensity": 10.666666666666666,
       "tile_count_expected": 1,
-      "sim_wall_clock_s": 0.112,
+      "sim_wall_clock_s": 0.226,
       "engines": {
         "pe_dma": {
           "occupancy_ns": 79.0,
@@ -100,7 +100,7 @@
       "bytes_hbm": 6144,
       "arith_intensity": 10.666666666666666,
       "tile_count_expected": 1,
-      "sim_wall_clock_s": 0.108,
+      "sim_wall_clock_s": 0.186,
       "engines": {
         "pe_dma": {
           "occupancy_ns": 71.0,
@@ -167,11 +167,11 @@
       "bytes_hbm": 6144,
       "arith_intensity": 10.666666666666666,
       "tile_count_expected": 1,
-      "sim_wall_clock_s": 0.105,
+      "sim_wall_clock_s": 0.181,
       "engines": {
         "pe_dma": {
           "occupancy_ns": 63.0,
-          "wall_ns": 63.0,
+          "wall_ns": 51.0,
           "record_count": 3
         },
         "pe_fetch_store": {
@@ -222,7 +222,7 @@
           "record_count": 0
         }
       },
-      "pe_window_ns": 183.394,
+      "pe_window_ns": 151.394,
       "composite_window_ns": 77.38400000000001
     },
     {
@@ -234,7 +234,7 @@
       "bytes_hbm": 10240,
       "arith_intensity": 12.8,
       "tile_count_expected": 1,
-      "sim_wall_clock_s": 0.169,
+      "sim_wall_clock_s": 0.314,
       "engines": {
         "pe_dma": {
           "occupancy_ns": 79.0,
@@ -301,7 +301,7 @@
       "bytes_hbm": 10240,
       "arith_intensity": 12.8,
       "tile_count_expected": 1,
-      "sim_wall_clock_s": 0.163,
+      "sim_wall_clock_s": 0.191,
       "engines": {
         "pe_dma": {
           "occupancy_ns": 79.0,
@@ -368,11 +368,11 @@
       "bytes_hbm": 10240,
       "arith_intensity": 12.8,
       "tile_count_expected": 1,
-      "sim_wall_clock_s": 0.104,
+      "sim_wall_clock_s": 0.186,
       "engines": {
         "pe_dma": {
-          "occupancy_ns": 79.0,
-          "wall_ns": 79.0,
+          "occupancy_ns": 86.0,
+          "wall_ns": 66.0,
           "record_count": 3
         },
         "pe_fetch_store": {
@@ -423,7 +423,7 @@
           "record_count": 0
         }
       },
-      "pe_window_ns": 191.394,
+      "pe_window_ns": 166.394,
       "composite_window_ns": 77.38400000000001
     },
     {
@@ -435,7 +435,7 @@
       "bytes_hbm": 18432,
       "arith_intensity": 14.222222222222221,
       "tile_count_expected": 2,
-      "sim_wall_clock_s": 0.162,
+      "sim_wall_clock_s": 0.318,
       "engines": {
         "pe_dma": {
           "occupancy_ns": 158.995,
@@ -502,7 +502,7 @@
       "bytes_hbm": 18432,
       "arith_intensity": 14.222222222222221,
       "tile_count_expected": 2,
-      "sim_wall_clock_s": 0.183,
+      "sim_wall_clock_s": 0.188,
       "engines": {
         "pe_dma": {
           "occupancy_ns": 139.995,
@@ -569,11 +569,11 @@
       "bytes_hbm": 18432,
       "arith_intensity": 14.222222222222221,
       "tile_count_expected": 2,
-      "sim_wall_clock_s": 0.106,
+      "sim_wall_clock_s": 0.186,
       "engines": {
         "pe_dma": {
-          "occupancy_ns": 112.0,
-          "wall_ns": 107.0,
+          "occupancy_ns": 134.0,
+          "wall_ns": 98.0,
           "record_count": 3
         },
         "pe_fetch_store": {
@@ -624,7 +624,7 @@
           "record_count": 0
         }
       },
-      "pe_window_ns": 223.77800000000002,
+      "pe_window_ns": 214.77800000000002,
       "composite_window_ns": 93.76800000000003
     },
     {
@@ -636,7 +636,7 @@
       "bytes_hbm": 49152,
       "arith_intensity": 21.333333333333332,
       "tile_count_expected": 8,
-      "sim_wall_clock_s": 0.177,
+      "sim_wall_clock_s": 0.308,
       "engines": {
         "pe_dma": {
           "occupancy_ns": 1798.8600000000001,
@@ -703,7 +703,7 @@
       "bytes_hbm": 49152,
       "arith_intensity": 21.333333333333332,
       "tile_count_expected": 8,
-      "sim_wall_clock_s": 0.22,
+      "sim_wall_clock_s": 0.263,
       "engines": {
         "pe_dma": {
           "occupancy_ns": 819.8600000000001,
@@ -770,11 +770,11 @@
       "bytes_hbm": 49152,
       "arith_intensity": 21.333333333333332,
       "tile_count_expected": 8,
-      "sim_wall_clock_s": 0.112,
+      "sim_wall_clock_s": 0.183,
       "engines": {
         "pe_dma": {
-          "occupancy_ns": 283.0,
-          "wall_ns": 238.07200000000012,
+          "occupancy_ns": 305.0,
+          "wall_ns": 229.07200000000012,
           "record_count": 6
         },
         "pe_fetch_store": {
@@ -825,7 +825,7 @@
           "record_count": 0
         }
       },
-      "pe_window_ns": 418.0820000000001,
+      "pe_window_ns": 409.0820000000001,
       "composite_window_ns": 192.07200000000012
     },
     {
@@ -837,7 +837,7 @@
       "bytes_hbm": 395264,
       "arith_intensity": 15.917098445595855,
       "tile_count_expected": 48,
-      "sim_wall_clock_s": 0.2,
+      "sim_wall_clock_s": 0.293,
       "engines": {
         "pe_dma": {
           "occupancy_ns": 56319.35999999987,
@@ -904,7 +904,7 @@
       "bytes_hbm": 395264,
       "arith_intensity": 15.917098445595855,
       "tile_count_expected": 48,
-      "sim_wall_clock_s": 0.237,
+      "sim_wall_clock_s": 0.292,
       "engines": {
         "pe_dma": {
           "occupancy_ns": 20236.359999999866,
@@ -912,13 +912,13 @@
           "record_count": 50
         },
         "pe_fetch_store": {
-          "occupancy_ns": 1543.9999999999998,
+          "occupancy_ns": 1543.9999999999995,
           "wall_ns": 791.9999999999998,
           "record_count": 49
         },
         "pe_gemm": {
-          "occupancy_ns": 1219.584000000017,
-          "wall_ns": 786.4320000000007,
+          "occupancy_ns": 1219.5840000000169,
+          "wall_ns": 786.4320000000005,
           "record_count": 48
         },
         "pe_math": {
@@ -939,7 +939,7 @@
           "record_count": 1
         },
         "FETCH": {
-          "occupancy_ns": 1535.9999999999998,
+          "occupancy_ns": 1535.9999999999995,
           "wall_ns": 783.9999999999998,
           "record_count": 48
         },
@@ -949,8 +949,8 @@
           "record_count": 1
         },
         "GEMM": {
-          "occupancy_ns": 1219.584000000017,
-          "wall_ns": 786.4320000000007,
+          "occupancy_ns": 1219.5840000000169,
+          "wall_ns": 786.4320000000005,
           "record_count": 48
         },
         "MATH": {
@@ -971,10 +971,10 @@
       "bytes_hbm": 395264,
       "arith_intensity": 15.917098445595855,
       "tile_count_expected": 48,
-      "sim_wall_clock_s": 0.208,
+      "sim_wall_clock_s": 0.237,
       "engines": {
         "pe_dma": {
-          "occupancy_ns": 2310.0,
+          "occupancy_ns": 2342.0,
           "wall_ns": 1569.9999999999998,
           "record_count": 3
         },
@@ -1038,7 +1038,7 @@
       "bytes_hbm": 36864,
       "arith_intensity": 7.111111111111111,
       "tile_count_expected": 8,
-      "sim_wall_clock_s": 0.089,
+      "sim_wall_clock_s": 0.259,
       "engines": {
         "pe_dma": {
           "occupancy_ns": 1798.8600000000001,
@@ -1105,7 +1105,7 @@
       "bytes_hbm": 36864,
       "arith_intensity": 7.111111111111111,
       "tile_count_expected": 8,
-      "sim_wall_clock_s": 0.155,
+      "sim_wall_clock_s": 0.313,
       "engines": {
         "pe_dma": {
           "occupancy_ns": 795.8600000000001,
@@ -1172,11 +1172,11 @@
       "bytes_hbm": 36864,
       "arith_intensity": 7.111111111111111,
       "tile_count_expected": 8,
-      "sim_wall_clock_s": 0.234,
+      "sim_wall_clock_s": 0.32,
       "engines": {
         "pe_dma": {
           "occupancy_ns": 258.0,
-          "wall_ns": 218.07200000000012,
+          "wall_ns": 206.07200000000012,
           "record_count": 6
         },
         "pe_fetch_store": {
@@ -1227,7 +1227,7 @@
           "record_count": 0
         }
       },
-      "pe_window_ns": 418.0820000000001,
+      "pe_window_ns": 386.0820000000001,
       "composite_window_ns": 192.07200000000012
     },
     {
@@ -1239,7 +1239,7 @@
       "bytes_hbm": 36864,
       "arith_intensity": 7.111111111111111,
       "tile_count_expected": 16,
-      "sim_wall_clock_s": 0.119,
+      "sim_wall_clock_s": 0.304,
       "engines": {
         "pe_dma": {
           "occupancy_ns": 7854.6320000000005,
@@ -1306,7 +1306,7 @@
       "bytes_hbm": 36864,
       "arith_intensity": 7.111111111111111,
       "tile_count_expected": 16,
-      "sim_wall_clock_s": 0.183,
+      "sim_wall_clock_s": 0.165,
       "engines": {
         "pe_dma": {
           "occupancy_ns": 3779.6320000000005,
@@ -1373,11 +1373,11 @@
       "bytes_hbm": 36864,
       "arith_intensity": 7.111111111111111,
       "tile_count_expected": 16,
-      "sim_wall_clock_s": 0.179,
+      "sim_wall_clock_s": 0.236,
       "engines": {
         "pe_dma": {
           "occupancy_ns": 858.0,
-          "wall_ns": 183.0,
+          "wall_ns": 171.0,
           "record_count": 18
         },
         "pe_fetch_store": {
@@ -1428,7 +1428,7 @@
           "record_count": 0
         }
       },
-      "pe_window_ns": 511.01,
+      "pe_window_ns": 479.01,
       "composite_window_ns": 405.0
     },
     {
@@ -1440,7 +1440,7 @@
       "bytes_hbm": 1572864,
       "arith_intensity": 170.66666666666666,
       "tile_count_expected": 2048,
-      "sim_wall_clock_s": 1.241,
+      "sim_wall_clock_s": 2.164,
       "engines": {
         "pe_dma": {
           "occupancy_ns": 104867178.97599977,
@@ -1507,7 +1507,7 @@
       "bytes_hbm": 1572864,
       "arith_intensity": 170.66666666666666,
       "tile_count_expected": 2048,
-      "sim_wall_clock_s": 0.742,
+      "sim_wall_clock_s": 1.398,
       "engines": {
         "pe_dma": {
           "occupancy_ns": 37760375.97599977,
@@ -1520,7 +1520,7 @@
           "record_count": 2304
         },
         "pe_gemm": {
-          "occupancy_ns": 36417.40799998891,
+          "occupancy_ns": 36417.407999988936,
           "wall_ns": 33554.43199999761,
           "record_count": 2048
         },
@@ -1547,12 +1547,12 @@
           "record_count": 2048
         },
         "STORE": {
-          "occupancy_ns": 130151.58400000146,
+          "occupancy_ns": 130151.58400000144,
           "wall_ns": 31914.784000000003,
           "record_count": 256
         },
         "GEMM": {
-          "occupancy_ns": 36417.40799998891,
+          "occupancy_ns": 36417.407999988936,
           "wall_ns": 33554.43199999761,
           "record_count": 2048
         },
@@ -1574,10 +1574,10 @@
       "bytes_hbm": 1572864,
       "arith_intensity": 170.66666666666666,
       "tile_count_expected": 2048,
-      "sim_wall_clock_s": 0.406,
+      "sim_wall_clock_s": 0.477,
       "engines": {
         "pe_dma": {
-          "occupancy_ns": 142065.0,
+          "occupancy_ns": 142097.0,
           "wall_ns": 6170.0,
           "record_count": 258
         },
@@ -1587,8 +1587,8 @@
           "record_count": 2304
         },
         "pe_gemm": {
-          "occupancy_ns": 838467.5839981049,
-          "wall_ns": 33554.43199999738,
+          "occupancy_ns": 838467.5839981119,
+          "wall_ns": 33554.431999997396,
           "record_count": 2048
         },
         "pe_math": {
@@ -1619,8 +1619,8 @@
           "record_count": 256
         },
         "GEMM": {
-          "occupancy_ns": 838467.5839981049,
-          "wall_ns": 33554.43199999738,
+          "occupancy_ns": 838467.5839981119,
+          "wall_ns": 33554.431999997396,
           "record_count": 2048
         },
         "MATH": {

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

@@ -0,0 +1,57 @@
+{
+  "S_kv_measured": 65536,
+  "S_kv_headline": 1048576,
+  "n_layers": 80,
+  "num_pes": 64,
+  "wo_ffn_per_token_bytes": 1310720,
+  "cases": {
+    "1": {
+      "label": "Case 1 (Cube-Repl x PE-repl)",
+      "total_ipcq_bytes_one_layer": 0,
+      "per_pe_partial_score_bytes_one_layer": 0,
+      "per_pe_mlo_bytes_one_layer": 0,
+      "per_pe_attn_bytes_per_token_at_1M": 0,
+      "per_pe_total_bytes_per_token_at_1M": 1310720
+    },
+    "2": {
+      "label": "Case 2 (Cube-SP   x PE-repl)",
+      "total_ipcq_bytes_one_layer": 14560,
+      "per_pe_partial_score_bytes_one_layer": 0,
+      "per_pe_mlo_bytes_one_layer": 227,
+      "per_pe_attn_bytes_per_token_at_1M": 18160,
+      "per_pe_total_bytes_per_token_at_1M": 1328880
+    },
+    "3": {
+      "label": "Case 3 (Cube-Repl x PE-SP)",
+      "total_ipcq_bytes_one_layer": 116480,
+      "per_pe_partial_score_bytes_one_layer": 0,
+      "per_pe_mlo_bytes_one_layer": 1820,
+      "per_pe_attn_bytes_per_token_at_1M": 145600,
+      "per_pe_total_bytes_per_token_at_1M": 1456320
+    },
+    "4": {
+      "label": "Case 4 (Cube-SP   x PE-TP d_head)",
+      "total_ipcq_bytes_one_layer": 14696192,
+      "per_pe_partial_score_bytes_one_layer": 131072,
+      "per_pe_mlo_bytes_one_layer": 98556,
+      "per_pe_attn_bytes_per_token_at_1M": 175656640,
+      "per_pe_total_bytes_per_token_at_1M": 176967360
+    },
+    "5": {
+      "label": "Case 5 (Cube-TP d_head x PE-SP)",
+      "total_ipcq_bytes_one_layer": 14696192,
+      "per_pe_partial_score_bytes_one_layer": 131072,
+      "per_pe_mlo_bytes_one_layer": 98556,
+      "per_pe_attn_bytes_per_token_at_1M": 175656640,
+      "per_pe_total_bytes_per_token_at_1M": 176967360
+    },
+    "6": {
+      "label": "Case 6 (Cube-SP   x PE-SP)  [*]",
+      "total_ipcq_bytes_one_layer": 131040,
+      "per_pe_partial_score_bytes_one_layer": 0,
+      "per_pe_mlo_bytes_one_layer": 2047,
+      "per_pe_attn_bytes_per_token_at_1M": 163760,
+      "per_pe_total_bytes_per_token_at_1M": 1474480
+    }
+  }
+}