report(gqa-short): draft results subsection (mode + composite-tier comparison, 4 figures)

test: fix CSV output paths in 3 sweep files (each variant now writes a unique CSV)
plot: merge plot scripts into single plot_short_context.py
2026-06-26 16:01:32 -07:00 · 2026-06-26 15:59:21 -07:00 · 2026-06-26 15:59:21 -07:00 · 2026-06-26 15:59:21 -07:00 · 2026-06-26 15:59:21 -07:00 · 2026-06-26 15:59:21 -07:00
39 changed files with 3295 additions and 872 deletions
@@ -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).
@@ -111,29 +111,144 @@ 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
+two phases. The \emph{prefill} kernel splits the query tile across the PEs
-shards around an inter-CUBE ring (``Ring KV''). The \emph{decode} kernel
+of a group and broadcasts each KV tile from the group's root PE over
-is head-replicated with a statically sharded KV cache and reduces partial
+intra-CUBE IPCQ, so the HBM K/V read is paid once per group instead of
-attention outputs through an M-fold intra-CUBE chain and, for multiple
+once per PE. The \emph{decode} kernel sequence-shards the KV cache across
-users, a two-level reduce-to-root. Two further primitives make long
+the same group of PEs, runs per-PE local attention, then chain-reduces the
-context practical: a \emph{lazy load} that issues the KV \textsf{DMA\_READ}
+partial $(m,\ell,O)$ triples back up to the group root. Two further
-and returns immediately, auto-waiting only at first use so that KV load
+primitives make long context practical: a \emph{lazy load} that issues the
-overlaps score computation; and per-tile \emph{scratch recycling} that
+KV \textsf{DMA\_READ} and returns immediately, auto-waiting only at first
-keeps the running softmax accumulators ($m,\ell,O$) in a persistent arena
+use so KV load overlaps score computation; and per-tile \emph{scratch
-while freeing per-tile temporaries, so the kernel fits the
+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
+The design question for short context is therefore not whether to fuse
-% (h_kv=8 KV heads -> 8 CUBEs, 1:1; intra-CUBE PE usage).
+softmax---that is settled---but how to distribute the eight KV heads
-% Bench code: src/kernbench/benches/gqa_helpers/short_ctx/
+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,
 HBM bandwidth pressure per CUBE grows linearly, the per-CUBE KV cache
 grows linearly, 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).
+A second axis is how each GEMM tile is issued. We isolate this with three
-% TODO: decode performance figure (latency, stage breakdown).
+\emph{composite tiers} applied to the same mapping: \textsf{without
-%       Bench output for short_ctx to be generated.
+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}.
 At every context length the \textsf{1-kv-per-cube} mapping is the
 fastest, and the $\textsf{8-kv-per-cube}/\textsf{1-kv-per-cube}$ ratio
 grows monotonically with $S_{kv}$ (decode: $0.98\times$ at $8$K,
 $1.06\times$ at $64$K; prefill follows the same trend). The wall gap is
 modest because short-context decode is dominated by setup overhead and
 its $M{=}G{=}8$ skinny shape leaves the MAC array idle most of the time;
 where the mappings really separate is per-CUBE pressure.
 \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). Prefill and decode follow the same ordering at
 every context length and the \textsf{1-kv-per-cube} mapping is the
 fastest throughout; the gap to \textsf{8-kv-per-cube} grows from
 $0.98\times$ at $8$K to $1.06\times$ at $64$K. The composite and
 softmax\_merge tiers track the baseline within a few percent at this
 shape (see Figure~\ref{fig:gqa-short-a1-variant}).}
 \label{fig:gqa-short-wall-baseline}
 \end{figure}
 Figure~\ref{fig:gqa-short-tradeoff} makes the per-CUBE trade-off
 explicit. Per-PE HBM bandwidth utilization is $\sim$\,$8\times$ higher on
 \textsf{8-kv-per-cube} than on \textsf{1-kv-per-cube} at every context
 length: putting all heads on a single CUBE saturates that CUBE's HBM
 channel while leaving seven CUBEs idle, whereas spreading one head per
 CUBE keeps each CUBE's read bandwidth well below the
 \SI{256}{\giga\byte\per\second} per-PE ceiling. IPCQ traffic moves in the
 opposite direction---\textsf{1-kv-per-cube} pays the broadcast or
 chain-reduce cost across all eight PEs, while \textsf{8-kv-per-cube} pays
 nothing---but absolute IPCQ volume stays under \SI{120}{\kibi\byte} per
 run, two orders of magnitude below the HBM traffic. Because per-PE HBM
 bandwidth and per-CUBE KV cache (which grows as $\text{kv\_per\_cube}\times
 S_{kv}\times d_{\text{head}}$, so $8\times$ more on
 \textsf{8-kv-per-cube}) both scale with the number of heads on a CUBE,
 the small wall advantage of \textsf{1-kv-per-cube} compounds with much
 larger headroom for longer contexts and for batched serving---which is
 why it is the choice carried into long context (\S\ref{sec:gqa-long}).
 \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;
 HBM and IPCQ totals are variant-invariant). Top row: per-PE HBM
 bandwidth utilization---\textsf{8-kv-per-cube} concentrates roughly
 $8\times$ more bandwidth on a single CUBE than \textsf{1-kv-per-cube}.
 Bottom row: IPCQ traffic per run on a log scale---\textsf{1-kv-per-cube}
 pays the broadcast/reduce cost across the largest group while
 \textsf{8-kv-per-cube} (single PE per head) pays zero.}
 \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 essentially
 unchanged across tiers ($< 0.5\%$ in either direction). 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 shows up in
 Figure~\ref{fig:gqa-short-gemm-util} as inflated GEMM engine utilization
 on the composite and softmax\_merge tiers (roughly $3$ and $6\times$ the
 baseline), which is bookkeeping cost rather than added useful work.
 Demonstrating fusion benefit therefore requires lifting $M$---either by
 batching multiple users (a separate batched-$M$ kernel, deferred to
 \S\ref{sec:future}) or by the long-context path, where Q-tile splitting
 yields a larger $M$ per PE. In the short-context regime the contribution
 of composite/softmax\_merge is a no-regression dispatch-cost isolation;
 the substantive design lever is the mapping, and
 \textsf{1-kv-per-cube} wins it.
 \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}: wall-clock
 latency is flat across the three tiers within $0.5\%$. With
 $M{=}G{=}8 < \textsf{TILE\_M}{=}32$, the composite scheduler pads the
 $M$ dimension $4\times$ and any fusion overlap is absorbed by the
 padding overhead. Fusion benefit is exercised only when $M$ fills the
 supertile---batched-$M$ inference or long-context Q-tile splitting.}
 \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 utilization on \textsf{1-kv-per-cube} across the
 three composite tiers. The $\sim 3$ and $\sim 6\times$ rise on the
 composite and softmax\_merge tiers reflects supertile padding (the
 $8\!\to\!32$ pad in $M$) being accounted as GEMM time, not added useful
 work. With actual $M=8$, the engine is below \SI{2}{\percent} busy in
 all tiers---the kernel is firmly setup- and KV-bandwidth-bound at this
 shape.}
 \label{fig:gqa-short-gemm-util}
 \end{figure}
 \subsection{Inference with Long-Context Length}
 \label{sec:gqa-long}
@@ -284,9 +399,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
+optimization for this memory-bound decode; it is a \emph{CPU-issue}
 single-rank caveat follows, Figure~\ref{fig:gqa-decode-stream}); 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 +410,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
