topology: set flit_bytes=4096 (was implicit 256 default)

Wire per-flit loop in engine._wire (ADR-0033 Phase 2c) iterates
nbytes/flit_bytes times per hop, firing ~2 SimPy events per flit.
At long-context decode, K/V tile loads are hundreds of KB per shard,
producing thousands of flit-iterations per hop. The default 256-byte
flit size made the wire loop the dominant wall-clock cost.

Raising flit_bytes to 4096 (still a multiple of hbm_ctrl.burst_bytes=256
per ADR-0033 D1) cuts flit iterations 16x with no meaningful change to
modeled latency. Wormhole pipelining and the available_at BW-share
chain in _wire behave identically at any flit size.

Measured impact on Case 4 (Cube-SP x PE-SP) decode, C=8, P=8, T_q=1,
LLaMA-3.1-70B single-KV-group, enable_data=False:

  S_kv     flit=256 wall    flit=4096 wall    speedup    latency drift
   128K       233.71 s          18.16 s        12.9x     +0.05% (0.25us)
   256K       485.75 s          25.31 s        19.2x     +0.01% (0.09us)
   512K      1045.64 s          48.45 s        21.6x     +0.01% (0.17us)
    1M      ~2250 s (proj)     104.75 s        ~21x      within extrap

Total sweep at flit=4096: ~3.3 min. Baseline would have been ~65 min.
1M decode iteration is now practical.

op_log records are byte-identical (same counts, same components); the
tiny latency drift (order of a few ns) comes from the last flit's
env.timeout landing on a slightly different granularity when flits are
larger. Well within noise for any test that asserts on structural
invariants (dma_write cubes, ipcq_copy counts).

Verified:
- tests/attention/test_milestone_gqa_decode_long_ctx_4cases.py: 18/18 pass
  (also 7x faster: 174s -> 24.83s per full run)

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

@@ -28,7 +28,7 @@
 \date{
 \small
 AGI Computing Lab, System Technology Group\\
-2026 H1 Report
+2026 Q1-Q3 Report
 }
 
 \begin{document}
@@ -40,8 +40,10 @@ AGI Computing Lab, System Technology Group\\
 \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/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/05-gqa.tex

@@ -111,29 +111,220 @@ redundant compute each one induces.
 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
+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.
 
-% 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/
+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).
 
-% TODO: prefill performance figure (latency, stage breakdown).
-% TODO: decode performance figure (latency, stage breakdown).
-%       Bench output for short_ctx to be generated.
+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}
@@ -205,22 +396,17 @@ absorb at link speed.}
 \label{fig:gqa-6cases-traffic}
 \end{figure}
 
-The measurements select the right placement for this regime. The fastest
-raw latency (Case~3, \SI{17.8}{\micro\second}) comes from replicating the
-full KV cache into every CUBE---which is exactly the placement the
-long-context memory budget forbids, and which the parallelism panel shows
-wastes 8$\times$ the compute. Restricting attention to the placements that
-fit production-context memory (the 64-way splits, Cases~4--6), the choice
-is Case~6~$\star$: it is the fastest of the three
-(\SI{30.6}{\micro\second}), and the op-count panel shows why---its
-softmax-state-only reduction charges 189 IPCQ copies and one DMA write
-against the 280 copies and 8 DMA writes the $d_{\text{head}}$-split
-Cases~4--5 pay for their partial-score all-reduce. For long-context
-decode, then, the appropriate data placement is the both-axes sequence
-shard (Case~6): it is the cheapest-communicating member of the only
-memory-feasible family, and the cross-PE softmax reduction it does pay is
-precisely the traffic the communication-side codesign of this report is
-built to move quickly.
+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}
@@ -284,9 +470,7 @@ 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 \emph{at the 64-way production scale} (a
-single-rank caveat follows, Figure~\ref{fig:gqa-decode-stream}); it is a
-\emph{CPU-issue}
+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
@@ -297,49 +481,8 @@ 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{Isolating the rank: the masked streaming win.} That neutrality
-is a property of the \emph{full 64-way} critical path, not of the local
-attention: at production scale the inter-CUBE $(m,\ell,O)$ reduce tail and
-shared-HBM contention set the wall clock, so a faster local attention does
-not surface. Stripping those away---a single rank, no cross-CUBE reduce,
-swept over the per-rank context $S_{kv}$ (so $S_{kv}{=}16$\,K here is the
-per-PE load of a 1M-token, 64-way-sharded decode)---exposes the local
-attention directly (Figure~\ref{fig:gqa-decode-stream}), and the composite
-\emph{does} win, by \SI{25}{}--\SI{28}{\percent}. The reason is the
-memory-bound mirror of prefill: its scheduler-streamed concurrent per-tile
-DMAs keep the HBM pipeline full and reach \SI{233}{\giga\byte\per\second}
----\SI{91}{\percent} of the per-rank \SI{256}{\giga\byte\per\second}
-roofline---whereas the primitive kernel's blocking \textsf{tl.dot}
-serializes one tile DMA at a time and plateaus at
-\SI{166}{\giga\byte\per\second}. So even for memory-bound decode the
-composite is not \emph{only} a CPU-issue optimization---it also extracts
-bandwidth---but that latency benefit materializes only when the local
-attention is on the critical path, which at 64-way production scale it is
-not.
-
-\begin{figure}[t]
-\centering
-\includegraphics[width=\linewidth]{gqa_decode_streaming.png}
-\caption{Single-rank memory-bound decode ($T_q{=}1$, $M{=}8$), three
-command forms, swept over per-rank context. \emph{Left:} end-to-end
-latency---the composite forms run \SI{25}{}--\SI{28}{\percent} below the
-primitive, a gap that widens with context. \emph{Right:} achieved HBM
-bandwidth against the per-rank \SI{256}{\giga\byte\per\second} roofline.
-The primitive's blocking load$\rightarrow$dot serializes the KV stream and
-plateaus at \SI{166}{\giga\byte\per\second}; the composite forms pipeline
-concurrent per-tile DMAs through the scheduler and reach
-\SI{233}{\giga\byte\per\second}. This is the memory-bound mirror of the
-prefill result (Figure~\ref{fig:gqa-prefill-cb}): there the composite
-approaches the MAC roofline, here the bandwidth roofline. Capped at
-$16$\,K---the plain composite materializes the full $(M,S_{kv})$ scores in
-TCM, so beyond that only the \textsf{softmax\_merge} recipe, which tiles
-the softmax, stays within scratch.}
-\label{fig:gqa-decode-stream}
-\end{figure}
-
-\paragraph{The compute-bound mirror: prefill.} Decode's \emph{production}
-verdict---command form is latency-neutral at 64-way scale---is a property
-of its regime, not of the
+\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
@@ -387,24 +530,24 @@ 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.
 
-\subsection{Comprehensive Analysis}
+\subsection{Summary}
 \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.
+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. What this implies for hardware investment
+across the report as a whole is taken up in the discussion.
 
 % TODO: cross-regime DP (data parallelism) applicability:
 %   - Does Case-4 long-context placement compose with batch-level DP

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

scripts/paper/paper_plot_gqa_decode_composite_breakdown.py

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

scripts/paper/paper_plot_gqa_decode_long_ctx_composite.py

@@ -45,11 +45,12 @@ _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", "composite", "composite_extended")
+_ORDER = ("primitive_tiled", "composite", "composite_extended")
 
 
 def _load() -> dict:

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