+\paragraph{The compute-bound mirror: prefill.} Decode's verdict---command
-is a property of the \emph{full 64-way} critical path, not of the local
+form is latency-neutral---is a property of its regime, not of the
 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
 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
@@ -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("/Users/lalicorne/Desktop/ahbm/kernbench2")
 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 = [
        ("hbm_bw_util", "HBM BW utilization (per PE)", False, "{:.3f}"),
        ("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: HBM BW + 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.")
@@ -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,146.95153450000072,64,0.0017838806576053131,656,0.0,0.013997812319543966,16.0625,8.0,57.75,2.0
 composite_fused,A1,1,8,16384,279.867342500003,64,0.007173570099630665,1296,13.903615999992937,0.0146676634841736,32.0625,8.0,57.75,4.0
 composite_fused,A1,1,8,32768,544.5763825000042,64,0.010097125355742134,2576,41.71084799996763,0.015059411798858053,64.0625,8.0,57.75,8.0
 composite_fused,A1,1,8,65536,1074.101862500392,64,0.011619635377527878,5136,97.32531200143696,0.015262053416273654,128.0625,8.0,57.75,16.0
 composite_fused,A2,2,4,8192,140.35490875000028,32,0.01430408111743636,624,6.951807999995537,0.029254409671653128,16.03125,8.0,24.75,4.0
 composite_fused,A2,2,4,16384,273.9919487500021,32,0.020068677291975947,1264,20.855423999989405,0.029935186188568382,32.03125,8.0,24.75,8.0
 composite_fused,A2,2,4,32768,541.3734287504049,32,0.0230537210304644,2544,48.66265599996224,0.030282239817052973,64.03125,8.0,24.75,16.0
 composite_fused,A2,2,4,65536,1076.851308751879,32,0.024557433125924955,5104,104.27712000153959,0.030438742780552714,128.03125,8.0,24.75,32.0
 composite_fused,A4,4,2,8192,140.2291675000001,16,0.03921192786066565,608,10.427711999993306,0.058504233792873325,16.015625,8.0,8.25,8.0
 composite_fused,A4,4,2,16384,276.4849475001752,16,0.04514051167338893,1248,24.33132799998764,0.05930160085835996,32.015625,8.0,8.25,16.0
 composite_fused,A4,4,2,32768,549.3539675010354,16,0.04813782290440163,2528,52.13855999995954,0.059670088757369746,64.015625,8.0,8.25,32.0
 composite_fused,A4,4,2,65536,1095.4730075019588,16,0.04963405545386165,5088,107.7530240015909,0.05983534012350623,128.015625,8.0,8.25,64.0
 composite_fused,B,8,1,8192,144.10828500005923,8,0.08660620726926088,600,12.16566399999219,0.11380331116974475,16.0078125,8.0,0.0,16.0
 composite_fused,B,8,1,16384,285.6728750004253,8,0.09256988084807763,1240,26.069279999986758,0.11476063311909189,32.0078125,8.0,0.0,32.0
 composite_fused,B,8,1,32768,568.9925550010347,8,0.09555971782205692,2520,53.8765119999582,0.1152071313127828,64.0078125,8.0,0.0,64.0
 composite_fused,B,8,1,65536,1135.6319150019874,8,0.09706392938508733,5080,109.49097600161657,0.11543176822374755,128.0078125,8.0,0.0,128.0
@@ -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,147.23658450000076,64,0.0057446592018384314,656,0.0,0.013970712557516501,16.0625,8.0,57.75,2.0
 composite,A1,1,8,16384,279.9533665000026,64,0.006042606385378475,1360,0.0,0.014663156408229735,32.0625,8.0,57.75,4.0
 composite,A1,1,8,32768,545.3869305000062,64,0.006203478321011016,2768,0.0,0.015037030668265908,64.0625,8.0,57.75,8.0
 composite,A1,1,8,65536,1076.3614585003581,64,0.006286542449928883,5584,0.0,0.015230013923798023,128.0625,8.0,57.75,16.0
 composite,A2,2,4,8192,140.44093275000077,32,0.012045263206857853,656,0.0,0.02923649052736712,16.03125,8.0,24.75,4.0
 composite,A2,2,4,16384,274.80249675000204,32,0.012311736756599198,1360,0.0,0.029846890392199246,32.03125,8.0,24.75,8.0
 composite,A2,2,4,32768,543.6330247504077,32,0.012446984807344498,2768,0.0,0.030156372504276757,64.03125,8.0,24.75,16.0
 composite,A2,2,4,65536,1082.0090007518106,32,0.012507459726542472,5584,0.0,0.030293648183356066,128.03125,8.0,24.75,32.0
 composite,A4,4,2,8192,141.0397155000009,16,0.023988250316491475,656,0.0,0.05816801296653173,16.015625,8.0,8.25,8.0
 composite,A4,4,2,16384,278.7445435001758,16,0.02427524469220757,1360,0.0,0.05882088235384476,32.015625,8.0,8.25,16.0
 composite,A4,4,2,32768,554.51165950104,16,0.024405589617097562,2768,0.0,0.059115077994024615,64.015625,8.0,8.25,32.0
 composite,A4,4,2,65536,1106.4268915018213,16,0.024462861675913012,5584,0.0,0.05924295631591859,128.015625,8.0,8.25,64.0
 composite,B,8,1,8192,146.3478710000607,8,0.0462363541999099,656,0.0,0.11206176002378058,16.0078125,8.0,0.0,16.0
 composite,B,8,1,16384,290.8105570004271,8,0.046536082251474285,1360,0.0,0.1127331838917796,32.0078125,8.0,0.0,32.0
 composite,B,8,1,32768,579.9264290010428,8,0.046672071914328234,2768,0.0,0.11303502775846438,64.0078125,8.0,0.0,64.0
 composite,B,8,1,65536,1158.158173001712,8,0.04674036523209764,5584,0.0,0.11318661220534877,128.0078125,8.0,0.0,128.0
@@ -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
@@ -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,146.95153450000072,64,0.0017838806576053131,656,0.013997812319543966,16.0625,8.0,57.75,2.0
 A1,1,8,16384,279.3832665000025,64,0.0018765905580817512,1360,0.014693077546933053,32.0625,8.0,57.75,4.0
 A1,1,8,32768,544.246730500006,64,0.0019266555796063503,2768,0.015068533333155061,64.0625,8.0,57.75,8.0
 A1,1,8,65536,1074.081058500367,64,0.0019525081309290412,5584,0.01526234902874828,128.0625,8.0,57.75,16.0
 A2,2,4,8192,139.87083275000066,32,0.003748372621310069,656,0.029355655637933423,16.03125,8.0,24.75,4.0
 A2,2,4,16384,273.6622967500018,32,0.0038316421825465713,1360,0.029971245938539923,32.03125,8.0,24.75,8.0
 A2,2,4,32768,541.3526247504072,32,0.0038739112070761647,2768,0.030283403553383558,64.03125,8.0,24.75,16.0
 A2,2,4,65536,1077.4482007518286,32,0.003892812663356761,5584,0.030421880121130614,128.03125,8.0,24.75,32.0
 A4,4,2,8192,139.89951550000075,16,0.007495208230366452,656,0.05864209015076937,16.015625,8.0,8.25,8.0
 A4,4,2,16384,276.4641435001754,16,0.007585620230706115,1360,0.05930606331952627,32.015625,8.0,8.25,16.0
 A4,4,2,32768,549.9508595010393,16,0.007626688689616412,2768,0.059605325518975856,64.015625,8.0,8.25,32.0
 A4,4,2,65536,1097.305291501857,16,0.0076447348472311214,5584,0.059735426874945555,128.015625,8.0,8.25,64.0
 B,8,1,8192,144.06747100006032,8,0.014556735017573057,656,0.11383555139933797,16.0078125,8.0,0.0,16.0
 B,8,1,16384,286.24975700042637,8,0.014652602831705778,1360,0.1145293548666704,32.0078125,8.0,0.0,32.0
 B,8,1,32768,570.8048290010414,8,0.014696105522939606,2768,0.11484135499470417,64.0078125,8.0,0.0,64.0
 B,8,1,65536,1139.9149730017834,8,0.014717953880200899,5584,0.11499805082373886,128.0078125,8.0,0.0,128.0
@@ -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,1101.4436329998505,64,0.011331194467560951,4800,97.32531200143696,0.0018739043362378583,16.0625,64.0,114688.0,2.0
 composite_fused,A1,1,8,16384,2178.44091796638,64,0.012139279878164511,9920,208.55424000307917,0.0018875884886695197,32.0625,64.0,229376.0,4.0
 composite_fused,A1,1,8,32768,4335.060767847123,64,0.012542561892233649,20160,431.01209600543973,0.0018933990639481287,64.0625,64.0,458752.0,8.0
 composite_fused,A1,1,8,65536,8649.804197613514,64,0.012743513431742019,40640,875.9278079059123,0.0018959966752224023,128.0625,64.0,917504.0,16.0
 composite_fused,A2,2,4,8192,1093.9720835019473,32,0.01853321150177446,2400,96.40870400077105,0.0037734052470386023,16.0625,64.0,49152.0,4.0
 composite_fused,A2,2,4,16384,2175.4760434340433,32,0.019695389489835093,4960,206.59008000165224,0.00378032202414797,32.0625,64.0,98304.0,8.0
 composite_fused,A2,2,4,32768,4338.482933349601,32,0.020281493171481216,10080,426.9528320015669,0.003783811127574436,64.0625,64.0,196608.0,16.0
 composite_fused,A2,2,4,65536,8664.499802991075,32,0.020575802877008553,20320,867.678335950613,0.0037855618611332994,128.0625,64.0,393216.0,32.0
 composite_fused,A4,4,2,8192,1102.9720360028073,16,0.03256574312686088,1200,92.28038400042057,0.007485230568419403,16.0625,64.0,16384.0,8.0
 composite_fused,A4,4,2,16384,2196.357595426575,16,0.03449854256839416,2480,197.74368000090124,0.007488762319145705,32.0625,64.0,32768.0,16.0
 composite_fused,A4,4,2,32768,4383.130075368471,16,0.03547127219906801,5040,408.6702720000148,0.00749053745507198,64.0625,64.0,65536.0,32.0
 composite_fused,A4,4,2,65536,8756.675034884938,16,0.03595923415220853,10160,830.5234559737444,0.007491427937962982,128.0625,64.0,131072.0,64.0
 composite_fused,B,8,1,8192,1125.1019960018205,8,0.12366810519999533,712,121.6138560003899,0.014676002761240575,16.0625,64.0,0.0,16.0
 composite_fused,B,8,1,16384,2241.8354359431773,8,0.13192705729555512,1480,260.6011200008355,0.0146736907948643,32.0625,64.0,0.0,32.0
 composite_fused,B,8,1,32768,4475.302315825972,8,0.1360796989728368,3016,538.575647996068,0.014672528326811126,64.0625,64.0,0.0,64.0
 composite_fused,B,8,1,65536,8942.236075603707,8,0.13816184916758104,6088,1094.5247039788662,0.014671945460927953,128.0625,64.0,0.0,128.0
@@ -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,1101.224596999844,64,0.006144606666764116,5248,0.0,0.0018742770599413814,16.0625,64.0,114688.0,2.0
 composite,A1,1,8,16384,2178.2218819663735,64,0.006212950164859459,10880,0.0,0.0018877782993750494,32.0625,64.0,229376.0,4.0
 composite,A1,1,8,32768,4334.841731847132,64,0.006243911466154239,22144,0.0,0.0018934947358510517,64.0625,64.0,458752.0,8.0
 composite,A1,1,8,65536,8649.585161613555,64,0.006258419909799432,44672,0.0,0.0018960446881062475,128.0625,64.0,917504.0,16.0
 composite,A2,2,4,8192,1094.5704665014161,32,0.007139940496760185,2624,0.0,0.0037713423907684606,16.0625,64.0,49152.0,4.0
 composite,A2,2,4,16384,2177.5659939366815,32,0.007177893135813299,5440,0.0,0.0037766938053309506,32.0625,64.0,98304.0,8.0
 composite,A2,2,4,32768,4341.645937350417,32,0.007200189155116074,11072,0.0,0.0037810545209999826,64.0625,64.0,196608.0,16.0
 composite,A2,2,4,65536,8671.42409207618,32,0.007210043393449057,22336,0.0,0.003782539021470782,128.0625,64.0,393216.0,32.0
 composite,A4,4,2,8192,1104.2310135028845,16,0.008976672344028575,1312,0.0,0.007476696360673657,16.0625,64.0,16384.0,8.0
 composite,A4,4,2,16384,2197.6415934265688,16,0.009020870400324496,2720,0.0,0.0074843869215062645,32.0625,64.0,32768.0,16.0
 composite,A4,4,2,32768,4384.414073368468,16,0.00904323344692278,5536,0.0,0.007488343813013936,64.0625,64.0,65536.0,32.0
 composite,A4,4,2,65536,8757.965852885032,16,0.009054449551299243,11168,0.0,0.0074903237922982055,128.0625,64.0,131072.0,64.0
 composite,B,8,1,8192,1125.5021760018114,8,0.020408994750147753,656,0.0,0.014670784608037422,16.0625,64.0,0.0,16.0
 composite,B,8,1,16384,2242.2356159431683,8,0.02048880843562789,1360,0.0,0.014671071927542597,32.0625,64.0,0.0,32.0
 composite,B,8,1,32768,4475.702495825993,8,0.020528949826292484,2768,0.0,0.014671216431663579,64.0625,64.0,0.0,64.0
 composite,B,8,1,65536,8942.63625560376,8,0.020549079566244906,5584,0.0,0.014671288896245289,128.0625,64.0,0.0,128.0
@@ -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
@@ -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,1100.677128999906,64,0.0019053289513740923,5248,0.001875209310359147,16.0625,64.0,114688.0,2.0
 A1,1,8,16384,2177.674413966317,64,0.0019260473343055045,10880,0.0018882528874050511,32.0625,64.0,229376.0,4.0
 A1,1,8,32768,4334.294263847072,64,0.001935403433484189,22144,0.0018937339046091134,64.0625,64.0,458752.0,8.0
 A1,1,8,65536,8649.037693613021,64,0.0019397783423315907,44672,0.0018961647042087428,128.0625,64.0,917504.0,16.0
 A2,2,4,8192,1093.6804565019731,32,0.003835036070239803,2624,0.003774411415563731,16.0625,64.0,49152.0,4.0
 A2,2,4,16384,2175.2205604341793,32,0.0038564401939640384,5440,0.0037807660287830624,32.0625,64.0,98304.0,8.0
 A2,2,4,32768,4338.30076834928,32,0.003867232102121915,11072,0.003783970009586559,64.0625,64.0,196608.0,16.0
 A2,2,4,65536,8664.461183990821,32,0.003872650738157816,22336,0.0037855787340364574,128.0625,64.0,393216.0,32.0
 A4,4,2,8192,1103.915289502884,16,0.0075989598837576024,1312,0.007478834724463186,16.0625,64.0,16384.0,8.0
 A4,4,2,16384,2197.52387042685,16,0.007634600117789289,2720,0.007484787865719574,32.0625,64.0,32768.0,16.0
 A4,4,2,32768,4384.131109368463,16,0.0076536105246211355,5536,0.007488827131524694,64.0625,64.0,65536.0,32.0
 A4,4,2,65536,8757.668308885037,16,0.007662868886230499,11168,0.007490578277947104,128.0625,64.0,131072.0,64.0
 B,8,1,8192,1125.2657520018113,8,0.014909558893223249,656,0.014673867013748253,16.0625,64.0,0.0,16.0
 B,8,1,16384,2241.999191943168,8,0.01496629977413786,1360,0.014672619026008048,32.0625,64.0,0.0,32.0
 B,8,1,32768,4475.466071825993,8,0.014994832476216706,2768,0.014671991463273241,64.0625,64.0,0.0,64.0
 B,8,1,65536,8942.39983160376,8,0.01500913966355881,5584,0.014671676783710771,128.0625,64.0,0.0,128.0
@@ -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()
@@ -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
    }
  ]
 }
@@ -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()
@@ -1,37 +1,74 @@
-"""GQA fused-attention decode kernel — short context (ADR-0060 §B.split.2).
+"""GQA decode kernel: short context, attention only, multi-tile per PE (1).
-Short context (``S_kv < 256K``): each CUBE owns ``kv_per_cube`` whole
+Unified A1/A2/A4/B decode mapping per ADR-0070 (phase mirror of
-KV heads, with no S_kv sharding across CUBEs and no inter-CUBE reduce.
+``_gqa_attention_prefill_short.py``). Mode selected at launch via
-PE-SP within each CUBE: the ``P`` PEs split into ``kv_per_cube`` groups
+``kv_per_cube ∈ {1, 2, 4, 8}``:
 of ``P/kv_per_cube`` PEs each; each group does PE-SP across the group
 for one owned head, then the group's root PE stores its head's output.
-The local attention uses an S_kv-axis tile sweep (ADR-0063 §A.2) so
+    Mode  kv_per_cube  C        group_size  Reduce topology
-per-rank scratch is bounded by ``TILE_S_KV``.
+    ----  -----------  -------  ----------  -----------------------------
    A1            1    h_kv      P (=8)     row 0 + col bridge + row 1
    A2            2    h_kv/2    P/2 (=4)   row chain only
    A4            4    h_kv/4    P/4 (=2)   single intra_W hop
    B             8    1         1          (no reduce, single PE)