@@ -1,124 +0,0 @@
-"""Comparative figure for the memory-bound decode-streaming composite study.
-
-Reads sweep_decode_streaming.json (emitted by milestone-1h-gqa, sweep
-``decode_streaming``) and writes one two-panel PNG:
-
-  gqa_decode_streaming.png
-    Left  — end-to-end single-rank decode latency (µs) vs per-rank context.
-    Right — achieved HBM bandwidth (GB/s) vs context, against the
-            256 GB/s per-rank roofline.
-
-The memory-bound mirror of the compute-bound prefill figure. With T_q=1
-the GEMMs are skinny (M=8) and the kernel is bound by streaming the KV
-cache. Isolating a single rank (no inter-CUBE reduce) reveals what the
-64-way Case-6 decode masks: the composite command still wins, not by
-feeding the MAC array but by keeping the DMA pipeline full — its
-scheduler-streamed concurrent tile DMAs extract ~230 GB/s (near the
-256 GB/s roofline) while the primitive kernel's blocking tl.dot serializes
-one tile DMA at a time and plateaus at ~166 GB/s. That bandwidth gap is a
-~25-28 % latency win that grows nowhere near prefill's compute-bound
-margin but is decidedly not zero.
-
-Run (after the bench):
-    GQA_1H_RUN=1 GQA_1H_SWEEPS=decode_streaming python -m kernbench.cli.main \\
-        run --bench milestone-1h-gqa --topology topology.yaml
-    python scripts/paper/paper_plot_gqa_decode_streaming.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_streaming.json"
-_PAPER_FIG_DIR = (
-    _REPO_ROOT / "docs" / "report" / "1H-codesign-paper" / "figures"
-)
-
-# Per-rank HBM roofline: 8 pseudo-channels × 32 GB/s (topology.yaml
-# hbm_ctrl.num_pcs / pc_bw_gbs; = pe_dma_to_noc_bw_gbs).
-_PEAK_HBM_GBS = 256.0
-
-_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["s_kv"], 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["s_kv_points"]
-
-    fig, (ax_lat, ax_bw) = 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, bw = _series(rows, v, "achieved_bw_gbs")
-        ax_bw.plot(xs, bw, marker=marker, color=color, label=label, lw=2)
-
-    for ax in (ax_lat, ax_bw):
-        ax.set_xscale("log", base=2)
-        ax.set_xticks(ctxs)
-        ax.set_xticklabels([_ctx_label(c) for c in ctxs])
-        ax.set_xlabel(r"per-rank context length $S_{kv}$ ($T_q{=}1$)")
-        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("Single-rank memory-bound decode latency per command form")
-
-    ax_bw.set_ylabel("achieved HBM bandwidth (GB/s)")
-    ax_bw.set_title(
-        "HBM bandwidth — composite keeps the DMA pipe full; primitive plateaus"
-    )
-    ax_bw.axhline(_PEAK_HBM_GBS, color="#888", ls="--", lw=1, alpha=0.7)
-    ax_bw.text(ctxs[0], _PEAK_HBM_GBS - 8, "256 GB/s roofline",
-               fontsize=8, color="#555", va="top")
-    ax_bw.set_ylim(0, _PEAK_HBM_GBS * 1.08)
-
-    fig.suptitle(
-        "Memory-bound decode streaming — use of composite commands\n"
-        "single-rank, GQA single-KV-head group ($h_q{=}8$, $d_{\\text{head}}"
-        "{=}128$); $M{=}8$ skinny, KV-streaming-bound",
-        fontsize=11,
-    )
-    fig.tight_layout(rect=(0, 0, 1, 0.92))
-
-    out = _FIG_DIR / "gqa_decode_streaming.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()

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

@@ -1,7 +1,7 @@
 {
   "version": 2,
   "variants": [
-    "primitive",
+    "primitive_tiled",
     "composite",
     "composite_extended"
   ],
@@ -21,7 +21,7 @@
   ],
   "rows": [
     {
-      "variant": "primitive",
+      "variant": "primitive_tiled",
       "S_kv": 8192,
       "C": 8,
       "P": 8,
@@ -29,9 +29,9 @@
       "d_head": 128,
       "h_q": 8,
       "h_kv": 1,
-      "pe_cpu_cmd_count": 96,
-      "pe_cpu_dispatch_cycles": 930,
-      "latency_ns": 30646.00300000079
+      "pe_cpu_cmd_count": 414,
+      "pe_cpu_dispatch_cycles": 3918,
+      "latency_ns": 63285.01999999998
     },
     {
       "variant": "composite",
@@ -60,7 +60,7 @@
       "latency_ns": 30483.359500000362
     },
     {
-      "variant": "primitive",
+      "variant": "primitive_tiled",
       "S_kv": 65536,
       "C": 8,
       "P": 8,
@@ -68,9 +68,9 @@
       "d_head": 128,
       "h_q": 8,
       "h_kv": 1,
-      "pe_cpu_cmd_count": 96,
-      "pe_cpu_dispatch_cycles": 930,
-      "latency_ns": 231579.37900000165
+      "pe_cpu_cmd_count": 2654,
+      "pe_cpu_dispatch_cycles": 24974,
+      "latency_ns": 491824.02999999997
     },
     {
       "variant": "composite",
@@ -99,7 +99,7 @@
       "latency_ns": 231211.1740000015
     },
     {
-      "variant": "primitive",
+      "variant": "primitive_tiled",
       "S_kv": 131072,
       "C": 8,
       "P": 8,
@@ -107,9 +107,9 @@
       "d_head": 128,
       "h_q": 8,
       "h_kv": 1,
-      "pe_cpu_cmd_count": 118,
-      "pe_cpu_dispatch_cycles": 1145,
-      "latency_ns": 461119.51900000183
+      "pe_cpu_cmd_count": 5235,
+      "pe_cpu_dispatch_cycles": 49242,
+      "latency_ns": 959456.0549999998
     },
     {
       "variant": "composite",
@@ -138,7 +138,7 @@
       "latency_ns": 460552.9805000025
     },
     {
-      "variant": "primitive",
+      "variant": "primitive_tiled",
       "S_kv": 262144,
       "C": 8,
       "P": 8,
@@ -146,8 +146,8 @@
       "d_head": 128,
       "h_q": 8,
       "h_kv": 1,
-      "pe_cpu_cmd_count": 162,
-      "pe_cpu_dispatch_cycles": 1575,
+      "pe_cpu_cmd_count": 10397,
+      "pe_cpu_dispatch_cycles": 97778,
       "latency_ns": null
     },
     {
@@ -177,7 +177,7 @@
       "latency_ns": null
     },
     {
-      "variant": "primitive",
+      "variant": "primitive_tiled",
       "S_kv": 524288,
       "C": 8,
       "P": 8,
@@ -185,8 +185,8 @@
       "d_head": 128,
       "h_q": 8,
       "h_kv": 1,
-      "pe_cpu_cmd_count": 250,
-      "pe_cpu_dispatch_cycles": 2435,
+      "pe_cpu_cmd_count": 20721,
+      "pe_cpu_dispatch_cycles": 194850,
       "latency_ns": null
     },
     {
@@ -216,7 +216,7 @@
       "latency_ns": null
     },
     {
-      "variant": "primitive",
+      "variant": "primitive_tiled",
       "S_kv": 1048576,
       "C": 8,
       "P": 8,
@@ -224,8 +224,8 @@
       "d_head": 128,
       "h_q": 8,
       "h_kv": 1,
-      "pe_cpu_cmd_count": 426,
-      "pe_cpu_dispatch_cycles": 4155,
+      "pe_cpu_cmd_count": 41369,
+      "pe_cpu_dispatch_cycles": 388994,
       "latency_ns": null
     },
     {

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

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

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

@@ -1,172 +0,0 @@
-{
-  "version": 1,
-  "variants": [
-    "primitive",
-    "composite",
-    "composite_extended"
-  ],
-  "s_kv_points": [
-    2048,
-    4096,
-    8192,
-    16384
-  ],
-  "rows": [
-    {
-      "variant": "primitive",
-      "s_kv": 2048,
-      "M": 8,
-      "latency_ns": 6188.437999999816,
-      "gemm_busy_ns": 1048.576000000001,
-      "dma_busy_ns": 6322.0,
-      "kv_bytes": 1048576.0,
-      "achieved_bw_gbs": 169.44114169036374,
-      "achieved_tflops": 1.3555291335229098,
-      "mac_util": 0.16944114169036373,
-      "dma_occupancy": 1.021582505957107
-    },
-    {
-      "variant": "composite",
-      "s_kv": 2048,
-      "M": 8,
-      "latency_ns": 4836.115999999989,
-      "gemm_busy_ns": 6629.632000000123,
-      "dma_busy_ns": 267480.720000013,
-      "kv_bytes": 1048576.0,
-      "achieved_bw_gbs": 216.82192900253062,
-      "achieved_tflops": 1.734575432020245,
-      "mac_util": 0.2168219290025306,
-      "dma_occupancy": 55.30899589671001
-    },
-    {
-      "variant": "composite_extended",
-      "s_kv": 2048,
-      "M": 8,
-      "latency_ns": 4737.751999999986,
-      "gemm_busy_ns": 8481.408000000116,
-      "dma_busy_ns": 251398.7400000122,
-      "kv_bytes": 1048576.0,
-      "achieved_bw_gbs": 221.32353065335693,
-      "achieved_tflops": 1.7705882452268553,
-      "mac_util": 0.22132353065335691,
-      "dma_occupancy": 53.06287454472352
-    },
-    {
-      "variant": "primitive",
-      "s_kv": 4096,
-      "M": 8,
-      "latency_ns": 12510.006000000496,
-      "gemm_busy_ns": 2097.152000000002,
-      "dma_busy_ns": 12602.0,
-      "kv_bytes": 2097152.0,
-      "achieved_bw_gbs": 167.6379691584414,
-      "achieved_tflops": 1.3411037532675312,
-      "mac_util": 0.1676379691584414,
-      "dma_occupancy": 1.0073536335633653
-    },
-    {
-      "variant": "composite",
-      "s_kv": 4096,
-      "M": 8,
-      "latency_ns": 9360.883999999554,
-      "gemm_busy_ns": 19520.799999939867,
-      "dma_busy_ns": 1059043.5999999335,
-      "kv_bytes": 2097152.0,
-      "achieved_bw_gbs": 224.03354213128802,
-      "achieved_tflops": 1.792268337050304,
-      "mac_util": 0.224033542131288,
-      "dma_occupancy": 113.13499878857424
-    },
-    {
-      "variant": "composite_extended",
-      "s_kv": 4096,
-      "M": 8,
-      "latency_ns": 9198.135999999562,
-      "gemm_busy_ns": 39531.583999941715,
-      "dma_busy_ns": 1026582.7399999356,
-      "kv_bytes": 2097152.0,
-      "achieved_bw_gbs": 227.99749862364504,
-      "achieved_tflops": 1.8239799889891604,
-      "mac_util": 0.22799749862364505,
-      "dma_occupancy": 111.60769312390951
-    },
-    {
-      "variant": "primitive",
-      "s_kv": 8192,
-      "M": 8,
-      "latency_ns": 25153.141999997388,
-      "gemm_busy_ns": 4194.304000000004,
-      "dma_busy_ns": 25162.0,
-      "kv_bytes": 4194304.0,
-      "achieved_bw_gbs": 166.7506985807354,
-      "achieved_tflops": 1.334005588645883,
-      "mac_util": 0.16675069858073538,
-      "dma_occupancy": 1.0003521627637062
-    },
-    {
-      "variant": "composite",
-      "s_kv": 8192,
-      "M": 8,
-      "latency_ns": 18419.187999999034,
-      "gemm_busy_ns": 64177.11999976149,
-      "dma_busy_ns": 4214541.840000685,
-      "kv_bytes": 4194304.0,
-      "achieved_bw_gbs": 227.713838416776,
-      "achieved_tflops": 1.821710707334208,
-      "mac_util": 0.227713838416776,
-      "dma_occupancy": 228.81257523409317
-    },
-    {
-      "variant": "composite_extended",
-      "s_kv": 8192,
-      "M": 8,
-      "latency_ns": 18128.439999999027,
-      "gemm_busy_ns": 169650.68799976518,
-      "dma_busy_ns": 4149323.2200006745,
-      "kv_bytes": 4194304.0,
-      "achieved_bw_gbs": 231.365964197704,
-      "achieved_tflops": 1.850927713581632,
-      "mac_util": 0.231365964197704,
-      "dma_occupancy": 228.88473691067168
-    },
-    {
-      "variant": "primitive",
-      "s_kv": 16384,
-      "M": 8,
-      "latency_ns": 50439.414000008954,
-      "gemm_busy_ns": 8388.607999999076,
-      "dma_busy_ns": 50282.0,
-      "kv_bytes": 8388608.0,
-      "achieved_bw_gbs": 166.31057609032712,
-      "achieved_tflops": 1.330484608722617,
-      "mac_util": 0.16631057609032712,
-      "dma_occupancy": 0.9968791469304357
-    },
-    {
-      "variant": "composite",
-      "s_kv": 16384,
-      "M": 8,
-      "latency_ns": 36535.79600000568,
-      "gemm_busy_ns": 228978.46400288073,
-      "dma_busy_ns": 16815028.239995122,
-      "kv_bytes": 8388608.0,
-      "achieved_bw_gbs": 229.59970545047648,
-      "achieved_tflops": 1.8367976436038118,
-      "mac_util": 0.22959970545047648,
-      "dma_occupancy": 460.2343477063565
-    },
-    {
-      "variant": "composite_extended",
-      "s_kv": 16384,
-      "M": 8,
-      "latency_ns": 35989.048000005685,
-      "gemm_busy_ns": 701990.3680028584,
-      "dma_busy_ns": 16684294.09999516,
-      "kv_bytes": 8388608.0,
-      "achieved_bw_gbs": 233.0877993771515,
-      "achieved_tflops": 1.864702395017212,
-      "mac_util": 0.2330877993771515,
-      "dma_occupancy": 463.5936493789034
-    }
-  ]
-}

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

@@ -37,9 +37,9 @@ Topology / SFR:
 from __future__ import annotations
 
 from kernbench.benches.gqa_helpers.long_ctx._gqa_mlo_reduce import (
-    _ROOT_CUBE,
     _merge_running,
     reduce_mlo,
+    root_cube_for,
 )
 
 TILE_S_KV = 1024  # ADR-0063 §A.2 S_kv-axis tile sweep (per-tile width).
@@ -108,9 +108,9 @@ def gqa_attention_decode_long_ctx_cube_sp_pe_sp_kernel(
             tl.copy_to(O_local, O_new)
 
     # ── Two-level (m, ℓ, O) reduce-to-root (shared helper) ──
-    reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, P, tl=tl)
+    reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, C, P, tl=tl)
 
-    # ── Final normalise + store (root only: PE 0 of CUBE 6) ──
-    if pe_id == 0 and cube_id == _ROOT_CUBE:
+    # ── Final normalise + store (root only: PE 0 of the C-derived root cube) ──
+    if pe_id == 0 and cube_id == root_cube_for(C):
         O_final = O_local / l_local
         tl.store(o_ptr, O_final)

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

@@ -18,8 +18,8 @@ under the ADR-0064 Rev2 structural model (the CPU-offload win).
 from __future__ import annotations
 
 from kernbench.benches.gqa_helpers.long_ctx._gqa_mlo_reduce import (
-    _ROOT_CUBE,
     reduce_mlo,
+    root_cube_for,
 )
 
 
@@ -63,9 +63,9 @@ def gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_kernel(
     tl.composite(op="gemm", a=exp_scores, b=V, out=O_local)  # P·V, one command
 
     # ── Two-level (m, ℓ, O) reduce-to-root (shared helper) ──
-    reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, P, tl=tl)
+    reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, C, P, tl=tl)
 
-    # ── Final normalise + store (root only: PE 0 of CUBE 6) ──
-    if pe_id == 0 and cube_id == _ROOT_CUBE:
+    # ── Final normalise + store (root only: PE 0 of the C-derived root cube) ──
+    if pe_id == 0 and cube_id == root_cube_for(C):
         O_final = O_local / l_local
         tl.store(o_ptr, O_final)

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

@@ -27,8 +27,8 @@ numeric parity of the recipe's 8 MATH ops is a separate follow-up
 from __future__ import annotations
 
 from kernbench.benches.gqa_helpers.long_ctx._gqa_mlo_reduce import (
-    _ROOT_CUBE,
     reduce_mlo,
+    root_cube_for,
 )
 
 # Small primitive seed slice that establishes the running (m, ℓ, O) the
@@ -86,9 +86,9 @@ def gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_ext_kernel(
         )
 
     # ── Two-level (m, ℓ, O) reduce-to-root (shared helper) ──
-    reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, P, tl=tl)
+    reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, C, P, tl=tl)
 
-    # ── Final normalise + store (root only: PE 0 of CUBE 6) ──
-    if pe_id == 0 and cube_id == _ROOT_CUBE:
+    # ── Final normalise + store (root only: PE 0 of the C-derived root cube) ──
+    if pe_id == 0 and cube_id == root_cube_for(C):
         O_final = O_local / l_local
         tl.store(o_ptr, O_final)

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

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

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

@@ -1,129 +0,0 @@
-"""GQA decode kernel — Case 6, **primitive-TILED** (16×16×16 MAC blocking).
-
-Same Case-6 placement and (m, ℓ, O) reduce as the primitive baseline
-(``_gqa_attention_decode_long_ctx_cube_sp_pe_sp``); the only difference is
-that each local-attention matmul is *hand-blocked into 16×16×16 GemmCmds*
-(mac=16) instead of one coarse ``tl.dot`` per tile. This models a kernel
-that issues the MAC-array fan-out from PE_CPU itself: the per-block GemmCmd
-count is ``ceil(M/16)·ceil(K/16)·ceil(N/16)`` per matmul, charging the full
-ADR-0064 dispatch cost for every block — the "dispatch explosion" the
-composite form offloads to PE_SCHEDULER.
-
-FLOPs are conserved (each 16³ GemmCmd carries the TFLOPS-model compute of
-its block; the blocks sum to the full matmul), so end-to-end compute time
-is unchanged vs the coarse primitive — only the PE_CPU command count and
-its dispatch cycles grow. Inputs are zero (decode bench convention), so the
-blocked accumulation is identically zero; the kernel returns a single
-zeroed ``(M, N)`` output handle that the downstream softmax consumes — no
-per-block accumulation handle needed for this zero-input study.
-"""
-from __future__ import annotations
-
-from math import ceil
-
-from kernbench.common.pe_commands import GemmCmd
-from kernbench.benches.gqa_helpers.long_ctx._gqa_mlo_reduce import (
-    _ROOT_CUBE,
-    _merge_running,
-    reduce_mlo,
-)
-
-TILE_S_KV = 1024  # ADR-0063 §A.2 S_kv-axis tile sweep (per-tile width).
-MAC = 16          # 16×16×16 MAC-array blocking granularity.
-
-
-def _blocked_dot(A, B, *, tl):
-    """``A @ B`` issued as one GemmCmd per 16×16×16 block.
-
-    ``A``:(M, K), ``B``:(K, N) → out:(M, N). Emits
-    ``ceil(M/16)·ceil(K/16)·ceil(N/16)`` GemmCmds (each charged the full
-    ADR-0064 dispatch cost via ``tl.dot``'s emit path). **All blocks write
-    the same ``(M, N)`` output handle** (``out``), and that handle is
-    returned — so downstream softmax ops depend on it exactly like the
-    coarse ``tl.dot`` path (the engine tracks the producer by output handle
-    id, and ``GemmCmd`` is a blocking PE_CPU command, so the block GEMMs
-    serialize on the critical path ahead of the consuming ``tl.max``/
-    ``tl.exp``). K is innermost so each (mi, ni) output tile accumulates
-    across the K blocks into the same handle.
-    """
-    M, K = A.shape[-2], A.shape[-1]
-    K2, N = B.shape[-2], B.shape[-1]
-    assert K == K2, f"blocked_dot shape mismatch K={K} != {K2}"
-    out = tl._make_compute_out(shape=(M, N), dtype="f16")
-    # One GemmCmd per 16³ block, all writing the shared `out` handle. A/B
-    # reuse the full operand handles at block dims (m,k,n = 16, or the
-    # ragged tail) — operands are TCM-resident (pinned), so this charges
-    # dispatch + the block's TFLOPS-model compute. K innermost = accumulate
-    # the K blocks into each (mi, ni) output tile of the shared handle.
-    for _mi in range(ceil(M / MAC)):
-        bm = min(MAC, M - _mi * MAC)
-        for _ni in range(ceil(N / MAC)):
-            bn = min(MAC, N - _ni * MAC)
-            for _ki in range(ceil(K / MAC)):
-                bk = min(MAC, K - _ki * MAC)
-                tl._emit(GemmCmd(a=A, b=B, out=out, m=bm, k=bk, n=bn))
-    return out
-
-
-def gqa_attention_decode_long_ctx_cube_sp_pe_sp_tiled_kernel(
-    q_ptr: int,
-    k_ptr: int,
-    v_ptr: int,
-    o_ptr: int,
-    T_q: int,
-    S_kv: int,
-    h_q: int,
-    h_kv: int,
-    d_head: int,
-    C: int,
-    P: int,
-    *,
-    tl,
-) -> None:
-    """Case-6 decode, primitive-TILED (16×16×16 GemmCmd blocking)."""
-    G = h_q // h_kv
-    n_ranks = C * P
-    S_local = S_kv // n_ranks
-    pe_id = tl.program_id(axis=0)
-    cube_id = tl.program_id(axis=1)
-
-    Q = tl.load(q_ptr, shape=(G * T_q, d_head), dtype="f16")
-    n_tiles = (S_local + TILE_S_KV - 1) // TILE_S_KV
-    KV_ROW_BYTES = d_head * 2  # f16
-
-    tile_s0 = min(TILE_S_KV, S_local)
-    K_T = tl.load(k_ptr, shape=(d_head, tile_s0), dtype="f16")
-    V = tl.load(v_ptr, shape=(tile_s0, d_head), dtype="f16")
-    scores = _blocked_dot(Q, K_T, tl=tl)
-    m_local = tl.max(scores, axis=-1)
-    centered = scores - m_local
-    exp_scores = tl.exp(centered)
-    l_local = tl.sum(exp_scores, axis=-1)
-    O_local = _blocked_dot(exp_scores, V, tl=tl)
-
-    for tile_idx in range(1, n_tiles):
-        tile_start = tile_idx * TILE_S_KV
-        tile_s = min(TILE_S_KV, S_local - tile_start)
-        with tl.scratch_scope():
-            K_T_t = tl.load(k_ptr + tile_start * KV_ROW_BYTES,
-                            shape=(d_head, tile_s), dtype="f16")
-            V_t = tl.load(v_ptr + tile_start * KV_ROW_BYTES,
-                          shape=(tile_s, d_head), dtype="f16")
-            scores_t = _blocked_dot(Q, K_T_t, tl=tl)
-            m_tile = tl.max(scores_t, axis=-1)
-            centered_t = scores_t - m_tile
-            exp_scores_t = tl.exp(centered_t)
-            l_tile = tl.sum(exp_scores_t, axis=-1)
-            O_tile = _blocked_dot(exp_scores_t, V_t, tl=tl)
-            m_new, l_new, O_new = _merge_running(
-                m_local, l_local, O_local, m_tile, l_tile, O_tile, tl=tl,
-            )
-            tl.copy_to(m_local, m_new)
-            tl.copy_to(l_local, l_new)
-            tl.copy_to(O_local, O_new)
-
-    reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, P, tl=tl)
-
-    if pe_id == 0 and cube_id == _ROOT_CUBE:
-        O_final = O_local / l_local
-        tl.store(o_ptr, O_final)

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

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

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

@@ -24,15 +24,15 @@ from __future__ import annotations
 import json
 from pathlib import Path
 
-from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_sp import (  # noqa: E501
-    gqa_attention_decode_long_ctx_cube_sp_pe_sp_kernel,
-)
 from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite import (  # noqa: E501
     gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_kernel,
 )
 from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_ext import (  # noqa: E501
     gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_ext_kernel,
 )
+from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_sp_hand_tiled_16x16x16 import (  # noqa: E501
+    gqa_attention_decode_long_ctx_cube_sp_pe_sp_hand_tiled_16x16x16_kernel,
+)
 from kernbench.benches.gqa_helpers.shared._gqa_panel_helpers import _ccl_cfg
 from kernbench.ccl.sfr_config import configure_sfr_intercube_multisip
 from kernbench.policy.placement.dp import DPPolicy
@@ -50,12 +50,12 @@ _SWEEP_JSON = _OUTPUT_DIR / "sweep_decode_composite.json"
 # Each kernel implements the same Case-6 placement; only the per-tile
 # command form differs (see module docstring).
 _VARIANT_KERNELS = {
-    "primitive": gqa_attention_decode_long_ctx_cube_sp_pe_sp_kernel,
+    "primitive_tiled": gqa_attention_decode_long_ctx_cube_sp_pe_sp_hand_tiled_16x16x16_kernel,
     "composite": gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_kernel,
     "composite_extended":
         gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_ext_kernel,
 }
-_VARIANTS = ("primitive", "composite", "composite_extended")
+_VARIANTS = ("primitive_tiled", "composite", "composite_extended")
 
 # Op-count (PE_CPU dispatch) is computed at emit time (exact, instant), so
 # it spans up to the 1M production point (S_local = 1M/64 = 16384, 16

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

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

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

@@ -1,168 +0,0 @@
-"""milestone-1h-gqa: memory-bound decode streaming composite-command study.
-
-Single-rank companion to the compute-bound prefill study
-(``gqa_prefill_compute_bound``). Same three command-form kernels
-(``_gqa_prefill_compute_bound``), same single-rank (C=P=1) harness, but
-driven with **T_q=1** (decode) and swept over a large S_kv. With T_q=1 the
-score / context GEMMs are skinny (M = G·T_q = 8), so the MAC array is
-barely fed and the kernel is **memory-bound** — streaming the KV cache out
-of HBM dominates. This is the regime mirror of prefill: command form is
-*latency-neutral* here, because the bottleneck is data movement, not
-issue.
-
-Running single-rank (no cross-CUBE (m,ℓ,O) reduce) isolates the
-local-attention streaming vs. issue trade-off cleanly — unlike the 64-way
-Case-6 decode sweep, whose small-S_kv end-to-end latency is masked by the
-inter-CUBE reduce tail. S_kv here is the *per-rank* context, so S_kv=64K
-single-rank is the per-PE load of a 4M-token, 64-way-sharded decode.
-
-Records per (variant, S_kv) the end-to-end latency, GEMM/DMA engine busy
-time, MAC utilization (achieved ÷ peak), and DMA occupancy (dma_busy ÷
-e2e), so the comparative plot can show all three command forms landing on
-the same latency curve while the MAC array stays floored and the DMA
-channel stays saturated.
-
-Runs in data mode (engine latency). Gated via the umbrella
-``GQA_1H_SWEEPS=decode_streaming``.
-"""
-from __future__ import annotations
-
-import json
-from pathlib import Path
-
-from kernbench.benches.gqa_helpers.long_ctx._gqa_prefill_compute_bound import (
-    gqa_prefill_composite_ext_kernel,
-    gqa_prefill_composite_kernel,
-    gqa_prefill_primitive_kernel,
-)
-from kernbench.policy.placement.dp import DPPolicy
-
-_OUTPUT_DIR = (
-    Path(__file__).resolve().parents[2]
-    / "1H_milestone_output" / "gqa" / "long_ctx"
-)
-_SWEEP_JSON = _OUTPUT_DIR / "sweep_decode_streaming.json"
-
-# Same kernels as the prefill study — they are general attention kernels;
-# the regime is set by T_q (=1 here → memory-bound).
-_VARIANT_KERNELS = {
-    "primitive": gqa_prefill_primitive_kernel,
-    "composite": gqa_prefill_composite_kernel,
-    "composite_extended": gqa_prefill_composite_ext_kernel,
-}
-_VARIANTS = ("primitive", "composite", "composite_extended")
-
-# Per-rank context length. T_q=1, so M = G·T_q = 8 (skinny, memory-bound)
-# at every point. Capped at 16K: the *plain* composite materializes the
-# full (M, S_kv) scores in TCM scratch (~48·S_kv B incl. the exp transients),
-# which overflows the 1 MB kernel scratch beyond ~21K — itself a sign that
-# the recipe (composite_extended), which tiles the softmax, is what long
-# context actually needs. S_kv is per-rank, so 16K single-rank already
-# equals the per-PE load of a 1M-token, 64-way-sharded decode.
-_S_KV_POINTS = (2048, 4096, 8192, 16384)
-
-_T_Q = 1
-_H_Q, _H_KV, _D_HEAD = 8, 1, 128
-_PEAK_TFLOPS = 8.0  # per-PE f16 GEMM peak (topology.yaml pe_gemm.peak_tflops_f16)
-
-
-def _run_panel_fn(variant: str, s_kv: int):
-    kernel = _VARIANT_KERNELS[variant]
-    panel = f"decode_stream_{variant}_s{s_kv}"
-
-    def _bench_fn(ctx):
-        dp = DPPolicy(cube="replicate", pe="replicate",
-                      num_cubes=1, num_pes=1)
-        q = ctx.zeros((_T_Q, _H_Q * _D_HEAD),
-                      dtype="f16", dp=dp, name=f"{panel}_q")
-        k = ctx.zeros((s_kv, _H_KV * _D_HEAD),
-                      dtype="f16", dp=dp, name=f"{panel}_k")
-        v = ctx.zeros((s_kv, _H_KV * _D_HEAD),
-                      dtype="f16", dp=dp, name=f"{panel}_v")
-        o = ctx.empty((_T_Q, _H_Q * _D_HEAD),
-                      dtype="f16", dp=dp, name=f"{panel}_o")
-        ctx.launch(panel, kernel, q, k, v, o,
-                   _T_Q, s_kv, _H_Q, _H_KV, _D_HEAD, 1, 1,
-                   _auto_dim_remap=False)
-
-    return _bench_fn
-
-
-def _end_to_end_ns(op_log) -> float:
-    if not op_log:
-        return 0.0
-    return max(r.t_end for r in op_log) - min(r.t_start for r in op_log)
-
-
-def _engine_busy_ns(op_log, suffix: str) -> float:
-    return sum(r.t_end - r.t_start
-               for r in op_log if r.component_id.endswith("." + suffix))
-
-
-def _run_panel(variant: str, s_kv: int, topology: str) -> dict:
-    from kernbench.runtime_api.bench_runner import run_bench
-    from kernbench.runtime_api.types import resolve_device
-    from kernbench.sim_engine.engine import GraphEngine
-    from kernbench.topology.builder import resolve_topology
-
-    topo = resolve_topology(topology)
-    result = run_bench(
-        topology=topo, bench_fn=_run_panel_fn(variant, s_kv),
-        device=resolve_device(None),
-        engine_factory=lambda t, d: GraphEngine(
-            getattr(t, "topology_obj", t), enable_data=True,
-        ),
-    )
-    if not result.completion.ok:
-        raise RuntimeError(
-            f"gqa-decode-streaming {variant}@{s_kv} failed: {result.completion}"
-        )
-    op_log = result.engine.op_log
-    e2e = _end_to_end_ns(op_log)
-    gemm = _engine_busy_ns(op_log, "pe_gemm")
-    dma = _engine_busy_ns(op_log, "pe_dma")
-    G = _H_Q // _H_KV
-    M = G * _T_Q
-    # Useful attention flops (Q·Kᵀ + P·V), single rank.
-    useful_flops = 4.0 * M * _D_HEAD * s_kv
-    # KV-cache bytes streamed from HBM (K + V, f16) — the memory-bound
-    # denominator. achieved_bw = bytes / e2e (GB/s) measures how close the
-    # command form gets to the HBM roofline (the memory-bound mirror of
-    # MAC utilization in the compute-bound prefill study).
-    kv_bytes = 2.0 * s_kv * _D_HEAD * 2.0
-    return {
-        "variant": variant,
-        "s_kv": s_kv,
-        "M": M,
-        "latency_ns": e2e,
-        "gemm_busy_ns": gemm,
-        "dma_busy_ns": dma,
-        "kv_bytes": kv_bytes,
-        "achieved_bw_gbs": (kv_bytes / e2e) if e2e > 0 else 0.0,
-        "achieved_tflops": (useful_flops / e2e / 1e3) if e2e > 0 else 0.0,
-        "mac_util": (useful_flops / e2e / 1e3 / _PEAK_TFLOPS) if e2e > 0 else 0.0,
-        "dma_occupancy": (dma / e2e) if e2e > 0 else 0.0,
-    }
-
-
-def run_sweep(topology: str = "topology.yaml") -> int:
-    """Drive all (variant, S_kv) decode-streaming panels; write sweep.json."""
-    _OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
-    rows = [
-        _run_panel(variant, s_kv, topology)
-        for s_kv in _S_KV_POINTS
-        for variant in _VARIANTS
-    ]
-    sweep = {
-        "version": 1,
-        "variants": list(_VARIANTS),
-        "s_kv_points": list(_S_KV_POINTS),
-        "rows": rows,
-    }
-    _SWEEP_JSON.write_text(json.dumps(sweep, indent=2))
-    print(f"  gqa-decode-streaming: {len(rows)} rows -> {_SWEEP_JSON}")
-    return len(rows)
-
-
-if __name__ == "__main__":
-    run_sweep()

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

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

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

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

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

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

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

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

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

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

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

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

src/kernbench/benches/milestone_1h_gqa.py

@@ -28,9 +28,6 @@ from kernbench.benches.gqa_helpers.long_ctx.gqa_decode_long_ctx_4cases import (
 from kernbench.benches.gqa_helpers.long_ctx.gqa_decode_long_ctx_composite import (
     run_sweep as _run_composite_sweep,
 )
-from kernbench.benches.gqa_helpers.long_ctx.gqa_decode_streaming import (
-    run_sweep as _run_decode_streaming_sweep,
-)
 from kernbench.benches.gqa_helpers.long_ctx.gqa_prefill_compute_bound import (
     run_sweep as _run_prefill_cb_sweep,
 )
@@ -69,7 +66,6 @@ def run(torch) -> None:
         "decode": _run_decode_sweep,
         "composite": _run_composite_sweep,
         "prefill_cb": _run_prefill_cb_sweep,
-        "decode_streaming": _run_decode_streaming_sweep,
     }
     unknown = [s for s in sweeps if s not in runners]
     if unknown:

src/kernbench/components/builtin/pe_cpu.py

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

src/kernbench/policy/routing/router.py

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

src/kernbench/runtime_api/context.py

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

src/kernbench/sim_engine/engine.py

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

src/kernbench/triton_emu/tl_context.py

@@ -697,6 +697,10 @@ class TLContext:
                 nbytes=self._nbytes(shape, dtype),
                 data=data,
                 space=slot_space,
+                # On-chip (tcm/sram) IPCQ slot: read in place as a composite
+                # operand, not DMA_READ of its non-HBM slot address. An HBM
+                # slot is a valid PA, so it stays unpinned (DMA streams it).
+                pinned=slot_space != "hbm",
             )
         return self._make_handle(addr=0, shape=shape, dtype=dtype)
 
@@ -738,6 +742,7 @@ class TLContext:
                 nbytes=self._nbytes(shape, dtype),
                 data=None,
                 space=slot_space,
+                pinned=slot_space != "hbm",
             )
         return self._make_handle(addr=0, shape=shape, dtype=dtype)
 
@@ -1105,6 +1110,7 @@ class TLContext:
                 nbytes=self._nbytes(handle.cmd.shape, handle.cmd.dtype),
                 data=data,
                 space=slot_space,
+                pinned=slot_space != "hbm",
             )
             handle.resolved = True
             handle.result = th

tests/analytical_visualization/autosuggest.py

@@ -0,0 +1,106 @@
+"""Auto-suggest (CP, TP, PP) that fits the model in the per-PE budget.
+
+Strategy: iterate over candidate (CP, TP, PP) triples in order of
+increasing total PE count and return the first one that satisfies
+memory + kernel-support constraints with >10% slack. If none fits, return
+the best-effort configuration and flag over-budget.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, replace
+from typing import Iterable
+
+from .model_config import FullConfig, ModelConfig, TopologyConfig, MachineParams
+from .memory_layout import compute_memory
+
+
+# Candidate parallelism dimensions (powers of 2 mostly, up to sensible limits).
+_TP_OPTIONS = (1, 2, 4, 8, 16, 32)
+_CP_OPTIONS = (1, 2, 4, 8, 16, 32, 64, 96)
+_PP_OPTIONS = (1, 2, 4, 8, 16, 32)
+
+
+@dataclass
+class Suggestion:
+    cp: int
+    tp: int
+    pp: int
+    weights_gb: float
+    kv_gb: float
+    transient_gb: float
+    slack_gb: float
+    fits: bool
+    pes_used: int
+    sips_used: int
+    reason: str = ""
+
+
+def _iter_candidates(model: ModelConfig) -> Iterable[tuple[int, int, int]]:
+    """Yield (CP, TP, PP) in order of increasing PE count."""
+    triples: list[tuple[int, int, int]] = []
+    for tp in _TP_OPTIONS:
+        for cp in _CP_OPTIONS:
+            for pp in _PP_OPTIONS:
+                # PP must not exceed layer count.
+                if pp > model.layers:
+                    continue
+                # Skip TP > H_q * some factor (unrealistic).
+                if tp > model.h_q * 4:
+                    continue
+                triples.append((cp, tp, pp))
+    # Sort by pe_count then by (pp, tp, cp) — prefer smaller PP first
+    # (avoids pipeline bubbles), then smaller TP (avoids head-dim split).
+    triples.sort(key=lambda t: (t[0] * t[1] * t[2], t[2], t[1], t[0]))
+    for t in triples:
+        yield t
+
+
+def _score_candidate(cp: int, tp: int, pp: int,
+                     model: ModelConfig, machine: MachineParams,
+                     s_kv: int, mode: str,
+                     slack_frac: float = 0.10) -> Suggestion:
+    topo = TopologyConfig(cp=cp, tp=tp, pp=pp, s_kv=s_kv, mode=mode)
+    cfg = FullConfig(model=model, topo=topo, machine=machine)
+    mem = compute_memory(cfg)
+
+    fits = (not mem.over_budget
+            and mem.slack_bytes >= slack_frac * mem.budget_bytes)
+    reason = ""
+    if mem.over_budget:
+        reason = (f"weights+KV+transient ({mem.used_bytes/1e9:.2f} GB) "
+                  f"exceeds budget ({mem.budget_bytes/1e9:.2f} GB)")
+    elif not fits:
+        reason = f"slack ({mem.slack_bytes/1e9:.2f} GB) below 10% of budget"
+    return Suggestion(
+        cp=cp, tp=tp, pp=pp,
+        weights_gb=mem.weights_bytes / 1e9,
+        kv_gb=mem.kv_cache_bytes / 1e9,
+        transient_gb=mem.transient_bytes / 1e9,
+        slack_gb=mem.slack_bytes / 1e9,
+        fits=fits,
+        pes_used=topo.total_pes,
+        sips_used=topo.sips_used,
+        reason=reason,
+    )
+
+
+def auto_suggest(model: ModelConfig, machine: MachineParams,
+                 s_kv: int, mode: str = "decode",
+                 slack_frac: float = 0.10) -> Suggestion:
+    """Return the smallest-PE-count (CP, TP, PP) that fits.
+
+    If no candidate fits, returns the best-effort one (highest slack,
+    even if negative) with fits=False.
+    """
+    best_fit: Suggestion | None = None
+    best_effort: Suggestion | None = None
+
+    for cp, tp, pp in _iter_candidates(model):
+        s = _score_candidate(cp, tp, pp, model, machine, s_kv, mode, slack_frac)
+        if s.fits and best_fit is None:
+            best_fit = s
+            break  # candidates are pre-sorted by PE count
+        if best_effort is None or s.slack_gb > best_effort.slack_gb:
+            best_effort = s
+
+    return best_fit or best_effort  # type: ignore

tests/analytical_visualization/memory_layout.py

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

tests/analytical_visualization/model_config.py

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

tests/analytical_visualization/model_presets.py

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

tests/analytical_visualization/optimization_report.py

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

tests/analytical_visualization/pe_weight_layout.py

@@ -0,0 +1,221 @@
+"""Draw a PE-level view of one CP group (one TP group of cubes).
+
+Shows how weights + KV cache are distributed across the PEs of ONE CP
+rank. Attention weights are sharded by head (H_q, H_kv split across TP);
+FFN is sharded by dim (ffn_dim / TP / EP); KV cache is sharded by CP
+(sequence) × TP (head).
+
+If TP > 8 the group spans multiple cubes — layout draws all of them.
+"""
+from __future__ import annotations
+
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
+
+from .model_config import FullConfig
+
+
+PE_COLS = 4
+PE_ROWS = 2
+PES_PER_CUBE = PE_COLS * PE_ROWS
+
+
+def layers_per_stage_val(cfg: FullConfig) -> int:
+    return (cfg.model.layers + cfg.topo.pp - 1) // cfg.topo.pp
+
+
+def _q_heads_for_pe(pe_id_in_group: int, tp: int, h_q: int) -> list[int]:
+    """Return the list of Q-head indices assigned to this PE."""
+    heads_per_pe = h_q / tp
+    start = int(round(pe_id_in_group * heads_per_pe))
+    end = int(round((pe_id_in_group + 1) * heads_per_pe))
+    return list(range(start, end))
+
+
+def _kv_heads_for_pe(pe_id_in_group: int, tp: int, h_kv: int,
+                      kv_shard_mode: str) -> tuple[list[int], str]:
+    """KV heads assigned to this PE and a note about replication/split."""
+    if h_kv >= tp:
+        heads_per_pe = h_kv // tp
+        start = pe_id_in_group * heads_per_pe
+        end = start + heads_per_pe
+        return list(range(start, end)), ""
+    # TP > H_kv
+    if kv_shard_mode == "replicate":
+        rep_factor = tp // h_kv
+        head_idx = pe_id_in_group // rep_factor
+        return [head_idx], f"replicated x{rep_factor}"
+    # split mode
+    split_factor = tp // h_kv
+    head_idx = pe_id_in_group // split_factor
+    part = pe_id_in_group % split_factor
+    return [head_idx], f"head-split ({part+1}/{split_factor})"
+
+
+def _per_pe_bytes(cfg: FullConfig, pe_id_in_group: int) -> dict:
+    """Return per-tensor bytes for one PE within a TP group."""
+    m = cfg.model
+    tp = cfg.topo.tp
+    pp = cfg.topo.pp
+    ep = max(1, cfg.topo.ep)
+    b = m.bytes_per_elem
+    layers_per_stage = (m.layers + pp - 1) // pp
+
+    hq_per_pe = m.h_q / tp
+    if cfg.topo.kv_shard_mode == "replicate":
+        hkv_per_pe = max(1.0, m.h_kv / tp)
+    else:
+        hkv_per_pe = m.h_kv / tp
+    ffn_div = cfg.ffn_shard_divisor * ep
+    ffn_per_pe = m.ffn_dim // ffn_div
+
+    per_layer = {
+        "W_Q": int(m.hidden * hq_per_pe * m.d_head * b),
+        "W_K": int(m.hidden * hkv_per_pe * m.d_head * b),
+        "W_V": int(m.hidden * hkv_per_pe * m.d_head * b),
+        "W_O": int(hq_per_pe * m.d_head * m.hidden * b),
+        "W_gate": int(m.hidden * ffn_per_pe * b),
+        "W_up": int(m.hidden * ffn_per_pe * b),
+        "W_down": int(ffn_per_pe * m.hidden * b),
+    }
+    kv_per_layer = int(2 * cfg.topo.s_local * hkv_per_pe * m.d_head * b)
+
+    all_layers = {k: v * layers_per_stage for k, v in per_layer.items()}
+    all_layers["KV cache"] = kv_per_layer * layers_per_stage
+    # transient estimate
+    if cfg.topo.mode == "decode":
+        all_layers["Transient"] = int(4 * (m.hidden + hq_per_pe * m.d_head) * b)
+    else:
+        TILE = 1024
+        all_layers["Transient"] = int(2 * hq_per_pe * cfg.topo.T_q * TILE * b
+                                       + cfg.topo.T_q * m.hidden * b)
+    return all_layers
+
+
+def draw_pe_layout(cfg: FullConfig, ax=None):
+    """Draw one CP group's PEs with per-tensor weight + KV breakdown.
+
+    Each PE shows: PE id, Q heads, KV heads, W_Q / W_K / W_V / W_O in
+    MB, FFN (gate+up+down) in MB, KV cache in GB, total in GB.
+
+    TP=8 -> 1 cube. TP=16 -> 2 cubes. TP=32 -> 4 cubes.
+    Under placement, n_cubes reflects cubes-per-group (may include CP spill).
+    """
+    tp = cfg.topo.tp
+    # cubes per one cube-level group (= intra-cube dims / PEs-per-cube, rounded up).
+    n_cubes = max(1, (cfg.topo.intra_cube_dims + PES_PER_CUBE - 1) // PES_PER_CUBE)
+
+    if ax is None:
+        fig, ax = plt.subplots(figsize=(7.5 * n_cubes, 8.5))
+    else:
+        fig = ax.figure
+
+    cube_w = 7.0
+    cube_h = 6.0
+    cube_gap = 0.6
+
+    for cube_idx in range(n_cubes):
+        x0 = cube_idx * (cube_w + cube_gap)
+        y0 = 0
+
+        rect = patches.FancyBboxPatch(
+            (x0, y0), cube_w, cube_h,
+            boxstyle="round,pad=0.05",
+            facecolor="#f8f9fa", edgecolor="#212529", linewidth=1.5,
+        )
+        ax.add_patch(rect)
+        ax.text(x0 + cube_w / 2, y0 + cube_h + 0.1,
+                f"Cube {cube_idx}  (PEs {cube_idx*8}-{cube_idx*8+7})",
+                ha="center", va="bottom", fontsize=11, fontweight="bold")
+
+        pe_pad_x = 0.15
+        pe_pad_y = 0.2
+        pe_gap = 0.10
+        inner_w = cube_w - 2 * pe_pad_x
+        inner_h = cube_h - 2 * pe_pad_y
+        pe_w = (inner_w - (PE_COLS - 1) * pe_gap) / PE_COLS
+        pe_h = (inner_h - (PE_ROWS - 1) * pe_gap) / PE_ROWS
+
+        for pr in range(PE_ROWS):
+            for pc in range(PE_COLS):
+                pe_local = pr * PE_COLS + pc
+                pe_id_in_group = cube_idx * PES_PER_CUBE + pe_local
+                if pe_id_in_group >= cfg.topo.intra_cube_dims:
+                    continue
+                px = x0 + pe_pad_x + pc * (pe_w + pe_gap)
+                py = y0 + pe_pad_y + (PE_ROWS - 1 - pr) * (pe_h + pe_gap)
+
+                pe_rect = patches.Rectangle(
+                    (px, py), pe_w, pe_h,
+                    facecolor="#f0f7ff", edgecolor="#3a86ff", linewidth=1.0,
+                )
+                ax.add_patch(pe_rect)
+
+                q_heads = _q_heads_for_pe(pe_id_in_group, tp, cfg.model.h_q)
+                kv_heads, kv_note = _kv_heads_for_pe(
+                    pe_id_in_group, tp, cfg.model.h_kv, cfg.topo.kv_shard_mode,
+                )
+                bytes_ = _per_pe_bytes(cfg, pe_id_in_group)
+                weights_gb = sum(v for k, v in bytes_.items()
+                                 if k not in ("KV cache", "Transient")) / 1e9
+                kv_gb = bytes_["KV cache"] / 1e9
+                trans_mb = bytes_["Transient"] / 1e6
+                total_gb = sum(bytes_.values()) / 1e9
+
+                q_str = (f"{q_heads[0]}-{q_heads[-1]}"
+                         if len(q_heads) > 1 else str(q_heads[0])
+                         if q_heads else "-")
+                kv_str = str(kv_heads[0]) if kv_heads else "-"
+                if kv_note:
+                    kv_str += f"({kv_note})"
+
+                ffn_mb = (bytes_["W_gate"] + bytes_["W_up"]
+                          + bytes_["W_down"]) / 1e6
+                # Header (PE id and heads) bigger, then per-tensor breakdown
+                header = (f"PE {pe_id_in_group}\n"
+                          f"Q heads: {q_str}\n"
+                          f"KV head: {kv_str}")
+                ax.text(px + 0.08, py + pe_h - 0.05, header,
+                        ha="left", va="top",
+                        fontsize=7.0, fontweight="bold",
+                        family="monospace", color="#0d47a1")
+
+                # Per-tensor bytes with explicit division divisors.
+                # Divisor is always the PRODUCT of the parallel dims listed
+                # (TP*CP means TP times CP, i.e. multiplication).
+                tp_d = cfg.topo.tp
+                ffn_div = cfg.ffn_shard_divisor
+                ffn_scope = cfg.topo.ffn_shard_scope.replace("+", "*")
+                kv_div = cfg.topo.cp * cfg.topo.tp
+                L = layers_per_stage_val(cfg)
+                detail = (
+                    f"weights (all {L} layers):\n"
+                    f" W_Q /TP={tp_d} : {bytes_['W_Q']/1e6:7.1f}MB\n"
+                    f" W_K /TP={tp_d} : {bytes_['W_K']/1e6:7.1f}MB\n"
+                    f" W_V /TP={tp_d} : {bytes_['W_V']/1e6:7.1f}MB\n"
+                    f" W_O /TP={tp_d} : {bytes_['W_O']/1e6:7.1f}MB\n"
+                    f" FFN /({ffn_scope})={ffn_div}: {ffn_mb:7.1f}MB\n"
+                    f"runtime:\n"
+                    f" KV /(CP*TP)={kv_div}: {kv_gb*1000:7.1f}MB\n"
+                    f" trans     : {trans_mb:7.1f}MB\n"
+                    f"W total: {weights_gb*1000:7.1f}MB\n"
+                    f"TOTAL  : {total_gb*1000:7.1f}MB"
+                )
+                ax.text(px + pe_w / 2, py + 0.06, detail,
+                        ha="center", va="bottom",
+                        fontsize=5.8, family="monospace", color="#1a1a1a")
+
+    total_w = n_cubes * cube_w + (n_cubes - 1) * cube_gap
+    ax.set_xlim(-0.2, total_w + 0.2)
+    ax.set_ylim(-0.3, cube_h + 0.6)
+    ax.set_aspect("equal")
+    ax.axis("off")
+
+    title = (
+        f"Per-PE layout of one CP group  "
+        f"(TP={tp}, H_q={cfg.model.h_q}, H_kv={cfg.model.h_kv})"
+        f" | KV mode: {cfg.topo.kv_shard_mode}"
+    )
+    ax.set_title(title, fontsize=11, fontweight="bold")
+
+    return fig

tests/analytical_visualization/pipeline_diagram.py

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

tests/analytical_visualization/stage_latencies.py

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

tests/analytical_visualization/tensor_sharding.py

@@ -0,0 +1,330 @@
+"""Pictorial view of how each weight/KV tensor is sharded across ranks.
+
+Every tensor is drawn as a rectangle labelled with its (rows, cols)
+shape. Grid lines show shard boundaries; one shard (this PE's share)
+is highlighted so it's obvious which dim the split is along.
+
+Sharding pattern:
+  - W_Q  (hidden, H_q*d_head) : split on **output cols** across TP
+  - W_K  (hidden, H_kv*d_head): split on **output cols** across TP
+  - W_V  (hidden, H_kv*d_head): split on **output cols** across TP
+  - W_O  (H_q*d_head, hidden) : split on **input rows** across TP
+  - W_gate/up (hidden, ffn_dim): split on **output cols** across FFN scope
+  - W_down (ffn_dim, hidden)   : split on **input rows** across FFN scope
+  - K/V cache (S_kv, H_kv*d_head): split on **rows (S axis)** by CP AND
+    **cols (head axis)** by TP -> 2D shard grid
+"""
+from __future__ import annotations
+
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
+
+from .model_config import FullConfig
+
+
+_UNSHARDED_FACE = "#e9ecef"
+_UNSHARDED_EDGE = "#adb5bd"
+_MY_SHARD_FACE = "#3a86ff"
+_MY_SHARD_EDGE = "#0057b3"
+_TP_LINE_COLOR = "#d90429"     # TP split lines
+_CP_LINE_COLOR = "#0f9d58"     # CP split lines
+_FFN_LINE_COLOR = "#e37400"    # FFN-scope split lines
+
+
+def _draw_tensor(ax, x, y, w, h, name, shape_str,
+                  row_splits: int = 1, col_splits: int = 1,
+                  my_row_shard: int = 0, my_col_shard: int = 0,
+                  row_line_color: str = _TP_LINE_COLOR,
+                  col_line_color: str = _TP_LINE_COLOR,
+                  row_label: str = "", col_label: str = "",
+                  note: str = "",
+                  cell_annot=None):
+    """Draw a tensor with sharding grid.
+
+    cell_annot: optional callable (row, col) -> str returning the physical
+    label to render inside shard cell (r, c). If it returns "", cell is not
+    annotated.
+    """
+    """Draw a tensor as a rectangle with grid lines showing shards."""
+    # Base rectangle (unsharded background)
+    rect = patches.Rectangle(
+        (x, y), w, h,
+        facecolor=_UNSHARDED_FACE, edgecolor=_UNSHARDED_EDGE, linewidth=1.0,
+    )
+    ax.add_patch(rect)
+
+    # Highlight this PE's shard
+    shard_w = w / max(1, col_splits)
+    shard_h = h / max(1, row_splits)
+    sx = x + my_col_shard * shard_w
+    sy = y + (row_splits - 1 - my_row_shard) * shard_h
+    highlight = patches.Rectangle(
+        (sx, sy), shard_w, shard_h,
+        facecolor=_MY_SHARD_FACE, edgecolor=_MY_SHARD_EDGE, linewidth=1.5,
+        alpha=0.8,
+    )
+    ax.add_patch(highlight)
+
+    # Horizontal grid lines (row splits)
+    for r in range(1, row_splits):
+        yy = y + r * shard_h
+        ax.plot([x, x + w], [yy, yy],
+                color=row_line_color, linewidth=1.2, linestyle="--")
+    # Vertical grid lines (col splits)
+    for c in range(1, col_splits):
+        xx = x + c * shard_w
+        ax.plot([xx, xx], [y, y + h],
+                color=col_line_color, linewidth=1.2, linestyle="--")
+
+    # Per-cell physical annotations (e.g. "cube3 PE5") if provided.
+    if cell_annot is not None:
+        for r in range(row_splits):
+            for c in range(col_splits):
+                text = cell_annot(r, c)
+                if not text:
+                    continue
+                cx = x + c * shard_w + shard_w / 2
+                cy = y + (row_splits - 1 - r) * shard_h + shard_h / 2
+                ax.text(cx, cy, text,
+                        ha="center", va="center",
+                        fontsize=6.0, color="#0d1b2a",
+                        alpha=0.85)
+
+    # Stacked labels ABOVE rectangle: note (italic, top), shape, name (bold, closest)
+    if note:
+        ax.text(x + w / 2, y + h + 0.55, note,
+                ha="center", va="bottom", fontsize=6.5,
+                color="#555", style="italic")
+    ax.text(x + w / 2, y + h + 0.30, shape_str,
+            ha="center", va="bottom", fontsize=7, color="#666")
+    ax.text(x + w / 2, y + h + 0.08, name,
+            ha="center", va="bottom",
+            fontsize=10, fontweight="bold")
+
+    # Split labels BELOW rectangle (stacked if both present)
+    if row_label and col_label:
+        ax.text(x + w / 2, y - 0.10, row_label,
+                ha="center", va="top", fontsize=6.5,
+                color=row_line_color, fontweight="bold")
+        ax.text(x + w / 2, y - 0.35, col_label,
+                ha="center", va="top", fontsize=6.5,
+                color=col_line_color, fontweight="bold")
+    elif row_label:
+        ax.text(x + w / 2, y - 0.10, row_label,
+                ha="center", va="top", fontsize=6.5,
+                color=row_line_color, fontweight="bold")
+    elif col_label:
+        ax.text(x + w / 2, y - 0.10, col_label,
+                ha="center", va="top", fontsize=6.5,
+                color=col_line_color, fontweight="bold")
+
+
+def _physical_label(cfg: FullConfig, dim: str, rank: int) -> str:
+    """Where does rank `rank` of dim (`tp`|`cp`|`ffn`) live physically?
+    Returns a short "cubeX PEy" style label based on the placement config.
+    """
+    topo = cfg.topo
+    if dim == "tp":
+        if topo.tp_placement == "pe":
+            # TP fills PEs within one cube (spilling if tp > 8)
+            pe_in_cube = rank % topo.pes_per_cube_hw
+            cube = rank // topo.pes_per_cube_hw
+            return f"c{cube}p{pe_in_cube}"
+        else:  # tp on cube
+            return f"cube{rank}"
+    if dim == "cp":
+        if topo.cp_placement == "pe":
+            pe_in_cube = rank % topo.pes_per_cube_hw
+            cube = rank // topo.pes_per_cube_hw
+            return f"c{cube}p{pe_in_cube}"
+        else:  # cp on cube
+            return f"cube{rank}"
+    return ""
+
+
+def _kv_cell_annot(cfg: FullConfig):
+    """Return an annotation function for KV cache cells: (cp_r, tp_r) -> label."""
+    def _ann(r, c):
+        cp_lbl = _physical_label(cfg, "cp", r)
+        tp_lbl = _physical_label(cfg, "tp", c)
+        # Two lines if both are non-trivial; otherwise a single line
+        if cp_lbl and tp_lbl:
+            return f"{cp_lbl}\n{tp_lbl}"
+        return cp_lbl or tp_lbl
+    return _ann
+
+
+def _tp_cell_annot(cfg: FullConfig):
+    """Col-parallel tensors (W_Q/W_K/W_V/W_gate/W_up): each col = one TP rank."""
+    def _ann(r, c):
+        return _physical_label(cfg, "tp", c)
+    return _ann
+
+
+def _tp_row_cell_annot(cfg: FullConfig):
+    """Row-parallel tensors split by TP (W_O): each row = one TP rank."""
+    def _ann(r, c):
+        return _physical_label(cfg, "tp", r)
+    return _ann
+
+
+def draw_tensor_sharding(cfg: FullConfig, ax=None, my_pe: int = 0,
+                          show_physical: bool = False):
+    """Draw all tensors with their sharding overlays for PE `my_pe`.
+
+    show_physical: annotate each shard cell with the owning cube/PE.
+    """
+    m = cfg.model
+    tp = cfg.topo.tp
+    cp = cfg.topo.cp
+
+    # PE's assigned shard index for each dim
+    q_shard = my_pe % tp  # W_Q col shard
+    kv_shard_col = min(my_pe % tp, m.h_kv - 1) if m.h_kv >= tp else 0
+    ffn_shard = my_pe % cfg.ffn_shard_divisor
+    kv_row_shard = 0  # per-CP; showing CP=0 view for clarity
+
+    # Enlarged mode gives more room for physical annotations.
+    _fig_w, _fig_h = (22, 13) if show_physical else (15, 9)
+    _cell_w = 3.4 if show_physical else 2.6
+    _cell_h = 3.0 if show_physical else 2.2
+    _gap_x = 0.8 if show_physical else 0.55
+    _gap_y = 2.0 if show_physical else 1.7
+
+    if ax is None:
+        fig, ax = plt.subplots(figsize=(_fig_w, _fig_h))
+    else:
+        fig = ax.figure
+
+    # Grid layout: 2 rows x 4 cols of tensor rects
+    # Attention row: W_Q, W_K, W_V, W_O
+    # FFN + KV row: W_gate, W_up, W_down, KV cache
+    cell_w = _cell_w
+    cell_h = _cell_h
+    gap_x = _gap_x
+    gap_y = _gap_y
+
+    _kv_ann = _kv_cell_annot(cfg) if show_physical else None
+    _tp_col_ann = _tp_cell_annot(cfg) if show_physical else None
+    _tp_row_ann = _tp_row_cell_annot(cfg) if show_physical else None
+
+    positions = [
+        (0, 1),  # W_Q at row 0
+        (1, 1),  # W_K
+        (2, 1),  # W_V
+        (3, 1),  # W_O
+        (0, 0),  # W_gate at row 1
+        (1, 0),  # W_up
+        (2, 0),  # W_down
+        (3, 0),  # KV cache
+    ]
+
+    # ── attention tensors ────────────────────────
+    tp_txt = f"TP={tp}"
+    ffn_txt = f"{cfg.topo.ffn_shard_scope}={cfg.ffn_shard_divisor}"
+    cp_txt = f"CP={cp}"
+
+    def px(col): return col * (cell_w + gap_x)
+    def py(row): return row * (cell_h + gap_y)
+
+    # W_Q: (hidden, H_q*d_head), split on COL by TP
+    _draw_tensor(ax, px(0), py(1), cell_w, cell_h,
+                  "W_Q", f"({m.hidden}, {m.h_q * m.d_head})",
+                  row_splits=1, col_splits=tp,
+                  my_col_shard=q_shard,
+                  col_line_color=_TP_LINE_COLOR,
+                  col_label=f"cols split by {tp_txt}",
+                  note="col-parallel (output dim)",
+                  cell_annot=_tp_col_ann)
+
+    # W_K, W_V: (hidden, H_kv*d_head)
+    for (name, col_idx) in [("W_K", 1), ("W_V", 2)]:
+        splits = min(tp, m.h_kv) if m.h_kv >= tp else tp
+        _draw_tensor(ax, px(col_idx), py(1), cell_w, cell_h,
+                      name, f"({m.hidden}, {m.h_kv * m.d_head})",
+                      row_splits=1, col_splits=splits,
+                      my_col_shard=min(q_shard, splits - 1),
+                      col_line_color=_TP_LINE_COLOR,
+                      col_label=f"cols split by {tp_txt}",
+                      note="col-parallel (heads)",
+                      cell_annot=_tp_col_ann)
+
+    # W_O: (H_q*d_head, hidden), split on ROW by TP
+    _draw_tensor(ax, px(3), py(1), cell_w, cell_h,
+                  "W_O", f"({m.h_q * m.d_head}, {m.hidden})",
+                  row_splits=tp, col_splits=1,
+                  my_row_shard=q_shard,
+                  row_line_color=_TP_LINE_COLOR,
+                  row_label=f"rows split by {tp_txt}",
+                  note="row-parallel (input dim)",
+                  cell_annot=_tp_row_ann)
+
+    # ── FFN tensors ──────────────────────────────
+    # W_gate, W_up: (hidden, ffn_dim), split on COL by FFN scope
+    for (name, col_idx) in [("W_gate", 0), ("W_up", 1)]:
+        _draw_tensor(ax, px(col_idx), py(0), cell_w, cell_h,
+                      name, f"({m.hidden}, {m.ffn_dim})",
+                      row_splits=1, col_splits=cfg.ffn_shard_divisor,
+                      my_col_shard=ffn_shard,
+                      col_line_color=_FFN_LINE_COLOR,
+                      col_label=f"cols split by {ffn_txt}",
+                      note="col-parallel")
+
+    # W_down: (ffn_dim, hidden), split on ROW by FFN scope
+    _draw_tensor(ax, px(2), py(0), cell_w, cell_h,
+                  "W_down", f"({m.ffn_dim}, {m.hidden})",
+                  row_splits=cfg.ffn_shard_divisor, col_splits=1,
+                  my_row_shard=ffn_shard,
+                  row_line_color=_FFN_LINE_COLOR,
+                  row_label=f"rows split by {ffn_txt}",
+                  note="row-parallel")
+
+    # ── KV cache: (S_kv, H_kv*d_head), 2D shard ─
+    kv_row_splits = cp
+    kv_col_splits = min(tp, m.h_kv) if m.h_kv >= tp else tp
+    _draw_tensor(ax, px(3), py(0), cell_w, cell_h,
+                  "KV cache (K and V)",
+                  f"({m.h_kv * m.d_head}, S_kv={cfg.topo.s_kv:,})",
+                  row_splits=kv_row_splits,
+                  col_splits=kv_col_splits if show_physical else 1,
+                  my_row_shard=kv_row_shard,
+                  row_line_color=_CP_LINE_COLOR,
+                  row_label=f"rows: S split by {cp_txt}",
+                  col_label=f"cols: heads split by {tp_txt}",
+                  col_line_color=_TP_LINE_COLOR,
+                  note="2D shard: seq axis (CP) x head axis (TP)",
+                  cell_annot=_kv_ann)
+    # In non-physical mode, KV was drawn with col_splits=1; overlay TP col
+    # splits as dotted lines to hint at the 2D nature. In physical mode
+    # the _draw_tensor call already drew proper vertical dashes.
+    if not show_physical:
+        kv_x = px(3)
+        kv_y = py(0)
+        shard_col_w = cell_w / max(1, kv_col_splits)
+        for c in range(1, kv_col_splits):
+            xx = kv_x + c * shard_col_w
+            ax.plot([xx, xx], [kv_y, kv_y + cell_h],
+                    color=_TP_LINE_COLOR, linewidth=1.2, linestyle=":")
+
+    ax.set_xlim(-0.3, 4 * (cell_w + gap_x) - gap_x + 0.3)
+    ax.set_ylim(-0.9, 2 * (cell_h + gap_y) - gap_y + 0.9)
+    ax.set_aspect("auto")
+    ax.axis("off")
+
+    ax.set_title(
+        f"Tensor sharding view (PE {my_pe})  |  TP={tp}, CP={cp}, "
+        f"FFN scope={cfg.topo.ffn_shard_scope}  |  "
+        f"blue = this PE's shard",
+        fontsize=11, fontweight="bold",
+    )
+    # Legend for line colors
+    ax.plot([], [], color=_TP_LINE_COLOR, linewidth=1.2, linestyle="--",
+            label=f"TP split ({tp} ranks)")
+    ax.plot([], [], color=_CP_LINE_COLOR, linewidth=1.2, linestyle="--",
+            label=f"CP split ({cp} ranks, sequence dim)")
+    ax.plot([], [], color=_FFN_LINE_COLOR, linewidth=1.2, linestyle="--",
+            label=f"FFN scope split ({cfg.ffn_shard_divisor} ranks)")
+    ax.legend(loc="lower center", bbox_to_anchor=(0.5, -0.05),
+              ncol=3, fontsize=8, frameon=False)
+
+    return fig

tests/analytical_visualization/topology_map.py

@@ -0,0 +1,611 @@
+"""Draw the SIP topology diagram, color-coded by parallelism group.
+
+Color scheme:
+  - CP ring:   orange arrows (unchanged)
+  - Inter-SIP: red arrows across SIP boundaries
+  - PP stage:  cube border color (each PP stage gets a distinct hue)
+  - TP group:  PE fill color (each TP group within a cube gets one hue,
+               matches its cube's PP stage but slightly desaturated)
+  - EP group:  a small badge on the cube (for MoE)
+  - DP replica: cubes of each DP replica get a hatched pattern
+"""
+from __future__ import annotations
+
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
+from matplotlib.lines import Line2D
+
+from .model_config import FullConfig
+
+
+MESH_W = 4
+MESH_H = 4
+PE_COLS = 4
+PE_ROWS = 2
+
+
+def _group_label(cfg, group_idx: int) -> str:
+    """Label for a cube-level group under the current placement.
+    - cp_placement=cube -> CP rank (or CP/TP if both on cube)
+    - cp_placement=pe, tp_placement=cube -> TP rank
+    - both on pe -> single group ("PE")
+    """
+    parts = []
+    if cfg.topo.cp_placement == "cube":
+        cp_r = group_idx % max(1, cfg.topo.cp)
+        parts.append(f"CP{cp_r}")
+    if cfg.topo.tp_placement == "cube":
+        div = cfg.topo.cp if cfg.topo.cp_placement == "cube" else 1
+        tp_r = (group_idx // max(1, div)) % max(1, cfg.topo.tp)
+        parts.append(f"TP{tp_r}")
+    return "/".join(parts) if parts else "PE"
+
+# Fully-distinct palette for (PP stage x CP rank) — 24 unique colors.
+# Each cube's (pp, cp) pair maps to one entry; wraps if you exceed 24 groups.
+_GROUP_PALETTE = [
+    "#e6194b", "#3cb44b", "#ffe119", "#4363d8",
+    "#f58231", "#911eb4", "#42d4f4", "#f032e6",
+    "#bfef45", "#fabed4", "#469990", "#dcbeff",
+    "#9a6324", "#fffac8", "#800000", "#aaffc3",
+    "#808000", "#ffd8b1", "#000075", "#a9a9a9",
+    "#004d40", "#c62828", "#4527a0", "#00695c",
+]
+
+# Kept for backward-compat / PP-stage arrows
+_PP_HUES = _GROUP_PALETTE[:8]
+
+
+def _shade(hue_hex: str, t: float) -> str:
+    """Mix hue_hex toward white by t in [0, 1] (0 = original, 1 = white)."""
+    h = hue_hex.lstrip("#")
+    r, g, b = int(h[0:2], 16), int(h[2:4], 16), int(h[4:6], 16)
+    rn = int(r + (255 - r) * t)
+    gn = int(g + (255 - g) * t)
+    bn = int(b + (255 - b) * t)
+    return f"#{rn:02x}{gn:02x}{bn:02x}"
+
+
+def _cp_color(pp_stage: int, cp_rank: int, cp_size: int) -> tuple[str, str]:
+    """Return (border_color, pe_fill_color) for a cube at (pp_stage, cp_rank).
+
+    Each unique (pp, cp) pair gets a **fully distinct color** from the
+    24-entry palette. Border is the pure color; PE fill is a lightened
+    version so the cube border stays visible around the PE grid.
+    """
+    idx = pp_stage * max(1, cp_size) + cp_rank
+    color = _GROUP_PALETTE[idx % len(_GROUP_PALETTE)]
+    pe_fill = _shade(color, 0.30)   # slightly lightened for PE fill
+    return color, pe_fill
+
+# Inactive (unused) colors.
+_INACTIVE_CUBE_FACE = "#f7f7f7"
+_INACTIVE_CUBE_EDGE = "#c8c8c8"
+_INACTIVE_PE = "#eaeaea"
+
+
+def _snake_path(n: int, mesh_w: int, mesh_h: int) -> list[int]:
+    path: list[int] = []
+    for r in range(mesh_h):
+        cols = range(mesh_w) if r % 2 == 0 else range(mesh_w - 1, -1, -1)
+        for c in cols:
+            path.append(r * mesh_w + c)
+            if len(path) == n:
+                return path
+    return path[:n]
+
+
+def _rect_shape(n: int, max_w: int, max_h: int) -> tuple[int, int]:
+    """Return (rows, cols) for n cubes packed as square as possible in max_wxmax_h."""
+    if n <= 0:
+        return (0, 0)
+    # Prefer near-square shapes; cap at mesh dims.
+    best = None
+    for cols in range(1, min(n, max_w) + 1):
+        rows = (n + cols - 1) // cols
+        if rows > max_h:
+            continue
+        # Score: penalise non-squareness + wasted cells
+        waste = rows * cols - n
+        aspect = abs(rows - cols)
+        score = aspect * 10 + waste
+        if best is None or score < best[0]:
+            best = (score, rows, cols)
+    if best is None:
+        # Fallback: single row
+        return (1, min(n, max_w))
+    return best[1], best[2]
+
+
+def _pack_groups_2d(items_per_group: int, n_groups: int,
+                     mesh_w: int, mesh_h: int,
+                     layout_mode: str = "compact") -> list[tuple[int, int]]:
+    """Return list of (row, col) positions for n_groups × items_per_group
+    cubes in a mesh_w x mesh_h mesh.
+
+    layout_mode:
+      - "compact": each group is a near-square rectangle; groups are also
+        tiled 2D as near-square. Ideal for 2x2, 2x4 etc.
+      - "linear": each group is a single row of cubes; groups tile down
+        as rows (original sequential layout).
+    """
+    if n_groups == 0 or items_per_group == 0:
+        return []
+
+    if layout_mode == "linear":
+        # Each group is a single horizontal row of cubes.
+        tp_rows, tp_cols = 1, min(items_per_group, mesh_w)
+    else:  # compact
+        tp_rows, tp_cols = _rect_shape(items_per_group, mesh_w, mesh_h)
+    if tp_rows == 0:
+        return []
+
+    # Group grid: how many groups per row/col of GROUPS
+    max_group_cols = max(1, mesh_w // tp_cols)
+    max_group_rows = max(1, mesh_h // tp_rows)
+    if layout_mode == "linear":
+        # Row-major: fill each mesh row with groups, then wrap down.
+        g_cols = max_group_cols
+        g_rows = max_group_rows
+    else:
+        g_rows, g_cols = _rect_shape(n_groups, max_group_cols, max_group_rows)
+
+    positions: list[tuple[int, int]] = []
+    for g in range(n_groups):
+        gr = g // g_cols
+        gc = g % g_cols
+        base_row = gr * tp_rows
+        base_col = gc * tp_cols
+        for i in range(items_per_group):
+            r = i // tp_cols
+            c = i % tp_cols
+            row = base_row + r
+            col = base_col + c
+            if row >= mesh_h or col >= mesh_w:
+                fallback = (g * items_per_group + i)
+                row = fallback // mesh_w
+                col = fallback % mesh_w
+            positions.append((row, col))
+    return positions
+
+
+def _draw_one_sip(ax, cfg: FullConfig, sip_idx: int,
+                   sip_x0: float, sip_y0: float,
+                   sip_width: float, sip_height: float,
+                   cube_pp_cp: dict[int, tuple[int, int]],
+                   ring_paths: dict[int, list[int]],
+                   tp_group_cubes: dict[int, list[int]]):
+    """Draw one SIP with PP hue + CP shade + TP group boundaries."""
+    cube_size = min(sip_width / MESH_W, sip_height / MESH_H) * 0.85
+    cube_gap = cube_size * 0.15
+    total_w = MESH_W * (cube_size + cube_gap) - cube_gap
+    total_h = MESH_H * (cube_size + cube_gap) - cube_gap
+    ox = sip_x0 + (sip_width - total_w) / 2
+    oy = sip_y0 + (sip_height - total_h) / 2
+
+    sip_rect = patches.FancyBboxPatch(
+        (sip_x0, sip_y0), sip_width, sip_height,
+        boxstyle="round,pad=0.02", facecolor="#f8f9fa",
+        edgecolor="#212529", linewidth=1.5,
+    )
+    ax.add_patch(sip_rect)
+    ax.text(sip_x0 + sip_width / 2, sip_y0 + sip_height + 0.05,
+            f"SIP {sip_idx}", ha="center", va="bottom",
+            fontsize=10, fontweight="bold")
+
+    def _cube_center(cube_local: int) -> tuple[float, float]:
+        r = cube_local // MESH_W
+        c = cube_local % MESH_W
+        x = ox + c * (cube_size + cube_gap)
+        y = oy + (MESH_H - 1 - r) * (cube_size + cube_gap)
+        return x, y
+
+    for cube_local in range(MESH_W * MESH_H):
+        x, y = _cube_center(cube_local)
+        info = cube_pp_cp.get(cube_local, None)
+        is_used = info is not None
+
+        if is_used:
+            pp_stage, cp_rank = info
+            cube_edge, pe_fill = _cp_color(
+                pp_stage, cp_rank,
+                max(1, cfg.topo.inter_cube_dims),
+            )
+            cube_face = _shade(cube_edge, 0.85)
+            cube_lw = 2.0
+        else:
+            cube_edge = _INACTIVE_CUBE_EDGE
+            cube_face = _INACTIVE_CUBE_FACE
+            cube_lw = 0.6
+            pe_fill = _INACTIVE_PE
+
+        rect = patches.FancyBboxPatch(
+            (x, y), cube_size, cube_size,
+            boxstyle="round,pad=0.01",
+            facecolor=cube_face, edgecolor=cube_edge, linewidth=cube_lw,
+        )
+        ax.add_patch(rect)
+        if is_used:
+            pp_stage, cp_rank = info
+            # Small cube ID above
+            ax.text(x + cube_size / 2, y + cube_size + 0.01,
+                    f"cube {cube_local}", ha="center", va="bottom",
+                    fontsize=5.5, color=cube_edge, fontweight="bold")
+            # Group label (CP/TP depending on placement)
+            _lbl = _group_label(cfg, cp_rank)
+            ax.text(x + 0.03, y + cube_size - 0.03,
+                    _lbl, ha="left", va="top",
+                    fontsize=11, color=cube_edge, fontweight="bold",
+                    bbox=dict(boxstyle="round,pad=0.15",
+                              facecolor="white", edgecolor=cube_edge,
+                              linewidth=1.0, alpha=0.9))
+            if cfg.topo.pp > 1:
+                ax.text(x + cube_size - 0.03, y + cube_size - 0.03,
+                        f"PP{pp_stage}", ha="right", va="top",
+                        fontsize=9, color=cube_edge, fontweight="bold",
+                        bbox=dict(boxstyle="round,pad=0.1",
+                                  facecolor="white", edgecolor=cube_edge,
+                                  linewidth=0.8, alpha=0.9))
+        else:
+            ax.text(x + cube_size / 2, y + cube_size + 0.008,
+                    f"{cube_local}", ha="center", va="bottom",
+                    fontsize=6, color=cube_edge)
+
+        # PEs
+        pe_pad = cube_size * 0.08
+        pe_gap = cube_size * 0.02
+        inner_w = cube_size - 2 * pe_pad
+        inner_h = cube_size - 2 * pe_pad
+        pe_w = (inner_w - (PE_COLS - 1) * pe_gap) / PE_COLS
+        pe_h = (inner_h - (PE_ROWS - 1) * pe_gap) / PE_ROWS
+        pes_used = cfg.topo.pes_per_cube_used if is_used else 0
+        for pr in range(PE_ROWS):
+            for pc in range(PE_COLS):
+                pe_id = pr * PE_COLS + pc
+                px = x + pe_pad + pc * (pe_w + pe_gap)
+                py = y + pe_pad + (PE_ROWS - 1 - pr) * (pe_h + pe_gap)
+                fill = pe_fill if pe_id < pes_used else _INACTIVE_PE
+                pe_rect = patches.Rectangle(
+                    (px, py), pe_w, pe_h,
+                    facecolor=fill, edgecolor="#666", linewidth=0.3,
+                )
+                ax.add_patch(pe_rect)
+
+        if is_used and cfg.topo.ep > 1:
+            ax.text(x + cube_size - 0.02, y + cube_size - 0.02,
+                    f"EP{cfg.topo.ep}", ha="right", va="top",
+                    fontsize=6, color="#5f0f40", fontweight="bold")
+
+    # (CP ring arrows removed per user request)
+
+    # Draw TP group bounding boxes (RED dashed) when TP > 8
+    for tp_g, cubes in tp_group_cubes.items():
+        if len(cubes) <= 1:
+            continue  # TP fits in one cube, redundant with cube border
+        # Bounding box around all cubes in this TP group
+        xs, ys = [], []
+        for cube_local in cubes:
+            x, y = _cube_center(cube_local)
+            xs.extend([x, x + cube_size])
+            ys.extend([y, y + cube_size])
+        pad = 0.03
+        bbox_x = min(xs) - pad
+        bbox_y = min(ys) - pad
+        bbox_w = (max(xs) - min(xs)) + 2 * pad
+        bbox_h = (max(ys) - min(ys)) + 2 * pad
+        tp_rect = patches.Rectangle(
+            (bbox_x, bbox_y), bbox_w, bbox_h,
+            fill=False, edgecolor="#d90429", linewidth=1.8,
+            linestyle="--",
+        )
+        ax.add_patch(tp_rect)
+        ax.text(bbox_x + bbox_w / 2, bbox_y - 0.05,
+                f"TP group {tp_g}", ha="center", va="top",
+                color="#d90429", fontsize=7, fontweight="bold")
+
+
+def _sip_grid_layout(n_sips: int, topology: str) -> tuple[int, int]:
+    """(rows, cols) for arranging SIPs on the page.
+    - ring: horizontal chain (1 x N)
+    - mesh2d / torus2d: near-square grid
+    """
+    if n_sips <= 0:
+        return 0, 0
+    if topology == "ring" or n_sips <= 2:
+        return 1, n_sips
+    import math
+    cols = int(math.ceil(math.sqrt(n_sips)))
+    rows = (n_sips + cols - 1) // cols
+    return rows, cols
+
+
+def _draw_sip_links(ax, sip_xy: list[tuple[float, float]],
+                     topology: str, sip_w: float, sip_h: float,
+                     grid_rows: int, grid_cols: int):
+    """Draw inter-SIP interconnect lines.
+
+    Solid lines: adjacent (grid-neighbor) links.
+    Dashed lines: wrap-around links (ring, torus).
+    """
+    n = len(sip_xy)
+    if n <= 1:
+        return
+    color = "#c62828"
+    lw = 2.0
+
+    def center(i):
+        x, y = sip_xy[i]
+        return x + sip_w / 2, y + sip_h / 2
+
+    if topology == "ring":
+        # 1D chain (adjacent solid)
+        for i in range(n - 1):
+            x1, y1 = center(i)
+            x2, y2 = center(i + 1)
+            ax.plot([x1 + sip_w / 2, x2 - sip_w / 2], [y1, y2],
+                    color=color, linewidth=lw)
+        # wrap-around from last back to first (dashed arc below)
+        if n > 2:
+            x_first, y_first = center(0)
+            x_last, y_last = center(n - 1)
+            arc_y = min(y_first, y_last) - sip_h / 2 - 0.35
+            ax.plot([x_first, x_first], [y_first - sip_h / 2, arc_y],
+                    color=color, linewidth=lw, linestyle="--")
+            ax.plot([x_first, x_last], [arc_y, arc_y],
+                    color=color, linewidth=lw, linestyle="--")
+            ax.plot([x_last, x_last], [arc_y, y_last - sip_h / 2],
+                    color=color, linewidth=lw, linestyle="--")
+        elif n == 2:
+            # 2 SIPs: single link
+            x1, y1 = center(0)
+            x2, y2 = center(1)
+            # wrap link: dashed loop below
+            arc_y = y1 - sip_h / 2 - 0.35
+            ax.plot([x1, x1], [y1 - sip_h / 2, arc_y],
+                    color=color, linewidth=lw, linestyle="--")
+            ax.plot([x1, x2], [arc_y, arc_y],
+                    color=color, linewidth=lw, linestyle="--")
+            ax.plot([x2, x2], [arc_y, y2 - sip_h / 2],
+                    color=color, linewidth=lw, linestyle="--")
+        return
+
+    # mesh2d / torus2d: grid neighbours
+    def idx(r, c):
+        return r * grid_cols + c
+
+    for r in range(grid_rows):
+        for c in range(grid_cols):
+            i = idx(r, c)
+            if i >= n:
+                continue
+            # right
+            if c + 1 < grid_cols and idx(r, c + 1) < n:
+                x1, y1 = center(i)
+                x2, y2 = center(idx(r, c + 1))
+                ax.plot([x1 + sip_w / 2, x2 - sip_w / 2], [y1, y2],
+                        color=color, linewidth=lw)
+            # down (visually — grid row+1 is drawn below)
+            if r + 1 < grid_rows and idx(r + 1, c) < n:
+                x1, y1 = center(i)
+                x2, y2 = center(idx(r + 1, c))
+                ax.plot([x1, x2], [y1 - sip_h / 2, y2 + sip_h / 2],
+                        color=color, linewidth=lw)
+
+    if topology == "torus2d":
+        # row wrap: last col -> first col in each row (dashed arc BELOW that row)
+        for r in range(grid_rows):
+            left_i = idx(r, 0)
+            right_i = idx(r, grid_cols - 1)
+            if left_i >= n or right_i >= n or left_i == right_i:
+                continue
+            x1, y1 = center(left_i)
+            x2, y2 = center(right_i)
+            arc_y = min(y1, y2) - sip_h / 2 - 0.35
+            ax.plot([x1, x1], [y1 - sip_h / 2, arc_y],
+                    color=color, linewidth=lw, linestyle="--")
+            ax.plot([x1, x2], [arc_y, arc_y],
+                    color=color, linewidth=lw, linestyle="--")
+            ax.plot([x2, x2], [arc_y, y2 - sip_h / 2],
+                    color=color, linewidth=lw, linestyle="--")
+        # col wrap: last row -> first row in each col (dashed arc to the RIGHT)
+        for c in range(grid_cols):
+            top_i = idx(0, c)
+            bot_i = idx(grid_rows - 1, c)
+            if top_i >= n or bot_i >= n or top_i == bot_i:
+                continue
+            x1, y1 = center(top_i)   # top row (higher y)
+            x2, y2 = center(bot_i)   # bottom row (lower y)
+            arc_x = max(x1, x2) + sip_w / 2 + 0.35
+            ax.plot([x1 + sip_w / 2, arc_x], [y1, y1],
+                    color=color, linewidth=lw, linestyle="--")
+            ax.plot([arc_x, arc_x], [y1, y2],
+                    color=color, linewidth=lw, linestyle="--")
+            ax.plot([arc_x, x2 + sip_w / 2], [y2, y2],
+                    color=color, linewidth=lw, linestyle="--")
+
+
+def draw_topology(cfg: FullConfig, ax=None, layout_mode: str = "linear"):
+    """Draw one or more SIPs. Each (PP,CP) group gets a unique color.
+
+    SIP grid arrangement + inter-SIP links depend on cfg.topo.sip_topology:
+      "ring"    -> 1D horizontal chain, wrap arc below
+      "mesh2d"  -> near-square grid, no wrap
+      "torus2d" -> grid + row/col wrap arcs
+    """
+    sips_used = max(1, cfg.topo.sips_used)
+    total_pes = cfg.topo.total_pes
+    pes_per_stage = cfg.topo.pes_per_stage    # CP*TP
+    n_replicas = cfg.topo.dp
+    n_pp_stages = cfg.topo.pp
+    pes_per_cube = cfg.topo.pes_per_cube_hw
+    # Placement-aware: how many cubes per "group" (was: cubes-per-TP).
+    # A group is one cube-level unit (one CP rank if cp on cube, or one TP
+    # rank if only tp on cube, etc.). Cubes per group == intra-cube spill.
+    cubes_per_group = max(
+        1, (cfg.topo.intra_cube_dims + pes_per_cube - 1) // pes_per_cube
+    )
+    n_groups_per_stage = max(1, cfg.topo.inter_cube_dims)
+    cubes_per_stage = cfg.topo.cubes_per_stage
+    cubes_used = cfg.topo.cubes_used
+    cubes_per_sip = MESH_W * MESH_H
+
+    sip_topo = cfg.topo.sip_topology
+    grid_rows, grid_cols = _sip_grid_layout(sips_used, sip_topo)
+
+    sip_w = 4.0
+    sip_h = 4.0
+    sip_gap = 0.6
+
+    fig_w = grid_cols * sip_w + (grid_cols - 1) * sip_gap + 1.5
+    fig_h = grid_rows * sip_h + (grid_rows - 1) * sip_gap + 2.0
+    if ax is None:
+        fig, ax = plt.subplots(figsize=(max(6, fig_w), max(6, fig_h)))
+    else:
+        fig = ax.figure
+
+    # 2D packing: each group's cubes form a rectangle inside the SIP mesh.
+    # If more groups than fit in one SIP, spill to next SIP.
+    tp_rows, tp_cols = _rect_shape(cubes_per_group, MESH_W, MESH_H)
+    if tp_rows == 0:
+        tp_rows, tp_cols = 1, 1
+    groups_per_sip = max(1, (MESH_W // tp_cols) * (MESH_H // tp_rows))
+
+    cube_pp_cp_global: dict[int, tuple[int, int]] = {}
+    ring_globals_by_pp: dict[int, list[int]] = {}
+    tp_group_cubes_global: dict[int, list[int]] = {}
+
+    current_sip = 0
+    for rep in range(n_replicas):
+        for pp_s in range(n_pp_stages):
+            # Chunk cube-level groups into per-SIP batches.
+            for chunk_start in range(0, n_groups_per_stage, groups_per_sip):
+                chunk = min(groups_per_sip, n_groups_per_stage - chunk_start)
+                positions = _pack_groups_2d(cubes_per_group, chunk,
+                                             MESH_W, MESH_H,
+                                             layout_mode=layout_mode)
+                for local_gi in range(chunk):
+                    group_idx = chunk_start + local_gi
+                    tp_gid = ((rep * n_pp_stages) + pp_s) * n_groups_per_stage + group_idx
+                    cubes_here: list[int] = []
+                    for i in range(cubes_per_group):
+                        idx_in_group = local_gi * cubes_per_group + i
+                        r, c = positions[idx_in_group]
+                        local_cube = r * MESH_W + c
+                        gc = current_sip * cubes_per_sip + local_cube
+                        cube_pp_cp_global[gc] = (pp_s, group_idx)
+                        cubes_here.append(gc)
+                    tp_group_cubes_global[tp_gid] = cubes_here
+                    ring_globals_by_pp.setdefault(pp_s, []).append(cubes_here[0])
+                current_sip += 1  # next SIP for the next chunk (or (rep, pp))
+
+    # Partition (pp, cp) map by SIP using cube_positions_global
+    cubes_by_sip: list[dict[int, tuple[int, int]]] = [dict() for _ in range(sips_used)]
+    for gc, ppcp in cube_pp_cp_global.items():
+        sip_idx = gc // cubes_per_sip
+        local_cube = gc % cubes_per_sip
+        if sip_idx < sips_used:
+            cubes_by_sip[sip_idx][local_cube] = ppcp
+
+    # Per-SIP snake-path CP rings (organized by pp_stage)
+    ring_by_sip: list[dict[int, list[int]]] = [dict() for _ in range(sips_used)]
+    for pp_s, gc_list in ring_globals_by_pp.items():
+        by_sip_local: dict[int, list[int]] = {}
+        for gc in gc_list:
+            si = gc // cubes_per_sip
+            lc = gc % cubes_per_sip
+            by_sip_local.setdefault(si, []).append(lc)
+        for si, locals_list in by_sip_local.items():
+            if si < sips_used:
+                # Snake through the packed positions in order
+                ring_by_sip[si][pp_s] = locals_list
+
+    # TP group cubes by SIP
+    tp_by_sip: list[dict[int, list[int]]] = [dict() for _ in range(sips_used)]
+    for tp_g, gc_list in tp_group_cubes_global.items():
+        by_sip_local: dict[int, list[int]] = {}
+        for gc in gc_list:
+            si = gc // cubes_per_sip
+            lc = gc % cubes_per_sip
+            by_sip_local.setdefault(si, []).append(lc)
+        for si, locals_list in by_sip_local.items():
+            if si < sips_used:
+                tp_by_sip[si][tp_g] = locals_list
+
+    # Compute (x, y) top-left for each SIP based on grid layout.
+    # Row 0 is drawn at the TOP visually, so y decreases with row index.
+    sip_xy: list[tuple[float, float]] = []
+    for sip_idx in range(sips_used):
+        r = sip_idx // grid_cols
+        c = sip_idx % grid_cols
+        x = c * (sip_w + sip_gap)
+        y = (grid_rows - 1 - r) * (sip_h + sip_gap)
+        sip_xy.append((x, y))
+
+    for sip_idx in range(sips_used):
+        x0, y0 = sip_xy[sip_idx]
+        _draw_one_sip(
+            ax, cfg, sip_idx,
+            sip_x0=x0, sip_y0=y0,
+            sip_width=sip_w, sip_height=sip_h,
+            cube_pp_cp=cubes_by_sip[sip_idx],
+            ring_paths=ring_by_sip[sip_idx],
+            tp_group_cubes=tp_by_sip[sip_idx],
+        )
+
+    # Inter-SIP interconnect links per user-selected topology.
+    _draw_sip_links(ax, sip_xy, sip_topo, sip_w, sip_h, grid_rows, grid_cols)
+
+    total_w = grid_cols * sip_w + (grid_cols - 1) * sip_gap
+    total_h = grid_rows * sip_h + (grid_rows - 1) * sip_gap
+    # Extra room: bottom for wrap arcs, right for torus col-wrap arcs, top for SIP labels
+    ax.set_xlim(-0.3, total_w + 0.9)
+    ax.set_ylim(-0.9, total_h + 0.7)
+    ax.set_aspect("equal")
+    ax.axis("off")
+
+    # Legend — one entry per unique (PP, cube-group) + TP boundary + SIP links
+    legend_handles = []
+    for pp_s in range(n_pp_stages):
+        for gi in range(n_groups_per_stage):
+            color, _ = _cp_color(pp_s, gi, n_groups_per_stage)
+            legend_handles.append(Line2D(
+                [], [], color=color, marker="s", linestyle="",
+                markersize=10,
+                label=(f"PP{pp_s} " if n_pp_stages > 1 else "")
+                       + _group_label(cfg, gi),
+            ))
+    if cfg.topo.tp > cfg.topo.pes_per_cube_hw:
+        legend_handles.append(Line2D(
+            [], [], color="#d90429", marker="s", linestyle="--",
+            markersize=10, markerfacecolor="none", markeredgewidth=1.5,
+            label="TP group (dashed red)",
+        ))
+    if cfg.topo.ep > 1:
+        legend_handles.append(Line2D(
+            [], [], color="#5f0f40", marker="s", linestyle="",
+            markersize=8, label=f"EP={cfg.topo.ep}",
+        ))
+    if sips_used > 1:
+        _topo_label = {
+            "ring": "Inter-SIP: ring (chain + wrap)",
+            "mesh2d": "Inter-SIP: 2D mesh (no wrap)",
+            "torus2d": "Inter-SIP: 2D torus (grid + wrap)",
+        }.get(sip_topo, f"Inter-SIP: {sip_topo}")
+        legend_handles.append(Line2D(
+            [], [], color="#c62828", linewidth=2.0,
+            label=_topo_label,
+        ))
+    # Legend below the figure; wrap ncols so it's readable
+    n_items = len(legend_handles)
+    ncol = min(6, max(3, n_items))
+    ax.legend(handles=legend_handles, loc="upper center",
+              bbox_to_anchor=(0.5, -0.02), ncol=ncol,
+              fontsize=8, frameon=False)
+
+    title = (
+        f"Layout: {total_pes} PEs = {n_replicas} DP replica(s) x "
+        f"{n_pp_stages} PP stage(s) x {cubes_per_stage} cubes (CP={cfg.topo.cp}) "
+        f"x TP={cfg.topo.tp} | cubes: {cubes_used}, SIPs: {sips_used}"
+    )
+    ax.set_title(title, fontsize=10)
+
+    return fig

tests/attention/test_gqa_decode_batch_sweep.py

@@ -0,0 +1,193 @@
+"""Batch GQA decode sweep — concurrent independent users on ONE SIP.
+
+Fixes the target to a single 16-cube SIP and scales the batch B = number
+of concurrent decode users, each with its own (Q, K, V) placed on a
+disjoint cube group ``[u*C, u*C + C)`` via ``DPPolicy.cube_start`` and
+addressed by the kernel's ``cube_base`` local-index parameter. All B
+launches are deferred (``_defer_wait=True``) so they overlap on the shared
+discrete-event clock; ``run_bench`` drains them together.
+
+Per mapping the SIP holds ``16 // C`` users at once (A1=2, A2=4, A4=8,
+B=16). We sweep B up to that capacity and record aggregate latency,
+throughput (users / latency), and mean per-user latency, to see which
+mapping's throughput scales best with batch.
+
+Metric convention matches the corrected short-context sweeps:
+``wall = max(t_end) - min(t_start)`` over the op log (deploy excluded).
+"""
+from __future__ import annotations
+
+import csv
+from pathlib import Path
+
+import pytest
+
+from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config
+from kernbench.ccl.sfr_config import configure_sfr_intercube_multisip
+from kernbench.policy.placement.dp import DPPolicy
+from kernbench.runtime_api.bench_runner import run_bench
+from kernbench.runtime_api.types import resolve_device
+from kernbench.sim_engine.engine import GraphEngine
+from kernbench.topology.builder import resolve_topology
+
+TOPOLOGY_DEFAULT = Path(__file__).resolve().parents[2] / "topology.yaml"
+
+D_HEAD = 64
+DTYPE = "f16"
+TILE_S_KV = 1024
+P = 8
+H_KV = 8
+CUBES_PER_SIP = 16
+
+# (mode, kv_per_cube, C); users-per-SIP capacity = CUBES_PER_SIP // C.
+MODES = [("A1", 1, 8), ("A2", 2, 4), ("A4", 4, 2), ("B", 8, 1)]
+CONTEXT_LENGTHS = [8 * 1024]
+
+CSV_OUT = (Path(__file__).resolve().parents[2]
+           / "docs" / "sweeps" / "decode_batch_sweep.csv")
+
+
+def _ccl_cfg():
+    return resolve_algorithm_config(
+        load_ccl_config(), name="lrab_hierarchical_allreduce",
+    )
+
+
+def _engine_factory(t, d):
+    return GraphEngine(getattr(t, "topology_obj", t), enable_data=True)
+
+
+def _run_batch(*, kv_per_cube: int, C: int, S_kv: int, B: int, h_q: int = 64):
+    """Launch B concurrent decode users on disjoint cube groups of SIP 0."""
+    from kernbench.benches.gqa_helpers.short_ctx._gqa_attention_decode_short import (
+        _validate_config as _validate_decode_config,
+        gqa_attention_decode_short_kernel,
+    )
+    T_q = 1
+    _validate_decode_config(kv_per_cube=kv_per_cube, T_q=T_q, P=P, C=C,
+                            h_q=h_q, h_kv=H_KV, S_kv=S_kv)
+    Q_ROWS = kv_per_cube * T_q
+    Q_COLS = (h_q * D_HEAD) // kv_per_cube
+    k_rows = (H_KV * S_kv * D_HEAD) // TILE_S_KV
+    topo = resolve_topology(str(TOPOLOGY_DEFAULT))
+
+    def _bench_fn(ctx):
+        configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
+        # Deploy ALL users' tensors first. Each deploy blocks (advances the
+        # clock), so if a launch were already submitted a later deploy would
+        # drive it to completion and serialize the batch. Creating every
+        # tensor before any launch keeps the deploys from touching kernels.
+        users = []
+        for u in range(B):
+            cs = u * C
+            dp_q = DPPolicy(cube="column_wise", pe="replicate",
+                            num_cubes=C, num_pes=P, cube_start=cs)
+            dp_kv = DPPolicy(cube="row_wise", pe="row_wise",
+                             num_cubes=C, num_pes=P, cube_start=cs)
+            dp_o = DPPolicy(cube="column_wise", pe="replicate",
+                            num_cubes=C, num_pes=P, cube_start=cs)
+            q = ctx.zeros((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_q,
+                          name=f"q_u{u}")
+            k = ctx.zeros((k_rows, TILE_S_KV), dtype=DTYPE, dp=dp_kv,
+                          name=f"k_u{u}")
+            v = ctx.zeros((H_KV * S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv,
+                          name=f"v_u{u}")
+            o = ctx.empty((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_o,
+                          name=f"o_u{u}")
+            users.append((u, cs, q, k, v, o))
+        # Now submit every launch deferred (no blocking wait between them),
+        # then drain together so all B run concurrently on the shared clock.
+        pending = []
+        for u, cs, q, k, v, o in users:
+            pending += ctx.launch(
+                f"gqa_decode_u{u}",
+                gqa_attention_decode_short_kernel,
+                q, k, v, o,
+                T_q, S_kv, h_q, H_KV, D_HEAD, C, P, kv_per_cube, cs,
+                _auto_dim_remap=False,
+                _defer_wait=True,
+            )
+        for h, _sip_id, meta in pending:
+            ctx.wait(h, _meta=meta)
+
+    return run_bench(
+        topology=topo, bench_fn=_bench_fn,
+        device=resolve_device("sip:0"), engine_factory=_engine_factory,
+    )
+
+
+def _cube_of(rec) -> int:
+    """Physical cube index of an op record, or -1 if not a PE op."""
+    cid = rec.component_id or ""
+    if ".cube" in cid:
+        try:
+            return int(cid.split(".cube")[1].split(".")[0])
+        except (ValueError, IndexError):
+            return -1
+    return -1
+
+
+def _metrics(r, *, mode, kv_per_cube, C, S_kv, B):
+    op_log = r.engine.op_log
+    wall_ns = max(rec.t_end for rec in op_log) - min(rec.t_start for rec in op_log)
+
+    # Per-user latency: group ops by which user's cube band they touch.
+    per_user = []
+    for u in range(B):
+        lo, hi = u * C, u * C + C
+        starts = [rec.t_start for rec in op_log if lo <= _cube_of(rec) < hi]
+        ends = [rec.t_end for rec in op_log if lo <= _cube_of(rec) < hi]
+        if starts and ends:
+            per_user.append(max(ends) - min(starts))
+    mean_user_ns = sum(per_user) / len(per_user) if per_user else 0.0
+
+    agg_us = wall_ns / 1000
+    return {
+        "mode": mode,
+        "kv_per_cube": kv_per_cube,
+        "C": C,
+        "S_kv": S_kv,
+        "B": B,
+        "capacity": CUBES_PER_SIP // C,
+        "agg_latency_us": agg_us,
+        "mean_user_latency_us": mean_user_ns / 1000,
+        "throughput_users_per_us": (B / agg_us) if agg_us else 0.0,
+    }
+
+
+@pytest.mark.slow
+def test_decode_batch_sweep():
+    """Sweep 4 modes × contexts × B ∈ {1, 2, capacity}; dump batch CSV.
+
+    Throughput is near-linear in B (aggregate latency stays within a few
+    percent of the single-user latency across the batch), so the endpoints
+    plus a low point capture the scaling; the high-B runs of the dense
+    mappings are the expensive ones and full 1..capacity granularity is not
+    needed for the trade-off.
+    """
+    rows = []
+    for mode, kv_per_cube, C in MODES:
+        capacity = CUBES_PER_SIP // C
+        batch_sizes = sorted({1, 2, capacity})
+        for S_kv in CONTEXT_LENGTHS:
+            for B in batch_sizes:
+                print(f">>> BATCH {mode} S_kv={S_kv//1024}K B={B}/{capacity}",
+                      flush=True)
+                r = _run_batch(kv_per_cube=kv_per_cube, C=C, S_kv=S_kv, B=B)
+                assert r.completion.ok, (
+                    f"BATCH {mode} S_kv={S_kv} B={B}: {r.completion}"
+                )
+                m = _metrics(r, mode=mode, kv_per_cube=kv_per_cube,
+                             C=C, S_kv=S_kv, B=B)
+                rows.append(m)
+                print(f"    agg={m['agg_latency_us']:.2f}us  "
+                      f"user={m['mean_user_latency_us']:.2f}us  "
+                      f"tput={m['throughput_users_per_us']:.3f} u/us",
+                      flush=True)
+
+    CSV_OUT.parent.mkdir(parents=True, exist_ok=True)
+    with CSV_OUT.open("w", newline="") as f:
+        w = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
+        w.writeheader()
+        w.writerows(rows)
+    print(f"\nWrote {CSV_OUT} ({len(rows)} rows)")

tests/attention/test_gqa_decode_long_ctx_tiled.py

@@ -0,0 +1,87 @@
+"""Command-count checks for the Case-6 primitive-TILED (streamed) decode kernel.
+
+The tiled kernel hand-blocks each Q·Kᵀ / P·V matmul into 16×16×16
+GEMMs, with per-block DMAs of the HBM operand slices and a deferred
+K-inner sum outside the K loop (ADR-0064 Rev2 D8 single-op FIXED path
+exercised across DMA, GEMM, and MATH engines).
+"""
+from __future__ import annotations
+
+from math import ceil
+
+from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_sp import (  # noqa: E501
+    gqa_attention_decode_long_ctx_cube_sp_pe_sp_kernel as _primitive,
+)
+from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_sp_hand_tiled_16x16x16 import (  # noqa: E501
+    gqa_attention_decode_long_ctx_cube_sp_pe_sp_hand_tiled_16x16x16_kernel as _tiled,
+)
+from kernbench.common.pe_commands import (
+    DmaReadCmd, GemmCmd, PeCpuOverheadCmd,
+)
+from kernbench.common.pe_cost_model import DEFAULT_PE_COST_MODEL
+from kernbench.triton_emu.tl_context import TLContext, run_kernel
+
+_C, _P = 8, 8
+_S_KV = 131_072
+_D_HEAD, _H_Q, _H_KV, _T_Q = 128, 8, 1, 1
+_TILE_S_KV = 1024
+_MAC = 16
+
+
+def _run(kernel, S_kv: int):
+    tl = TLContext(
+        pe_id=0, num_programs=_P, cost_model=DEFAULT_PE_COST_MODEL,
+        cube_id=6, num_cubes=_C, scratch_base=0x200000, scratch_size=1 << 20,
+    )
+    run_kernel(
+        kernel, tl, 0x1000, 0x2000, 0x3000, 0x4000,
+        _T_Q, S_kv, _H_Q, _H_KV, _D_HEAD, _C, _P,
+    )
+    return tl.commands
+
+
+def _block_counts():
+    """Return (n_tiles, qk_blocks, pv_blocks) for the fixed test params."""
+    G = _H_Q // _H_KV
+    S_local = _S_KV // (_C * _P)
+    n_tiles = (S_local + _TILE_S_KV - 1) // _TILE_S_KV
+    tile_s = min(_TILE_S_KV, S_local)
+    M = G * _T_Q
+    qk = ceil(M / _MAC) * ceil(tile_s / _MAC) * ceil(_D_HEAD / _MAC)
+    pv = ceil(M / _MAC) * ceil(_D_HEAD / _MAC) * ceil(tile_s / _MAC)
+    return n_tiles, qk, pv
+
+
+def test_tiled_gemm_count_matches_block_formula():
+    """One GemmCmd per 16³ block across both matmuls."""
+    n_tiles, qk, pv = _block_counts()
+    expected = n_tiles * (qk + pv)
+
+    cmds = _run(_tiled, _S_KV)
+    got = sum(1 for c in cmds if isinstance(c, GemmCmd))
+    assert got == expected, f"GemmCmd count: got {got}, expected {expected}"
+
+
+def test_tiled_dma_count_matches_streamed_formula():
+    """Streamed operand DMAs — Q·Kᵀ blocks pay 2 DMAs (A+B slice), P·V
+    blocks pay 1 DMA (V slice only; the P side is TCM-resident post-softmax).
+    Each matmul also carries full-shape coarse loads for the GemmCmd
+    operand handles (2 for Q·Kᵀ: Q + K_T, 1 for P·V: V only)."""
+    n_tiles, qk, pv = _block_counts()
+    expected = n_tiles * (2 * qk + 1 * pv + 3)
+
+    cmds = _run(_tiled, _S_KV)
+    got = sum(1 for c in cmds if isinstance(c, DmaReadCmd))
+    assert got == expected, f"DmaReadCmd count: got {got}, expected {expected}"
+
+
+def test_tiled_amplifies_dispatch_vs_primitive():
+    """The streamed tiled kernel emits ≫40× more PE_CPU dispatches than
+    the coarse primitive — the whole point of adding this variant."""
+    n_tiled = sum(1 for c in _run(_tiled, _S_KV)
+                  if isinstance(c, PeCpuOverheadCmd))
+    n_prim = sum(1 for c in _run(_primitive, _S_KV)
+                 if isinstance(c, PeCpuOverheadCmd))
+    assert n_tiled > 40 * n_prim, (
+        f"expected >40× dispatch amplification, got {n_tiled} vs {n_prim}"
+    )

tests/attention/test_gqa_short_context.py

@@ -1,31 +1,17 @@
-"""Phase 1 spec test for Phase D: short-context GQA kernels
-(ADR-0060 §B "Items from the long/short context split", item B.split.2).
+"""Short-context GQA kernel tests — unified A1/A2/A4/B (ADR-0070).
 
-The long-context kernels (§5.2 decode / §5.5 prefill) shard each KV
-head row-wise across all CUBEs. At short context (S_kv < 256K), that
-shard is too thin to feed the engines and the cube-level collective
-overhead dominates. The short-context kernels drop cube-SP entirely:
-each CUBE owns ``kv_per_cube`` *whole* KV heads, no S_kv sharding across
-CUBEs.
+Both ``_gqa_attention_prefill_short.py`` and
+``_gqa_attention_decode_short.py`` are single unified kernels selected
+at launch via ``kv_per_cube ∈ {1, 2, 4, 8}``:
 
-Design (per AskUserQuestion answers in this session):
-  - PE-parallel heads: P PEs split into ``kv_per_cube`` groups, each
-    group does PE-SP across (P/kv_per_cube) PEs for one owned head.
-  - scratch_scope + tl.copy_to discipline mirrors the long kernels.
-  - One short kernel handles kv_per_cube ∈ {1, 2, 4, 8} via a parameter.
+    Mode  kv_per_cube  C        group_size
+    ----  -----------  -------  ----------
+    A1            1    h_kv      P (=8)
+    A2            2    h_kv/2    P/2 (=4)
+    A4            4    h_kv/4    P/4 (=2)
+    B             8    1         1
 
-Group layout on the 2×4 PE grid:
-  kv_per_cube=1, group=8 PEs (full 2×4): row chain + col bridge.
-  kv_per_cube=2, group=4 PEs (one row):  row chain only.
-  kv_per_cube=4, group=2 PEs (adj cols): 1-step chain.
-  kv_per_cube=8, group=1 PE:             no chain.
-
-After chain reduce-to-group-root, the group's root PE writes its
-owned head's output. No inter-CUBE reduce.
-
-Phase 1 (this commit): tests only — production code lands in Phase 2.
-All tests fail because the short kernels (``_gqa_attention_decode_short.py`` and
-``_gqa_attention_prefill_short.py``) do not exist yet.
+Tests are mode-parametrized over the four (kv_per_cube, C) tuples.
 """
 from __future__ import annotations
 
@@ -45,6 +31,17 @@ TOPOLOGY_DEFAULT = Path(__file__).resolve().parents[2] / "topology.yaml"
 
 D_HEAD = 64
 DTYPE = "f16"
+TILE_S_KV = 1024
+P = 8
+H_KV = 8
+
+# (kv_per_cube, C) for the four modes.
+MODES = [
+    pytest.param(1, 8, id="A1"),
+    pytest.param(2, 4, id="A2"),
+    pytest.param(4, 2, id="A4"),
+    pytest.param(8, 1, id="B"),
+]
 
 
 def _ccl_cfg():
@@ -61,166 +58,41 @@ def _count(op_log, name: str) -> int:
     return sum(1 for r in op_log if r.op_name == name)
 
 
-# ── Decode short-context kernel ──────────────────────────────────────
+# ── Prefill ──────────────────────────────────────────────────────────
 
 
-def _run_decode_short(*, h_q: int, h_kv: int, kv_per_cube: int,
-                     C: int, P: int, S_kv: int):
-    """Run the short-context decode kernel with PE-parallel heads.
-
-    Layout (after design iteration — see Phase D failure-recovery):
-      Q: (T_q, h_q·D_HEAD) replicated; kernel reshapes byte-conservingly.
-      K, V: (h_kv·S_kv, D_HEAD) head-stacked, ``cube=row_wise, pe=row_wise``
-          so each PE gets a contiguous (S_local, D_HEAD) chunk at its own
-          addressable shard (no partial reads needed).
-      O: replicated; each group root writes the full byte-conserving
-          (h_q·T_q, D_HEAD) result.
-    """
-    from kernbench.benches.gqa_helpers.short_ctx._gqa_attention_decode_short import gqa_attention_decode_short_kernel  # Phase 2
-
+def _run_prefill(*, kv_per_cube: int, C: int, T_q: int, S_kv: int,
+                 h_q: int = 8):
+    """Run the unified prefill kernel in the given mode."""
+    from kernbench.benches.gqa_helpers.short_ctx._gqa_attention_prefill_short import (
+        _validate_config as _validate_prefill_config,
+        gqa_attention_prefill_short_kernel,
+    )
+    _validate_prefill_config(kv_per_cube=kv_per_cube, T_q=T_q, P=P, C=C,
+                             h_q=h_q, h_kv=H_KV, S_kv=S_kv)
+    n_tiles = S_kv // TILE_S_KV
+    Q_ROWS = kv_per_cube * T_q
+    Q_COLS = (h_q * D_HEAD) // kv_per_cube
     topo = resolve_topology(str(TOPOLOGY_DEFAULT))
 
     def _bench_fn(ctx):
         configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
-        dp_full = DPPolicy(cube="replicate", pe="replicate",
-                           num_cubes=C, num_pes=P)
-        # Head-stacked KV with row_wise sharding so each PE's chunk is
-        # contiguous and exactly (S_local, D_HEAD) addressable.
-        dp_kv = DPPolicy(cube="row_wise", pe="row_wise",
-                         num_cubes=C, num_pes=P)
-        q = ctx.zeros((1, h_q * D_HEAD),
-                      dtype=DTYPE, dp=dp_full, name=f"q_short_{kv_per_cube}_{C}")
-        k = ctx.zeros((h_kv * S_kv, D_HEAD),
-                      dtype=DTYPE, dp=dp_kv, name=f"k_short_{kv_per_cube}_{C}")
-        v = ctx.zeros((h_kv * S_kv, D_HEAD),
-                      dtype=DTYPE, dp=dp_kv, name=f"v_short_{kv_per_cube}_{C}")
-        o = ctx.empty((1, h_q * D_HEAD),
-                      dtype=DTYPE, dp=dp_full, name=f"o_short_{kv_per_cube}_{C}")
+        dp_q = DPPolicy(cube="column_wise", pe="replicate", num_cubes=C, num_pes=P)
+        dp_kv = DPPolicy(cube="row_wise", pe="replicate", num_cubes=C, num_pes=P)
+        dp_o = DPPolicy(cube="column_wise", pe="replicate", num_cubes=C, num_pes=P)
+        q = ctx.zeros((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_q,
+                      name=f"q_pre_kv{kv_per_cube}")
+        k = ctx.zeros((H_KV * n_tiles * D_HEAD, TILE_S_KV),
+                      dtype=DTYPE, dp=dp_kv, name=f"k_pre_kv{kv_per_cube}")
+        v = ctx.zeros((H_KV * S_kv, D_HEAD),
+                      dtype=DTYPE, dp=dp_kv, name=f"v_pre_kv{kv_per_cube}")
+        o = ctx.empty((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_o,
+                      name=f"o_pre_kv{kv_per_cube}")
         ctx.launch(
-            f"gqa_decode_short_{kv_per_cube}_{C}",
-            gqa_attention_decode_short_kernel,
-            q, k, v, o,
-            1, S_kv, h_q, h_kv, D_HEAD, C, P, kv_per_cube,
-            _auto_dim_remap=False,
-        )
-
-    return run_bench(
-        topology=topo, bench_fn=_bench_fn,
-        device=resolve_device(None), engine_factory=_engine_factory,
-    )
-
-
-def test_short_decode_smoke_kv_per_cube_2_C_4():
-    """ADR-0060 §B.split.2 headline: kv_per_cube=2, C=4 — each CUBE
-    owns 2 heads (8 heads / 4 CUBEs). PE-parallel heads splits the 8
-    PEs into 2 groups of 4, each group does PE-SP for one owned head.
-    Smoke: kernel completes."""
-    result = _run_decode_short(
-        h_q=8, h_kv=8, kv_per_cube=2, C=4, P=8, S_kv=64,
-    )
-    assert result.completion.ok, (
-        f"short decode kv_per_cube=2 must complete; got {result.completion}"
-    )
-
-
-def test_short_decode_smoke_kv_per_cube_4_C_2():
-    """kv_per_cube=4, C=2 — half the CUBEs participate, each owns 4
-    heads, 4 PE groups of 2 PEs each."""
-    result = _run_decode_short(
-        h_q=8, h_kv=8, kv_per_cube=4, C=2, P=8, S_kv=64,
-    )
-    assert result.completion.ok, (
-        f"short decode kv_per_cube=4 must complete; got {result.completion}"
-    )
-
-
-def test_short_decode_smoke_kv_per_cube_8_C_1():
-    """kv_per_cube=8, C=1 — all heads on one CUBE. 8 PE groups of 1 PE
-    each → no chain reduce, each PE writes its head's output."""
-    result = _run_decode_short(
-        h_q=8, h_kv=8, kv_per_cube=8, C=1, P=8, S_kv=64,
-    )
-    assert result.completion.ok, (
-        f"short decode kv_per_cube=8 must complete; got {result.completion}"
-    )
-
-
-def test_short_decode_dma_write_count_equals_h_kv():
-    """ADR-0060 §B.split.2: each owned head produces exactly one
-    output (no inter-CUBE reduce). Total dma_writes = h_kv across all
-    participating CUBEs and groups.
-
-    For h_kv=8, kv_per_cube=2, C=4: each CUBE writes 2 outputs →
-    4 × 2 = 8 dma_writes total.
-    """
-    result = _run_decode_short(
-        h_q=8, h_kv=8, kv_per_cube=2, C=4, P=8, S_kv=64,
-    )
-    assert result.completion.ok
-    n_writes = _count(result.engine.op_log, "dma_write")
-    assert n_writes == 8, (
-        f"short decode: expected 8 dma_writes (h_kv); got {n_writes}"
-    )
-
-
-def test_short_decode_no_inter_cube_traffic():
-    """ADR-0060 §B.split.2: each head is fully owned by one CUBE → no
-    inter-CUBE reduce. The kernel must not invoke CUBE-level E/W IPCQ.
-
-    Today's long-context kernel at C=4 emits ~12 inter-CUBE ipcq_copy
-    via direction "E"/"W". The short kernel must emit zero of those,
-    keeping all IPCQ traffic on the ``intra_*`` directions.
-    """
-    result = _run_decode_short(
-        h_q=8, h_kv=8, kv_per_cube=2, C=4, P=8, S_kv=64,
-    )
-    assert result.completion.ok
-    # Count IPCQ traffic that targeted the CUBE-level "E"/"W" directions.
-    inter_cube_ipcq = sum(
-        1 for r in result.engine.op_log
-        if r.op_name == "ipcq_copy"
-        and r.params.get("direction") in ("E", "W")
-    )
-    assert inter_cube_ipcq == 0, (
-        f"short decode must have no inter-CUBE E/W IPCQ; got {inter_cube_ipcq}"
-    )
-
-
-# ── Prefill short-context kernel ─────────────────────────────────────
-
-
-def _run_prefill_short(*, h_kv: int, kv_per_cube: int,
-                      C: int, P: int, T_q: int, S_kv: int):
-    """Run the short-context prefill kernel with PE-parallel heads.
-
-    Layout: same head-stacked K/V scheme as decode short, with Q/O
-    replicated and byte-conserving reshape inside the kernel.
-    """
-    from kernbench.benches.gqa_helpers.short_ctx._gqa_attention_prefill_short import gqa_attention_prefill_short_kernel  # Phase 2
-
-    topo = resolve_topology(str(TOPOLOGY_DEFAULT))
-
-    def _bench_fn(ctx):
-        configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
-        dp_q = DPPolicy(cube="replicate", pe="replicate",
-                        num_cubes=C, num_pes=P)
-        dp_kv = DPPolicy(cube="row_wise", pe="row_wise",
-                         num_cubes=C, num_pes=P)
-        dp_o = DPPolicy(cube="replicate", pe="replicate",
-                        num_cubes=C, num_pes=P)
-        q = ctx.zeros((T_q, h_kv * D_HEAD),
-                      dtype=DTYPE, dp=dp_q, name=f"q_pre_short_{kv_per_cube}_{C}")
-        k = ctx.zeros((h_kv * S_kv, D_HEAD),
-                      dtype=DTYPE, dp=dp_kv, name=f"k_pre_short_{kv_per_cube}_{C}")
-        v = ctx.zeros((h_kv * S_kv, D_HEAD),
-                      dtype=DTYPE, dp=dp_kv, name=f"v_pre_short_{kv_per_cube}_{C}")
-        o = ctx.empty((T_q, h_kv * D_HEAD),
-                      dtype=DTYPE, dp=dp_o, name=f"o_pre_short_{kv_per_cube}_{C}")
-        ctx.launch(
-            f"gqa_prefill_short_{kv_per_cube}_{C}",
+            f"gqa_prefill_kv{kv_per_cube}",
             gqa_attention_prefill_short_kernel,
             q, k, v, o,
-            T_q, S_kv, h_kv, D_HEAD, C, P, kv_per_cube,
+            T_q, S_kv, h_q, H_KV, D_HEAD, C, P, kv_per_cube,
             _auto_dim_remap=False,
         )
 
@@ -230,31 +102,507 @@ def _run_prefill_short(*, h_kv: int, kv_per_cube: int,
     )
 
 
-def test_short_prefill_smoke_kv_per_cube_2_C_4():
-    """Short prefill kv_per_cube=2, C=4. Smoke: completes."""
-    result = _run_prefill_short(
-        h_kv=8, kv_per_cube=2, C=4, P=8, T_q=4, S_kv=64,
-    )
-    assert result.completion.ok, (
-        f"short prefill kv_per_cube=2 must complete; got {result.completion}"
+@pytest.mark.parametrize("kv_per_cube,C", MODES)
+def test_prefill_smoke(kv_per_cube, C):
+    """All four prefill modes complete on a single-tile mini config."""
+    r = _run_prefill(kv_per_cube=kv_per_cube, C=C, T_q=8, S_kv=1024)
+    assert r.completion.ok, (
+        f"prefill kv_per_cube={kv_per_cube}: {r.completion}"
     )
 
 
-def test_short_prefill_no_ring_KV_traffic():
-    """ADR-0060 §B.split.2: short prefill DOES NOT use Ring KV — each
-    CUBE owns its KV heads fully, no rotation. The kernel must not
-    emit any KV-rotation IPCQ traffic at the CUBE level.
+@pytest.mark.parametrize("kv_per_cube,C", MODES)
+def test_prefill_qtile_split_dma_writes(kv_per_cube, C):
+    """Every PE in every active group stores its q-tile output:
+    dma_writes = group_size · kv_per_cube · C."""
+    r = _run_prefill(kv_per_cube=kv_per_cube, C=C, T_q=8, S_kv=1024)
+    assert r.completion.ok
+    group_size = P // kv_per_cube
+    expected = group_size * kv_per_cube * C
+    n_writes = _count(r.engine.op_log, "dma_write")
+    assert n_writes == expected, (
+        f"prefill kv_per_cube={kv_per_cube}: expected {expected} "
+        f"dma_writes; got {n_writes}"
+    )
+
+
+@pytest.mark.parametrize("kv_per_cube,C", MODES)
+def test_prefill_kv_hbm_read_collapsed(kv_per_cube, C):
+    """IPCQ KV broadcast collapses HBM K/V reads to one per group per tile:
+    dma_reads = (P · C) Q-reads + (kv_per_cube · 2 · n_tiles · C) KV-reads.
     """
-    result = _run_prefill_short(
-        h_kv=8, kv_per_cube=2, C=4, P=8, T_q=4, S_kv=64,
+    r = _run_prefill(kv_per_cube=kv_per_cube, C=C, T_q=8, S_kv=1024)
+    assert r.completion.ok
+    n_tiles = 1024 // TILE_S_KV
+    expected = (P * C) + (kv_per_cube * 2 * n_tiles * C)
+    n_reads = _count(r.engine.op_log, "dma_read")
+    assert n_reads == expected, (
+        f"prefill kv_per_cube={kv_per_cube}: expected {expected} dma_reads "
+        f"(Q + IPCQ-broadcasted KV); got {n_reads}"
     )
-    assert result.completion.ok
+
+
+def test_prefill_no_ring_KV_traffic():
+    """Short prefill never uses Ring KV — no inter-CUBE E/W IPCQ."""
+    r = _run_prefill(kv_per_cube=2, C=4, T_q=8, S_kv=1024)
+    assert r.completion.ok
     inter_cube_ipcq = sum(
-        1 for r in result.engine.op_log
-        if r.op_name == "ipcq_copy"
-        and r.params.get("direction") in ("E", "W")
+        1 for rec in r.engine.op_log
+        if rec.op_name == "ipcq_copy"
+        and rec.params.get("direction") in ("E", "W")
     )
     assert inter_cube_ipcq == 0, (
-        f"short prefill must have no Ring KV (no inter-CUBE E/W IPCQ); "
+        f"short prefill must have no inter-CUBE E/W IPCQ; "
         f"got {inter_cube_ipcq}"
     )
+
+
+@pytest.mark.parametrize("kv_per_cube,C", MODES)
+def test_prefill_multitile_scaling(kv_per_cube, C):
+    """Prefill dma_read scales linearly with n_tiles in every mode.
+
+    Per-mode formula (IPCQ broadcast: cube's group-root PE reads K/V):
+        dma_reads = P·C + 2·kv_per_cube·C·n_tiles
+                  = (Q: 1 per PE)  +  (K + V: 1 per group per tile)
+    """
+    for S_kv in (1024, 2048, 4096):
+        r = _run_prefill(kv_per_cube=kv_per_cube, C=C, T_q=8, S_kv=S_kv)
+        assert r.completion.ok, f"kv={kv_per_cube} S_kv={S_kv}"
+        n_tiles = S_kv // TILE_S_KV
+        expected = (P * C) + (2 * kv_per_cube * C * n_tiles)
+        n_reads = _count(r.engine.op_log, "dma_read")
+        assert n_reads == expected, (
+            f"prefill kv={kv_per_cube} n_tiles={n_tiles}: expected "
+            f"{expected} dma_reads; got {n_reads}"
+        )
+
+
+# ── Decode ───────────────────────────────────────────────────────────
+
+
+def _run_decode(*, kv_per_cube: int, C: int, S_kv: int, h_q: int = 8):
+    """Run the unified decode kernel in the given mode."""
+    from kernbench.benches.gqa_helpers.short_ctx._gqa_attention_decode_short import (
+        _validate_config as _validate_decode_config,
+        gqa_attention_decode_short_kernel,
+    )
+    T_q = 1
+    _validate_decode_config(kv_per_cube=kv_per_cube, T_q=T_q, P=P, C=C,
+                            h_q=h_q, h_kv=H_KV, S_kv=S_kv)
+    Q_ROWS = kv_per_cube * T_q
+    Q_COLS = (h_q * D_HEAD) // kv_per_cube
+    # K total elements = h_kv · S_kv · d_head; mode-invariant deploy.
+    k_rows = (H_KV * S_kv * D_HEAD) // TILE_S_KV
+    topo = resolve_topology(str(TOPOLOGY_DEFAULT))
+
+    def _bench_fn(ctx):
+        configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
+        dp_q = DPPolicy(cube="column_wise", pe="replicate", num_cubes=C, num_pes=P)
+        dp_kv = DPPolicy(cube="row_wise", pe="row_wise", num_cubes=C, num_pes=P)
+        dp_o = DPPolicy(cube="column_wise", pe="replicate", num_cubes=C, num_pes=P)
+        q = ctx.zeros((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_q,
+                      name=f"q_dec_kv{kv_per_cube}")
+        k = ctx.zeros((k_rows, TILE_S_KV), dtype=DTYPE, dp=dp_kv,
+                      name=f"k_dec_kv{kv_per_cube}")
+        v = ctx.zeros((H_KV * S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv,
+                      name=f"v_dec_kv{kv_per_cube}")
+        o = ctx.empty((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_o,
+                      name=f"o_dec_kv{kv_per_cube}")
+        ctx.launch(
+            f"gqa_decode_kv{kv_per_cube}",
+            gqa_attention_decode_short_kernel,
+            q, k, v, o,
+            T_q, S_kv, h_q, H_KV, D_HEAD, C, P, kv_per_cube,
+            _auto_dim_remap=False,
+        )
+
+    return run_bench(
+        topology=topo, bench_fn=_bench_fn,
+        device=resolve_device(None), engine_factory=_engine_factory,
+    )
+
+
+@pytest.mark.parametrize("kv_per_cube,C", MODES)
+def test_decode_smoke(kv_per_cube, C):
+    """All four decode modes complete at S_kv = 8K (min multi-tile-aligned
+    config that satisfies every mode's group-size constraint)."""
+    r = _run_decode(kv_per_cube=kv_per_cube, C=C, S_kv=8192)
+    assert r.completion.ok, (
+        f"decode kv_per_cube={kv_per_cube}: {r.completion}"
+    )
+
+
+@pytest.mark.parametrize("kv_per_cube,C", MODES)
+def test_decode_chain_reduce_store_count(kv_per_cube, C):
+    """Decode chain-reduce: only group root (PE 0 per group) stores —
+    dma_writes = kv_per_cube · C (one per group root × cubes)."""
+    r = _run_decode(kv_per_cube=kv_per_cube, C=C, S_kv=8192)
+    assert r.completion.ok
+    expected = kv_per_cube * C
+    n_writes = _count(r.engine.op_log, "dma_write")
+    assert n_writes == expected, (
+        f"decode kv_per_cube={kv_per_cube}: expected {expected} dma_writes "
+        f"(1 per group root); got {n_writes}"
+    )
+
+
+@pytest.mark.parametrize("kv_per_cube,C", MODES)
+def test_decode_multitile_per_pe(kv_per_cube, C):
+    """Decode multi-tile sweep path exercised in every mode.
+
+    S_kv is chosen so each PE owns 2 tiles (n_tiles_per_pe == 2),
+    forcing the post-tile-0 sweep loop to run.
+    """
+    group_size = P // kv_per_cube
+    S_kv = group_size * 2 * TILE_S_KV
+    r = _run_decode(kv_per_cube=kv_per_cube, C=C, S_kv=S_kv)
+    assert r.completion.ok, (
+        f"decode kv={kv_per_cube} S_kv={S_kv}: {r.completion}"
+    )
+
+
+# ── ADR-0011 D-VA1 contract regression tests ───────────────────────
+
+
+def _per_cube_dma_busy(op_log):
+    """Sum of dma_read t_end-t_start grouped by cube index."""
+    from collections import defaultdict
+    busy = defaultdict(float)
+    for rec in op_log:
+        if rec.op_name != "dma_read":
+            continue
+        cid = rec.component_id or ""
+        for part in cid.split("."):
+            if part.startswith("cube"):
+                try:
+                    busy[int(part[4:])] += rec.t_end - rec.t_start
+                except ValueError:
+                    pass
+                break
+    return dict(busy)
+
+
+def _per_cube_large_read_addrs(op_log, *, min_nbytes: int = 1024):
+    """Group dma_read src_addrs (>= min_nbytes) by issuing cube."""
+    from collections import defaultdict
+    per_cube: dict[int, set[int]] = defaultdict(set)
+    for rec in op_log:
+        if rec.op_name != "dma_read":
+            continue
+        if rec.params.get("nbytes", 0) < min_nbytes:
+            continue
+        cid = rec.component_id or ""
+        for part in cid.split("."):
+            if part.startswith("cube"):
+                try:
+                    per_cube[int(part[4:])].add(rec.params["src_addr"])
+                except ValueError:
+                    pass
+                break
+    return per_cube
+
+
+def _assert_cubes_disjoint(per_cube_addrs, label):
+    seen = sorted(per_cube_addrs.items())
+    for i, (ci, ai) in enumerate(seen):
+        for cj, aj in seen[i + 1:]:
+            assert ai.isdisjoint(aj), (
+                f"{label}: cube{ci} and cube{cj} share {len(ai & aj)} "
+                f"dma_read src_addrs; cubes must target disjoint HBM regions."
+            )
+
+
+def _assert_cube_balanced(busy, label, max_ratio=1.1):
+    if len(busy) <= 1:
+        return  # single-cube modes auto-pass
+    vals = list(busy.values())
+    ratio = max(vals) / min(vals) if min(vals) > 0 else 0
+    assert ratio < max_ratio, (
+        f"{label}: per-cube DMA_READ unbalanced: max/min = {ratio:.2f}× "
+        f"(expected < {max_ratio}×). Per-cube busy (μs): "
+        f"{[f'cube{c}={v/1000:.1f}' for c,v in sorted(busy.items())]}"
+    )
+
+
+@pytest.mark.parametrize("kv_per_cube,C", MODES)
+def test_decode_per_cube_disjoint_src_addrs(kv_per_cube, C):
+    """Decode: cubes must read from disjoint HBM regions (per ADR-0011
+    D-VA1, kernel computes cube/PE shard offset from program_id).
+
+    Regression guard: pre-fix all 64 PEs read offset 0 → all cubes
+    share the same src_addrs.
+    """
+    r = _run_decode(kv_per_cube=kv_per_cube, C=C, S_kv=8192)
+    assert r.completion.ok
+    per_cube = _per_cube_large_read_addrs(r.engine.op_log)
+    assert len(per_cube) == C, (
+        f"decode kv={kv_per_cube}: expected {C} cubes; got {len(per_cube)}"
+    )
+    _assert_cubes_disjoint(per_cube, f"decode kv={kv_per_cube}")
+
+
+@pytest.mark.parametrize("kv_per_cube,C", MODES)
+def test_decode_per_cube_dma_balanced(kv_per_cube, C):
+    """Decode: per-cube DMA_READ busy must be ~equal (cubes operate on
+    independent local HBM). max/min < 1.1× for multi-cube modes.
+
+    Regression guard: pre-fix 11.5× ratio (cross-cube traffic to cube0).
+    """
+    r = _run_decode(kv_per_cube=kv_per_cube, C=C, S_kv=8192)
+    assert r.completion.ok
+    busy = _per_cube_dma_busy(r.engine.op_log)
+    assert len(busy) == C, (
+        f"decode kv={kv_per_cube}: expected {C} cubes; got {sorted(busy.keys())}"
+    )
+    _assert_cube_balanced(busy, f"decode kv={kv_per_cube}")
+
+
+@pytest.mark.parametrize("kv_per_cube,C", MODES)
+def test_prefill_per_cube_disjoint_src_addrs(kv_per_cube, C):
+    """Prefill: cubes must read from disjoint HBM regions."""
+    r = _run_prefill(kv_per_cube=kv_per_cube, C=C, T_q=8, S_kv=1024)
+    assert r.completion.ok
+    per_cube = _per_cube_large_read_addrs(r.engine.op_log)
+    assert len(per_cube) == C, (
+        f"prefill kv={kv_per_cube}: expected {C} cubes; got {len(per_cube)}"
+    )
+    _assert_cubes_disjoint(per_cube, f"prefill kv={kv_per_cube}")
+
+
+@pytest.mark.parametrize("kv_per_cube,C", MODES)
+def test_prefill_per_cube_dma_balanced(kv_per_cube, C):
+    """Prefill: per-cube DMA_READ busy must be ~equal (cube-parallel)."""
+    r = _run_prefill(kv_per_cube=kv_per_cube, C=C, T_q=8, S_kv=1024)
+    assert r.completion.ok
+    busy = _per_cube_dma_busy(r.engine.op_log)
+    assert len(busy) == C, (
+        f"prefill kv={kv_per_cube}: expected {C} cubes; got {sorted(busy.keys())}"
+    )
+    _assert_cube_balanced(busy, f"prefill kv={kv_per_cube}")
+
+
+# ── Composite (second-level) variant smoke tests ────────────────────
+
+
+def _run_prefill_composite(*, kv_per_cube: int, C: int, T_q: int,
+                           S_kv: int, h_q: int = 8):
+    """Same helper as _run_prefill but for the composite-GEMM kernel."""
+    from kernbench.benches.gqa_helpers.short_ctx._gqa_attention_prefill_short_composite import (
+        _validate_config as _validate_prefill_cmp_config,
+        gqa_attention_prefill_short_composite_kernel,
+    )
+    _validate_prefill_cmp_config(kv_per_cube=kv_per_cube, T_q=T_q, P=P,
+                                 C=C, h_q=h_q, h_kv=H_KV, S_kv=S_kv)
+    n_tiles = S_kv // TILE_S_KV
+    Q_ROWS = kv_per_cube * T_q
+    Q_COLS = (h_q * D_HEAD) // kv_per_cube
+    topo = resolve_topology(str(TOPOLOGY_DEFAULT))
+
+    def _bench_fn(ctx):
+        configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
+        dp_q = DPPolicy(cube="column_wise", pe="replicate", num_cubes=C, num_pes=P)
+        dp_kv = DPPolicy(cube="row_wise", pe="replicate", num_cubes=C, num_pes=P)
+        dp_o = DPPolicy(cube="column_wise", pe="replicate", num_cubes=C, num_pes=P)
+        q = ctx.zeros((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_q,
+                      name=f"q_pre_cmp_kv{kv_per_cube}")
+        k = ctx.zeros((H_KV * n_tiles * D_HEAD, TILE_S_KV),
+                      dtype=DTYPE, dp=dp_kv, name=f"k_pre_cmp_kv{kv_per_cube}")
+        v = ctx.zeros((H_KV * S_kv, D_HEAD),
+                      dtype=DTYPE, dp=dp_kv, name=f"v_pre_cmp_kv{kv_per_cube}")
+        o = ctx.empty((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_o,
+                      name=f"o_pre_cmp_kv{kv_per_cube}")
+        ctx.launch(
+            f"gqa_prefill_cmp_kv{kv_per_cube}",
+            gqa_attention_prefill_short_composite_kernel,
+            q, k, v, o,
+            T_q, S_kv, h_q, H_KV, D_HEAD, C, P, kv_per_cube,
+            _auto_dim_remap=False,
+        )
+
+    return run_bench(
+        topology=topo, bench_fn=_bench_fn,
+        device=resolve_device(None), engine_factory=_engine_factory,
+    )
+
+
+def _run_decode_composite(*, kv_per_cube: int, C: int, S_kv: int,
+                          h_q: int = 8):
+    """Same helper as _run_decode but for the composite-GEMM kernel."""
+    from kernbench.benches.gqa_helpers.short_ctx._gqa_attention_decode_short_composite import (
+        _validate_config as _validate_decode_cmp_config,
+        gqa_attention_decode_short_composite_kernel,
+    )
+    T_q = 1
+    _validate_decode_cmp_config(kv_per_cube=kv_per_cube, T_q=T_q, P=P,
+                                C=C, h_q=h_q, h_kv=H_KV, S_kv=S_kv)
+    Q_ROWS = kv_per_cube * T_q
+    Q_COLS = (h_q * D_HEAD) // kv_per_cube
+    k_rows = (H_KV * S_kv * D_HEAD) // TILE_S_KV
+    topo = resolve_topology(str(TOPOLOGY_DEFAULT))
+
+    def _bench_fn(ctx):
+        configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
+        dp_q = DPPolicy(cube="column_wise", pe="replicate", num_cubes=C, num_pes=P)
+        dp_kv = DPPolicy(cube="row_wise", pe="row_wise", num_cubes=C, num_pes=P)
+        dp_o = DPPolicy(cube="column_wise", pe="replicate", num_cubes=C, num_pes=P)
+        q = ctx.zeros((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_q,
+                      name=f"q_dec_cmp_kv{kv_per_cube}")
+        k = ctx.zeros((k_rows, TILE_S_KV), dtype=DTYPE, dp=dp_kv,
+                      name=f"k_dec_cmp_kv{kv_per_cube}")
+        v = ctx.zeros((H_KV * S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv,
+                      name=f"v_dec_cmp_kv{kv_per_cube}")
+        o = ctx.empty((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_o,
+                      name=f"o_dec_cmp_kv{kv_per_cube}")
+        ctx.launch(
+            f"gqa_decode_cmp_kv{kv_per_cube}",
+            gqa_attention_decode_short_composite_kernel,
+            q, k, v, o,
+            T_q, S_kv, h_q, H_KV, D_HEAD, C, P, kv_per_cube,
+            _auto_dim_remap=False,
+        )
+
+    return run_bench(
+        topology=topo, bench_fn=_bench_fn,
+        device=resolve_device(None), engine_factory=_engine_factory,
+    )
+
+
+def _run_prefill_composite_fused(*, kv_per_cube: int, C: int, T_q: int,
+                                 S_kv: int, h_q: int = 8):
+    """Variant (3): GEMM-only composite + softmax_merge prologue fusion."""
+    from kernbench.benches.gqa_helpers.short_ctx._gqa_attention_prefill_short_composite_fused import (
+        _validate_config as _validate_prefill_fuse_config,
+        gqa_attention_prefill_short_composite_fused_kernel,
+    )
+    _validate_prefill_fuse_config(kv_per_cube=kv_per_cube, T_q=T_q, P=P,
+                                  C=C, h_q=h_q, h_kv=H_KV, S_kv=S_kv)
+    n_tiles = S_kv // TILE_S_KV
+    Q_ROWS = kv_per_cube * T_q
+    Q_COLS = (h_q * D_HEAD) // kv_per_cube
+    topo = resolve_topology(str(TOPOLOGY_DEFAULT))
+
+    def _bench_fn(ctx):
+        configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
+        dp_q = DPPolicy(cube="column_wise", pe="replicate", num_cubes=C, num_pes=P)
+        dp_kv = DPPolicy(cube="row_wise", pe="replicate", num_cubes=C, num_pes=P)
+        dp_o = DPPolicy(cube="column_wise", pe="replicate", num_cubes=C, num_pes=P)
+        q = ctx.zeros((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_q,
+                      name=f"q_pre_fuse_kv{kv_per_cube}")
+        k = ctx.zeros((H_KV * n_tiles * D_HEAD, TILE_S_KV),
+                      dtype=DTYPE, dp=dp_kv, name=f"k_pre_fuse_kv{kv_per_cube}")
+        v = ctx.zeros((H_KV * S_kv, D_HEAD),
+                      dtype=DTYPE, dp=dp_kv, name=f"v_pre_fuse_kv{kv_per_cube}")
+        o = ctx.empty((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_o,
+                      name=f"o_pre_fuse_kv{kv_per_cube}")
+        ctx.launch(
+            f"gqa_prefill_fuse_kv{kv_per_cube}",
+            gqa_attention_prefill_short_composite_fused_kernel,
+            q, k, v, o,
+            T_q, S_kv, h_q, H_KV, D_HEAD, C, P, kv_per_cube,
+            _auto_dim_remap=False,
+        )
+
+    return run_bench(
+        topology=topo, bench_fn=_bench_fn,
+        device=resolve_device(None), engine_factory=_engine_factory,
+    )
+
+
+def _run_decode_composite_fused(*, kv_per_cube: int, C: int, S_kv: int,
+                                h_q: int = 8):
+    """Variant (3): GEMM-only composite + softmax_merge prologue fusion."""
+    from kernbench.benches.gqa_helpers.short_ctx._gqa_attention_decode_short_composite_fused import (
+        _validate_config as _validate_decode_fuse_config,
+        gqa_attention_decode_short_composite_fused_kernel,
+    )
+    T_q = 1
+    _validate_decode_fuse_config(kv_per_cube=kv_per_cube, T_q=T_q, P=P,
+                                 C=C, h_q=h_q, h_kv=H_KV, S_kv=S_kv)
+    Q_ROWS = kv_per_cube * T_q
+    Q_COLS = (h_q * D_HEAD) // kv_per_cube
+    k_rows = (H_KV * S_kv * D_HEAD) // TILE_S_KV
+    topo = resolve_topology(str(TOPOLOGY_DEFAULT))
+
+    def _bench_fn(ctx):
+        configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
+        dp_q = DPPolicy(cube="column_wise", pe="replicate", num_cubes=C, num_pes=P)
+        dp_kv = DPPolicy(cube="row_wise", pe="row_wise", num_cubes=C, num_pes=P)
+        dp_o = DPPolicy(cube="column_wise", pe="replicate", num_cubes=C, num_pes=P)
+        q = ctx.zeros((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_q,
+                      name=f"q_dec_fuse_kv{kv_per_cube}")
+        k = ctx.zeros((k_rows, TILE_S_KV), dtype=DTYPE, dp=dp_kv,
+                      name=f"k_dec_fuse_kv{kv_per_cube}")
+        v = ctx.zeros((H_KV * S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv,
+                      name=f"v_dec_fuse_kv{kv_per_cube}")
+        o = ctx.empty((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_o,
+                      name=f"o_dec_fuse_kv{kv_per_cube}")
+        ctx.launch(
+            f"gqa_decode_fuse_kv{kv_per_cube}",
+            gqa_attention_decode_short_composite_fused_kernel,
+            q, k, v, o,
+            T_q, S_kv, h_q, H_KV, D_HEAD, C, P, kv_per_cube,
+            _auto_dim_remap=False,
+        )
+
+    return run_bench(
+        topology=topo, bench_fn=_bench_fn,
+        device=resolve_device(None), engine_factory=_engine_factory,
+    )
+
+
+@pytest.mark.parametrize("kv_per_cube,C", MODES)
+def test_prefill_composite_smoke(kv_per_cube, C):
+    """Variant (2) GEMM-only composite prefill — all 4 modes."""
+    r = _run_prefill_composite(kv_per_cube=kv_per_cube, C=C, T_q=8, S_kv=1024)
+    assert r.completion.ok, (
+        f"prefill-composite kv_per_cube={kv_per_cube}: {r.completion}"
+    )
+
+
+@pytest.mark.parametrize("kv_per_cube,C", MODES)
+def test_decode_composite_smoke(kv_per_cube, C):
+    """Variant (2) GEMM-only composite decode — all 4 modes."""
+    r = _run_decode_composite(kv_per_cube=kv_per_cube, C=C, S_kv=8192)
+    assert r.completion.ok, (
+        f"decode-composite kv_per_cube={kv_per_cube}: {r.completion}"
+    )
+
+
+@pytest.mark.parametrize("kv_per_cube,C", MODES)
+def test_decode_composite_fused_smoke(kv_per_cube, C):
+    """Variant (3) composite + softmax_merge fused decode — all 4 modes."""
+    r = _run_decode_composite_fused(kv_per_cube=kv_per_cube, C=C, S_kv=8192)
+    assert r.completion.ok, (
+        f"decode-composite-fused kv_per_cube={kv_per_cube}: {r.completion}"
+    )
+
+
+@pytest.mark.parametrize("kv_per_cube,C", MODES)
+def test_decode_composite_fused_multitile(kv_per_cube, C):
+    """Fused decode with n_tiles_per_pe == 2 in EVERY mode, so the tile-1+
+    fused composite loop actually runs.
+
+    The smoke test at S_kv=8192 gives A1 (group_size=8) n_tiles_per_pe=1,
+    which never enters the fused path. S_kv = group_size·2·TILE_S_KV forces
+    2 tiles per PE for all modes (mirrors test_decode_multitile_per_pe)."""
+    group_size = P // kv_per_cube
+    S_kv = group_size * 2 * TILE_S_KV
+    r = _run_decode_composite_fused(kv_per_cube=kv_per_cube, C=C, S_kv=S_kv)
+    assert r.completion.ok, (
+        f"decode-composite-fused-multitile kv={kv_per_cube} S_kv={S_kv}: "
+        f"{r.completion}"
+    )
+
+
+@pytest.mark.parametrize("kv_per_cube,C", MODES)
+def test_prefill_composite_fused_smoke(kv_per_cube, C):
+    """Variant (3) composite + softmax_merge fused prefill — all 4 modes.
+
+    S_kv=2048 (n_tiles=2) so the tile-1+ fused composite loop actually
+    runs; S_kv=1024 gives n_tiles=1 and never enters the fused path.
+    Multi-cube (A1/A2/A4) works now that recv'd K/V slots are pinned."""
+    r = _run_prefill_composite_fused(kv_per_cube=kv_per_cube, C=C,
+                                     T_q=8, S_kv=2048)
+    assert r.completion.ok, (
+        f"prefill-composite-fused kv_per_cube={kv_per_cube}: {r.completion}"
+    )

tests/attention/test_gqa_short_context_sweep_decode.py

@@ -0,0 +1,131 @@
+"""Decode mode-comparison sweep across context lengths.
+
+Measures per (mode ∈ {A1, A2, A4, B}) × (S_kv ∈ {8K, 16K, 32K, 64K}):
+  1. Wall clock latency (μs)
+  2. Hardware utilization
+       - GEMM engine utilization (Σ gemm time / wall × n_pe)
+       - MATH op count
+  3. HBM bandwidth utilization
+       (Σ dma_read+write bytes / (wall × n_pe × peak_pe_hbm_bw))
+  4. Communication cost (IPCQ bytes — chain reduce)
+  5. KV cache space cost (per-cube bytes)
+
+Output: ``docs/sweeps/short_context_decode_sweep.csv`` (16 rows).
+
+Run: ``pytest tests/attention/test_gqa_short_context_sweep_decode.py -m slow``
+"""
+from __future__ import annotations
+
+import csv
+from pathlib import Path
+
+import pytest
+
+from tests.attention.test_gqa_short_context import (
+    D_HEAD, H_KV, P, TILE_S_KV, _run_decode,
+)
+
+PEAK_PE_HBM_BPS = 8 * 32.0  # 8 channels/PE x 32 GB/s = 256 bytes/ns per PE
+
+MATH_OPS = {"max", "sum", "exp", "sub", "mul_bcast", "add", "div",
+            "maximum", "minimum", "neg", "rmax", "rsum", "exp_diff",
+            "max_elem"}
+
+MODES = [
+    ("A1", 1, 8),
+    ("A2", 2, 4),
+    ("A4", 4, 2),
+    ("B",  8, 1),
+]
+
+CONTEXT_LENGTHS = [8 * 1024, 16 * 1024, 32 * 1024, 64 * 1024]
+
+CSV_OUT = (Path(__file__).resolve().parents[2]
+           / "docs" / "sweeps" / "short_context_decode_sweep.csv")
+
+
+def _compute_metrics(r, *, mode, kv_per_cube, C, S_kv):
+    op_log = r.engine.op_log
+    wall_ns = max(rec.t_end for rec in op_log) - min(rec.t_start for rec in op_log)
+
+    active_pes = set()
+    for rec in op_log:
+        cid = rec.component_id or ""
+        if cid and ".cube" in cid and ".pe" in cid:
+            parts = cid.split(".")
+            active_pes.add(".".join(parts[:3]))
+    n_pe = len(active_pes)
+
+    gemm_time_ns = 0.0
+    math_count = 0
+    dma_read_bytes = 0
+    dma_write_bytes = 0
+    ipcq_bytes = 0
+    for rec in op_log:
+        name = rec.op_name
+        nbytes = rec.params.get("nbytes", 0) or 0
+        if name == "gemm_f16":
+            gemm_time_ns += rec.t_end - rec.t_start
+        elif name in MATH_OPS:
+            math_count += 1
+        elif name == "dma_read":
+            dma_read_bytes += nbytes
+        elif name == "dma_write":
+            dma_write_bytes += nbytes
+        elif name == "ipcq_copy":
+            ipcq_bytes += nbytes
+
+    gemm_util = gemm_time_ns / (wall_ns * n_pe) if n_pe else 0.0
+    hbm_bw_util = (
+        (dma_read_bytes + dma_write_bytes)
+        / (wall_ns * n_pe * PEAK_PE_HBM_BPS)
+    ) if n_pe else 0.0
+
+    kv_cache_per_cube_bytes = kv_per_cube * S_kv * D_HEAD * 2 * 2
+
+    return {
+        "mode": mode,
+        "kv_per_cube": kv_per_cube,
+        "C": C,
+        "S_kv": S_kv,
+        "wall_us": wall_ns / 1000,
+        "n_pe": n_pe,
+        "gemm_util": gemm_util,
+        "math_count": math_count,
+        "hbm_bw_util": hbm_bw_util,
+        "hbm_read_mb": dma_read_bytes / (1024 * 1024),
+        "hbm_write_kb": dma_write_bytes / 1024,
+        "ipcq_kb": ipcq_bytes / 1024,
+        "kv_cache_per_cube_mb": kv_cache_per_cube_bytes / (1024 * 1024),
+    }
+
+
+@pytest.mark.slow
+def test_decode_mode_context_sweep():
+    """Run 4 modes × 4 context lengths, dump 5 metrics to CSV."""
+    rows = []
+    for mode, kv_per_cube, C in MODES:
+        for S_kv in CONTEXT_LENGTHS:
+            print(f">>> DECODE {mode} S_kv={S_kv//1024}K "
+                  f"(kv_per_cube={kv_per_cube}, C={C})", flush=True)
+            r = _run_decode(kv_per_cube=kv_per_cube, C=C, S_kv=S_kv, h_q=64)
+            assert r.completion.ok, (
+                f"DECODE {mode} S_kv={S_kv}: {r.completion}"
+            )
+            m = _compute_metrics(r, mode=mode, kv_per_cube=kv_per_cube,
+                                 C=C, S_kv=S_kv)
+            rows.append(m)
+            print(f"    wall={m['wall_us']:.1f}μs  "
+                  f"gemm_util={m['gemm_util']:.3f}  "
+                  f"hbm_bw_util={m['hbm_bw_util']:.3f}  "
+                  f"ipcq={m['ipcq_kb']:.1f}KB  "
+                  f"kv_cache={m['kv_cache_per_cube_mb']:.1f}MB/cube",
+                  flush=True)
+
+    CSV_OUT.parent.mkdir(parents=True, exist_ok=True)
+    with CSV_OUT.open("w", newline="") as f:
+        w = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
+        w.writeheader()
+        for row in rows:
+            w.writerow(row)
+    print(f"\nWrote {len(rows)} rows to {CSV_OUT}")

tests/attention/test_gqa_short_context_sweep_decode_composite.py

@@ -0,0 +1,143 @@
+"""Decode (variant 2) GEMM-only composite mode-comparison sweep.
+
+Per (mode ∈ {A1, A2, A4, B}) × (S_kv ∈ {8K, 16K, 32K, 64K}) the
+same 5 metrics as the first-level sweep are collected:
+  1. Wall clock latency (μs)
+  2. Hardware utilization (GEMM util, MATH op count)
+  3. HBM bandwidth utilization
+  4. Communication cost (IPCQ bytes — chain reduce)
+  5. KV cache space cost (per-cube bytes)
+
+Kernel: ``_gqa_attention_decode_short_composite.py`` (variant 2,
+GEMM-only composite; Q·Kᵀ via tl.composite, P·V via tl.dot).
+
+Output: ``docs/sweeps/short_context_decode_composite_sweep.csv``.
+
+Run: ``pytest tests/attention/test_gqa_short_context_sweep_decode_composite.py -m slow``
+"""
+from __future__ import annotations
+
+import csv
+from pathlib import Path
+
+import pytest
+
+from tests.attention.test_gqa_short_context import (
+    D_HEAD, H_KV, P, TILE_S_KV, _run_decode_composite,
+)
+
+PEAK_PE_HBM_BPS = 8 * 32.0  # 8 channels/PE x 32 GB/s = 256 bytes/ns per PE
+
+MATH_OPS = {"max", "sum", "exp", "sub", "mul_bcast", "add", "div",
+            "maximum", "minimum", "neg", "rmax", "rsum", "exp_diff",
+            "max_elem"}
+
+MODES = [
+    ("A1", 1, 8),
+    ("A2", 2, 4),
+    ("A4", 4, 2),
+    ("B",  8, 1),
+]
+
+CONTEXT_LENGTHS = [8 * 1024, 16 * 1024, 32 * 1024, 64 * 1024]
+
+CSV_OUT = (Path(__file__).resolve().parents[2]
+           / "docs" / "sweeps" / "short_context_decode_composite_sweep.csv")
+
+
+def _compute_metrics(r, *, mode, kv_per_cube, C, S_kv):
+    op_log = r.engine.op_log
+    wall_ns = max(rec.t_end for rec in op_log) - min(rec.t_start for rec in op_log)
+
+    active_pes = set()
+    for rec in op_log:
+        cid = rec.component_id or ""
+        if cid and ".cube" in cid and ".pe" in cid:
+            parts = cid.split(".")
+            active_pes.add(".".join(parts[:3]))
+    n_pe = len(active_pes)
+
+    gemm_time_ns = 0.0
+    math_count = 0
+    math_pipeline_time_ns = 0.0
+    dma_read_bytes = 0
+    dma_write_bytes = 0
+    ipcq_bytes = 0
+    for rec in op_log:
+        name = rec.op_name
+        nbytes = rec.params.get("nbytes", 0) or 0
+        # Composite/fused variants emit GEMM/MATH via two paths:
+        #   - non-pipeline ``tl.dot`` / primitives → ``gemm_f16`` / MATH_OPS
+        #   - pipeline ``tl.composite("gemm")`` → ``TileToken/GEMM`` (and
+        #     ``TileToken/MATH`` when a recipe prologue/epilogue is fused)
+        # Count both so GEMM util / MATH op count reflect the real work.
+        if name in ("gemm_f16", "TileToken/GEMM"):
+            gemm_time_ns += rec.t_end - rec.t_start
+        elif name == "TileToken/MATH":
+            math_count += 1
+            math_pipeline_time_ns = math_pipeline_time_ns + (rec.t_end - rec.t_start)
+        elif name in MATH_OPS:
+            math_count += 1
+        elif name == "dma_read":
+            dma_read_bytes += nbytes
+        elif name == "dma_write":
+            dma_write_bytes += nbytes
+        elif name == "ipcq_copy":
+            ipcq_bytes += nbytes
+
+    gemm_util = gemm_time_ns / (wall_ns * n_pe) if n_pe else 0.0
+    hbm_bw_util = (
+        (dma_read_bytes + dma_write_bytes)
+        / (wall_ns * n_pe * PEAK_PE_HBM_BPS)
+    ) if n_pe else 0.0
+    kv_cache_per_cube_bytes = kv_per_cube * S_kv * D_HEAD * 2 * 2
+
+    return {
+        "variant": "composite",
+        "mode": mode,
+        "kv_per_cube": kv_per_cube,
+        "C": C,
+        "S_kv": S_kv,
+        "wall_us": wall_ns / 1000,
+        "n_pe": n_pe,
+        "gemm_util": gemm_util,
+        "math_count": math_count,
+        "math_pipeline_us": math_pipeline_time_ns / 1000,
+        "hbm_bw_util": hbm_bw_util,
+        "hbm_read_mb": dma_read_bytes / (1024 * 1024),
+        "hbm_write_kb": dma_write_bytes / 1024,
+        "ipcq_kb": ipcq_bytes / 1024,
+        "kv_cache_per_cube_mb": kv_cache_per_cube_bytes / (1024 * 1024),
+    }
+
+
+@pytest.mark.slow
+def test_decode_composite_mode_context_sweep():
+    """Run 4 modes × 5 context lengths (variant 2), dump 5 metrics to CSV."""
+    rows = []
+    for mode, kv_per_cube, C in MODES:
+        for S_kv in CONTEXT_LENGTHS:
+            print(f">>> DECODE-cmp {mode} S_kv={S_kv//1024}K "
+                  f"(kv_per_cube={kv_per_cube}, C={C})", flush=True)
+            r = _run_decode_composite(kv_per_cube=kv_per_cube, C=C,
+                                      S_kv=S_kv, h_q=64)
+            assert r.completion.ok, (
+                f"DECODE-cmp {mode} S_kv={S_kv}: {r.completion}"
+            )
+            m = _compute_metrics(r, mode=mode, kv_per_cube=kv_per_cube,
+                                 C=C, S_kv=S_kv)
+            rows.append(m)
+            print(f"    wall={m['wall_us']:.1f}μs  "
+                  f"gemm_util={m['gemm_util']:.3f}  "
+                  f"hbm_bw_util={m['hbm_bw_util']:.3f}  "
+                  f"ipcq={m['ipcq_kb']:.1f}KB  "
+                  f"kv_cache={m['kv_cache_per_cube_mb']:.1f}MB/cube",
+                  flush=True)
+
+    CSV_OUT.parent.mkdir(parents=True, exist_ok=True)
+    with CSV_OUT.open("w", newline="") as f:
+        w = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
+        w.writeheader()
+        for row in rows:
+            w.writerow(row)
+    print(f"\nWrote {len(rows)} rows to {CSV_OUT}")

tests/attention/test_gqa_short_context_sweep_decode_composite_fused.py

@@ -0,0 +1,143 @@
+"""Decode (variant 3) composite + softmax_merge fused mode-comparison sweep.
+
+Per (mode ∈ {A1, A2, A4, B}) × (S_kv ∈ {8K, 16K, 32K, 64K}) the
+same 5 metrics as the first-level sweep are collected:
+  1. Wall clock latency (μs)
+  2. Hardware utilization (GEMM util, MATH op count)
+  3. HBM bandwidth utilization
+  4. Communication cost (IPCQ bytes — chain reduce)
+  5. KV cache space cost (per-cube bytes)
+
+Kernel: ``_gqa_attention_decode_short_composite.py`` (variant 2,
+GEMM-only composite; Q·Kᵀ via tl.composite, P·V via tl.dot).
+
+Output: ``docs/sweeps/short_context_decode_composite_fused_sweep.csv``.
+
+Run: ``pytest tests/attention/test_gqa_short_context_sweep_decode_composite_fused.py -m slow``
+"""
+from __future__ import annotations
+
+import csv
+from pathlib import Path
+
+import pytest
+
+from tests.attention.test_gqa_short_context import (
+    D_HEAD, H_KV, P, TILE_S_KV, _run_decode_composite_fused,
+)
+
+PEAK_PE_HBM_BPS = 8 * 32.0  # 8 channels/PE x 32 GB/s = 256 bytes/ns per PE
+
+MATH_OPS = {"max", "sum", "exp", "sub", "mul_bcast", "add", "div",
+            "maximum", "minimum", "neg", "rmax", "rsum", "exp_diff",
+            "max_elem"}
+
+MODES = [
+    ("A1", 1, 8),
+    ("A2", 2, 4),
+    ("A4", 4, 2),
+    ("B",  8, 1),
+]
+
+CONTEXT_LENGTHS = [8 * 1024, 16 * 1024, 32 * 1024, 64 * 1024]
+
+CSV_OUT = (Path(__file__).resolve().parents[2]
+           / "docs" / "sweeps" / "short_context_decode_composite_fused_sweep.csv")
+
+
+def _compute_metrics(r, *, mode, kv_per_cube, C, S_kv):
+    op_log = r.engine.op_log
+    wall_ns = max(rec.t_end for rec in op_log) - min(rec.t_start for rec in op_log)
+
+    active_pes = set()
+    for rec in op_log:
+        cid = rec.component_id or ""
+        if cid and ".cube" in cid and ".pe" in cid:
+            parts = cid.split(".")
+            active_pes.add(".".join(parts[:3]))
+    n_pe = len(active_pes)
+
+    gemm_time_ns = 0.0
+    math_count = 0
+    math_pipeline_time_ns = 0.0
+    dma_read_bytes = 0
+    dma_write_bytes = 0
+    ipcq_bytes = 0
+    for rec in op_log:
+        name = rec.op_name
+        nbytes = rec.params.get("nbytes", 0) or 0
+        # Composite/fused variants emit GEMM/MATH via two paths:
+        #   - non-pipeline ``tl.dot`` / primitives → ``gemm_f16`` / MATH_OPS
+        #   - pipeline ``tl.composite("gemm")`` → ``TileToken/GEMM`` (and
+        #     ``TileToken/MATH`` when a recipe prologue/epilogue is fused)
+        # Count both so GEMM util / MATH op count reflect the real work.
+        if name in ("gemm_f16", "TileToken/GEMM"):
+            gemm_time_ns += rec.t_end - rec.t_start
+        elif name == "TileToken/MATH":
+            math_count += 1
+            math_pipeline_time_ns = math_pipeline_time_ns + (rec.t_end - rec.t_start)
+        elif name in MATH_OPS:
+            math_count += 1
+        elif name == "dma_read":
+            dma_read_bytes += nbytes
+        elif name == "dma_write":
+            dma_write_bytes += nbytes
+        elif name == "ipcq_copy":
+            ipcq_bytes += nbytes
+
+    gemm_util = gemm_time_ns / (wall_ns * n_pe) if n_pe else 0.0
+    hbm_bw_util = (
+        (dma_read_bytes + dma_write_bytes)
+        / (wall_ns * n_pe * PEAK_PE_HBM_BPS)
+    ) if n_pe else 0.0
+    kv_cache_per_cube_bytes = kv_per_cube * S_kv * D_HEAD * 2 * 2
+
+    return {
+        "variant": "composite_fused",
+        "mode": mode,
+        "kv_per_cube": kv_per_cube,
+        "C": C,
+        "S_kv": S_kv,
+        "wall_us": wall_ns / 1000,
+        "n_pe": n_pe,
+        "gemm_util": gemm_util,
+        "math_count": math_count,
+        "math_pipeline_us": math_pipeline_time_ns / 1000,
+        "hbm_bw_util": hbm_bw_util,
+        "hbm_read_mb": dma_read_bytes / (1024 * 1024),
+        "hbm_write_kb": dma_write_bytes / 1024,
+        "ipcq_kb": ipcq_bytes / 1024,
+        "kv_cache_per_cube_mb": kv_cache_per_cube_bytes / (1024 * 1024),
+    }
+
+
+@pytest.mark.slow
+def test_decode_composite_fused_mode_context_sweep():
+    """Run 4 modes × 5 context lengths (variant 2), dump 5 metrics to CSV."""
+    rows = []
+    for mode, kv_per_cube, C in MODES:
+        for S_kv in CONTEXT_LENGTHS:
+            print(f">>> DECODE-fuse {mode} S_kv={S_kv//1024}K "
+                  f"(kv_per_cube={kv_per_cube}, C={C})", flush=True)
+            r = _run_decode_composite_fused(kv_per_cube=kv_per_cube, C=C,
+                                      S_kv=S_kv, h_q=64)
+            assert r.completion.ok, (
+                f"DECODE-fuse {mode} S_kv={S_kv}: {r.completion}"
+            )
+            m = _compute_metrics(r, mode=mode, kv_per_cube=kv_per_cube,
+                                 C=C, S_kv=S_kv)
+            rows.append(m)
+            print(f"    wall={m['wall_us']:.1f}μs  "
+                  f"gemm_util={m['gemm_util']:.3f}  "
+                  f"hbm_bw_util={m['hbm_bw_util']:.3f}  "
+                  f"ipcq={m['ipcq_kb']:.1f}KB  "
+                  f"kv_cache={m['kv_cache_per_cube_mb']:.1f}MB/cube",
+                  flush=True)
+
+    CSV_OUT.parent.mkdir(parents=True, exist_ok=True)
+    with CSV_OUT.open("w", newline="") as f:
+        w = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
+        w.writeheader()
+        for row in rows:
+            w.writerow(row)
+    print(f"\nWrote {len(rows)} rows to {CSV_OUT}")

tests/attention/test_gqa_short_context_sweep_prefill.py

@@ -0,0 +1,145 @@
+"""Prefill mode-comparison sweep across context lengths.
+
+Measures per (mode ∈ {A1, A2, A4, B}) × (S_kv ∈ {8K, 16K, 32K, 64K}):
+  1. Wall clock latency (μs)
+  2. Hardware utilization
+       - GEMM engine utilization (Σ gemm time / wall × n_pe)
+       - MATH op count
+  3. HBM bandwidth utilization
+       (Σ dma_read+write bytes / (wall × n_pe × peak_pe_hbm_bw))
+  4. Communication cost (IPCQ bytes — broadcast)
+  5. KV cache space cost (per-cube bytes)
+
+Sliced-prefill convention: T_q = 8.
+
+Output: ``docs/sweeps/prefill_sweep.csv`` (16 rows).
+
+Run: ``pytest tests/attention/test_gqa_short_context_sweep_prefill.py -m slow``
+"""
+from __future__ import annotations
+
+import csv
+from pathlib import Path
+
+import pytest
+
+from tests.attention.test_gqa_short_context import (
+    D_HEAD, H_KV, P, TILE_S_KV, _run_prefill,
+)
+
+# Peak HBM BW per PE — derived from topology (cube total / PE count).
+# topology.yaml: hbm_total_bw_gbs = 1024.0 per cube, 8 PE per cube.
+PEAK_PE_HBM_BPS = 8 * 32.0  # 8 channels/PE x 32 GB/s = 256 bytes/ns per PE
+
+MATH_OPS = {"max", "sum", "exp", "sub", "mul_bcast", "add", "div",
+            "maximum", "minimum", "neg", "rmax", "rsum", "exp_diff",
+            "max_elem"}
+
+MODES = [
+    ("A1", 1, 8),
+    ("A2", 2, 4),
+    ("A4", 4, 2),
+    ("B",  8, 1),
+]
+
+CONTEXT_LENGTHS = [8 * 1024, 16 * 1024, 32 * 1024, 64 * 1024]
+
+T_Q_PREFILL = 8
+
+CSV_OUT = (Path(__file__).resolve().parents[2]
+           / "docs" / "sweeps" / "short_context_prefill_sweep.csv")
+
+
+def _compute_metrics(r, *, mode, kv_per_cube, C, S_kv):
+    op_log = r.engine.op_log
+    wall_ns = max(rec.t_end for rec in op_log) - min(rec.t_start for rec in op_log)
+
+    # Per-PE detection (cube.pe)
+    active_pes = set()
+    for rec in op_log:
+        cid = rec.component_id or ""
+        if cid and ".cube" in cid and ".pe" in cid:
+            parts = cid.split(".")
+            active_pes.add(".".join(parts[:3]))
+    n_pe = len(active_pes)
+
+    # Stage time / count / bytes
+    gemm_time_ns = 0.0
+    math_count = 0
+    dma_read_bytes = 0
+    dma_write_bytes = 0
+    ipcq_bytes = 0
+    for rec in op_log:
+        name = rec.op_name
+        nbytes = rec.params.get("nbytes", 0) or 0
+        if name == "gemm_f16":
+            gemm_time_ns += rec.t_end - rec.t_start
+        elif name in MATH_OPS:
+            math_count += 1
+        elif name == "dma_read":
+            dma_read_bytes += nbytes
+        elif name == "dma_write":
+            dma_write_bytes += nbytes
+        elif name == "ipcq_copy":
+            ipcq_bytes += nbytes
+
+    gemm_util = gemm_time_ns / (wall_ns * n_pe) if n_pe else 0.0
+    hbm_bw_util = (
+        (dma_read_bytes + dma_write_bytes)
+        / (wall_ns * n_pe * PEAK_PE_HBM_BPS)
+    ) if n_pe else 0.0
+
+    # KV cache space — per-cube bytes (K + V, f16)
+    # Per cube owns kv_per_cube heads × S_kv tokens × d_head × 2 (K+V) × 2 (f16)
+    kv_cache_per_cube_bytes = kv_per_cube * S_kv * D_HEAD * 2 * 2
+
+    return {
+        "mode": mode,
+        "kv_per_cube": kv_per_cube,
+        "C": C,
+        "S_kv": S_kv,
+        "wall_us": wall_ns / 1000,
+        "n_pe": n_pe,
+        "gemm_util": gemm_util,
+        "math_count": math_count,
+        "hbm_bw_util": hbm_bw_util,
+        "hbm_read_mb": dma_read_bytes / (1024 * 1024),
+        "hbm_write_kb": dma_write_bytes / 1024,
+        "ipcq_kb": ipcq_bytes / 1024,
+        "kv_cache_per_cube_mb": kv_cache_per_cube_bytes / (1024 * 1024),
+    }
+
+
+@pytest.mark.slow
+def test_prefill_mode_context_sweep():
+    """Run 4 modes × 4 context lengths, dump 5 metrics to CSV.
+
+    Smoke-asserts each run completes; the measurement output is the CSV.
+    """
+    rows = []
+    for mode, kv_per_cube, C in MODES:
+        for S_kv in CONTEXT_LENGTHS:
+            print(f">>> PREFILL {mode} S_kv={S_kv//1024}K "
+                  f"(kv_per_cube={kv_per_cube}, C={C})", flush=True)
+            r = _run_prefill(kv_per_cube=kv_per_cube, C=C,
+                             T_q=T_Q_PREFILL, S_kv=S_kv, h_q=64)
+            assert r.completion.ok, (
+                f"PREFILL {mode} S_kv={S_kv}: {r.completion}"
+            )
+            m = _compute_metrics(r, mode=mode, kv_per_cube=kv_per_cube,
+                                 C=C, S_kv=S_kv)
+            rows.append(m)
+            print(f"    wall={m['wall_us']:.1f}μs  "
+                  f"gemm_util={m['gemm_util']:.3f}  "
+                  f"hbm_bw_util={m['hbm_bw_util']:.3f}  "
+                  f"ipcq={m['ipcq_kb']:.1f}KB  "
+                  f"kv_cache={m['kv_cache_per_cube_mb']:.1f}MB/cube",
+                  flush=True)
+
+    CSV_OUT.parent.mkdir(parents=True, exist_ok=True)
+    with CSV_OUT.open("w", newline="") as f:
+        w = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
+        w.writeheader()
+        for row in rows:
+            w.writerow(row)
+    print(f"\nWrote {len(rows)} rows to {CSV_OUT}")

tests/attention/test_gqa_short_context_sweep_prefill_composite.py

@@ -0,0 +1,143 @@
+"""Prefill (variant 2) GEMM-only composite mode-comparison sweep.
+
+Per (mode ∈ {A1, A2, A4, B}) × (S_kv ∈ {8K, 16K, 32K, 64K}) the
+same 5 metrics as the first-level sweep are collected:
+  1. Wall clock latency (μs)
+  2. Hardware utilization (GEMM util, MATH op count)
+  3. HBM bandwidth utilization
+  4. Communication cost (IPCQ bytes — broadcast)
+  5. KV cache space cost (per-cube bytes)
+
+Kernel: ``_gqa_attention_prefill_short_composite.py`` (variant 2,
+GEMM-only composite; Q·Kᵀ via tl.composite, P·V via tl.dot).
+
+Output: ``docs/sweeps/short_context_prefill_composite_sweep.csv``.
+
+Run: ``pytest tests/attention/test_gqa_short_context_sweep_prefill_composite.py -m slow``
+"""
+from __future__ import annotations
+
+import csv
+from pathlib import Path
+
+import pytest
+
+from tests.attention.test_gqa_short_context import (
+    D_HEAD, H_KV, P, TILE_S_KV, _run_prefill_composite,
+)
+
+PEAK_PE_HBM_BPS = 8 * 32.0  # 8 channels/PE x 32 GB/s = 256 bytes/ns per PE
+
+MATH_OPS = {"max", "sum", "exp", "sub", "mul_bcast", "add", "div",
+            "maximum", "minimum", "neg", "rmax", "rsum", "exp_diff",
+            "max_elem"}
+
+MODES = [
+    ("A1", 1, 8),
+    ("A2", 2, 4),
+    ("A4", 4, 2),
+    ("B",  8, 1),
+]
+
+CONTEXT_LENGTHS = [8 * 1024, 16 * 1024, 32 * 1024, 64 * 1024]
+
+CSV_OUT = (Path(__file__).resolve().parents[2]
+           / "docs" / "sweeps" / "short_context_prefill_composite_sweep.csv")
+
+
+def _compute_metrics(r, *, mode, kv_per_cube, C, S_kv):
+    op_log = r.engine.op_log
+    wall_ns = max(rec.t_end for rec in op_log) - min(rec.t_start for rec in op_log)
+
+    active_pes = set()
+    for rec in op_log:
+        cid = rec.component_id or ""
+        if cid and ".cube" in cid and ".pe" in cid:
+            parts = cid.split(".")
+            active_pes.add(".".join(parts[:3]))
+    n_pe = len(active_pes)
+
+    gemm_time_ns = 0.0
+    math_count = 0
+    math_pipeline_time_ns = 0.0
+    dma_read_bytes = 0
+    dma_write_bytes = 0
+    ipcq_bytes = 0
+    for rec in op_log:
+        name = rec.op_name
+        nbytes = rec.params.get("nbytes", 0) or 0
+        # Composite/fused variants emit GEMM/MATH via two paths:
+        #   - non-pipeline ``tl.dot`` / primitives → ``gemm_f16`` / MATH_OPS
+        #   - pipeline ``tl.composite("gemm")`` → ``TileToken/GEMM`` (and
+        #     ``TileToken/MATH`` when a recipe prologue/epilogue is fused)
+        # Count both so GEMM util / MATH op count reflect the real work.
+        if name in ("gemm_f16", "TileToken/GEMM"):
+            gemm_time_ns += rec.t_end - rec.t_start
+        elif name == "TileToken/MATH":
+            math_count += 1
+            math_pipeline_time_ns = math_pipeline_time_ns + (rec.t_end - rec.t_start)
+        elif name in MATH_OPS:
+            math_count += 1
+        elif name == "dma_read":
+            dma_read_bytes += nbytes
+        elif name == "dma_write":
+            dma_write_bytes += nbytes
+        elif name == "ipcq_copy":
+            ipcq_bytes += nbytes
+
+    gemm_util = gemm_time_ns / (wall_ns * n_pe) if n_pe else 0.0
+    hbm_bw_util = (
+        (dma_read_bytes + dma_write_bytes)
+        / (wall_ns * n_pe * PEAK_PE_HBM_BPS)
+    ) if n_pe else 0.0
+    kv_cache_per_cube_bytes = kv_per_cube * S_kv * D_HEAD * 2 * 2
+
+    return {
+        "variant": "composite",
+        "mode": mode,
+        "kv_per_cube": kv_per_cube,
+        "C": C,
+        "S_kv": S_kv,
+        "wall_us": wall_ns / 1000,
+        "n_pe": n_pe,
+        "gemm_util": gemm_util,
+        "math_count": math_count,
+        "math_pipeline_us": math_pipeline_time_ns / 1000,
+        "hbm_bw_util": hbm_bw_util,
+        "hbm_read_mb": dma_read_bytes / (1024 * 1024),
+        "hbm_write_kb": dma_write_bytes / 1024,
+        "ipcq_kb": ipcq_bytes / 1024,
+        "kv_cache_per_cube_mb": kv_cache_per_cube_bytes / (1024 * 1024),
+    }
+
+
+@pytest.mark.slow
+def test_prefill_composite_mode_context_sweep():
+    """Run 4 modes × 5 context lengths (variant 2 prefill), dump 5 metrics to CSV."""
+    rows = []
+    for mode, kv_per_cube, C in MODES:
+        for S_kv in CONTEXT_LENGTHS:
+            print(f">>> PREFILL-cmp {mode} S_kv={S_kv//1024}K "
+                  f"(kv_per_cube={kv_per_cube}, C={C})", flush=True)
+            r = _run_prefill_composite(kv_per_cube=kv_per_cube, C=C,
+                                       T_q=8, S_kv=S_kv, h_q=64)
+            assert r.completion.ok, (
+                f"PREFILL-cmp {mode} S_kv={S_kv}: {r.completion}"
+            )
+            m = _compute_metrics(r, mode=mode, kv_per_cube=kv_per_cube,
+                                 C=C, S_kv=S_kv)
+            rows.append(m)
+            print(f"    wall={m['wall_us']:.1f}μs  "
+                  f"gemm_util={m['gemm_util']:.3f}  "
+                  f"hbm_bw_util={m['hbm_bw_util']:.3f}  "
+                  f"ipcq={m['ipcq_kb']:.1f}KB  "
+                  f"kv_cache={m['kv_cache_per_cube_mb']:.1f}MB/cube",
+                  flush=True)
+
+    CSV_OUT.parent.mkdir(parents=True, exist_ok=True)
+    with CSV_OUT.open("w", newline="") as f:
+        w = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
+        w.writeheader()
+        for row in rows:
+            w.writerow(row)
+    print(f"\nWrote {len(rows)} rows to {CSV_OUT}")

tests/attention/test_gqa_short_context_sweep_prefill_composite_fused.py

@@ -0,0 +1,143 @@
+"""Prefill (variant 3) composite + softmax_merge fused mode-comparison sweep.
+
+Per (mode ∈ {A1, A2, A4, B}) × (S_kv ∈ {8K, 16K, 32K, 64K}) the
+same 5 metrics as the first-level sweep are collected:
+  1. Wall clock latency (μs)
+  2. Hardware utilization (GEMM util, MATH op count)
+  3. HBM bandwidth utilization
+  4. Communication cost (IPCQ bytes — broadcast)
+  5. KV cache space cost (per-cube bytes)
+
+Kernel: ``_gqa_attention_prefill_short_composite_fused.py`` (variant 2,
+GEMM-only composite; Q·Kᵀ via tl.composite, P·V via tl.dot).
+
+Output: ``docs/sweeps/short_context_prefill_composite_fused_sweep.csv``.
+
+Run: ``pytest tests/attention/test_gqa_short_context_sweep_prefill_composite_fused.py -m slow``
+"""
+from __future__ import annotations
+
+import csv
+from pathlib import Path
+
+import pytest
+
+from tests.attention.test_gqa_short_context import (
+    D_HEAD, H_KV, P, TILE_S_KV, _run_prefill_composite_fused,
+)
+
+PEAK_PE_HBM_BPS = 8 * 32.0  # 8 channels/PE x 32 GB/s = 256 bytes/ns per PE
+
+MATH_OPS = {"max", "sum", "exp", "sub", "mul_bcast", "add", "div",
+            "maximum", "minimum", "neg", "rmax", "rsum", "exp_diff",
+            "max_elem"}
+
+MODES = [
+    ("A1", 1, 8),
+    ("A2", 2, 4),
+    ("A4", 4, 2),
+    ("B",  8, 1),
+]
+
+CONTEXT_LENGTHS = [8 * 1024, 16 * 1024, 32 * 1024, 64 * 1024]
+
+CSV_OUT = (Path(__file__).resolve().parents[2]
+           / "docs" / "sweeps" / "short_context_prefill_composite_fused_sweep.csv")
+
+
+def _compute_metrics(r, *, mode, kv_per_cube, C, S_kv):
+    op_log = r.engine.op_log
+    wall_ns = max(rec.t_end for rec in op_log) - min(rec.t_start for rec in op_log)
+
+    active_pes = set()
+    for rec in op_log:
+        cid = rec.component_id or ""
+        if cid and ".cube" in cid and ".pe" in cid:
+            parts = cid.split(".")
+            active_pes.add(".".join(parts[:3]))
+    n_pe = len(active_pes)
+
+    gemm_time_ns = 0.0
+    math_count = 0
+    math_pipeline_time_ns = 0.0
+    dma_read_bytes = 0
+    dma_write_bytes = 0
+    ipcq_bytes = 0
+    for rec in op_log:
+        name = rec.op_name
+        nbytes = rec.params.get("nbytes", 0) or 0
+        # Composite/fused variants emit GEMM/MATH via two paths:
+        #   - non-pipeline ``tl.dot`` / primitives → ``gemm_f16`` / MATH_OPS
+        #   - pipeline ``tl.composite("gemm")`` → ``TileToken/GEMM`` (and
+        #     ``TileToken/MATH`` when a recipe prologue/epilogue is fused)
+        # Count both so GEMM util / MATH op count reflect the real work.
+        if name in ("gemm_f16", "TileToken/GEMM"):
+            gemm_time_ns += rec.t_end - rec.t_start
+        elif name == "TileToken/MATH":
+            math_count += 1
+            math_pipeline_time_ns = math_pipeline_time_ns + (rec.t_end - rec.t_start)
+        elif name in MATH_OPS:
+            math_count += 1
+        elif name == "dma_read":
+            dma_read_bytes += nbytes
+        elif name == "dma_write":
+            dma_write_bytes += nbytes
+        elif name == "ipcq_copy":
+            ipcq_bytes += nbytes
+
+    gemm_util = gemm_time_ns / (wall_ns * n_pe) if n_pe else 0.0
+    hbm_bw_util = (
+        (dma_read_bytes + dma_write_bytes)
+        / (wall_ns * n_pe * PEAK_PE_HBM_BPS)
+    ) if n_pe else 0.0
+    kv_cache_per_cube_bytes = kv_per_cube * S_kv * D_HEAD * 2 * 2
+
+    return {
+        "variant": "composite_fused",
+        "mode": mode,
+        "kv_per_cube": kv_per_cube,
+        "C": C,
+        "S_kv": S_kv,
+        "wall_us": wall_ns / 1000,
+        "n_pe": n_pe,
+        "gemm_util": gemm_util,
+        "math_count": math_count,
+        "math_pipeline_us": math_pipeline_time_ns / 1000,
+        "hbm_bw_util": hbm_bw_util,
+        "hbm_read_mb": dma_read_bytes / (1024 * 1024),
+        "hbm_write_kb": dma_write_bytes / 1024,
+        "ipcq_kb": ipcq_bytes / 1024,
+        "kv_cache_per_cube_mb": kv_cache_per_cube_bytes / (1024 * 1024),
+    }
+
+
+@pytest.mark.slow
+def test_prefill_composite_fused_mode_context_sweep():
+    """Run 4 modes × 5 context lengths (variant 2 prefill), dump 5 metrics to CSV."""
+    rows = []
+    for mode, kv_per_cube, C in MODES:
+        for S_kv in CONTEXT_LENGTHS:
+            print(f">>> PREFILL-fuse {mode} S_kv={S_kv//1024}K "
+                  f"(kv_per_cube={kv_per_cube}, C={C})", flush=True)
+            r = _run_prefill_composite_fused(kv_per_cube=kv_per_cube, C=C,
+                                       T_q=8, S_kv=S_kv, h_q=64)
+            assert r.completion.ok, (
+                f"PREFILL-fuse {mode} S_kv={S_kv}: {r.completion}"
+            )
+            m = _compute_metrics(r, mode=mode, kv_per_cube=kv_per_cube,
+                                 C=C, S_kv=S_kv)
+            rows.append(m)
+            print(f"    wall={m['wall_us']:.1f}μs  "
+                  f"gemm_util={m['gemm_util']:.3f}  "
+                  f"hbm_bw_util={m['hbm_bw_util']:.3f}  "
+                  f"ipcq={m['ipcq_kb']:.1f}KB  "
+                  f"kv_cache={m['kv_cache_per_cube_mb']:.1f}MB/cube",
+                  flush=True)
+
+    CSV_OUT.parent.mkdir(parents=True, exist_ok=True)
+    with CSV_OUT.open("w", newline="") as f:
+        w = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
+        w.writeheader()
+        for row in rows:
+            w.writerow(row)
+    print(f"\nWrote {len(rows)} rows to {CSV_OUT}")

tests/test_composite_recv_operand.py

@@ -0,0 +1,75 @@
+"""Regression: a recv handle used as a composite ``b`` operand must be read
+in place (on-chip), NOT DMA_READ from its non-HBM slot address.
+
+Before the recv-pinning fix, recv handles were unpinned → the composite
+scheduler emitted DMA_READ(B) from the IPCQ slot address (non-HBM) →
+PhysAddr.decode raised PhysAddrError. On-chip (TCM/SRAM) recv slots are
+now pinned (ADR-0065 D4), so a composite operand reads them in place.
+
+Isolates operand ``b``: ``a`` is a pinned ``tl.load``; only ``b`` is the
+recv handle. Mirrors the decode row-chain IPCQ directions (intra_E recv /
+intra_W send) so the recv actually lands without deadlock.
+"""
+from __future__ import annotations
+
+from pathlib import Path
+
+from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config
+from kernbench.ccl.sfr_config import configure_sfr_intercube_multisip
+from kernbench.policy.placement.dp import DPPolicy
+from kernbench.runtime_api.bench_runner import run_bench
+from kernbench.runtime_api.types import resolve_device
+from kernbench.sim_engine.engine import GraphEngine
+from kernbench.topology.builder import resolve_topology
+
+TOPOLOGY_DEFAULT = Path(__file__).resolve().parents[1] / "topology.yaml"
+P = 8
+
+
+def _ccl_cfg():
+    return resolve_algorithm_config(
+        load_ccl_config(), name="lrab_hierarchical_allreduce",
+    )
+
+
+def _engine_factory(t, d):
+    return GraphEngine(getattr(t, "topology_obj", t), enable_data=True)
+
+
+def _recv_operand_kernel(a_ptr, b_ptr, o_ptr, M, K, N, P, *, tl):
+    pe_id = tl.program_id(axis=0)
+    pe_in_group = pe_id % P
+    group_cols = min(4, P)
+    pe_col = pe_in_group % group_cols
+
+    A = tl.load(a_ptr, shape=(M, K), dtype="f16")          # pinned (no DMA)
+    if pe_col < group_cols - 1:
+        B = tl.recv(dir="intra_E", shape=(K, N), dtype="f16")  # recv handle
+        out = tl.composite("gemm", a=A, b=B)                   # ← b = recv slot
+        tl.store(o_ptr, out)
+    if pe_col > 0:
+        B_send = tl.load(b_ptr, shape=(K, N), dtype="f16")
+        tl.send(dir="intra_W", src=B_send)
+
+
+def test_recv_handle_as_composite_operand_reads_in_place():
+    M = K = N = 64
+    topo = resolve_topology(str(TOPOLOGY_DEFAULT))
+
+    def _bench_fn(ctx):
+        configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
+        dp = DPPolicy(cube="replicate", pe="replicate", num_cubes=1, num_pes=P)
+        a = ctx.zeros((M, K), dtype="f16", dp=dp, name="recv_operand_a")
+        b = ctx.zeros((K, N), dtype="f16", dp=dp, name="recv_operand_b")
+        o = ctx.empty((M, N), dtype="f16", dp=dp, name="recv_operand_o")
+        ctx.launch(
+            "recv_operand_composite", _recv_operand_kernel,
+            a, b, o, M, K, N, P,
+            _auto_dim_remap=False,
+        )
+
+    r = run_bench(
+        topology=topo, bench_fn=_bench_fn,
+        device=resolve_device(None), engine_factory=_engine_factory,
+    )
+    assert r.completion.ok, f"recv-operand composite completion: {r.completion}"

topology.yaml

@@ -1,6 +1,7 @@
 
 system:
   ns_per_mm: 0.01   # wire propagation delay: 10 ps/mm (on-chip silicon)
+  flit_bytes: 4096  # wire chunk size (ADR-0033 D1: multiple of hbm_ctrl.burst_bytes=256)
 
   sips:
     count: 2