-Group layout on the 2×4 PE grid:
+Cube ``h`` owns ``kv_per_cube`` whole KV heads. Within each cube the
-  kv_per_cube=1, group=8 PEs (full 2×4): row chain + col bridge.
+8 PEs split into ``kv_per_cube`` groups of ``group_size = P/kv_per_cube``
-  kv_per_cube=2, group=4 PEs (one row):  row chain only.
+PEs; each PE owns ``S_local = S_kv/group_size`` KV tokens (sequence
-  kv_per_cube=4, group=2 PEs (adj cols): 1-step chain.
+shard). FA2 fuses ``G = h_q/h_kv`` Q heads into the M dim of one GEMM
-  kv_per_cube=8, group=1 PE:             no chain — direct write.
+per tile. After ``n_tiles_per_pe = S_local/TILE_S_KV`` tiles, partials
 ``(m, ℓ, O)`` chain-reduce up to group root (PE 0), which normalizes
 and stores the cube's O slab.
-Layout caveats:
+Per ADR-0011 D-VA1 the kernel computes its own shard base offset from
-  - K, V: ``(h_kv·S_kv, d_head)`` head-stacked, deployed with
+``program_id(axis=0)`` (PE id) and ``program_id(axis=1)`` (cube id);
-    ``dp = (cube=row_wise, pe=row_wise)`` so each PE's chunk is
+deploy places shards in HBM but the kernel must address them.
-    contiguously ``(S_local, d_head)`` at its own shard. K loaded as
+
-    ``(d_head, S_local)`` via byte-conserving reshape (ADR-0060 §3).
+Tensor layouts (host-side, mode-invariant byte totals):
-  - Q: replicated ``(T_q, h_q·d_head)``, reshaped byte-conservingly to
+  - Q: ``(kv_per_cube·T_q, h_q·d_head/kv_per_cube)``, T_q=1.
-    ``(h_q·T_q, d_head)``. Kernel computes attention for ALL Q rows
+    dp=(cube=column_wise, pe=replicate). Caller pre-scales by
-    against the group's owned K head; only the group's owned head rows
+    ``1/sqrt(d_head)``.
-    are semantically meaningful (correct for zero / symmetric inputs).
+  - K: ``(h_kv·S_kv·d_head/TILE_S_KV, TILE_S_KV)`` tile-major.
-  - O: replicated; each group root writes its head's
+    dp=(cube=row_wise, pe=row_wise).
-    ``(h_q·T_q, d_head)`` result at disjoint PE-local addresses.
+  - V: ``(h_kv·S_kv, d_head)`` native. Same dp as K.
-  - GEMMs use ``tl.dot`` (no composite epilogue / ``softmax_scale``).
+  - O: same dp as Q; only group root (pe_in_group==0) stores.
 Configuration constraints:
  kv_per_cube ∈ {1, 2, 4, 8},  T_q == 1,  P == 8,
  C == h_kv/kv_per_cube,  h_q % h_kv == 0,
  S_kv % (group_size·TILE_S_KV) == 0.
 Out of scope: causal mask (decode is auto-causal), f32 accumulator.
 """
 from __future__ import annotations
-TILE_S_KV = 1024  # ADR-0063 §A.2 S_kv-axis tile sweep (per-tile width).
+TILE_S_KV = 1024
 def _validate_config(*, kv_per_cube: int, T_q: int, P: int, C: int,
                     h_q: int, h_kv: int, S_kv: int) -> None:
    """Validate decode kernel configuration before the run.
    """
    if kv_per_cube not in (1, 2, 4, 8):
        raise ValueError(f"kv_per_cube must be in {{1,2,4,8}}; got {kv_per_cube}")
    if T_q != 1:
        raise ValueError(f"decode requires T_q == 1; got {T_q}")
    if P != 8:
        raise ValueError(f"2x4 mesh chain-reduce geometry requires P == 8; got {P}")
    if h_q % h_kv != 0:
        raise ValueError(f"GQA group G = h_q/h_kv must be integer; got h_q={h_q}, h_kv={h_kv}")
    if C != h_kv // kv_per_cube:
        raise ValueError(
            f"C must equal h_kv/kv_per_cube = {h_kv // kv_per_cube}; "
            f"got C={C} (h_kv={h_kv}, kv_per_cube={kv_per_cube})"
        )
    group_size = P // kv_per_cube
    if S_kv % (group_size * TILE_S_KV) != 0:
        raise ValueError(
            f"S_kv ({S_kv}) must be a multiple of group_size·TILE_S_KV "
            f"({group_size * TILE_S_KV}); each PE's sequence shard must "
            f"be a whole number of tiles"
        )
 def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl):
@@ -60,32 +97,40 @@ def gqa_attention_decode_short_kernel(
    *,
    tl,
 ) -> None:
-    """Short-context GQA decode with PE-parallel heads + intra-group PE-SP."""
+    """Unified decode: sequence-shard + chain reduce + FA2 (ADR-0070).
    """
    group_size = P // kv_per_cube
    pe_id = tl.program_id(axis=0)
    pe_in_group = pe_id % group_size
    S_local = S_kv // group_size
    n_tiles_per_pe = S_local // TILE_S_KV
-    # ── Local attention (S_kv-axis tile sweep, ADR-0063 §A.2) ──
+    G = h_q // h_kv
-    Q = tl.load(q_ptr, shape=(h_q * T_q, d_head), dtype="f16")
+    pe_id = tl.program_id(axis=0)
-    n_tiles = (S_local + TILE_S_KV - 1) // TILE_S_KV
+    cube_id = tl.program_id(axis=1)
-    KV_ROW_BYTES = d_head * 2  # f16
+    pe_in_group = pe_id % group_size
    group_id_in_cube = pe_id // group_size
-    # Tile 0: establishes persistent (m_local, l_local, O_local).
+    KV_ROW_BYTES = d_head * 2
-    #
+    K_TILE_BYTES = d_head * TILE_S_KV * 2
-    # Cannot be folded into the Tiles 1..N loop (kernbench-only limitation):
+    Q_ROW_BYTES = G * d_head * 2
-    #   - persistent (m, ℓ, O) must live OUTSIDE ``tl.scratch_scope``,
+    K_HEAD_BYTES = S_kv * d_head * 2
-    #     otherwise scope teardown discards them before the next tile's
+    V_HEAD_BYTES = S_kv * KV_ROW_BYTES
-    #     merge can read them;
+
-    #   - kernbench has no scratch-backed initializer — ``tl.zeros`` /
+    # Global VA per ADR-0011 D-VA1: kernel computes its own shard base.
-    #     ``tl.full`` return addr=0 handles with no backing storage, so
+    q_base = q_ptr + cube_id * kv_per_cube * Q_ROW_BYTES + group_id_in_cube * Q_ROW_BYTES
-    #     they cannot be overwritten via ``tl.copy_to`` to seed (-inf, 0, 0).
+    k_shard_base = (k_ptr
-    # So Tile 0 computes the initial running state directly; Tiles 1..N
+                    + cube_id * kv_per_cube * K_HEAD_BYTES
-    # fold into it. Triton port: limitation does not apply (SSA tensors
+                    + group_id_in_cube * K_HEAD_BYTES
-    # stay live across iterations) — a single unified loop suffices.
+                    + pe_in_group * n_tiles_per_pe * K_TILE_BYTES)
-    tile_s0 = min(TILE_S_KV, S_local)
+    v_shard_base = (v_ptr
-    K_T = tl.load(k_ptr, shape=(d_head, tile_s0), dtype="f16")
+                    + cube_id * kv_per_cube * V_HEAD_BYTES
-    V = tl.load(v_ptr, shape=(tile_s0, d_head), dtype="f16")
+                    + group_id_in_cube * V_HEAD_BYTES
                    + pe_in_group * n_tiles_per_pe * TILE_S_KV * KV_ROW_BYTES)
    Q = tl.load(q_base, shape=(G, d_head), dtype="f16")
    # ── Tile 0: establish (m, ℓ, O) ──────────────────────────────────
    K_T = tl.load(k_shard_base, shape=(d_head, TILE_S_KV), dtype="f16")
    V = tl.load(v_shard_base, shape=(TILE_S_KV, d_head), dtype="f16")
    scores = tl.dot(Q, K_T)
    m_local = tl.max(scores, axis=-1)
    centered = scores - m_local
@@ -93,17 +138,13 @@ def gqa_attention_decode_short_kernel(
    l_local = tl.sum(exp_scores, axis=-1)
    O_local = tl.dot(exp_scores, V)
-    # Tiles 1..n_tiles-1: fold into running state via online-softmax merge.
+    # ── Tiles 1..n_tiles_per_pe-1: sweep this PE's sequence shard ──
-    # Triton port: drop the ``with tl.scratch_scope():`` line and replace
+    for tile_idx in range(1, n_tiles_per_pe):
    # each ``copy_to`` with a Python rebind.
    for tile_idx in range(1, n_tiles):
        tile_start = tile_idx * TILE_S_KV
        tile_s = min(TILE_S_KV, S_local - tile_start)
        with tl.scratch_scope():
-            K_T_t = tl.load(k_ptr + tile_start * KV_ROW_BYTES,
+            K_T_t = tl.load(k_shard_base + tile_idx * K_TILE_BYTES,
-                            shape=(d_head, tile_s), dtype="f16")
+                            shape=(d_head, TILE_S_KV), dtype="f16")
-            V_t = tl.load(v_ptr + tile_start * KV_ROW_BYTES,
+            V_t = tl.load(v_shard_base + tile_idx * TILE_S_KV * KV_ROW_BYTES,
-                          shape=(tile_s, d_head), dtype="f16")
+                          shape=(TILE_S_KV, d_head), dtype="f16")
            scores_t = tl.dot(Q, K_T_t)
            m_tile = tl.max(scores_t, axis=-1)
            centered_t = scores_t - m_tile
@@ -117,13 +158,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 +182,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 +200,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_id * kv_per_cube * Q_ROW_BYTES
                  + group_id_in_cube * Q_ROW_BYTES)
        tl.store(o_base, O_final)
@@ -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)
@@ -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)
@@ -1,15 +1,45 @@
-"""GQA fused-attention prefill kernel — short context (ADR-0060 §B.split.2).
+"""GQA prefill kernel: short context, attention only, multi-tile (1).
-Prefill analogue of ``_gqa_decode_short.py`` — same CUBE/PE layout
+Unified A1/A2/A4/B prefill mapping per ADR-0070 (supersedes
-(``kv_per_cube`` heads per CUBE, group-PE-SP, within-group chain reduce,
+ADR-0060 §B.split.2 prefill-short clause). Mode selected at launch
-no inter-CUBE reduce). The only structural difference from short decode:
+via ``kv_per_cube ∈ {1, 2, 4, 8}``:
 ``T_q`` may be > 1 (prefill processes multiple query tokens) and Q is
 shaped ``(T_q, h_kv·d_head)`` — one Q head per KV head, no GQA M-fold.
-The local attention uses an S_kv-axis tile sweep (ADR-0063 §A.2) so
+    Mode  kv_per_cube  C        group_size  Broadcast topology
-per-rank scratch is bounded by ``TILE_S_KV``.
+    ----  -----------  -------  ----------  ---------------------------
    A1            1    h_kv      P (=8)     row 0 + col bridge + row 1
    A2            2    h_kv/2    P/2 (=4)   row chain only
    A4            4    h_kv/4    P/4 (=2)   single intra_E hop
    B             8    1         1          (no broadcast, single PE)
-No Ring KV here — each owned head is fully resident at its CUBE.
+Cube ``h`` owns ``kv_per_cube`` whole KV heads. Within each cube the
 8 PEs split into ``kv_per_cube`` groups of ``group_size = P/kv_per_cube``
 PEs; the group's T_q rows are floor-balanced across its PEs (Q-tile
 split). FA2 fuses ``G = h_q/h_kv`` Q heads into one batched GEMM per
 PE per tile. Group root (pe_in_group==0) loads K/V from HBM and IPCQ-
 broadcasts each tile to the rest of its group; q-tiles are
 independent so no intra-group reduce — every PE writes its own
 ``(T_q_pe·G, d_head)`` slab to disjoint rows of O.
 Per ADR-0011 D-VA1 the kernel computes its own shard base offset from
 ``program_id(axis=0)`` (PE id) and ``program_id(axis=1)`` (cube id);
 deploy places shards in HBM but the kernel must address them.
 Tensor layouts (host-side, mode-invariant byte totals):
  - Q: ``(kv_per_cube·T_q, h_q·d_head/kv_per_cube)``.
    dp=(cube=column_wise, pe=replicate). Caller pre-scales by
    ``1/sqrt(d_head)``.
  - K: ``(h_kv·n_tiles·d_head, TILE_S_KV)`` tile-major.
    dp=(cube=row_wise, pe=replicate).
  - V: ``(h_kv·S_kv, d_head)`` native. Same dp as K.
  - O: same dp as Q; every PE stores.
 Configuration constraints:
  kv_per_cube ∈ {1, 2, 4, 8},  P == 8,
  C == h_kv/kv_per_cube,  h_q % h_kv == 0,
  T_q >= group_size,  S_kv % TILE_S_KV == 0.
 Out of scope: causal mask, f32 accumulator. See ADR-0070 §Known
 Limitations.
 """
 from __future__ import annotations
@@ -17,6 +47,33 @@ from __future__ import annotations
 TILE_S_KV = 1024  # ADR-0063 §A.2 S_kv-axis tile sweep (per-tile width).
 def _validate_config(*, kv_per_cube: int, T_q: int, P: int, C: int,
                     h_q: int, h_kv: int, S_kv: int) -> None:
    """Validate prefill kernel configuration before the run.
    """
    if kv_per_cube not in (1, 2, 4, 8):
        raise ValueError(f"kv_per_cube must be in {{1,2,4,8}}; got {kv_per_cube}")
    if P != 8:
        raise ValueError(f"2x4 mesh broadcast geometry requires P == 8; got {P}")
    if h_q % h_kv != 0:
        raise ValueError(f"GQA group G = h_q/h_kv must be integer; got h_q={h_q}, h_kv={h_kv}")
    if C != h_kv // kv_per_cube:
        raise ValueError(
            f"C must equal h_kv/kv_per_cube = {h_kv // kv_per_cube}; "
            f"got C={C} (h_kv={h_kv}, kv_per_cube={kv_per_cube})"
        )
    group_size = P // kv_per_cube
    if T_q < group_size:
        raise ValueError(
            f"T_q ({T_q}) must be >= group_size ({group_size}) so every PE "
            f"has q-tile work for the broadcast chain"
        )
    if S_kv % TILE_S_KV != 0:
        raise ValueError(
            f"S_kv ({S_kv}) must be a multiple of TILE_S_KV ({TILE_S_KV})"
        )
 def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl):
    """Online-softmax merge of two partial ``(m, ℓ, O)`` triples."""
    m_new = tl.maximum(m_local, m_other)
@@ -34,6 +91,7 @@ def gqa_attention_prefill_short_kernel(
    o_ptr: int,
    T_q: int,
    S_kv: int,
    h_q: int,
    h_kv: int,
    d_head: int,
    C: int,
@@ -42,32 +100,79 @@ def gqa_attention_prefill_short_kernel(
    *,
    tl,
 ) -> None:
-    """Short-context prefill with PE-parallel heads + intra-group PE-SP."""
+    """Unified prefill: Q-tile split + FA2 + IPCQ KV broadcast (ADR-0070)."""
    group_size = P // kv_per_cube
    G = h_q // h_kv
    pe_id = tl.program_id(axis=0)
    cube_id = tl.program_id(axis=1)
    pe_in_group = pe_id % group_size
-    S_local = S_kv // group_size
+    group_id_in_cube = pe_id // group_size
-    # ── Local attention (S_kv-axis tile sweep, ADR-0063 §A.2) ──
+    # Floor-balanced Q-tile partition.
-    Q = tl.load(q_ptr, shape=(h_kv * T_q, d_head), dtype="f16")
+    q_start = (T_q * pe_in_group) // group_size
-    n_tiles = (S_local + TILE_S_KV - 1) // TILE_S_KV
+    q_end = (T_q * (pe_in_group + 1)) // group_size
-    KV_ROW_BYTES = d_head * 2  # f16
+    T_q_pe = q_end - q_start
    n_tiles = S_kv // TILE_S_KV
    Q_ROW_BYTES = G * d_head * 2
    KV_ROW_BYTES = d_head * 2
    K_TILE_BYTES = d_head * TILE_S_KV * 2
    K_HEAD_BYTES = n_tiles * K_TILE_BYTES
    V_HEAD_BYTES = S_kv * KV_ROW_BYTES
    # Global VA per ADR-0011 D-VA1: kernel computes its own shard base.
    q_base = (q_ptr
              + cube_id * kv_per_cube * T_q * Q_ROW_BYTES
              + (group_id_in_cube * T_q + q_start) * Q_ROW_BYTES)
    k_head_shard_base = (k_ptr
                         + cube_id * kv_per_cube * K_HEAD_BYTES
                         + group_id_in_cube * K_HEAD_BYTES)
    v_head_shard_base = (v_ptr
                         + cube_id * kv_per_cube * V_HEAD_BYTES
                         + group_id_in_cube * V_HEAD_BYTES)
    Q = tl.load(q_base, shape=(T_q_pe * G, d_head), dtype="f16")
    # 2x4 mesh broadcast geometry within the group.
    group_cols = min(4, group_size)
    group_rows = (group_size + group_cols - 1) // group_cols
    pe_col_in_group = pe_in_group % group_cols
    pe_row_in_group = pe_in_group // group_cols
    # ──────────────────────────────────────────────────────────
    # Tile 0: KV broadcast + establish persistent (m, ℓ, O).
    # Persistent state lives OUTSIDE scratch_scope (ADR-0063 §A.3).
    # ──────────────────────────────────────────────────────────
    if pe_in_group == 0:
        K_T = tl.load(k_head_shard_base,
                      shape=(d_head, TILE_S_KV), dtype="f16")
        V = tl.load(v_head_shard_base,
                    shape=(TILE_S_KV, d_head), dtype="f16")
        if group_cols > 1:
            tl.send(dir="intra_E", src=K_T)
            tl.send(dir="intra_E", src=V)
        if group_rows > 1:
            tl.send(dir="intra_S", src=K_T)
            tl.send(dir="intra_S", src=V)
    elif pe_col_in_group == 0 and pe_row_in_group > 0:
        K_T = tl.recv(dir="intra_N",
                      shape=(d_head, TILE_S_KV), dtype="f16")
        V = tl.recv(dir="intra_N",
                    shape=(TILE_S_KV, d_head), dtype="f16")
        if group_cols > 1:
            tl.send(dir="intra_E", src=K_T)
            tl.send(dir="intra_E", src=V)
    else:
        K_T = tl.recv(dir="intra_W",
                      shape=(d_head, TILE_S_KV), dtype="f16")
        V = tl.recv(dir="intra_W",
                    shape=(TILE_S_KV, d_head), dtype="f16")
        if pe_col_in_group < group_cols - 1:
            tl.send(dir="intra_E", src=K_T)
            tl.send(dir="intra_E", src=V)
    # Tile 0: establishes persistent (m_local, l_local, O_local).
    #
    # Cannot be folded into the Tiles 1..N loop (kernbench-only limitation):
    #   - persistent (m, ℓ, O) must live OUTSIDE ``tl.scratch_scope``,
    #     otherwise scope teardown discards them before the next tile's
    #     merge can read them;
    #   - kernbench has no scratch-backed initializer — ``tl.zeros`` /
    #     ``tl.full`` return addr=0 handles with no backing storage, so
    #     they cannot be overwritten via ``tl.copy_to`` to seed (-inf, 0, 0).
    # So Tile 0 computes the initial running state directly; Tiles 1..N
    # fold into it. Triton port: limitation does not apply (SSA tensors
    # stay live across iterations) — a single unified loop suffices.
    tile_s0 = min(TILE_S_KV, S_local)
    K_T = tl.load(k_ptr, shape=(d_head, tile_s0), dtype="f16")
    V = tl.load(v_ptr, shape=(tile_s0, d_head), dtype="f16")
    scores = tl.dot(Q, K_T)
    m_local = tl.max(scores, axis=-1)
    centered = scores - m_local
@@ -75,17 +180,38 @@ def gqa_attention_prefill_short_kernel(
    l_local = tl.sum(exp_scores, axis=-1)
    O_local = tl.dot(exp_scores, V)
-    # Tiles 1..n_tiles-1: fold into running state via online-softmax merge.
+    # ──────────────────────────────────────────────────────────
-    # Triton port: drop the ``with tl.scratch_scope():`` line and replace
+    # Tiles 1..n_tiles-1: broadcast + fold into running state.
-    # each ``copy_to`` with a Python rebind.
+    # ──────────────────────────────────────────────────────────
    for tile_idx in range(1, n_tiles):
        tile_start = tile_idx * TILE_S_KV
        tile_s = min(TILE_S_KV, S_local - tile_start)
        with tl.scratch_scope():
-            K_T_t = tl.load(k_ptr + tile_start * KV_ROW_BYTES,
+            if pe_in_group == 0:
-                            shape=(d_head, tile_s), dtype="f16")
+                K_T_t = tl.load(k_head_shard_base + tile_idx * K_TILE_BYTES,
-            V_t = tl.load(v_ptr + tile_start * KV_ROW_BYTES,
+                                shape=(d_head, TILE_S_KV), dtype="f16")
-                          shape=(tile_s, d_head), dtype="f16")
+                V_t = tl.load(v_head_shard_base + tile_idx * TILE_S_KV * KV_ROW_BYTES,
                              shape=(TILE_S_KV, d_head), dtype="f16")
                if group_cols > 1:
                    tl.send(dir="intra_E", src=K_T_t)
                    tl.send(dir="intra_E", src=V_t)
                if group_rows > 1:
                    tl.send(dir="intra_S", src=K_T_t)
                    tl.send(dir="intra_S", src=V_t)
            elif pe_col_in_group == 0 and pe_row_in_group > 0:
                K_T_t = tl.recv(dir="intra_N",
                                shape=(d_head, TILE_S_KV), dtype="f16")
                V_t = tl.recv(dir="intra_N",
                              shape=(TILE_S_KV, d_head), dtype="f16")
                if group_cols > 1:
                    tl.send(dir="intra_E", src=K_T_t)
                    tl.send(dir="intra_E", src=V_t)
            else:
                K_T_t = tl.recv(dir="intra_W",
                                shape=(d_head, TILE_S_KV), dtype="f16")
                V_t = tl.recv(dir="intra_W",
                              shape=(TILE_S_KV, d_head), dtype="f16")
                if pe_col_in_group < group_cols - 1:
                    tl.send(dir="intra_E", src=K_T_t)
                    tl.send(dir="intra_E", src=V_t)
            scores_t = tl.dot(Q, K_T_t)
            m_tile = tl.max(scores_t, axis=-1)
            centered_t = scores_t - m_tile
@@ -99,49 +225,10 @@ def gqa_attention_prefill_short_kernel(
            tl.copy_to(l_local, l_new)
            tl.copy_to(O_local, O_new)
-    # ── Communication: within-group chain reduce-to-root (Level-2 only) ──
+    # No intra-group reduce — q-tiles independent. Each PE writes its own
-    group_cols = min(4, group_size)
+    # slab to disjoint rows of the cube's column-wise O slab.
-    group_rows = (group_size + group_cols - 1) // group_cols
+    O_final = O_local / l_local
-    pe_col_in_group = pe_in_group % group_cols
+    o_base = (o_ptr
-    pe_row_in_group = pe_in_group // group_cols
+              + cube_id * kv_per_cube * T_q * Q_ROW_BYTES
-
+              + (group_id_in_cube * T_q + q_start) * Q_ROW_BYTES)
-    # Row chain (within group's row, along intra_W, leftward).
+    tl.store(o_base, O_final)
    if group_cols > 1:
        if pe_col_in_group < group_cols - 1:
            with tl.scratch_scope():
                m_other = tl.recv(dir="intra_E", shape=m_local.shape, dtype="f16")
                l_other = tl.recv(dir="intra_E", shape=l_local.shape, dtype="f16")
                O_other = tl.recv(dir="intra_E", shape=O_local.shape, dtype="f16")
                m_new, l_new, O_new = _merge_running(
                    m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
                )
                tl.copy_to(m_local, m_new)
                tl.copy_to(l_local, l_new)
                tl.copy_to(O_local, O_new)
        if pe_col_in_group > 0:
            tl.send(dir="intra_W", src=m_local)
            tl.send(dir="intra_W", src=l_local)
            tl.send(dir="intra_W", src=O_local)
    # Col bridge (within group, along intra_N, row-1 → row-0).
    if pe_col_in_group == 0 and group_rows > 1:
        if pe_row_in_group < group_rows - 1:
            with tl.scratch_scope():
                m_other = tl.recv(dir="intra_S", shape=m_local.shape, dtype="f16")
                l_other = tl.recv(dir="intra_S", shape=l_local.shape, dtype="f16")
                O_other = tl.recv(dir="intra_S", shape=O_local.shape, dtype="f16")
                m_new, l_new, O_new = _merge_running(
                    m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
                )
                tl.copy_to(m_local, m_new)
                tl.copy_to(l_local, l_new)
                tl.copy_to(O_local, O_new)
        if pe_row_in_group > 0:
            tl.send(dir="intra_N", src=m_local)
            tl.send(dir="intra_N", src=l_local)
            tl.send(dir="intra_N", src=O_local)
    # ── Final normalise + store (group root only) ──
    if pe_in_group == 0:
        O_final = O_local / l_local
        tl.store(o_ptr, O_final)
@@ -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)
@@ -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)
@@ -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:
@@ -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
@@ -1,31 +1,17 @@
-"""Phase 1 spec test for Phase D: short-context GQA kernels
+"""Short-context GQA kernel tests — unified A1/A2/A4/B (ADR-0070).
 (ADR-0060 §B "Items from the long/short context split", item B.split.2).
-The long-context kernels (§5.2 decode / §5.5 prefill) shard each KV
+Both ``_gqa_attention_prefill_short.py`` and
-head row-wise across all CUBEs. At short context (S_kv < 256K), that
+``_gqa_attention_decode_short.py`` are single unified kernels selected
-shard is too thin to feed the engines and the cube-level collective
+at launch via ``kv_per_cube ∈ {1, 2, 4, 8}``:
 overhead dominates. The short-context kernels drop cube-SP entirely:
 each CUBE owns ``kv_per_cube`` *whole* KV heads, no S_kv sharding across
 CUBEs.
-Design (per AskUserQuestion answers in this session):
+    Mode  kv_per_cube  C        group_size
-  - 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.
+    A1            1    h_kv      P (=8)
-  - scratch_scope + tl.copy_to discipline mirrors the long kernels.
+    A2            2    h_kv/2    P/2 (=4)
-  - One short kernel handles kv_per_cube ∈ {1, 2, 4, 8} via a parameter.
+    A4            4    h_kv/4    P/4 (=2)
    B             8    1         1
-Group layout on the 2×4 PE grid:
+Tests are mode-parametrized over the four (kv_per_cube, C) tuples.
  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.
 """
 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,
+def _run_prefill(*, kv_per_cube: int, C: int, T_q: int, S_kv: int,
-                     C: int, P: int, S_kv: int):
+                 h_q: int = 8):
-    """Run the short-context decode kernel with PE-parallel heads.
+    """Run the unified prefill kernel in the given mode."""
-
+    from kernbench.benches.gqa_helpers.short_ctx._gqa_attention_prefill_short import (
-    Layout (after design iteration — see Phase D failure-recovery):
+        _validate_config as _validate_prefill_config,
-      Q: (T_q, h_q·D_HEAD) replicated; kernel reshapes byte-conservingly.
+        gqa_attention_prefill_short_kernel,
-      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
+    _validate_prefill_config(kv_per_cube=kv_per_cube, T_q=T_q, P=P, C=C,
-          addressable shard (no partial reads needed).
+                             h_q=h_q, h_kv=H_KV, S_kv=S_kv)
-      O: replicated; each group root writes the full byte-conserving
+    n_tiles = S_kv // TILE_S_KV
-          (h_q·T_q, D_HEAD) result.
+    Q_ROWS = kv_per_cube * T_q
-    """
+    Q_COLS = (h_q * D_HEAD) // kv_per_cube
    from kernbench.benches.gqa_helpers.short_ctx._gqa_attention_decode_short import gqa_attention_decode_short_kernel  # Phase 2
    topo = resolve_topology(str(TOPOLOGY_DEFAULT))
    def _bench_fn(ctx):
        configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
-        dp_full = DPPolicy(cube="replicate", pe="replicate",
+        dp_q = DPPolicy(cube="column_wise", pe="replicate", num_cubes=C, num_pes=P)
-                           num_cubes=C, num_pes=P)
+        dp_kv = DPPolicy(cube="row_wise", pe="replicate", num_cubes=C, num_pes=P)
-        # Head-stacked KV with row_wise sharding so each PE's chunk is
+        dp_o = DPPolicy(cube="column_wise", pe="replicate", num_cubes=C, num_pes=P)
-        # contiguous and exactly (S_local, D_HEAD) addressable.
+        q = ctx.zeros((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_q,
-        dp_kv = DPPolicy(cube="row_wise", pe="row_wise",
+                      name=f"q_pre_kv{kv_per_cube}")
-                         num_cubes=C, num_pes=P)
+        k = ctx.zeros((H_KV * n_tiles * D_HEAD, TILE_S_KV),
-        q = ctx.zeros((1, h_q * D_HEAD),
+                      dtype=DTYPE, dp=dp_kv, name=f"k_pre_kv{kv_per_cube}")
-                      dtype=DTYPE, dp=dp_full, name=f"q_short_{kv_per_cube}_{C}")
+        v = ctx.zeros((H_KV * S_kv, D_HEAD),
-        k = ctx.zeros((h_kv * S_kv, D_HEAD),
+                      dtype=DTYPE, dp=dp_kv, name=f"v_pre_kv{kv_per_cube}")
-                      dtype=DTYPE, dp=dp_kv, name=f"k_short_{kv_per_cube}_{C}")
+        o = ctx.empty((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_o,
-        v = ctx.zeros((h_kv * S_kv, D_HEAD),
+                      name=f"o_pre_kv{kv_per_cube}")
                      dtype=DTYPE, dp=dp_kv, name=f"v_short_{kv_per_cube}_{C}")
        o = ctx.empty((1, h_q * D_HEAD),
                      dtype=DTYPE, dp=dp_full, name=f"o_short_{kv_per_cube}_{C}")
        ctx.launch(
-            f"gqa_decode_short_{kv_per_cube}_{C}",
+            f"gqa_prefill_kv{kv_per_cube}",
            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}",
            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():
+@pytest.mark.parametrize("kv_per_cube,C", MODES)
-    """Short prefill kv_per_cube=2, C=4. Smoke: completes."""
+def test_prefill_smoke(kv_per_cube, C):
-    result = _run_prefill_short(
+    """All four prefill modes complete on a single-tile mini config."""
-        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, (
-    assert result.completion.ok, (
+        f"prefill kv_per_cube={kv_per_cube}: {r.completion}"
        f"short prefill kv_per_cube=2 must complete; got {result.completion}"
    )
-def test_short_prefill_no_ring_KV_traffic():
+@pytest.mark.parametrize("kv_per_cube,C", MODES)
-    """ADR-0060 §B.split.2: short prefill DOES NOT use Ring KV — each
+def test_prefill_qtile_split_dma_writes(kv_per_cube, C):
-    CUBE owns its KV heads fully, no rotation. The kernel must not
+    """Every PE in every active group stores its q-tile output:
-    emit any KV-rotation IPCQ traffic at the CUBE level.
+    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(
+    r = _run_prefill(kv_per_cube=kv_per_cube, C=C, T_q=8, S_kv=1024)
-        h_kv=8, kv_per_cube=2, C=4, P=8, T_q=4, S_kv=64,
+    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
+        1 for rec in r.engine.op_log
-        if r.op_name == "ipcq_copy"
+        if rec.op_name == "ipcq_copy"
-        and r.params.get("direction") in ("E", "W")
+        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}"
    )
@@ -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 = 1024.0 / 8  # 128 bytes/ns
 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)
    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}")
@@ -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 = 1024.0 / 8  # 128 bytes/ns
 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)
    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}")
@@ -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 = 1024.0 / 8  # 128 bytes/ns
 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)
    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}")
@@ -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 = 1024.0 / 8  # 128 bytes/ns
 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)
    # 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}")
@@ -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 = 1024.0 / 8  # 128 bytes/ns
 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)
    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}")
@@ -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 = 1024.0 / 8  # 128 bytes/ns
 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)
    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}")
@@ -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}"
Author	SHA1	Message	Date
eusonice	f8caf2e16c	report(gqa-short): draft results subsection (mode + composite-tier comparison, 4 figures)	2026-06-26 16:01:32 -07:00
eusonice	09e70bd516	test: fix CSV output paths in 3 sweep files (each variant now writes a unique CSV)	2026-06-26 15:59:21 -07:00
eusonice	e20e1f7362	plot: merge plot scripts into single plot_short_context.py	2026-06-26 15:59:21 -07:00
eusonice	3437d7e079	sweep: 6 short-context CSVs (3 variants x prefill+decode) + 6 wall/util figures	2026-06-26 15:59:21 -07:00
eusonice	3da116d40d	adr-0070: gqa short context attention prefill/decode + 3 composite tiers	2026-06-26 15:59:21 -07:00
eusonice	f7aa8134b9	test: 3-variant sweep (baseline + composite + composite-fused) for prefill/decode	2026-06-26 15:59:21 -07:00
eusonice	a455ae61b4	sim: pin compute outputs and on-chip recv slots for composite operands	2026-06-26 15:59:21 -07:00
eusonice	1971ac731c	add three gqa attention kernel variants for short context length	2026-06-26 15:59:21 -07:00
eusonice	a98e93110b	short context length composite command	2026-06-26 15:59:21 -07:00
eusonice	f08dda3bd4	test short context length attention kernel	2026-06-26 15:59:21 -07:00
eusonice	7d90d53d4d	short context length attention kernel	2026-06-26 15:59:21 -07:00