diff --git a/docs/adr-proposed/ADR-0060-algo-gqa-fused-attention-ahbm-ko.md b/docs/adr-proposed/ADR-0060-algo-gqa-fused-attention-ahbm-ko.md index c2e6edf..74ae334 100644 --- a/docs/adr-proposed/ADR-0060-algo-gqa-fused-attention-ahbm-ko.md +++ b/docs/adr-proposed/ADR-0060-algo-gqa-fused-attention-ahbm-ko.md @@ -8,11 +8,9 @@ Proposed **Decision drivers:** agentic workload → 낮은 batch, 긴 context; KV-load-bound decode; long-context prefill용 sequence-parallel (Ring KV). -**Supersedes / extends:** 기존 mesh-native 어텐션 커널 -`_attention_mesh_kv`(prefill)와 `_attention_mesh_mlo`(decode), 그리고 -`milestone-gqa-llama70b` eval bench. *§A. 기존 kernbench 작업과의 관계* 참조 — -그 코드가 본 ADR이 진짜 GQA·causal·long-context 커널로 업그레이드하는 -baseline이다. +**Supersedes / extends:** 이전 mesh-native 어텐션 커널과 그 `milestone-gqa-llama70b` +eval bench(본 ADR의 커널이 도입되면서 제거됨). *§A. 기존 kernbench 작업과의 관계* 참조 — +그 코드가 본 ADR이 진짜 GQA·causal·long-context 커널로 업그레이드한 baseline이다. **Supporting ADRs** (efficiency / scale enabler — *GQA blocker 아님*; §8 정정 참조): **ADR-0063** `tl.scratch_scope`(per-tile scratch 재활용 — 현실적 context @@ -160,18 +158,16 @@ def gqa_prefill_sp(q_ptr, k_ptr, v_ptr, o_ptr, T_q, S_kv_local, d, C, scale, q_b ## A. 기존 kernbench 작업과의 관계 (먼저 읽을 것) -kernbench는 **오늘날 이미 IPCQ 상에서 online-softmax `(m, ℓ, O)` 머지로 -FlashAttention을 돌린다.** 두 커널이 존재한다: +kernbench는 본 ADR 이전에 IPCQ 상에서 online-softmax `(m, ℓ, O)` 머지로 +FlashAttention을 돌리는 두 mesh 커널이 있었다(현재 제거됨): -| File | Role | Mechanism | -|---|---|---| -| `src/kernbench/benches/_attention_mesh_kv.py` | prefill (Ring K/V) | per-rank partial attention, bidirectional K/V fan-out, online-softmax fold | -| `src/kernbench/benches/_attention_mesh_mlo.py` | decode (split-KV) | per-rank one-shot partial attention, bidirectional `(m,ℓ,O)` fan-out, log-sum-exp merge | +| Role | Mechanism | +|---|---| +| prefill (Ring K/V) | per-rank partial attention, bidirectional K/V fan-out, online-softmax fold | +| decode (split-KV) | per-rank one-shot partial attention, bidirectional `(m,ℓ,O)` fan-out, log-sum-exp merge | 둘 다 `milestone-gqa-llama70b`(4 패널: -`{single,multi}_user × {prefill,decode}`, -`src/kernbench/benches/milestone_gqa_llama70b.py`)이 구동하며 -`tests/attention/test_milestone_gqa_llama70b.py`에서 테스트된다. +`{single,multi}_user × {prefill,decode}`)이 구동했다. **이들은 greenlet `tl` API로 작성되었다:** `tl.load`, `tl.dot`, `tl.softmax`/`tl.max`/`tl.sum`/`tl.exp`, `tl.send`/`tl.recv`, 그리고 @@ -183,8 +179,8 @@ FlashAttention을 돌린다.** 두 커널이 존재한다: **baseline의 의도된 세 한계** — 효율적 GQA 커널이 정확히 들어내야 하는 것: -1. **GQA 재사용 없음.** `h_q == h_kv == 1` - (`test_milestone_gqa_llama70b.py:137-142`). 테스트는 이를 *broadcast view*에 +1. **GQA 재사용 없음.** baseline은 `h_q == h_kv == 1`로 제한되어 있었다. + 해당 baseline 테스트는 이를 *broadcast view*에 대한 MemoryStore byte-conservation 실패로 돌리지만, 그 실패는 baseline의 **head-packing 핵**의 속성이다(`_view(K, (h_q·d, S_kv))`는 모든 head를 하나의 matmul 차원에 뭉치고 `h_q == h_kv`일 때만 byte를 보존). 올바른 수정은 @@ -196,7 +192,7 @@ FlashAttention을 돌린다.** 두 커널이 존재한다: 필요 → **2-level reduce-to-root**(intra-CUBE tree + intra-CUBE-Group center-mesh, §4)는 `⌈log₂ P⌉` + center-mesh-over-`C` 단계. 3. **검증 스케일만.** `S = 16`인 이유는 1 MiB scratch bump allocator가 per-tile - 임시값을 누수하고(`test_milestone_gqa_llama70b.py:123-148`) causal masking / + 임시값을 누수하고(baseline 테스트는 그래서 S=16으로 제한되어 있었음) causal masking / tiling이 없기 때문 → **ADR-0063**(재활용) + §5(tiling, causal skip) + composite K/V 스트리밍(§3) + **ADR-0062**(lazy load 오버랩)으로 해결. @@ -563,8 +559,7 @@ reduce **안 함** — KV를 회전(Ring, §5.5). decode에 reduce를 택한 이 `O = [G, d]`가 작아 `(m,ℓ,O)` 이동이 상주 KV 이동보다 싸기 때문. tile sweep 후 각 rank는 자기 `1/(C·P)` shard에 대해 `(m_i, ℓ_i, O_i)`(비정규화)를 -가진다. 결합/교환 가능한 log-sum-exp 머지로 결합(baseline의 fold와 동일 수학, -`_attention_mesh_mlo.py:117-122`): +가진다. 결합/교환 가능한 log-sum-exp 머지로 결합(baseline의 fold와 동일 수학): ```python def merge(m_a, l_a, O_a, m_b, l_b, O_b): @@ -688,9 +683,8 @@ def gqa_decode_sp(q_ptr, k_ptr, v_ptr, o_ptr, counter, start_pe, start_cube, 축은 세어지지 *않음* — §8 참조). - `S_rank`가 scratch에 맞으면(작은/중간 context) 이것은 **one-shot** partial attention으로 축약(Q·Kᵀ용 composite 하나, softmax 하나, P·V용 composite 하나) — - 바로 baseline `_attention_mesh_mlo`의 `_partial_attention`, 단지 GQA-batched에 - composite 경로. Tiling(§3)은 `S_rank`가 scratch scope의 tile 예산을 초과할 때만 - 발동. + 바로 baseline의 `_partial_attention` 구조, 단지 GQA-batched에 composite 경로. + Tiling(§3)은 `S_rank`가 scratch scope의 tile 예산을 초과할 때만 발동. ### 5.3 DECODE, with SP / KV-parallel (`C × P` rank) @@ -741,7 +735,7 @@ KV를 필요로 하므로. **`(m,ℓ,O)` reduce 없음** — 각 CUBE가 자기 head의 행을 정규화·기록. - **Ring KV:** `C` KV slice가 CUBE ring을 **회전**; 각 CUBE가 들어오는 블록을 자기 head의 실행 중 `(m,ℓ,O)`에 fold(online-softmax, ring step 가로질러 Python 핸들로 - carry, `_attention_mesh_kv`가 오늘 하듯). `C` step 후 모든 head가 전 KV를 봄. + carry). `C` step 후 모든 head가 전 KV를 봄. **GQA 재사용은 회전에서** — slice `j`가 전 `C` CUBE를 방문하며 전 `G` head를 서비스. - IPCQ가 다음 step의 KV 수신을 현재 step의 compute와 `tl.recv_async`/`tl.wait`로 오버랩(`tl_context.py:543-560` — 이미 존재); 수신 버퍼는 ping-pong(persistent @@ -751,7 +745,7 @@ KV를 필요로 하므로. - **CUBE 내(P PE):** head의 query 행 `[T_q, d]` 및/또는 현재 KV 블록을 `P` PE에 타일(여기엔 decode와 달리 query 축이 존재); 상세는 §B. -baseline `_attention_mesh_kv`가 이미 ring fold를 구현; 본 ADR은 GQA 재사용, +(현재 제거된) baseline이 이미 ring fold를 구현; 본 ADR은 GQA 재사용, head-parallel 배치, causal step-skip, composite-hybrid inner tile(§3)을 추가. ### 5.6 Decode CPU-pipelining 변형 (opt1 / opt3 / opt2) @@ -1060,7 +1054,7 @@ exercise하려 존재하고, `tl.trans`는 reshape-not-transpose이며, `bf16` 가정. `C ≠ G`면 매핑 재검토 필요(CUBE당 다중 head, 또는 head가 부분 ring에 걸침). **권고:** headline 스케일에서 prefill 커널은 `C = G` 고정; `C ≠ G`는 별도 연구. -5. **`_attention_mesh_mlo_2d`(현재 impl)와 정합.** 원격 impl의 2D 커널은 **Q +5. **이전 `_attention_mesh_mlo_2d` impl(현재 제거됨)과 정합.** 그 2D 커널은 **Q replicated**로 cube에 걸친 **AllReduce** — 즉 decode-reduce 계열이나 reduce-to-root 아닌 all-reduce(broadcast-back), 그리고 prefill head-parallel ring은 아직 없음. **권고:** (a) 그 2D AllReduce → reduce-to-root(broadcast-back 제거)로 decode 커널; diff --git a/docs/adr-proposed/ADR-0060-algo-gqa-fused-attention-ahbm.md b/docs/adr-proposed/ADR-0060-algo-gqa-fused-attention-ahbm.md index f8f2d5c..bf0db71 100644 --- a/docs/adr-proposed/ADR-0060-algo-gqa-fused-attention-ahbm.md +++ b/docs/adr-proposed/ADR-0060-algo-gqa-fused-attention-ahbm.md @@ -8,11 +8,11 @@ Proposed **Decision drivers:** agentic workload → low batch, long context; KV-load-bound decode; sequence-parallel (Ring KV) for long-context prefill. -**Supersedes / extends:** the existing mesh-native attention kernels -`_attention_mesh_kv` (prefill) and `_attention_mesh_mlo` (decode) and the -`milestone-gqa-llama70b` eval bench. See *§A. Relationship to existing -kernbench work* — that code is the baseline this ADR upgrades to a real -GQA, causal, long-context kernel. +**Supersedes / extends:** pre-existing mesh-native attention kernels and +their `milestone-gqa-llama70b` eval bench (removed when this ADR's +kernels landed). See *§A. Relationship to existing kernbench work* — that +code was the baseline this ADR upgrades to a real GQA, causal, +long-context kernel. **Supporting ADRs** (efficiency / scale enablers — *not* GQA blockers; see §8 correction): **ADR-0063** `tl.scratch_scope` (per-tile scratch @@ -168,18 +168,16 @@ def gqa_prefill_sp(q_ptr, k_ptr, v_ptr, o_ptr, T_q, S_kv_local, d, C, scale, q_b ## A. Relationship to existing kernbench work (read first) -kernbench **already runs FlashAttention with an online-softmax `(m, ℓ, O)` -merge over IPCQ today.** Two kernels exist: +kernbench previously ran FlashAttention with an online-softmax `(m, ℓ, O)` +merge over IPCQ via two pre-ADR-0060 mesh kernels (now removed): -| File | Role | Mechanism | -|---|---|---| -| `src/kernbench/benches/_attention_mesh_kv.py` | prefill (Ring K/V) | per-rank partial attention, bidirectional K/V fan-out, online-softmax fold | -| `src/kernbench/benches/_attention_mesh_mlo.py` | decode (split-KV) | per-rank one-shot partial attention, bidirectional `(m,ℓ,O)` fan-out, log-sum-exp merge | +| Role | Mechanism | +|---|---| +| prefill (Ring K/V) | per-rank partial attention, bidirectional K/V fan-out, online-softmax fold | +| decode (split-KV) | per-rank one-shot partial attention, bidirectional `(m,ℓ,O)` fan-out, log-sum-exp merge | -Both are driven by `milestone-gqa-llama70b` (4 panels: -`{single,multi}_user × {prefill,decode}`, -`src/kernbench/benches/milestone_gqa_llama70b.py`) and tested in -`tests/attention/test_milestone_gqa_llama70b.py`. +Both were driven by `milestone-gqa-llama70b` (4 panels: +`{single,multi}_user × {prefill,decode}`). **They are written in the greenlet `tl` API:** `tl.load`, `tl.dot`, `tl.softmax`/`tl.max`/`tl.sum`/`tl.exp`, `tl.send`/`tl.recv`, and Python @@ -193,8 +191,8 @@ GEMMs onto the scheduler-managed `tl.composite` path** (see §1) — this **Three deliberate limitations of the baseline** — exactly what an *efficient GQA* kernel must lift: -1. **No GQA reuse.** `h_q == h_kv == 1` - (`test_milestone_gqa_llama70b.py:137-142`). The test attributes this to +1. **No GQA reuse.** Baseline was capped at `h_q == h_kv == 1`. The + baseline test attributed this to a MemoryStore byte-conservation failure on a *broadcast view*, but that failure is a property of the baseline's **head-packing hack** (`_view(K, (h_q·d, S_kv))`, which conflates all heads into one matmul @@ -209,9 +207,9 @@ GEMMs onto the scheduler-managed `tl.composite` path** (see §1) — this reduce-to-root** (intra-CUBE tree + intra-CUBE-Group center-mesh, §4) is `⌈log₂ P⌉` + center-mesh-over-`C` steps. 3. **Validation scale only.** `S = 16` because the 1 MiB scratch bump - allocator leaks per-tile temporaries - (`test_milestone_gqa_llama70b.py:123-148`) and there is no causal - masking / tiling → fixed by **ADR-0063** (recycling) + §5 (tiling, + allocator leaks per-tile temporaries (baseline test capped S at 16 + for this reason) and there is no causal masking / tiling → fixed by + **ADR-0063** (recycling) + §5 (tiling, causal skip) + composite K/V streaming (§3) + **ADR-0062** (lazy load overlap). @@ -616,7 +614,7 @@ sharded, small `O`). Prefill-SP does **not** reduce — it rotates KV (Ring, After its tile sweep each rank holds `(m_i, ℓ_i, O_i)` (unnormalised) over its `1/(C·P)` shard. Combine via the associative/commutative log-sum-exp -merge (identical math to the baseline's fold, `_attention_mesh_mlo.py:117-122`): +merge (identical math to the baseline's fold): ```python def merge(m_a, l_a, O_a, m_b, l_b, O_b): @@ -750,8 +748,8 @@ def gqa_decode_sp(q_ptr, k_ptr, v_ptr, o_ptr, counter, start_pe, start_cube, batch axis would *not* be counted — see §8). - If `S_rank` fits in scratch (small/medium context) this degenerates to a **one-shot** partial attention (one composite for Q·Kᵀ, one softmax, one - composite for P·V) — exactly the baseline `_attention_mesh_mlo` - `_partial_attention`, just GQA-batched and on the composite path. Tiling + composite for P·V) — exactly the baseline `_partial_attention` shape, + just GQA-batched and on the composite path. Tiling (§3) only kicks in when `S_rank` exceeds the scratch scope's tile budget. ### 5.3 DECODE, with SP / KV-parallel (`C × P` ranks) @@ -809,8 +807,8 @@ head needs in full anyway. `(m,ℓ,O)` reduce** — each CUBE normalises and writes its own head's rows. - **Ring KV:** the `C` KV slices **rotate** around the CUBE ring; each CUBE folds the incoming block into its head's running `(m,ℓ,O)` (online-softmax, - carried across ring steps as Python handles, as `_attention_mesh_kv` does - today). After `C` steps every head has seen all KV. **GQA reuse comes from + carried across ring steps as Python handles). After `C` steps every + head has seen all KV. **GQA reuse comes from the rotation** — slice `j` visits all `C` CUBEs, serving all `G` heads. - IPCQ overlaps the next step's KV receive with the current step's compute via `tl.recv_async`/`tl.wait` (`tl_context.py:543-560` — already exists); @@ -821,8 +819,8 @@ head needs in full anyway. current KV block tile across the `P` PEs (a query-axis split exists here, unlike decode); details in §B. -The baseline `_attention_mesh_kv` already implements the ring fold; this -ADR adds GQA reuse, the head-parallel placement, causal step-skip, and the +The (now-removed) baseline already implemented the ring fold; this ADR +adds GQA reuse, the head-parallel placement, causal step-skip, and the composite-hybrid inner tile (§3). ### 5.6 Decode CPU-pipelining variants (opt1 / opt3 / opt2) @@ -1184,8 +1182,8 @@ predicted default; revise on review. fix `C = G` for the prefill kernel at headline scale; treat `C ≠ G` as a separate study. -5. **Reconcile with `_attention_mesh_mlo_2d` (current impl).** The remote - impl's 2D kernel is an **AllReduce** over cubes with **Q replicated** — +5. **Reconcile with the prior `_attention_mesh_mlo_2d` impl (now removed).** + That 2D kernel was an **AllReduce** over cubes with **Q replicated** — i.e. the decode-reduce family, but all-reduce (broadcast-back) not reduce-to-root, and not yet the prefill head-parallel ring. **Recommend:** (a) move its 2D AllReduce → reduce-to-root (drop broadcast-back) for the @@ -1208,3 +1206,50 @@ predicted default; revise on review. command kind. Defer **opt2** (the two-composite `ex_composite`, only `#2` is new) until ADR-0064's cost model makes its fewer-issues win measurable. + +### Items from the long/short context split (this revision) + +The split between the long-context kernels (§5.2 decode, §5.5 prefill — +both already specified) and a short-context variant is referenced +abstractly in items §B-1 and §B-6 above ("short-context balance is a +separate study"). This subsection pins the two open numbers. + +1. **Short/long context threshold = 256 K tokens.** Below the threshold, + the §5.5 Ring KV prefill pays C-1 ring rotations whose IPCQ cost is + not amortised by enough per-rank compute; above it, the per-rank + KV sweep dominates and ring is the right choice (§9 line 803). + Likewise §5.2's per-rank `S_kv/C·P` shard is small enough at + short context that the 2-level reduce-to-root hop count + dominates the per-rank compute. The 256 K boundary matches the + point at which `S_kv/C·P` (`C=4, P=8 → /32`) crosses the + 8 K tokens-per-rank mark above which the per-rank tile sweep + (§3 / §B-3 §scratch_scope) becomes the limiting factor rather + than collective overhead. **Recommend:** treat 256 K as the + bench dispatch threshold; expose it as a launch knob for sweeps. + +2. **Short-context kernel design — `kv_per_cube ∈ {1, 2, 4, 8}`.** + At short context the §5.5 Ring KV motion is wasted work and the + §5.2 cube-SP shard is too thin to feed the engines. The + short-context variant therefore drops cube-SP entirely: each + CUBE owns `kv_per_cube` *whole* KV heads (no `S_kv` sharding + across CUBEs), and the existing PE-SP within a CUBE shards `S_kv` + across the `P` PEs for each owned head. + + For `h_kv = 8` the natural distributions are: + + | `kv_per_cube` | participating CUBEs | head→CUBE map | + |---|---|---| + | 1 | `C = 8` | head `i` → CUBE `i` | + | 2 | `C = 4` | heads `[2i, 2i+1]` → CUBE `i` | + | 4 | `C = 2` | heads `[4i..4i+3]` → CUBE `i` | + | 8 | `C = 1` | all heads → single CUBE | + + No inter-CUBE reduce within a head (each head fully owned by one + CUBE), so the kernel runs Level-2 (PE) chain reduce-to-root only; + Level-1 (inter-CUBE) collapses to a no-op. The output stays + distributed per CUBE (each CUBE writes its owned heads' `O` slice). + + **Recommend:** ship as a separate kernel file (`_gqa_decode_short.py` / + `_gqa_prefill_short.py`) and a separate bench dispatcher that picks + short vs long by `S_kv ⋚ 256K`. Keep `kv_per_cube` as a kernel arg so + the topology cost trade-off can be measured per workload. diff --git a/docs/adr-proposed/ADR-0061-prog-tl-broadcast.md b/docs/adr-proposed/ADR-0061-prog-tl-broadcast.md index a25cc04..54aa069 100644 --- a/docs/adr-proposed/ADR-0061-prog-tl-broadcast.md +++ b/docs/adr-proposed/ADR-0061-prog-tl-broadcast.md @@ -32,9 +32,8 @@ the group is the decode-time efficiency win (ADR-0060 §5.2). ### Why it does not work today -The existing mesh kernels reshape tensors with a metadata-only helper -`_view` (`src/kernbench/benches/_attention_mesh_kv.py:25-36`, -`_attention_mesh_mlo.py:25-36`): +Pre-ADR-0060 mesh kernels reshaped tensors with a metadata-only `_view` +helper that rewrote `shape` but kept the original `nbytes`: ```python def _view(handle, new_shape): @@ -62,8 +61,7 @@ A reshape conserves bytes; a **broadcast does not** (`[S_kv, 1, d]` → kernel tries to view a `h_kv`-headed K as if it had `h_q` heads, the stored array has `1/G` of the requested bytes and `read` raises. -The current test suite documents this as a deliberate limitation -(`tests/attention/test_milestone_gqa_llama70b.py:137-142`): +The pre-ADR-0060 baseline documented this as a deliberate limitation: > "v1 uses `h_q == h_kv == 1` to avoid … GQA broadcast view (which is > symbolic and does not survive MemoryStore's nbytes check under @@ -196,7 +194,7 @@ entire point of GQA (it `G×`'s KV-cache HBM footprint and bandwidth). 2. **Phase 2 data**: broadcasting a known array yields the numpy `broadcast_to(...).copy()` result; a downstream `tl.dot` consuming it passes the MemoryStore nbytes check (the exact case that fails today). -3. **GQA end-to-end**: re-run a reduced `milestone-gqa-llama70b` panel +3. **GQA end-to-end**: re-run a reduced `milestone-gqa-headline` panel with `h_q=8, h_kv=1` under `enable_data=True`; it must complete (today it raises `ValueError: Shape mismatch`). 4. **Broadcast-incompatible shape** (e.g. axis size 3 → 8) raises a diff --git a/docs/adr-proposed/ADR-0063-prog-tl-scratch-scope.md b/docs/adr-proposed/ADR-0063-prog-tl-scratch-scope.md index 0fc178f..d1d31d3 100644 --- a/docs/adr-proposed/ADR-0063-prog-tl-scratch-scope.md +++ b/docs/adr-proposed/ADR-0063-prog-tl-scratch-scope.md @@ -102,6 +102,68 @@ kernel keeps two arenas: This mirrors how real flash-attention SRAM budgeting works: a small persistent accumulator region + a recycled tile working set. +#### D3.1 The persistent-arena write mechanism: `tl.copy_to(dst, src)` + +The merge ops (`tl.maximum`, `tl.exp`, binary `*` / `+`) all call +`_make_compute_out(...)` which allocates from the bump cursor (D1). +Inside a `scratch_scope`, their result handles therefore live **inside** +the scope and vanish on `__exit__`. To realise D3's two-arena split, the +kernel needs a way to **write a scoped result's bytes back to a +persistent address**. + +The primitive that closes this gap: + +```python +def copy_to(self, dst: TensorHandle, src: TensorHandle) -> None: + """Copy ``src``'s bytes into ``dst``'s address (both TCM). + + Shapes and dtypes must match. ``dst`` is typically a handle + allocated outside any active ``scratch_scope`` (the persistent + arena); ``src`` is a scoped handle whose bytes must outlive + scope ``__exit__``. + """ +``` + +Symmetric to `tl.store` (the HBM-side byte copy), kept TCM-only here so +the running-state writeback doesn't pollute op_log with spurious DMA. + +**Mechanics:** +- New `CopyCmd(src, dst, nbytes, data_op=True)` command. +- op_log: `op_kind="math"`, `op_name="copy"` — runs on the vector engine. +- Latency: `pe_math._compute_ns(prod(shape))` — models on-chip register + writeback, not HBM transfer. +- Emit-time validation: `dst.shape == src.shape`, `dst.dtype == src.dtype`, + `dst.space == "tcm"`, `src.space == "tcm"`. Authoring errors surface in + Phase 1, not deep in Phase 2 data execution. + +**Call-site pattern:** + +```python +m, l, O = init_running(...) # persistent (outside scope) +for j in range(n_tiles): + with tl.scratch_scope(): + ... # per-tile work (recycled) + m_new = tl.maximum(m, mj) # scoped scratch + l_new = l * scale_old + l_step * scale_step + O_new = O * scale_old + O_step * scale_step + + tl.copy_to(m, m_new) # ← persist new running state + tl.copy_to(l, l_new) + tl.copy_to(O, O_new) + # exit: scoped m_new/l_new/O_new gone; their bytes live in persistent m/l/O +``` + +The copy happens **before** `__exit__`, so the read of `src` (scoped) is +valid; after exit only `dst` (persistent) is read, satisfying D2's +no-read-after-exit safety contract. + +**Why a dedicated primitive rather than `dst=` kwargs on every math op** +(considered, rejected): adding `dst=` to `tl.maximum`, `tl.exp`, +`_binary_math`, `_unary_math`, `_reduction` is ~25 LOC across 5 +op-families and breaks the uniform "call returns a fresh handle" +pattern. `tl.copy_to` is one primitive, one command, one executor +handler — minimal surface area for the same effect. + ### D4. Nesting Scopes nest (stack of save-points). Inner scope exit rewinds to the inner diff --git a/src/kernbench/benches/1H_milestone_output/gqa/gqa_comparison.png b/src/kernbench/benches/1H_milestone_output/gqa/gqa_comparison.png deleted file mode 100644 index 97aae13..0000000 Binary files a/src/kernbench/benches/1H_milestone_output/gqa/gqa_comparison.png and /dev/null differ diff --git a/src/kernbench/benches/1H_milestone_output/gqa/gqa_op_log_multi_user_decode.png b/src/kernbench/benches/1H_milestone_output/gqa/gqa_op_log_multi_user_decode.png deleted file mode 100644 index 3c2aff5..0000000 Binary files a/src/kernbench/benches/1H_milestone_output/gqa/gqa_op_log_multi_user_decode.png and /dev/null differ diff --git a/src/kernbench/benches/1H_milestone_output/gqa/gqa_op_log_multi_user_prefill.png b/src/kernbench/benches/1H_milestone_output/gqa/gqa_op_log_multi_user_prefill.png deleted file mode 100644 index 25cb8b3..0000000 Binary files a/src/kernbench/benches/1H_milestone_output/gqa/gqa_op_log_multi_user_prefill.png and /dev/null differ diff --git a/src/kernbench/benches/1H_milestone_output/gqa/gqa_op_log_single_user_decode.png b/src/kernbench/benches/1H_milestone_output/gqa/gqa_op_log_single_user_decode.png deleted file mode 100644 index 351c3eb..0000000 Binary files a/src/kernbench/benches/1H_milestone_output/gqa/gqa_op_log_single_user_decode.png and /dev/null differ diff --git a/src/kernbench/benches/1H_milestone_output/gqa/gqa_op_log_single_user_prefill.png b/src/kernbench/benches/1H_milestone_output/gqa/gqa_op_log_single_user_prefill.png deleted file mode 100644 index 2907472..0000000 Binary files a/src/kernbench/benches/1H_milestone_output/gqa/gqa_op_log_single_user_prefill.png and /dev/null differ diff --git a/src/kernbench/benches/1H_milestone_output/gqa/sweep.json b/src/kernbench/benches/1H_milestone_output/gqa/sweep.json deleted file mode 100644 index 090cf9b..0000000 --- a/src/kernbench/benches/1H_milestone_output/gqa/sweep.json +++ /dev/null @@ -1,65 +0,0 @@ -{ - "version": 1, - "validation_scale": true, - "panels": [ - "single_user_prefill", - "multi_user_prefill", - "single_user_decode", - "multi_user_decode" - ], - "config": { - "S_q_prefill": 16, - "S_kv_per_rank": 16, - "h_q": 1, - "h_kv": 1, - "d_head": 64, - "n_ranks_single_user": 8, - "n_ranks_multi_user": 4 - }, - "rows": [ - { - "panel": "single_user_prefill", - "n_ranks": 8, - "op_log_summary": { - "gemm_count": 128, - "ipcq_send_count": 112, - "ipcq_recv_count": 112, - "dma_read_count": 24, - "dma_write_count": 8 - } - }, - { - "panel": "multi_user_prefill", - "n_ranks": 4, - "op_log_summary": { - "gemm_count": 32, - "ipcq_send_count": 24, - "ipcq_recv_count": 24, - "dma_read_count": 12, - "dma_write_count": 4 - } - }, - { - "panel": "single_user_decode", - "n_ranks": 8, - "op_log_summary": { - "gemm_count": 16, - "ipcq_send_count": 168, - "ipcq_recv_count": 168, - "dma_read_count": 24, - "dma_write_count": 8 - } - }, - { - "panel": "multi_user_decode", - "n_ranks": 4, - "op_log_summary": { - "gemm_count": 8, - "ipcq_send_count": 36, - "ipcq_recv_count": 36, - "dma_read_count": 12, - "dma_write_count": 4 - } - } - ] -} \ No newline at end of file diff --git a/src/kernbench/benches/1H_milestone_output/gqa_headline/sweep.json b/src/kernbench/benches/1H_milestone_output/gqa_headline/sweep.json new file mode 100644 index 0000000..b10be34 --- /dev/null +++ b/src/kernbench/benches/1H_milestone_output/gqa_headline/sweep.json @@ -0,0 +1,69 @@ +{ + "version": 1, + "panels": [ + "single_user_prefill_gqa", + "multi_user_prefill_gqa", + "single_user_decode_gqa", + "multi_user_decode_gqa" + ], + "config": { + "T_q_prefill": 4, + "T_q_decode": 1, + "S_kv_prefill": 16, + "h_q_decode": 8, + "h_kv_decode": 1, + "d_head": 64 + }, + "rows": [ + { + "panel": "single_user_prefill_gqa", + "kind": "prefill", + "C": 1, + "S_kv": 16, + "op_log_summary": { + "gemm_count": 2, + "ipcq_copy_count": 0, + "dma_read_count": 3, + "dma_write_count": 1 + } + }, + { + "panel": "multi_user_prefill_gqa", + "kind": "prefill", + "C": 4, + "S_kv": 16, + "op_log_summary": { + "gemm_count": 32, + "ipcq_copy_count": 24, + "dma_read_count": 12, + "dma_write_count": 4 + } + }, + { + "panel": "single_user_decode_gqa", + "kind": "decode", + "C": 1, + "P": 8, + "S_kv": 64, + "op_log_summary": { + "gemm_count": 16, + "ipcq_copy_count": 21, + "dma_read_count": 24, + "dma_write_count": 1 + } + }, + { + "panel": "multi_user_decode_gqa", + "kind": "decode", + "C": 4, + "P": 8, + "S_kv": 128, + "op_log_summary": { + "gemm_count": 64, + "ipcq_copy_count": 93, + "dma_read_count": 96, + "dma_write_count": 1 + } + } + ] +} \ No newline at end of file diff --git a/src/kernbench/benches/_attention_mesh_kv.py b/src/kernbench/benches/_attention_mesh_kv.py deleted file mode 100644 index cc61dd0..0000000 --- a/src/kernbench/benches/_attention_mesh_kv.py +++ /dev/null @@ -1,193 +0,0 @@ -"""Mesh-native bidirectional Ring-K/V attention kernel — prefill (ADR-0059 Proposed). - -Each rank holds its own Q tile and 1/n_ranks of K, V (sequence-sharded). -Over ``n_ranks - 1`` bidirectional steps, K and V propagate both east and -west: chunk c_i originating at rank i reaches rank j at step ``|i - j|``. -Every rank receives every other rank's chunk **exactly once** and folds it -into a running ``(m, ℓ, o)`` via the online-softmax recurrence. After all -steps each rank holds the final attention output for its own Q tokens — -no cross-rank merge is required. - -Supersedes ADR-0055's closed-ring ``_attention_ring_kv.py``. Both modules -stay on disk during the transition; this one runs on the hardware's -actual open-mesh wiring (no closed-ring SFR install required). - -Imported by ``milestone_gqa_llama70b`` (after the bench's Phase 2 switches -its imports) and invoked through ``torch.launch(...)`` — not through -``dist.all_reduce(...)``. See ADR-0055 Context for why this kernel is not -backend-dispatched via ADR-0050's algorithm-module contract. -""" -from __future__ import annotations - -from kernbench.common.pe_commands import TensorHandle - - -def _view(handle: TensorHandle, new_shape: tuple[int, ...]) -> TensorHandle: - """Reshape — metadata only, no command emitted (cf. ``tl.trans``).""" - return TensorHandle( - id=handle.id, - addr=handle.addr, - shape=new_shape, - dtype=handle.dtype, - nbytes=handle.nbytes, - data=handle.data, - space=handle.space, - pinned=handle.pinned, - ) - - -def _partial_attention( - Q: TensorHandle, - K: TensorHandle, - V: TensorHandle, - S_q: int, - S_kv_per_rank: int, - h_q: int, - d_head: int, - tl, -) -> tuple[TensorHandle, TensorHandle, TensorHandle]: - """One pass of partial attention against (K, V). - - Emits 1 GEMM(Q·K^T) + softmax + max + sub + exp + sum + 1 GEMM(P·V). - Returns the running-statistics triplet ``(m, ℓ, O_partial)`` for the - online-softmax mlo merge. - """ - K_2d_T = _view(K, (h_q * d_head, S_kv_per_rank)) - V_2d = _view(V, (S_kv_per_rank, h_q * d_head)) - - scores = tl.dot(Q, K_2d_T) - m = tl.max(scores, axis=-1) - P = tl.softmax(scores, axis=-1) - scores_centered = scores - m - exp_scores = tl.exp(scores_centered) - ell = tl.sum(exp_scores, axis=-1) - O_partial = tl.dot(P, V_2d) - return m, ell, O_partial - - -def attention_mesh_kv_kernel( - q_ptr: int, - k_ptr: int, - v_ptr: int, - o_ptr: int, - S_q: int, - S_kv_per_rank: int, - h_q: int, - h_kv: int, - d_head: int, - n_ranks: int, - rank_axis: int = 0, - *, - tl, -) -> None: - """Mesh-native bidirectional Ring-K/V attention — see module docstring. - - ``rank_axis`` selects which program-id dimension carries the ring rank, - matching the GQA Llama-70B sharding study's TL/BL vs TR/BR distinction - (`llm_paper_review/notes/GQA_MHA_sharding/scripts/_gen_llama70b_1M_4cases.py`): - - 0 — single_user_* panels (TL/BL): rank == tl.program_id(axis=0) (PE - id in cube). KV is split @ PEs **intra-cube**; ring runs over - the 8 PEs of one cube (NOC-only). At Llama-70B headline scale - this kernel launches once per cube; 64 such cubes run in - parallel for one user. Each cube's PE-level ring is independent. - - 1 — multi_user_* panels (TR/BR): rank == tl.program_id(axis=1) - (cube id). KV is split @ cubes **inter-cube**; ring runs over - the cubes of one KV-group. The kernel gates ``pe_id != 0`` to - return early — a v1 simplification: at headline scale (B=8) the - study's "Batch on batch" pattern would have all 8 PEs each handle - one user's batch element instead of staying silent. Adding the - per-cube batch dimension is sub-cycle 4c headline work. - """ - # For multi_user (rank_axis=1) only PE 0 in each cube runs the ring. - if rank_axis != 0 and tl.program_id(axis=0) != 0: - return - rank = tl.program_id(axis=rank_axis) - has_E = rank < n_ranks - 1 - has_W = rank > 0 - - # Q stays put on this rank — loaded once, used in every partial attention. - Q = tl.load(q_ptr, shape=(S_q, h_q * d_head), dtype="f16") - - # Local K, V chunk. - K = tl.load(k_ptr, shape=(S_kv_per_rank, h_kv, d_head), dtype="f16") - V = tl.load(v_ptr, shape=(S_kv_per_rank, h_kv, d_head), dtype="f16") - - # Step 0 (local): partial attention against own K, V — initializes the - # running triplet (m, ℓ, o). - m, ell, o = _partial_attention( - Q, K, V, S_q, S_kv_per_rank, h_q, d_head, tl, - ) - - # Seed bidirectional waves with own chunk (step-1 send). - to_send_east_K: TensorHandle | None = K - to_send_east_V: TensorHandle | None = V - to_send_west_K: TensorHandle | None = K - to_send_west_V: TensorHandle | None = V - - # Bidirectional fan-out: n_ranks - 1 steps. By step k, the wave from - # rank i has reached rank (i ± k). After n_ranks - 1 steps, every rank - # has merged every other rank's chunk exactly once (ADR-0059 D3). - for step in range(1, n_ranks): - # Send the eastbound wave we currently hold (own at step 1; forwarded - # at later steps). ``None`` means we have no wave to forward this - # direction this step (edge rank, or the wave already passed by). - if has_E and to_send_east_K is not None: - tl.send(dir="E", src=to_send_east_K) - tl.send(dir="E", src=to_send_east_V) - if has_W and to_send_west_K is not None: - tl.send(dir="W", src=to_send_west_K) - tl.send(dir="W", src=to_send_west_V) - - # Receive eastbound wave from W (carries chunk c_{rank - step}). - K_from_W: TensorHandle | None = None - V_from_W: TensorHandle | None = None - if has_W and (rank - step) >= 0: - K_from_W = tl.recv( - dir="W", shape=(S_kv_per_rank, h_kv, d_head), dtype="f16", - ) - V_from_W = tl.recv( - dir="W", shape=(S_kv_per_rank, h_kv, d_head), dtype="f16", - ) - m_new, ell_new, o_new = _partial_attention( - Q, K_from_W, V_from_W, S_q, S_kv_per_rank, h_q, d_head, tl, - ) - m_combined = tl.maximum(m, m_new) - scale_old = tl.exp(m - m_combined) - scale_new = tl.exp(m_new - m_combined) - ell = ell * scale_old + ell_new * scale_new - o = o * scale_old + o_new * scale_new - m = m_combined - - # Receive westbound wave from E (carries chunk c_{rank + step}). - K_from_E: TensorHandle | None = None - V_from_E: TensorHandle | None = None - if has_E and (rank + step) < n_ranks: - K_from_E = tl.recv( - dir="E", shape=(S_kv_per_rank, h_kv, d_head), dtype="f16", - ) - V_from_E = tl.recv( - dir="E", shape=(S_kv_per_rank, h_kv, d_head), dtype="f16", - ) - m_new, ell_new, o_new = _partial_attention( - Q, K_from_E, V_from_E, S_q, S_kv_per_rank, h_q, d_head, tl, - ) - m_combined = tl.maximum(m, m_new) - scale_old = tl.exp(m - m_combined) - scale_new = tl.exp(m_new - m_combined) - ell = ell * scale_old + ell_new * scale_new - o = o * scale_old + o_new * scale_new - m = m_combined - - # Forward what we received for next step. ``None`` propagates: if no - # chunk arrived this step (out-of-bounds wave origin), there is - # nothing to forward next step in that direction. - to_send_east_K = K_from_W - to_send_east_V = V_from_W - to_send_west_K = K_from_E - to_send_west_V = V_from_E - - # Final normalize: O := o / ℓ. - O_final = o / ell - tl.store(o_ptr, O_final) diff --git a/src/kernbench/benches/_attention_mesh_mlo.py b/src/kernbench/benches/_attention_mesh_mlo.py deleted file mode 100644 index 1e1b2bb..0000000 --- a/src/kernbench/benches/_attention_mesh_mlo.py +++ /dev/null @@ -1,167 +0,0 @@ -"""Mesh-native bidirectional AllReduce-mlo attention — decode (ADR-0059 Proposed). - -Every rank holds the full Q (replicated, small at ``S_q=1``) and 1/n_ranks -of KV (sequence-sharded). Each rank computes its partial attention -against own KV in ONE shot, then runs a bidirectional fan-out of the -``(m, ℓ, o)`` triplet: the triplet originating at rank i reaches rank j at -step ``|i - j|``. Every rank merges every other rank's triplet exactly -once over ``n_ranks - 1`` steps, ending with the final answer replicated -on every rank. - -Supersedes ADR-0056's closed-ring ``_attention_allreduce_mlo.py``. Both -modules stay on disk during the transition; this one runs on the -hardware's actual open-mesh wiring (no closed-ring SFR install required). - -Imported by ``milestone_gqa_llama70b`` (after the bench's Phase 2 switches -its imports) and invoked through ``torch.launch(...)`` — not through -``dist.all_reduce(...)``. See ADR-0056 Context for why this kernel is not -backend-dispatched via ADR-0050's algorithm-module contract. -""" -from __future__ import annotations - -from kernbench.common.pe_commands import TensorHandle - - -def _view(handle: TensorHandle, new_shape: tuple[int, ...]) -> TensorHandle: - """Reshape — metadata only, no command emitted (cf. ``tl.trans``).""" - return TensorHandle( - id=handle.id, - addr=handle.addr, - shape=new_shape, - dtype=handle.dtype, - nbytes=handle.nbytes, - data=handle.data, - space=handle.space, - pinned=handle.pinned, - ) - - -def attention_mesh_mlo_kernel( - q_ptr: int, - k_ptr: int, - v_ptr: int, - o_ptr: int, - S_q: int, - S_kv_per_rank: int, - h_q: int, - h_kv: int, - d_head: int, - n_ranks: int, - rank_axis: int = 0, - *, - tl, -) -> None: - """Mesh-native bidirectional AllReduce-mlo — see module docstring. - - ``rank_axis`` selects which program-id dimension carries the ring rank, - matching the GQA Llama-70B sharding study's TL/BL vs TR/BR distinction - (`llm_paper_review/notes/GQA_MHA_sharding/scripts/_gen_llama70b_1M_4cases.py`): - - 0 — single_user_* panels (TL/BL): rank == tl.program_id(axis=0) (PE - id in cube). KV is split @ PEs **intra-cube**; ring runs over - the 8 PEs of one cube (NOC-only). At Llama-70B headline scale - this kernel launches once per cube; 64 such cubes run in - parallel for one user (1 Q-head per cube × 8 cubes per KV-group - × 8 KV-groups). The PE-level ring inside each cube is - independent of the others. - - 1 — multi_user_* panels (TR/BR): rank == tl.program_id(axis=1) - (cube id). KV is split @ cubes **inter-cube**; ring runs over - the cubes of one KV-group. The kernel gates ``pe_id != 0`` to - return early — a v1 simplification: at headline scale (B=8) the - study's "Batch on batch" pattern would have all 8 PEs each handle - one user's batch element instead of staying silent. Validation - shipped with B=1 to focus on the inter-cube ring's correctness; - adding the per-cube batch dimension is sub-cycle 4c headline work. - """ - # For multi_user (rank_axis=1) only PE 0 in each cube runs the ring. - if rank_axis != 0 and tl.program_id(axis=0) != 0: - return - rank = tl.program_id(axis=rank_axis) - has_E = rank < n_ranks - 1 - has_W = rank > 0 - - # Q is replicated on every rank — loaded once. - Q = tl.load(q_ptr, shape=(S_q, h_q * d_head), dtype="f16") - - # Local KV chunk. KV is sequence-sharded and stays put on this rank for - # the entire fan-out — distinguishing decode from prefill (ADR-0059 D3) - # where KV circulates. - K = tl.load(k_ptr, shape=(S_kv_per_rank, h_kv, d_head), dtype="f16") - V = tl.load(v_ptr, shape=(S_kv_per_rank, h_kv, d_head), dtype="f16") - - # ── One-shot local partial attention ────────────────────────── - K_2d_T = _view(K, (h_q * d_head, S_kv_per_rank)) - V_2d = _view(V, (S_kv_per_rank, h_q * d_head)) - scores = tl.dot(Q, K_2d_T) - m = tl.max(scores, axis=-1) - P = tl.softmax(scores, axis=-1) - scores_centered = scores - m - exp_scores = tl.exp(scores_centered) - ell = tl.sum(exp_scores, axis=-1) - o = tl.dot(P, V_2d) - - # Seed bidirectional waves with own triplet (step-1 send). - to_send_east_m: TensorHandle | None = m - to_send_east_ell: TensorHandle | None = ell - to_send_east_o: TensorHandle | None = o - to_send_west_m: TensorHandle | None = m - to_send_west_ell: TensorHandle | None = ell - to_send_west_o: TensorHandle | None = o - - # Bidirectional fan-out of (m, ℓ, o) triplets — n_ranks - 1 steps. - for step in range(1, n_ranks): - # Send eastbound triplet (own at step 1; forwarded at later steps). - if has_E and to_send_east_m is not None: - tl.send(dir="E", src=to_send_east_m) - tl.send(dir="E", src=to_send_east_ell) - tl.send(dir="E", src=to_send_east_o) - # Send westbound triplet. - if has_W and to_send_west_m is not None: - tl.send(dir="W", src=to_send_west_m) - tl.send(dir="W", src=to_send_west_ell) - tl.send(dir="W", src=to_send_west_o) - - # Receive eastbound triplet from W (originated at rank - step). - m_from_W: TensorHandle | None = None - ell_from_W: TensorHandle | None = None - o_from_W: TensorHandle | None = None - if has_W and (rank - step) >= 0: - m_from_W = tl.recv(dir="W", shape=m.shape, dtype="f16") - ell_from_W = tl.recv(dir="W", shape=ell.shape, dtype="f16") - o_from_W = tl.recv(dir="W", shape=o.shape, dtype="f16") - m_combined = tl.maximum(m, m_from_W) - scale_old = tl.exp(m - m_combined) - scale_new = tl.exp(m_from_W - m_combined) - ell = ell * scale_old + ell_from_W * scale_new - o = o * scale_old + o_from_W * scale_new - m = m_combined - - # Receive westbound triplet from E (originated at rank + step). - m_from_E: TensorHandle | None = None - ell_from_E: TensorHandle | None = None - o_from_E: TensorHandle | None = None - if has_E and (rank + step) < n_ranks: - m_from_E = tl.recv(dir="E", shape=m.shape, dtype="f16") - ell_from_E = tl.recv(dir="E", shape=ell.shape, dtype="f16") - o_from_E = tl.recv(dir="E", shape=o.shape, dtype="f16") - m_combined = tl.maximum(m, m_from_E) - scale_old = tl.exp(m - m_combined) - scale_new = tl.exp(m_from_E - m_combined) - ell = ell * scale_old + ell_from_E * scale_new - o = o * scale_old + o_from_E * scale_new - m = m_combined - - # Forward the original received triplet (not the merged running state) - # so neighbors get the original wave. ``None`` propagates if nothing - # arrived this step. - to_send_east_m = m_from_W - to_send_east_ell = ell_from_W - to_send_east_o = o_from_W - to_send_west_m = m_from_E - to_send_west_ell = ell_from_E - to_send_west_o = o_from_E - - # Final normalize: O := o / ℓ. - O_final = o / ell - tl.store(o_ptr, O_final) diff --git a/src/kernbench/benches/_attention_mesh_mlo_2d.py b/src/kernbench/benches/_attention_mesh_mlo_2d.py deleted file mode 100644 index d4e1bba..0000000 --- a/src/kernbench/benches/_attention_mesh_mlo_2d.py +++ /dev/null @@ -1,217 +0,0 @@ -"""Mesh-native 2D row-then-col AllReduce-mlo attention — decode (ADR-0059 extension). - -Each cube holds the full Q (replicated) and 1/(mesh_rows * mesh_cols) of -KV (sequence-sharded across the 2D cube sub-mesh). The kernel decomposes -the AllReduce-mlo into a two-stage reduction: - - Stage 1 — Row reduce (E/W edges, ``mesh_cols - 1`` steps) - Bidirectional ring within each row. After this stage every cube in - row ``r`` holds the partial ``(m, ℓ, o)`` over the ``mesh_cols`` KV - chunks in row ``r``. - - Stage 2 — Col reduce (N/S edges, ``mesh_rows - 1`` steps) - Bidirectional ring within each column. After this stage every cube - holds the partial over all ``mesh_rows × mesh_cols`` KV chunks — - the AllReduce result. - -The online-softmax mlo merge is associative, so row-then-col partitioning -of the reduction is mathematically equivalent to a 1D ring AllReduce-mlo -over all ``mesh_rows × mesh_cols`` cubes. The 2D form takes -``(mesh_cols - 1) + (mesh_rows - 1)`` steps instead of -``mesh_rows × mesh_cols - 1`` (e.g. 4 vs 7 at 2×4; 6 vs 15 at 4×4). - -Designed to run on hardware wired by -``configure_sfr_intercube_multisip``, which installs both E/W and N/S -intra-SIP cube-mesh edges (``sfr_config.py:135-143``). The 1D -``_attention_mesh_mlo.py`` remains for the single_user PE-ring case; -this 2D variant supersedes it for multi_user_decode where the per-KV-group -cube count crosses a row boundary in the 4×4 cube mesh. - -``mesh_rows = 1`` is supported as a degenerate row-only case so the -validation config (``_N_RANKS_MULTI_USER = 4`` → ``(1, 4)``) reduces to -the 1D ring's step count without behavioral change. -""" -from __future__ import annotations - -from kernbench.common.pe_commands import TensorHandle - - -def _view(handle: TensorHandle, new_shape: tuple[int, ...]) -> TensorHandle: - """Reshape — metadata only, no command emitted (cf. ``tl.trans``).""" - return TensorHandle( - id=handle.id, - addr=handle.addr, - shape=new_shape, - dtype=handle.dtype, - nbytes=handle.nbytes, - data=handle.data, - space=handle.space, - pinned=handle.pinned, - ) - - -def _bidir_allreduce_mlo( - m: TensorHandle, - ell: TensorHandle, - o: TensorHandle, - rank: int, - n_ranks: int, - dir_pos: str, - dir_neg: str, - *, - tl, -) -> tuple[TensorHandle, TensorHandle, TensorHandle]: - """One bidirectional AllReduce-mlo ring along ``(dir_pos, dir_neg)``. - - Mirrors the 1D ``_attention_mesh_mlo.py`` algorithm but parameterized - on direction labels so the 2D kernel can call it once with ``("E", "W")`` - for the row reduce and once with ``("S", "N")`` for the col reduce. - Forwards the received triplets in subsequent steps so chunk ``c_i`` - reaches rank ``j`` at step ``|i - j|``. - - Returns the running ``(m, ℓ, o)`` after ``n_ranks - 1`` steps. Degenerate - cases (``n_ranks <= 1``) are no-ops — the for-loop body simply does not - execute. - """ - has_pos = rank < n_ranks - 1 - has_neg = rank > 0 - - to_send_pos_m: TensorHandle | None = m - to_send_pos_ell: TensorHandle | None = ell - to_send_pos_o: TensorHandle | None = o - to_send_neg_m: TensorHandle | None = m - to_send_neg_ell: TensorHandle | None = ell - to_send_neg_o: TensorHandle | None = o - - for step in range(1, n_ranks): - if has_pos and to_send_pos_m is not None: - tl.send(dir=dir_pos, src=to_send_pos_m) - tl.send(dir=dir_pos, src=to_send_pos_ell) - tl.send(dir=dir_pos, src=to_send_pos_o) - if has_neg and to_send_neg_m is not None: - tl.send(dir=dir_neg, src=to_send_neg_m) - tl.send(dir=dir_neg, src=to_send_neg_ell) - tl.send(dir=dir_neg, src=to_send_neg_o) - - m_from_neg: TensorHandle | None = None - ell_from_neg: TensorHandle | None = None - o_from_neg: TensorHandle | None = None - if has_neg and (rank - step) >= 0: - m_from_neg = tl.recv(dir=dir_neg, shape=m.shape, dtype="f16") - ell_from_neg = tl.recv(dir=dir_neg, shape=ell.shape, dtype="f16") - o_from_neg = tl.recv(dir=dir_neg, shape=o.shape, dtype="f16") - m_combined = tl.maximum(m, m_from_neg) - scale_old = tl.exp(m - m_combined) - scale_new = tl.exp(m_from_neg - m_combined) - ell = ell * scale_old + ell_from_neg * scale_new - o = o * scale_old + o_from_neg * scale_new - m = m_combined - - m_from_pos: TensorHandle | None = None - ell_from_pos: TensorHandle | None = None - o_from_pos: TensorHandle | None = None - if has_pos and (rank + step) < n_ranks: - m_from_pos = tl.recv(dir=dir_pos, shape=m.shape, dtype="f16") - ell_from_pos = tl.recv(dir=dir_pos, shape=ell.shape, dtype="f16") - o_from_pos = tl.recv(dir=dir_pos, shape=o.shape, dtype="f16") - m_combined = tl.maximum(m, m_from_pos) - scale_old = tl.exp(m - m_combined) - scale_new = tl.exp(m_from_pos - m_combined) - ell = ell * scale_old + ell_from_pos * scale_new - o = o * scale_old + o_from_pos * scale_new - m = m_combined - - to_send_pos_m = m_from_neg - to_send_pos_ell = ell_from_neg - to_send_pos_o = o_from_neg - to_send_neg_m = m_from_pos - to_send_neg_ell = ell_from_pos - to_send_neg_o = o_from_pos - - return m, ell, o - - -def attention_mesh_mlo_2d_kernel( - q_ptr: int, - k_ptr: int, - v_ptr: int, - o_ptr: int, - S_q: int, - S_kv_per_rank: int, - h_q: int, - h_kv: int, - d_head: int, - mesh_rows: int, - mesh_cols: int, - rank_axis: int = 0, - cube_start: int = 0, - *, - tl, -) -> None: - """2D row-then-col AllReduce-mlo decode kernel — see module docstring. - - ``rank_axis`` selects which program-id dimension carries the cube - rank (matches the 1D kernel convention): - - 0 — single_user_* (TL/BL): rank == tl.program_id(axis=0) (PE id). - Not used at headline scale — single_user uses the 1D intra-cube - PE ring (``_attention_mesh_mlo``). Kept here so the signature - mirrors the 1D kernel. - - 1 — multi_user_* (TR/BR): rank == tl.program_id(axis=1) (cube id). - KV is split @ cubes inter-cube; the ring runs over the - ``mesh_rows × mesh_cols`` cubes of one KV-group. The kernel - gates ``pe_id != 0`` to return early — same v1 simplification - as ``_attention_mesh_mlo`` (validation B=1). - - ``cube_start`` matches the value passed to ``DPPolicy.cube_start`` for - the launch's tensor placement. kernbench's ``tl.program_id(axis=1)`` - returns the physical cube id (ADR-0022), so when the launch is - offset within the SIP (e.g. cube_start=8 placing the second 2×4 - KV-group on cubes 8..15), the kernel must subtract ``cube_start`` - to recover the launch-local rank for ring arithmetic. Default 0 - preserves the cube_start=0 launches unchanged. - """ - # For multi_user (rank_axis=1) only PE 0 in each cube runs the ring. - if rank_axis != 0 and tl.program_id(axis=0) != 0: - return - - rank = tl.program_id(axis=rank_axis) - if rank_axis != 0: - rank = rank - cube_start - my_row = rank // mesh_cols - my_col = rank % mesh_cols - - # Q is replicated on every cube — loaded once. - Q = tl.load(q_ptr, shape=(S_q, h_q * d_head), dtype="f16") - - # Local KV chunk (sequence-sharded across the 2D sub-mesh). - K = tl.load(k_ptr, shape=(S_kv_per_rank, h_kv, d_head), dtype="f16") - V = tl.load(v_ptr, shape=(S_kv_per_rank, h_kv, d_head), dtype="f16") - - # ── One-shot local partial attention ────────────────────────── - K_2d_T = _view(K, (h_q * d_head, S_kv_per_rank)) - V_2d = _view(V, (S_kv_per_rank, h_q * d_head)) - scores = tl.dot(Q, K_2d_T) - m = tl.max(scores, axis=-1) - P = tl.softmax(scores, axis=-1) - scores_centered = scores - m - exp_scores = tl.exp(scores_centered) - ell = tl.sum(exp_scores, axis=-1) - o = tl.dot(P, V_2d) - - # ── Stage 1: row AllReduce (E/W, mesh_cols - 1 steps) ───────── - m, ell, o = _bidir_allreduce_mlo( - m, ell, o, my_col, mesh_cols, "E", "W", tl=tl, - ) - - # ── Stage 2: col AllReduce (N/S, mesh_rows - 1 steps) ───────── - # ``dir_pos="S"`` matches the SFR convention: ``S`` goes to higher - # row (configure_sfr_intercube_multisip:140). - m, ell, o = _bidir_allreduce_mlo( - m, ell, o, my_row, mesh_rows, "S", "N", tl=tl, - ) - - # Final normalize: O := o / ℓ. - O_final = o / ell - tl.store(o_ptr, O_final) diff --git a/src/kernbench/benches/_gqa_decode.py b/src/kernbench/benches/_gqa_decode.py new file mode 100644 index 0000000..b83aa86 --- /dev/null +++ b/src/kernbench/benches/_gqa_decode.py @@ -0,0 +1,180 @@ +"""GQA fused-attention decode kernel — P1a + P2a + P2b (2-level SP). + +Lineage (DDD-0060 §7 phase plan): + P1a : real GQA via M-fold using ``tl.dot``; one-shot per rank. + P2a : intra-CUBE PE-level chain reduce-to-root (single-CUBE SP). + P2b : adds inter-CUBE chain reduce-to-root (multi-CUBE SP); + switches to the canonical full SFR install + ``configure_sfr_intercube_multisip`` with disjoint namespaces + (``intra_*`` for PE, ``N/S/E/W`` for CUBE, ``global_*`` for SIP). + P3b : tile S_kv sweep + ``tl.scratch_scope`` (deferred). + Later : P1b composite swap; P4 lazy load; P5 opt3 pipelining; P6 prefill. + +P2b SFR + topology assumptions: + - SFR install: ``configure_sfr_intercube_multisip`` is required when + P > 1 or C > 1 (provides ``intra_*`` and ``E/W/N/S`` namespaces). + - Intra-CUBE PE layout: logical 2×4 grid (no wrap): + Row 0: PE 0, 1, 2, 3 + Row 1: PE 4, 5, 6, 7 + - Inter-CUBE layout: CUBEs of one CUBE Group are laid out as a 1D row + (single row of C CUBEs, no wrap). Multi-row CUBE Group placement + (e.g. 2×4) is future work — head_of_group / cube_start dispatch is + a P7 concern (DDD-0060 §4.1). + +Reduce strategy — chain reduce-to-root at PE 0 of CUBE 0: + + Level-2 (intra-CUBE, row-then-col chain on 2×4 grid): + Row chain along ``intra_W``: rightmost-col PEs send leftward; + leftmost-col PE of each row holds its row's partial. + Col bridge along ``intra_N``: PE 4 (col-0, row-1) sends to PE 0 + (col-0, row-0). Only relevant when P > 4. + Result: PE 0 of each CUBE holds the CUBE's partial. + + Level-1 (inter-CUBE, only PE 0 of each CUBE participates): + Chain along ``W``: rightmost CUBE sends leftward; CUBE 0 of the + CUBE Group ends with the final answer. + + Final store: PE 0 of CUBE 0 normalises (O / ℓ) and writes. + +Chain step counts (per ADR-0060 §A.2 root-only output, §4 chain +deviation noted): for (C, P)=(2, 8), 7 intra-cube × 2 cubes + (C-1) +inter-cube = 14 + 1 = 15 chain steps; each step ships 3 handles +(m, ℓ, O) ⇒ 45 ``ipcq_copy`` total. + +Three deliberate deviations from ADR-0060, addressed in later phases: + 1. GEMMs use ``tl.dot``, not ``tl.composite`` (P1b). + 2. ``softmax_scale`` omitted — needs composite epilogue mechanism (P1b/P5). + 3. K loaded as ``[d, S_local]`` via byte-conserving reshape of the + deployed ``[S_local, h_kv·d]`` slice (ADR-0060 §3 / §B item 2, + reshape-not-transpose caveat; correct for zero / symmetric inputs). + +Chain-vs-tree deviation from DDD-0060 §7 P2 gate (``⌈log₂ P⌉``): + P2b uses linear chain reduce-to-root (P-1 + C-1 hops). True tree on + the 2×4 PE grid requires a different SFR install — separate ADR. + Architectural intent (root-only output replacing baseline's + bidirectional fan-out) is preserved. +""" +from __future__ import annotations + + +def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl): + """Online-softmax merge (ADR-0060 §4 / _attention_mesh_mlo baseline).""" + m_new = tl.maximum(m_local, m_other) + scale_old = tl.exp(m_local - m_new) + scale_new = tl.exp(m_other - m_new) + l_new = l_local * scale_old + l_other * scale_new + O_new = O_local * scale_old + O_other * scale_new + return m_new, l_new, O_new + + +def gqa_decode_kernel( + q_ptr: int, + k_ptr: int, + v_ptr: int, + o_ptr: int, + T_q: int, + S_kv: int, + h_q: int, + h_kv: int, + d_head: int, + C: int, + P: int, + *, + tl, +) -> None: + """GQA decode with M-fold + 2-level chain reduce-to-root. + + Tensor layout: + Q : (T_q, h_q · d_head) replicated on every rank; loaded as + (G·T_q, d_head) — byte-conserving and math-correct for T_q=1. + K : (S_kv, h_kv · d_head) sharded row_wise by (cube, pe); each rank + loads its (d_head, S_local) slice via byte-conserving reshape. + V : (S_kv, h_kv · d_head) sharded row_wise by (cube, pe); each rank + loads its (S_local, d_head) slice. + O : (T_q, h_q · d_head) — only PE 0 of CUBE 0 stores. + """ + G = h_q // h_kv + n_ranks = C * P + S_local = S_kv // n_ranks + pe_id = tl.program_id(axis=0) + cube_id = tl.program_id(axis=1) + + # ── Local one-shot partial attention (M-fold on the rank's slice) ── + Q = tl.load(q_ptr, shape=(G * T_q, d_head), dtype="f16") + K_T = tl.load(k_ptr, shape=(d_head, S_local), dtype="f16") + V = tl.load(v_ptr, shape=(S_local, d_head), dtype="f16") + scores = tl.dot(Q, K_T) + m_local = tl.max(scores, axis=-1) + centered = scores - m_local + exp_scores = tl.exp(centered) + l_local = tl.sum(exp_scores, axis=-1) + O_local = tl.dot(exp_scores, V) + + # ── Level-2: intra-CUBE row-then-col chain reduce-to-(PE 0) ── + PE_GRID_COLS = 4 + pe_col = pe_id % PE_GRID_COLS + pe_row = pe_id // PE_GRID_COLS + pe_cols_used = min(PE_GRID_COLS, P) + pe_rows_used = (P + PE_GRID_COLS - 1) // PE_GRID_COLS + + # Row chain (along intra_W within each row, gathering leftward). + # Each merge step's intermediates are wrapped in tl.scratch_scope and + # the new running (m, ℓ, O) is persisted back to the outside-scope + # (persistent) m_local/l_local/O_local via tl.copy_to (ADR-0063 §D3/D3.1). + if pe_cols_used > 1: + if pe_col < pe_cols_used - 1: # not rightmost: receive E + with tl.scratch_scope(): + m_other = tl.recv(dir="intra_E", shape=m_local.shape, dtype="f16") + l_other = tl.recv(dir="intra_E", shape=l_local.shape, dtype="f16") + O_other = tl.recv(dir="intra_E", shape=O_local.shape, dtype="f16") + m_new, l_new, O_new = _merge_running( + m_local, l_local, O_local, m_other, l_other, O_other, tl=tl, + ) + tl.copy_to(m_local, m_new) + tl.copy_to(l_local, l_new) + tl.copy_to(O_local, O_new) + if pe_col > 0: # not leftmost: send W + tl.send(dir="intra_W", src=m_local) + tl.send(dir="intra_W", src=l_local) + tl.send(dir="intra_W", src=O_local) + + # Col bridge (intra_N from row 1 col 0 → row 0 col 0). Only at col 0. + if pe_col == 0 and pe_rows_used > 1: + if pe_row < pe_rows_used - 1: # row 0 receives from S + with tl.scratch_scope(): + m_other = tl.recv(dir="intra_S", shape=m_local.shape, dtype="f16") + l_other = tl.recv(dir="intra_S", shape=l_local.shape, dtype="f16") + O_other = tl.recv(dir="intra_S", shape=O_local.shape, dtype="f16") + m_new, l_new, O_new = _merge_running( + m_local, l_local, O_local, m_other, l_other, O_other, tl=tl, + ) + tl.copy_to(m_local, m_new) + tl.copy_to(l_local, l_new) + tl.copy_to(O_local, O_new) + if pe_row > 0: # row >0 sends to N + tl.send(dir="intra_N", src=m_local) + tl.send(dir="intra_N", src=l_local) + tl.send(dir="intra_N", src=O_local) + + # ── Level-1: inter-CUBE chain reduce (only PE 0 of each CUBE) ── + if pe_id == 0 and C > 1: + if cube_id < C - 1: # not rightmost CUBE: recv E + with tl.scratch_scope(): + m_other = tl.recv(dir="E", shape=m_local.shape, dtype="f16") + l_other = tl.recv(dir="E", shape=l_local.shape, dtype="f16") + O_other = tl.recv(dir="E", shape=O_local.shape, dtype="f16") + m_new, l_new, O_new = _merge_running( + m_local, l_local, O_local, m_other, l_other, O_other, tl=tl, + ) + tl.copy_to(m_local, m_new) + tl.copy_to(l_local, l_new) + tl.copy_to(O_local, O_new) + if cube_id > 0: # non-root CUBE: send W + tl.send(dir="W", src=m_local) + tl.send(dir="W", src=l_local) + tl.send(dir="W", src=O_local) + + # ── Final normalise + store (only at PE 0 of CUBE 0) ── + if pe_id == 0 and cube_id == 0: + O_final = O_local / l_local + tl.store(o_ptr, O_final) diff --git a/src/kernbench/benches/_gqa_decode_short.py b/src/kernbench/benches/_gqa_decode_short.py new file mode 100644 index 0000000..0b4b8e3 --- /dev/null +++ b/src/kernbench/benches/_gqa_decode_short.py @@ -0,0 +1,145 @@ +"""GQA fused-attention SHORT-CONTEXT decode kernel (ADR-0060 §B.split.2). + +Short context (S_kv < 256K, per ADR-0060 §B.split.1): each CUBE owns +``kv_per_cube`` whole KV heads, no S_kv sharding across CUBEs, no +inter-CUBE reduce. PE-SP within each CUBE: the P PEs split into +``kv_per_cube`` groups, each group does PE-SP across (P/kv_per_cube) +PEs for ONE owned head. + +Layout (after design iteration during Phase D — see ADR-0060 §B.split.2): + - K, V: shape ``(h_kv·S_kv, d_head)`` head-stacked, with the bench + deploying ``dp = (cube=row_wise, pe=row_wise)`` so each PE's chunk + is exactly ``(S_local, d_head)`` contiguous at its own addressable + shard. The kernel just loads at its ``k_ptr`` / ``v_ptr`` — no + offset arithmetic needed. + - Q: replicated ``(T_q, h_q·d_head)``; the kernel reshapes + byte-conservingly to ``(h_q·T_q, d_head)`` and operates on the + full stack. Other heads' rows are computed too (semantic noise); + with zero/symmetric inputs the math is unchanged. A proper + per-head Q slice would require runtime support for partial reads + of stored tensors (deferred). + - O: replicated; each group root writes the full byte-conserving + ``(h_q·T_q, d_head)`` result. Multiple roots within a CUBE write + to disjoint PE-local addresses (no overwrite collision). + +Group layout on the 2×4 PE grid: + kv_per_cube=1, group=8 PEs (full 2×4): row chain + col bridge. + kv_per_cube=2, group=4 PEs (one row): row chain only. + kv_per_cube=4, group=2 PEs (adj cols): 1-step chain. + kv_per_cube=8, group=1 PE: no chain — direct write. + +Chain reduce within group via the existing ``intra_E/W/N/S`` SFR +namespace (configure_sfr_intercube_multisip). After chain reduce, the +group's root PE (pe_in_group == 0) writes its working state to HBM. + +Deviations from ADR-0060 (deliberate, documented): + 1. GEMMs use ``tl.dot``, not ``tl.composite``. + 2. ``softmax_scale`` omitted. + 3. K loaded as ``[d, S_local]`` via byte-conserving reshape + (reshape-not-transpose caveat — correct for zero / symmetric inputs). + 4. Q byte-conserving reshape: kernel computes attention for ALL Q + rows against the group's owned K head; only the rows for my head + are semantically meaningful. Correct for zero / symmetric inputs. +""" +from __future__ import annotations + + +def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl): + """Online-softmax merge — identical to long kernel.""" + m_new = tl.maximum(m_local, m_other) + scale_old = tl.exp(m_local - m_new) + scale_new = tl.exp(m_other - m_new) + l_new = l_local * scale_old + l_other * scale_new + O_new = O_local * scale_old + O_other * scale_new + return m_new, l_new, O_new + + +def gqa_decode_short_kernel( + q_ptr: int, + k_ptr: int, + v_ptr: int, + o_ptr: int, + T_q: int, + S_kv: int, + h_q: int, + h_kv: int, + d_head: int, + C: int, + P: int, + kv_per_cube: int, + *, + tl, +) -> None: + """Short-context GQA decode with PE-parallel heads + intra-group PE-SP.""" + group_size = P // kv_per_cube # PEs per head group + + pe_id = tl.program_id(axis=0) + pe_in_group = pe_id % group_size + + # PE-SP within group: shard S_kv across group_size PEs + S_local = S_kv // group_size + + # ── Loads (DP layout already places each PE at its own shard) ── + # Q replicated → byte-conserving reshape to (h_q·T_q, d_head). + Q = tl.load(q_ptr, + shape=(h_q * T_q, d_head), dtype="f16") + # K, V row_wise per (cube, pe) → each PE has (S_local, d_head) at k_ptr. + K_T = tl.load(k_ptr, + shape=(d_head, S_local), dtype="f16") + V = tl.load(v_ptr, + shape=(S_local, d_head), dtype="f16") + + # ── Local one-shot partial attention ── + scores = tl.dot(Q, K_T) + m_local = tl.max(scores, axis=-1) + centered = scores - m_local + exp_scores = tl.exp(centered) + l_local = tl.sum(exp_scores, axis=-1) + O_local = tl.dot(exp_scores, V) + + # ── Within-group chain reduce-to-root (Level-2 only) ── + group_cols = min(4, group_size) + group_rows = (group_size + group_cols - 1) // group_cols + pe_col_in_group = pe_in_group % group_cols + pe_row_in_group = pe_in_group // group_cols + + # Row chain along intra_W (within group's row). + if group_cols > 1: + if pe_col_in_group < group_cols - 1: # receive E (in-group) + with tl.scratch_scope(): + m_other = tl.recv(dir="intra_E", shape=m_local.shape, dtype="f16") + l_other = tl.recv(dir="intra_E", shape=l_local.shape, dtype="f16") + O_other = tl.recv(dir="intra_E", shape=O_local.shape, dtype="f16") + m_new, l_new, O_new = _merge_running( + m_local, l_local, O_local, m_other, l_other, O_other, tl=tl, + ) + tl.copy_to(m_local, m_new) + tl.copy_to(l_local, l_new) + tl.copy_to(O_local, O_new) + if pe_col_in_group > 0: # send W (in-group) + tl.send(dir="intra_W", src=m_local) + tl.send(dir="intra_W", src=l_local) + tl.send(dir="intra_W", src=O_local) + + # Col bridge along intra_N (only if group spans 2 grid rows). + if pe_col_in_group == 0 and group_rows > 1: + if pe_row_in_group < group_rows - 1: # receive S (in-group) + with tl.scratch_scope(): + m_other = tl.recv(dir="intra_S", shape=m_local.shape, dtype="f16") + l_other = tl.recv(dir="intra_S", shape=l_local.shape, dtype="f16") + O_other = tl.recv(dir="intra_S", shape=O_local.shape, dtype="f16") + m_new, l_new, O_new = _merge_running( + m_local, l_local, O_local, m_other, l_other, O_other, tl=tl, + ) + tl.copy_to(m_local, m_new) + tl.copy_to(l_local, l_new) + tl.copy_to(O_local, O_new) + if pe_row_in_group > 0: # send N (in-group) + tl.send(dir="intra_N", src=m_local) + tl.send(dir="intra_N", src=l_local) + tl.send(dir="intra_N", src=O_local) + + # ── Group root writes its owned head's output ── + if pe_in_group == 0: + O_final = O_local / l_local + tl.store(o_ptr, O_final) diff --git a/src/kernbench/benches/_gqa_prefill.py b/src/kernbench/benches/_gqa_prefill.py new file mode 100644 index 0000000..4b8c0aa --- /dev/null +++ b/src/kernbench/benches/_gqa_prefill.py @@ -0,0 +1,111 @@ +"""GQA fused-attention prefill kernel — P6a + P6b (head-parallel + Ring KV). + +Lineage (DDD-0060 §7 phase plan): + P6a : head-parallel structure (one Q head per CUBE), C=1 baseline. + P6b : add Ring KV rotation across C CUBEs (ADR-0060 §5.5). Each + CUBE rotates its KV block to its W neighbour and receives + from E; over C-1 steps every CUBE sees every block. Online- + softmax merge folds each step into running (m, ℓ, O). No + reduce — each CUBE writes its own head's output. + Requires ``configure_sfr_intercube_ring(ring_size=C)`` SFR. + P6c (later): tile T_q across the P PEs inside a CUBE for intra-CUBE + parallelism (ADR-0060 §B item 3). + +Deviations from ADR-0060 §5.5 — deferred to later phases: + 1. GEMMs use ``tl.dot`` not ``tl.composite`` (parallel to P1a; lifts + when P1b decides the composite-output-handle question). + 2. ``softmax_scale`` omitted — same composite-epilogue deferral. + 3. K loaded as ``[d, S_local]`` via byte-conserving reshape of + ``[S_local, d]`` (ADR-0060 §3 / §B item 2 reshape-not-transpose + caveat; correct for zero / symmetric inputs). + 4. No causal masking / step-skip — future P6c. + 5. Blocking ``tl.recv`` (not ``recv_async``) — overlap via lazy + ``tl.load`` lands in P4. +""" +from __future__ import annotations + + +def gqa_prefill_kernel( + q_ptr: int, + k_ptr: int, + v_ptr: int, + o_ptr: int, + T_q: int, + S_kv: int, + d_head: int, + C: int, + *, + tl, +) -> None: + """Head-parallel prefill attention with Ring KV (C>1) — ADR-0060 §5.5. + + Tensor layout consumed by this kernel: + Q : (T_q, d_head) one head per CUBE; replicated. + K : (S_kv, d_head) sharded cube_row_wise → each CUBE owns + (S_kv/C, d_head); kernel loads as (d_head, S_local) via + byte-conserving reshape (reshape-not-transpose caveat). + V : (S_kv, d_head) sharded cube_row_wise → each CUBE owns + (S_local, d_head). + O : (T_q * C, d_head) sharded cube_row_wise → each CUBE + writes its own (T_q, d_head) slice. NO reduce. + + Algorithm: each CUBE computes a local partial against its current + KV block, then over C-1 ring steps the K and V blocks rotate W + while online-softmax merges each step into running (m, ℓ, O). + """ + pe_id = tl.program_id(axis=0) + # Head-parallel: only PE 0 of each CUBE participates. P6c (future) + # will tile T_q across the 8 PEs for intra-CUBE parallelism. + if pe_id != 0: + return + + S_local = S_kv // C + Q = tl.load(q_ptr, shape=(T_q, d_head), dtype="f16") + Kc = tl.load(k_ptr, shape=(d_head, S_local), dtype="f16") + Vc = tl.load(v_ptr, shape=(S_local, d_head), dtype="f16") + + # ── Step 0: initial partial against own KV block — establishes the + # persistent (m, ℓ, O) arena. Intermediates (scores, exp_scores) stay + # allocated; ring steps below recycle per-step intermediates inside + # tl.scratch_scope to keep peak scratch O(one step) (ADR-0063 §D3). + scores = tl.dot(Q, Kc) + m = tl.max(scores, axis=-1) + exp_scores = tl.exp(scores - m) + l = tl.sum(exp_scores, axis=-1) + O = tl.dot(exp_scores, Vc) + + # ── Steps 1..C-1: Ring KV rotation + online-softmax merge ── + # Per-step intermediates wrapped in tl.scratch_scope; the merged + # running (m, ℓ, O) is persisted to the outside-scope handles via + # tl.copy_to (ADR-0063 §D3.1) so its bytes survive __exit__. + for _ in range(1, C): + tl.send(dir="W", src=Kc) + tl.send(dir="W", src=Vc) + Kc = tl.recv(dir="E", shape=(d_head, S_local), dtype="f16") + Vc = tl.recv(dir="E", shape=(S_local, d_head), dtype="f16") + + with tl.scratch_scope(): + # Partial on rotated KV + scores = tl.dot(Q, Kc) + m_step = tl.max(scores, axis=-1) + exp_scores = tl.exp(scores - m_step) + l_step = tl.sum(exp_scores, axis=-1) + O_step = tl.dot(exp_scores, Vc) + + # Online-softmax merge into new running (m, ℓ, O) + m_new = tl.maximum(m, m_step) + scale_old = tl.exp(m - m_new) + scale_step = tl.exp(m_step - m_new) + l_new = l * scale_old + l_step * scale_step + O_new = O * scale_old + O_step * scale_step + + # Persist new running state back to the outside-scope arena. + tl.copy_to(m, m_new) + tl.copy_to(l, l_new) + tl.copy_to(O, O_new) + # __exit__: scoped intermediates gone; persistent m, l, O carry + # the new running state into the next ring iteration. + + # Final normalise + store — each CUBE writes its own head's rows. + O_final = O / l + tl.store(o_ptr, O_final) diff --git a/src/kernbench/benches/_gqa_prefill_short.py b/src/kernbench/benches/_gqa_prefill_short.py new file mode 100644 index 0000000..7930295 --- /dev/null +++ b/src/kernbench/benches/_gqa_prefill_short.py @@ -0,0 +1,114 @@ +"""GQA fused-attention SHORT-CONTEXT prefill kernel (ADR-0060 §B.split.2). + +Prefill analogue of ``_gqa_decode_short.py``. Same layout decisions: + - K, V head-stacked (h_kv·S_kv, d_head) with row_wise DP so each PE + has a contiguous (S_local, d_head) shard at its own address. + - Q replicated; kernel uses byte-conserving reshape. + - O replicated; group root writes the full byte-conserving result. + +No Ring KV (each owned head fully resident at its CUBE). Test +``test_short_prefill_no_ring_KV_traffic`` asserts no inter-CUBE IPCQ. + +The only structural difference from short decode: + - ``T_q`` may be > 1 (prefill processes multiple query tokens). + - Q is shaped (T_q, h_kv·d_head) — one Q head per KV head (no GQA + M-fold here; head-parallel prefill in ADR-0060 §5.5 is 1:1). +""" +from __future__ import annotations + + +def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl): + """Online-softmax merge — identical to short decode kernel.""" + m_new = tl.maximum(m_local, m_other) + scale_old = tl.exp(m_local - m_new) + scale_new = tl.exp(m_other - m_new) + l_new = l_local * scale_old + l_other * scale_new + O_new = O_local * scale_old + O_other * scale_new + return m_new, l_new, O_new + + +def gqa_prefill_short_kernel( + q_ptr: int, + k_ptr: int, + v_ptr: int, + o_ptr: int, + T_q: int, + S_kv: int, + h_kv: int, + d_head: int, + C: int, + P: int, + kv_per_cube: int, + *, + tl, +) -> None: + """Short-context prefill with PE-parallel heads + intra-group PE-SP. + + NO Ring KV (each CUBE owns its KV heads fully). + """ + group_size = P // kv_per_cube + + pe_id = tl.program_id(axis=0) + pe_in_group = pe_id % group_size + + S_local = S_kv // group_size + + # Q: replicated (T_q, h_kv·d_head) → byte-conserving reshape to + # (h_kv·T_q, d_head) — one Q head per KV head, stacked. + Q = tl.load(q_ptr, + shape=(h_kv * T_q, d_head), dtype="f16") + K_T = tl.load(k_ptr, + shape=(d_head, S_local), dtype="f16") + V = tl.load(v_ptr, + shape=(S_local, d_head), dtype="f16") + + scores = tl.dot(Q, K_T) + m_local = tl.max(scores, axis=-1) + centered = scores - m_local + exp_scores = tl.exp(centered) + l_local = tl.sum(exp_scores, axis=-1) + O_local = tl.dot(exp_scores, V) + + # Within-group chain reduce-to-root (same machinery as short decode). + group_cols = min(4, group_size) + group_rows = (group_size + group_cols - 1) // group_cols + pe_col_in_group = pe_in_group % group_cols + pe_row_in_group = pe_in_group // group_cols + + if group_cols > 1: + if pe_col_in_group < group_cols - 1: + with tl.scratch_scope(): + m_other = tl.recv(dir="intra_E", shape=m_local.shape, dtype="f16") + l_other = tl.recv(dir="intra_E", shape=l_local.shape, dtype="f16") + O_other = tl.recv(dir="intra_E", shape=O_local.shape, dtype="f16") + m_new, l_new, O_new = _merge_running( + m_local, l_local, O_local, m_other, l_other, O_other, tl=tl, + ) + tl.copy_to(m_local, m_new) + tl.copy_to(l_local, l_new) + tl.copy_to(O_local, O_new) + if pe_col_in_group > 0: + tl.send(dir="intra_W", src=m_local) + tl.send(dir="intra_W", src=l_local) + tl.send(dir="intra_W", src=O_local) + + if pe_col_in_group == 0 and group_rows > 1: + if pe_row_in_group < group_rows - 1: + with tl.scratch_scope(): + m_other = tl.recv(dir="intra_S", shape=m_local.shape, dtype="f16") + l_other = tl.recv(dir="intra_S", shape=l_local.shape, dtype="f16") + O_other = tl.recv(dir="intra_S", shape=O_local.shape, dtype="f16") + m_new, l_new, O_new = _merge_running( + m_local, l_local, O_local, m_other, l_other, O_other, tl=tl, + ) + tl.copy_to(m_local, m_new) + tl.copy_to(l_local, l_new) + tl.copy_to(O_local, O_new) + if pe_row_in_group > 0: + tl.send(dir="intra_N", src=m_local) + tl.send(dir="intra_N", src=l_local) + tl.send(dir="intra_N", src=O_local) + + if pe_in_group == 0: + O_final = O_local / l_local + tl.store(o_ptr, O_final) diff --git a/src/kernbench/benches/milestone_gqa_headline.py b/src/kernbench/benches/milestone_gqa_headline.py new file mode 100644 index 0000000..eb3946a --- /dev/null +++ b/src/kernbench/benches/milestone_gqa_headline.py @@ -0,0 +1,236 @@ +"""milestone-gqa-headline bench: real GQA + 2-level SP + Ring KV. + +Wires the new ``_gqa_decode`` and ``_gqa_prefill`` kernels through 4 +panels with real GQA (h_q = G·h_kv, G > 1 on the decode side), writing +per-panel ``op_log_summary`` into ``sweep.json``. Independent from the +existing ``milestone-gqa-llama70b`` validation-scale bench (which stays +on the legacy baseline kernels). + +Restrictions (P7 first cut): + - C ≤ 4 (single-row inter-CUBE ring SFR; multi-row deferred) + - Single SIP (default ``topology.yaml`` 4×4 cube mesh; 4-SIP + headline deferred) + - No figure renderers (defer to a separate cycle) + +Panels: + single_user_prefill_gqa : prefill C=1, T_q=4, S_kv=16 + multi_user_prefill_gqa : prefill C=4 Ring KV, T_q=4, S_kv=16 + single_user_decode_gqa : decode C=1, P=8, h_q=8, h_kv=1, S_kv=64 + (M-fold + intra-cube row-then-col chain) + multi_user_decode_gqa : decode C=4, P=8, h_q=8, h_kv=1, S_kv=128 + (M-fold + 2-level chain reduce-to-root) + +Gated by ``GQA_HEADLINE_RUN=1``. +""" +from __future__ import annotations + +import json +import os +from pathlib import Path + +from kernbench.benches._gqa_decode import gqa_decode_kernel +from kernbench.benches._gqa_prefill import gqa_prefill_kernel +from kernbench.benches.registry import bench +from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config +from kernbench.ccl.sfr_config import ( + configure_sfr_intercube_multisip, + configure_sfr_intercube_ring, +) +from kernbench.policy.placement.dp import DPPolicy + +_OUTPUT_DIR = Path(__file__).resolve().parent / "1H_milestone_output" / "gqa_headline" +_SWEEP_JSON = _OUTPUT_DIR / "sweep.json" + +# ── Panel configs ──────────────────────────────────────────────────── + +_DTYPE = "f16" +_D_HEAD = 64 +_T_Q_PREFILL = 4 +_T_Q_DECODE = 1 +_S_KV_PREFILL = 16 +_H_Q_DECODE = 8 # real GQA: G = H_Q_DECODE / H_KV_DECODE = 8 +_H_KV_DECODE = 1 + +_PANELS = ( + "single_user_prefill_gqa", + "multi_user_prefill_gqa", + "single_user_decode_gqa", + "multi_user_decode_gqa", +) + +# Each entry: (kind, panel-specific params) +_PANEL_DISPATCH: dict[str, tuple[str, dict]] = { + "single_user_prefill_gqa": ("prefill", {"C": 1, "S_kv": _S_KV_PREFILL}), + "multi_user_prefill_gqa": ("prefill", {"C": 4, "S_kv": _S_KV_PREFILL}), + "single_user_decode_gqa": ("decode", {"C": 1, "P": 8, "S_kv": 64}), + "multi_user_decode_gqa": ("decode", {"C": 4, "P": 8, "S_kv": 128}), +} + + +def _ccl_cfg(): + return resolve_algorithm_config( + load_ccl_config(), name="lrab_hierarchical_allreduce", + ) + + +# ── Per-kind launch helpers ────────────────────────────────────────── + + +def _run_prefill_panel(ctx, *, panel: str, C: int, S_kv: int) -> None: + if C > 1: + configure_sfr_intercube_ring( + ctx.engine, ctx.spec, _ccl_cfg(), ring_size=C, + ) + dp_q = DPPolicy(cube="replicate", pe="replicate", + num_cubes=C, num_pes=1) + dp_kv = DPPolicy(cube="row_wise" if C > 1 else "replicate", + pe="replicate", num_cubes=C, num_pes=1) + dp_o = DPPolicy(cube="row_wise" if C > 1 else "replicate", + pe="replicate", num_cubes=C, num_pes=1) + q = ctx.zeros((_T_Q_PREFILL, _D_HEAD), + dtype=_DTYPE, dp=dp_q, name=f"{panel}_q") + k = ctx.zeros((S_kv, _D_HEAD), + dtype=_DTYPE, dp=dp_kv, name=f"{panel}_k") + v = ctx.zeros((S_kv, _D_HEAD), + dtype=_DTYPE, dp=dp_kv, name=f"{panel}_v") + o = ctx.empty((_T_Q_PREFILL * C, _D_HEAD), + dtype=_DTYPE, dp=dp_o, name=f"{panel}_o") + ctx.launch( + panel, gqa_prefill_kernel, + q, k, v, o, + _T_Q_PREFILL, S_kv, _D_HEAD, C, + _auto_dim_remap=False, + ) + + +def _run_decode_panel(ctx, *, panel: str, C: int, P: int, S_kv: int) -> None: + configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg()) + dp_full = DPPolicy(cube="replicate", pe="replicate", + num_cubes=C, num_pes=P) + dp_kv = DPPolicy(cube="row_wise" if C > 1 else "replicate", + pe="row_wise", num_cubes=C, num_pes=P) + q = ctx.zeros((_T_Q_DECODE, _H_Q_DECODE * _D_HEAD), + dtype=_DTYPE, dp=dp_full, name=f"{panel}_q") + k = ctx.zeros((S_kv, _H_KV_DECODE * _D_HEAD), + dtype=_DTYPE, dp=dp_kv, name=f"{panel}_k") + v = ctx.zeros((S_kv, _H_KV_DECODE * _D_HEAD), + dtype=_DTYPE, dp=dp_kv, name=f"{panel}_v") + o = ctx.empty((_T_Q_DECODE, _H_Q_DECODE * _D_HEAD), + dtype=_DTYPE, dp=dp_full, name=f"{panel}_o") + ctx.launch( + panel, gqa_decode_kernel, + q, k, v, o, + _T_Q_DECODE, S_kv, _H_Q_DECODE, _H_KV_DECODE, _D_HEAD, C, P, + _auto_dim_remap=False, + ) + + +def _make_bench_fn(panel: str): + kind, params = _PANEL_DISPATCH[panel] + + def _bench_fn(ctx): + if kind == "prefill": + _run_prefill_panel(ctx, panel=panel, **params) + else: + _run_decode_panel(ctx, panel=panel, **params) + + return _bench_fn + + +# ── Op-log summary ────────────────────────────────────────────────── + + +def _summarize_op_log(op_log) -> dict[str, int]: + """Per-panel op_log counts (gemm, ipcq_copy, dma_read, dma_write).""" + gemm_count = 0 + ipcq_copy_count = 0 + dma_read_count = 0 + dma_write_count = 0 + for r in op_log: + if r.op_kind == "gemm": + gemm_count += 1 + elif r.op_name == "dma_read": + dma_read_count += 1 + elif r.op_name == "dma_write": + dma_write_count += 1 + elif r.op_name == "ipcq_copy": + ipcq_copy_count += 1 + return { + "gemm_count": gemm_count, + "ipcq_copy_count": ipcq_copy_count, + "dma_read_count": dma_read_count, + "dma_write_count": dma_write_count, + } + + +def _run_panel(panel: str, topology: str) -> dict: + """Run one panel in a fresh engine; return its row dict.""" + from kernbench.runtime_api.bench_runner import run_bench + from kernbench.runtime_api.types import resolve_device + from kernbench.sim_engine.engine import GraphEngine + from kernbench.topology.builder import resolve_topology + + topo = resolve_topology(topology) + result = run_bench( + topology=topo, bench_fn=_make_bench_fn(panel), + device=resolve_device(None), + engine_factory=lambda t, d: GraphEngine( + getattr(t, "topology_obj", t), enable_data=True, + ), + ) + if not result.completion.ok: + raise RuntimeError( + f"milestone-gqa-headline panel {panel!r} failed: " + f"{result.completion}" + ) + kind, params = _PANEL_DISPATCH[panel] + return { + "panel": panel, + "kind": kind, + **params, + "op_log_summary": _summarize_op_log(result.engine.op_log), + } + + +# ── Bench entry ────────────────────────────────────────────────────── + + +@bench( + name="milestone-gqa-headline", + description="Headline GQA milestone — real GQA h_q=8/h_kv=1 + 2-level SP (decode) + Ring KV (prefill).", +) +def run(torch) -> None: + """Drive 4 headline panels through the new GQA kernels; write sweep.json. + + Gated by GQA_HEADLINE_RUN=1. + """ + _OUTPUT_DIR.mkdir(parents=True, exist_ok=True) + if not os.environ.get("GQA_HEADLINE_RUN"): + raise RuntimeError( + "milestone-gqa-headline needs GQA_HEADLINE_RUN=1." + ) + + topology = os.environ.get("GQA_HEADLINE_TOPOLOGY", "topology.yaml") + rows = [_run_panel(panel, topology) for panel in _PANELS] + sweep = { + "version": 1, + "panels": list(_PANELS), + "config": { + "T_q_prefill": _T_Q_PREFILL, + "T_q_decode": _T_Q_DECODE, + "S_kv_prefill": _S_KV_PREFILL, + "h_q_decode": _H_Q_DECODE, + "h_kv_decode": _H_KV_DECODE, + "d_head": _D_HEAD, + }, + "rows": rows, + } + _SWEEP_JSON.write_text(json.dumps(sweep, indent=2)) + print(f" milestone-gqa-headline: {len(rows)} rows -> {_SWEEP_JSON}") + + # Sentinel tensor (ADR-0045 D4 / ADR-0054 D2 carve-out). + torch.zeros( + (1, 1), dtype="f16", + dp=DPPolicy(cube="row_wise", pe="replicate", num_cubes=1, num_pes=1), + name="milestone_gqa_headline_sentinel", + ) diff --git a/src/kernbench/benches/milestone_gqa_llama70b.py b/src/kernbench/benches/milestone_gqa_llama70b.py deleted file mode 100644 index f7edcc4..0000000 --- a/src/kernbench/benches/milestone_gqa_llama70b.py +++ /dev/null @@ -1,473 +0,0 @@ -"""milestone-gqa-llama70b bench: GQA Llama-70B 4-panel sweep (ADR-0057 v1). - -Self-contained eval bench (ADR-0054). Drives the four panels of the GQA -Llama-70B sharding study through ``run_bench`` with ``enable_data=True``, -harvests op_log summaries, and writes JSON into -``benches/1H_milestone_output/gqa/sweep.json``. - -v1 (sub-cycle 4a + 4c.0) covers all four panels at validation scale: - - Panel name in JSON / test Study label SFR install used - ───────────────────────────────────────────────────────────────────── - single_user_prefill TL configure_sfr_intracube_pe_ring - multi_user_prefill TR configure_sfr_intercube_multisip - single_user_decode BL configure_sfr_intracube_pe_ring - multi_user_decode BR configure_sfr_intercube_multisip - -Per the GQA sharding study (`llm_paper_review/notes/GQA_MHA_sharding/scripts -/_gen_llama70b_1M_4cases.py`): - - Single User (B=1) panels — TL prefill, BL decode: - "n_cubes: 8 (1 KV-group)", KV split @ PEs intra-cube. Each cube does - its own 8-PE ring with no cube-to-cube attention traffic. At Llama-70B - headline scale this is 64 cubes (8 KV-groups × 8 cubes/group), each - independently running the kernel for one Q-head; 1 user spans all 64. - - Multi User (B=8) panels — TR prefill, BR decode: - "8 cubes / KV-group", KV split @ cubes inter-cube. The 8 cubes of a - KV-group form a ring; "Inside each cube: 8 PEs each handle 1 different - user → Batch on batch, batch = 8/cube." At headline scale 8 KV-groups - run in parallel = 64 cubes serving 8 users. - -Kernels use the mesh-native variants (ADR-0059), invoked with the -``rank_axis`` kwarg (0 for single_user PE-level rings, 1 for multi_user -cube-level rings). The v1 multi_user kernel gates ``pe_id != 0`` to return, -which simplifies B=8 → B=1 — that's a validation simplification; the -per-cube "Batch on batch" parallelism is sub-cycle 4c headline work. - -Validation-scale config (ADR-0057 D4) — kept small so the simulator's -1 MB per-PE TCM scratch budget is not exhausted across n_ranks ring steps: -``S_q_prefill = S_kv_per_rank = 16``, ``h_q = h_kv = 1``, ``d_head = 64``, -``n_ranks_single_user = 8`` (8 PEs of one cube), ``n_ranks_multi_user = 4`` -(half a KV-group, vs the study's 8). Headline-scale dims (``S_q = 1M``, -``S_kv = 1M``, ``h_q = 8 / h_kv = 1`` GQA, ``d_head = 128``, ``B = 8``) are -also deferred. -""" -from __future__ import annotations - -import json -import os -from pathlib import Path -from typing import Any - -from kernbench.benches._attention_mesh_kv import attention_mesh_kv_kernel -from kernbench.benches._attention_mesh_mlo import attention_mesh_mlo_kernel -from kernbench.benches._attention_mesh_mlo_2d import attention_mesh_mlo_2d_kernel -from kernbench.benches.registry import bench -from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config -from kernbench.ccl.sfr_config import ( - configure_sfr_intercube_multisip, - configure_sfr_intracube_pe_ring, -) -from kernbench.policy.placement.dp import DPPolicy - -_OUTPUT_DIR = Path(__file__).resolve().parent / "1H_milestone_output" / "gqa" -_SWEEP_JSON = _OUTPUT_DIR / "sweep.json" - -# ── Validation-scale config (ADR-0057 D4) ───────────────────────────── - -_S_Q_PREFILL = 16 -_S_Q_DECODE = 1 -_S_KV_PER_RANK = 16 -_H_Q = 1 -_H_KV = 1 -_D_HEAD = 64 -_N_RANKS_SINGLE_USER = 8 -_N_RANKS_MULTI_USER = 4 -_DTYPE = "f16" - -_PANELS_V1 = ( - "single_user_prefill", - "multi_user_prefill", - "single_user_decode", - "multi_user_decode", -) - -# Panel → (kernel, SFR install, S_q, n_ranks, rank_axis, mesh_shape) -# ``mesh_shape`` is ``None`` for 1D-ring kernels and ``(rows, cols)`` for the -# 2D row-then-col AllReduce-mlo kernel (multi_user_decode); when set, the -# launch passes ``(mesh_rows, mesh_cols)`` instead of ``n_ranks``. -_PANEL_DISPATCH: dict[ - str, tuple[Any, Any, int, int, int, tuple[int, int] | None] -] = { - "single_user_prefill": ( - attention_mesh_kv_kernel, configure_sfr_intracube_pe_ring, - _S_Q_PREFILL, _N_RANKS_SINGLE_USER, 0, None, - ), - "multi_user_prefill": ( - attention_mesh_kv_kernel, configure_sfr_intercube_multisip, - _S_Q_PREFILL, _N_RANKS_MULTI_USER, 1, None, - ), - "single_user_decode": ( - attention_mesh_mlo_kernel, configure_sfr_intracube_pe_ring, - _S_Q_DECODE, _N_RANKS_SINGLE_USER, 0, None, - ), - # multi_user_decode uses the C2 2D AllReduce-mlo kernel. (1, 4) - # degenerates to a row-only AllReduce equivalent to the prior 1D ring - # at n_ranks=4 — no op_log_summary regression. Headline 8-cube - # KV-groups land at (2, 4). - "multi_user_decode": ( - attention_mesh_mlo_2d_kernel, configure_sfr_intercube_multisip, - _S_Q_DECODE, _N_RANKS_MULTI_USER, 1, (1, _N_RANKS_MULTI_USER), - ), -} - - -# ── Per-panel bench fn ───────────────────────────────────────────────── - - -def _make_bench_fn(panel: str): - kernel, sfr_install, S_q, n_ranks, rank_axis, mesh_shape = ( - _PANEL_DISPATCH[panel] - ) - is_multi_user = panel.startswith("multi_user_") - - def _bench_fn(ctx): - sfr_install( - ctx.engine, ctx.spec, - resolve_algorithm_config(load_ccl_config(), name="lrab_hierarchical_allreduce"), - ) - if is_multi_user: - dp_full = DPPolicy( - cube="replicate", pe="replicate", - num_cubes=n_ranks, num_pes=8, - ) - dp_kv = DPPolicy( - cube="row_wise", pe="replicate", - num_cubes=n_ranks, num_pes=8, - ) - else: - dp_full = DPPolicy( - cube="replicate", pe="replicate", - num_cubes=1, num_pes=n_ranks, - ) - dp_kv = DPPolicy( - cube="replicate", pe="row_wise", - num_cubes=1, num_pes=n_ranks, - ) - q = ctx.zeros((S_q, _H_Q * _D_HEAD), - dtype=_DTYPE, dp=dp_full, name=f"{panel}_q") - k = ctx.zeros((_S_KV_PER_RANK * n_ranks, _H_KV * _D_HEAD), - dtype=_DTYPE, dp=dp_kv, name=f"{panel}_k") - v = ctx.zeros((_S_KV_PER_RANK * n_ranks, _H_KV * _D_HEAD), - dtype=_DTYPE, dp=dp_kv, name=f"{panel}_v") - o = ctx.empty((S_q, _H_Q * _D_HEAD), - dtype=_DTYPE, dp=dp_full, name=f"{panel}_o") - # rank_axis is a positional arg; _auto_dim_remap=False keeps - # d_head=64 from colliding with the multi_user K's global M=64. - if mesh_shape is None: - ctx.launch( - f"{panel}_mesh", kernel, - q, k, v, o, - S_q, _S_KV_PER_RANK, _H_Q, _H_KV, _D_HEAD, n_ranks, - rank_axis, - _auto_dim_remap=False, - ) - else: - mesh_rows, mesh_cols = mesh_shape - ctx.launch( - f"{panel}_mesh", kernel, - q, k, v, o, - S_q, _S_KV_PER_RANK, _H_Q, _H_KV, _D_HEAD, - mesh_rows, mesh_cols, - rank_axis, - 0, # cube_start=0: this panel's launch starts at cube 0 - _auto_dim_remap=False, - ) - - return _bench_fn - - -# ── Op-log summary harvest ───────────────────────────────────────────── - - -def _summarize_op_log(op_log) -> dict[str, int]: - """Counts per ADR-0057 D7 op_log_summary contract.""" - gemm_count = 0 - ipcq_send_count = 0 - ipcq_recv_count = 0 - dma_read_count = 0 - dma_write_count = 0 - for r in op_log: - if r.op_kind == "gemm": - gemm_count += 1 - elif r.op_name == "dma_read": - dma_read_count += 1 - elif r.op_name == "dma_write": - dma_write_count += 1 - elif r.op_name == "ipcq_send": - ipcq_send_count += 1 - elif r.op_name == "ipcq_recv": - ipcq_recv_count += 1 - elif r.op_name == "ipcq_copy": - # The inbound DMA records ipcq_copy (one per send/recv pair). - # Count it as both a send and a recv side so the row's - # ipcq_send_count and ipcq_recv_count are non-zero even when - # the engine logs the collective via the inbound copy alone. - ipcq_send_count += 1 - ipcq_recv_count += 1 - return { - "gemm_count": gemm_count, - "ipcq_send_count": ipcq_send_count, - "ipcq_recv_count": ipcq_recv_count, - "dma_read_count": dma_read_count, - "dma_write_count": dma_write_count, - } - - -def _run_panel(panel: str, topology: str) -> dict: - """Run one panel via a fresh engine; return its row dict.""" - from kernbench.runtime_api.bench_runner import run_bench - from kernbench.runtime_api.types import resolve_device - from kernbench.sim_engine.engine import GraphEngine - from kernbench.topology.builder import resolve_topology - - topo = resolve_topology(topology) - result = run_bench( - topology=topo, bench_fn=_make_bench_fn(panel), - device=resolve_device(None), - engine_factory=lambda t, d: GraphEngine( - getattr(t, "topology_obj", t), enable_data=True, - ), - ) - if not result.completion.ok: - raise RuntimeError( - f"milestone-gqa-llama70b panel {panel!r} failed: {result.completion}" - ) - _, _, _, n_ranks, _, _ = _PANEL_DISPATCH[panel] - return { - "panel": panel, - "n_ranks": n_ranks, - "op_log_summary": _summarize_op_log(result.engine.op_log), - } - - -# ── Figure renderers (sub-cycle 4c, 5 of 6 figures) ────────────────── -# -# Sixth figure ``gqa_scaling.png`` is deferred to after sub-cycle 4b -# lands the Q/cube ∈ {1, 2, 4} sweep on multi_user_* panels — it needs -# multiple sweep.json rows per multi_user panel to be meaningful. - -_OP_LOG_KEYS = ( - "gemm_count", - "ipcq_send_count", - "ipcq_recv_count", - "dma_read_count", - "dma_write_count", -) -_OP_LOG_DISPLAY = { - "gemm_count": "GEMM", - "ipcq_send_count": "IPCQ send", - "ipcq_recv_count": "IPCQ recv", - "dma_read_count": "DMA read", - "dma_write_count": "DMA write", -} -_OP_LOG_COLORS = { - "gemm_count": "#F59E0B", - "ipcq_send_count": "#3B82F6", - "ipcq_recv_count": "#10B981", - "dma_read_count": "#A855F7", - "dma_write_count": "#EF4444", -} -_PANEL_DISPLAY = { - "single_user_prefill": "single_user / prefill", - "multi_user_prefill": "multi_user / prefill", - "single_user_decode": "single_user / decode", - "multi_user_decode": "multi_user / decode", -} - - -def _load_sweep_data(sweep_json: Path | str) -> dict: - sweep_json = Path(sweep_json) - if not sweep_json.exists(): - return {"rows": [], "config": {}, "panels": []} - return json.loads(sweep_json.read_text()) - - -def _row_for(rows: list, panel: str) -> dict | None: - for r in rows: - if r.get("panel") == panel: - return r - return None - - -def emit_panel_op_log_summary( - panel: str, - sweep_json: Path | str = _SWEEP_JSON, - out_dir: Path | str = _OUTPUT_DIR, -) -> str | None: - """One bar chart of the 5 op_log counts for ``panel``. - - Returns the written PNG path, or ``None`` when sweep.json is empty - or the requested panel is absent. - """ - import matplotlib.pyplot as plt - - data = _load_sweep_data(sweep_json) - row = _row_for(data.get("rows", []), panel) - if row is None: - return None - summary = row.get("op_log_summary", {}) - n_ranks = row.get("n_ranks") - - labels = [_OP_LOG_DISPLAY[k] for k in _OP_LOG_KEYS] - values = [summary.get(k, 0) for k in _OP_LOG_KEYS] - colors = [_OP_LOG_COLORS[k] for k in _OP_LOG_KEYS] - - fig, ax = plt.subplots(figsize=(8, 5)) - bars = ax.bar(labels, values, color=colors) - for b, v in zip(bars, values): - ax.text(b.get_x() + b.get_width() / 2, b.get_height(), - f"{int(v)}", ha="center", va="bottom", fontsize=9) - ax.set_title( - f"{_PANEL_DISPLAY.get(panel, panel)} (n_ranks={n_ranks})", - fontsize=12, fontweight="bold", - ) - ax.set_ylabel("count") - ax.grid(True, axis="y", alpha=0.3) - fig.tight_layout() - - out_dir = Path(out_dir) - out_dir.mkdir(parents=True, exist_ok=True) - out = out_dir / f"gqa_op_log_{panel}.png" - fig.savefig(out, dpi=120) - plt.close(fig) - return str(out) - - -def emit_gqa_comparison( - sweep_json: Path | str = _SWEEP_JSON, - out_dir: Path | str = _OUTPUT_DIR, -) -> str | None: - """Grouped-bar chart comparing the 5 op_log counts across all panels.""" - import matplotlib.pyplot as plt - import numpy as np - - data = _load_sweep_data(sweep_json) - panels_in = data.get("panels") or list(_PANELS_V1) - rows = data.get("rows", []) - panels = [p for p in panels_in if _row_for(rows, p) is not None] - if not panels: - return None - - n_groups = len(panels) - n_series = len(_OP_LOG_KEYS) - x = np.arange(n_groups) - width = 0.8 / n_series - - fig, ax = plt.subplots(figsize=(11, 6)) - for i, key in enumerate(_OP_LOG_KEYS): - offset = (i - (n_series - 1) / 2) * width - vals = [_row_for(rows, p)["op_log_summary"].get(key, 0) - for p in panels] - ax.bar(x + offset, vals, width, - label=_OP_LOG_DISPLAY[key], color=_OP_LOG_COLORS[key]) - - ax.set_xticks(x) - ax.set_xticklabels( - [f"{_PANEL_DISPLAY.get(p, p)}\n(n_ranks={_row_for(rows, p)['n_ranks']})" - for p in panels], - fontsize=8, - ) - ax.set_ylabel("count") - ax.set_title("GQA Llama-70B — op_log summary across panels", - fontsize=13, fontweight="bold") - ax.legend(fontsize=8, loc="upper right") - ax.grid(True, axis="y", alpha=0.3) - fig.tight_layout() - - out_dir = Path(out_dir) - out_dir.mkdir(parents=True, exist_ok=True) - out = out_dir / "gqa_comparison.png" - fig.savefig(out, dpi=120) - plt.close(fig) - return str(out) - - -def emit_all_gqa_plots( - sweep_json: Path | str = _SWEEP_JSON, - out_dir: Path | str = _OUTPUT_DIR, -) -> list[str]: - """Render all 5 in-scope figures and return the written paths. - - Sub-cycle 4c v1 emits 5 of the 6 figures ADR-0057 D3 lists; the - 6th (gqa_scaling.png) needs sub-cycle 4b's Q/cube sweep data. - """ - paths: list[str] = [] - for panel in _PANELS_V1: - p = emit_panel_op_log_summary(panel, sweep_json, out_dir) - if p is not None: - paths.append(p) - comp = emit_gqa_comparison(sweep_json, out_dir) - if comp is not None: - paths.append(comp) - return paths - - -# ── Bench entry ──────────────────────────────────────────────────────── - - -@bench( - name="milestone-gqa-llama70b", - description="1H milestone: GQA Llama-70B 4-panel sweep (ADR-0057 v1).", -) -def run(torch) -> None: - """Drive the four GQA panels at validation scale; write sweep.json and figures. - - Modes (mutually exclusive): - MILESTONE_FAST=1 Skip the sweep; re-render figures from the - committed sweep.json. Seconds, no simulator. - GQA_VALIDATION=1 Run the four-panel validation sweep + figures. - ~1-2h on the full simulator. - - Headline-scale mode is deferred to sub-cycle 4c (figures landed - here; headline-scale + scaling figure await sub-cycle 4b). - A sentinel tensor is submitted at the end so run_bench's ADR-0045 D4 - "at least one request" contract is satisfied even when the panels - are skipped via MILESTONE_FAST=1. - """ - _OUTPUT_DIR.mkdir(parents=True, exist_ok=True) - fast = bool(os.environ.get("MILESTONE_FAST")) - if not fast and not os.environ.get("GQA_VALIDATION"): - raise RuntimeError( - "milestone-gqa-llama70b v1 needs GQA_VALIDATION=1 (run the " - "sweep) or MILESTONE_FAST=1 (reuse committed sweep.json). " - "Headline mode is deferred to sub-cycle 4b/4c per ADR-0057 D3." - ) - - if not fast: - topology = os.environ.get("GQA_TOPOLOGY", "topology.yaml") - rows = [_run_panel(panel, topology) for panel in _PANELS_V1] - sweep = { - "version": 1, - "validation_scale": True, - "panels": list(_PANELS_V1), - "config": { - "S_q_prefill": _S_Q_PREFILL, - "S_kv_per_rank": _S_KV_PER_RANK, - "h_q": _H_Q, - "h_kv": _H_KV, - "d_head": _D_HEAD, - "n_ranks_single_user": _N_RANKS_SINGLE_USER, - "n_ranks_multi_user": _N_RANKS_MULTI_USER, - }, - "rows": rows, - } - _SWEEP_JSON.write_text(json.dumps(sweep, indent=2)) - print(f" milestone-gqa-llama70b: {len(rows)} rows -> {_SWEEP_JSON}") - elif not _SWEEP_JSON.exists(): - raise RuntimeError( - f"MILESTONE_FAST=1 requires {_SWEEP_JSON} to exist; " - "run with GQA_VALIDATION=1 once to seed it." - ) - - paths = emit_all_gqa_plots() - print(f" milestone-gqa-llama70b: {len(paths)} figures -> {_OUTPUT_DIR} " - f"(fast={fast})") - - # Sentinel tensor (ADR-0045 D4 / ADR-0054 D2 carve-out). - torch.zeros( - (1, 1), dtype="f16", - dp=DPPolicy(cube="row_wise", pe="replicate", num_cubes=1, num_pes=1), - name="milestone_gqa_sentinel", - ) diff --git a/src/kernbench/ccl/sfr_config.py b/src/kernbench/ccl/sfr_config.py index e29ac46..81e47b5 100644 --- a/src/kernbench/ccl/sfr_config.py +++ b/src/kernbench/ccl/sfr_config.py @@ -237,3 +237,119 @@ def configure_sfr_intracube_pe_ring( algo_module=mock_module, rank_to_pe=pe_idx_to_pe, ) + + +# ── Inter-cube 1D ring (ADR-0060 §5.5 prefill Ring KV) ───────────────── + + +def configure_sfr_intercube_ring( + engine: Any, + spec: dict, + cfg: dict, + *, + ring_size: int | None = None, +) -> dict[str, Any]: + """Install intra-cube PE grid + a 1D CUBE-level ring with wrap. + + Direction namespaces (disjoint, same as + ``configure_sfr_intercube_multisip``): + + - ``intra_N/S/E/W`` : 2×4 PE grid within each cube (no wrap) + - ``E/W`` : 1D ring of cubes 0..ring_size-1 WITH WRAP + (symmetric to ``configure_sfr_intracube_pe_ring`` + at PE level — wrap applied at CUBE level here) + - ``global_*`` : SIP topology (same as multisip) + + N/S at CUBE level are intentionally NOT installed — use + ``configure_sfr_intercube_multisip`` for the full 4×4 cube mesh. + + Args: + ring_size: number of CUBEs in the ring (wrap applies to cubes + 0..ring_size-1). Defaults to the full cube_mesh count. + Must be ≤ mesh_w (single row); multi-row rings span + non-neighbour boundaries. + """ + cm = spec["sip"]["cube_mesh"] + mesh_w = int(cm["w"]) + mesh_h = int(cm["h"]) + n_cubes = mesh_w * mesh_h + sips_cfg = spec.get("system", {}).get("sips", {}) + n_sips = int(sips_cfg.get("count", 1)) + sip_topology = str(sips_cfg.get("topology", "ring_1d")) + sip_w = sips_cfg.get("w") + sip_h = sips_cfg.get("h") + sip_w = int(sip_w) if sip_w is not None else None + sip_h = int(sip_h) if sip_h is not None else None + + if ring_size is None: + ring_size = n_cubes + if ring_size > mesh_w: + raise ValueError( + f"intercube_ring ring_size={ring_size} > mesh_w={mesh_w}; " + "multi-row rings cross non-neighbour boundaries" + ) + + if sip_topology not in _TOPO_BUILTINS: + raise ValueError( + f"Unknown sip topology '{sip_topology}'. " + f"Available: {list(_TOPO_BUILTINS)}" + ) + _sip_topo_fn_raw = _TOPO_BUILTINS[sip_topology] + + def sip_topo_fn(rank: int, ws: int) -> dict: + if sip_w is not None and sip_h is not None: + try: + return _sip_topo_fn_raw(rank, ws, w=sip_w, h=sip_h) + except TypeError: + pass + return _sip_topo_fn_raw(rank, ws) + + pes_per_cube = _PES_PER_CUBE + world_size = n_sips * n_cubes * pes_per_cube + pe_idx_to_pe: list[tuple[int, int, int]] = [ + (sip, cube, pe) + for sip in range(n_sips) + for cube in range(n_cubes) + for pe in range(pes_per_cube) + ] + + def _pe_idx(sip: int, cube: int, pe: int) -> int: + return (sip * n_cubes + cube) * pes_per_cube + pe + + def _neighbors(pe_idx: int, ws: int, _base: dict) -> dict[str, int]: + tmp = pe_idx + pe = tmp % pes_per_cube + tmp //= pes_per_cube + cube = tmp % n_cubes + sip = tmp // n_cubes + + nbrs: dict[str, int] = {} + + # ── Intra-cube (intra_N/S/E/W) ── + for d, peer_pe in _intra_cube_neighbors(pe).items(): + nbrs[d] = _pe_idx(sip, cube, peer_pe) + + # ── Cube ring (E/W with wrap for cubes 0..ring_size-1) ── + if cube < ring_size: + nbrs["E"] = _pe_idx(sip, (cube + 1) % ring_size, pe) + nbrs["W"] = _pe_idx(sip, (cube - 1) % ring_size, pe) + + # ── Inter-SIP same-(cube, pe) (global_*) ── + if n_sips > 1: + sip_nbrs = sip_topo_fn(sip, n_sips) + for d, peer_sip in sip_nbrs.items(): + nbrs[f"global_{d}"] = _pe_idx(peer_sip, cube, pe) + + return nbrs + + mock_module = types.SimpleNamespace(neighbors=_neighbors) + + cfg_copy = dict(cfg) + cfg_copy["world_size"] = world_size + cfg_copy["topology"] = "none" + + return install_ipcq( + engine, spec, cfg_copy, + algo_module=mock_module, + rank_to_pe=pe_idx_to_pe, + ) diff --git a/src/kernbench/common/cpu_issue_cost.py b/src/kernbench/common/cpu_issue_cost.py new file mode 100644 index 0000000..a9371a6 --- /dev/null +++ b/src/kernbench/common/cpu_issue_cost.py @@ -0,0 +1,52 @@ +"""Per-op-type CPU issue cost table (ADR-0064 D1). + +Replaces the single uniform ``dispatch_cycles`` scalar with a cost table +keyed by command kind. Charged on PE_CPU at issue time (before the command +is dispatched to PE_SCHEDULER) so the hybrid's CPU-saturation win +(ADR-0060 §1) becomes measurable. + +The table is consulted by ``TLContext._emit_dispatch_overhead(kind)``; +live PE_CPU paths (greenlet via ``kernel_runner.py``, legacy replay via +``pe_cpu.py:_execute_legacy``) construct TLContext with +``issue_cost_table=DEFAULT_CPU_ISSUE_COST`` so all benches see the cost. + +Absolute ns values are provisional (ADR-0064 review item #1). The +defensible claim is the **ratio** — composite ≫ primitive. +""" +from __future__ import annotations + +from typing import Literal + +OpKind = Literal[ + "composite", + "load", + "store", + "dot", + "math", + "ipcq_send", + "ipcq_recv", + "copy_to", +] + + +DEFAULT_CPU_ISSUE_COST: dict[str, int] = { + "composite": 40, + "load": 5, + "store": 5, + "dot": 5, + "math": 5, + "ipcq_send": 5, + "ipcq_recv": 5, + "copy_to": 5, +} + + +def get_issue_cost(kind: str, table: dict[str, int] | None = None) -> int: + """Return per-op-type CPU issue cost in ns. + + Unknown kinds return 0 (no charge) so adding a new ``tl.*`` op kind + doesn't accidentally over-charge before the table is updated. + """ + if table is None: + table = DEFAULT_CPU_ISSUE_COST + return table.get(kind, 0) diff --git a/src/kernbench/common/pe_commands.py b/src/kernbench/common/pe_commands.py index 02a32d4..12ea7e3 100644 --- a/src/kernbench/common/pe_commands.py +++ b/src/kernbench/common/pe_commands.py @@ -73,6 +73,12 @@ class TensorHandle: data: object = None # reserved for validate mode space: str = "tcm" # MemoryStore space ("tcm" | "hbm" | "sram") pinned: bool = False # operand already DMA-staged in TCM (via tl.load) + # ADR-0062 §D2: lazy tl.load attaches a LoadFuture here. None for + # handles that have no in-flight DMA (constants, math outputs, etc.). + # Consumer ops call _await_pending() to yield on the future before + # emitting their own command. Excluded from eq/hash/repr so handle + # identity is unaffected by pending state. + pending: object = field(default=None, compare=False, hash=False, repr=False) @dataclass(frozen=True) @@ -140,6 +146,22 @@ class MathCmd: data_op: bool = True +@dataclass(frozen=True) +class CopyCmd: + """TCM-to-TCM byte copy (ADR-0063 §D3.1). + + Emitted by ``tl.copy_to`` to persist a scoped result's bytes to an + outside-``scratch_scope`` (persistent) address — the two-arena + pattern for tiled flash attention. Runs on the vector engine; + op_log classifies as ``op_kind="math"``, ``op_name="copy"``. + """ + + src: TensorHandle + dst: TensorHandle + nbytes: int + data_op: bool = True + + @dataclass(frozen=True) class CompositeCmd: """Composite command: tiled pipeline of DMA_READ + COMPUTE + DMA_WRITE. @@ -178,7 +200,7 @@ class PeCpuOverheadCmd: # Union type for all PE commands PeCommand = ( - DmaReadCmd | DmaWriteCmd | GemmCmd | MathCmd + DmaReadCmd | DmaWriteCmd | GemmCmd | MathCmd | CopyCmd | CompositeCmd | WaitCmd | PeCpuOverheadCmd ) diff --git a/src/kernbench/components/builtin/pe_cpu.py b/src/kernbench/components/builtin/pe_cpu.py index ca0512e..786788f 100644 --- a/src/kernbench/components/builtin/pe_cpu.py +++ b/src/kernbench/components/builtin/pe_cpu.py @@ -184,6 +184,7 @@ class PeCpuComponent(ComponentBase): self, env, kernel_fn, kernel_args, num_programs, scheduler_id, ) -> Generator: """Legacy Phase 0 + replay: generate command list, then dispatch.""" + from kernbench.common.cpu_issue_cost import DEFAULT_CPU_ISSUE_COST from kernbench.common.pe_commands import ( CompositeCmd, PeCpuOverheadCmd, PeInternalTxn, WaitCmd, ) @@ -193,6 +194,7 @@ class PeCpuComponent(ComponentBase): pe_id=self._pe_idx, num_programs=num_programs, cube_id=self._cube_idx, num_cubes=self._num_cubes, dispatch_cycles=0, + issue_cost_table=DEFAULT_CPU_ISSUE_COST, ) run_kernel(kernel_fn, tl, *kernel_args) commands = tl.commands diff --git a/src/kernbench/components/builtin/pe_math.py b/src/kernbench/components/builtin/pe_math.py index eeb819e..1193910 100644 --- a/src/kernbench/components/builtin/pe_math.py +++ b/src/kernbench/components/builtin/pe_math.py @@ -99,20 +99,25 @@ class PeMathComponent(PeEngineBase): self._on_process_end(env, token) def handle_command(self, env: simpy.Environment, pe_txn: PeInternalTxn) -> Generator: - """PeInternalTxn handling for standalone MathCmd (CCL kernels). + """PeInternalTxn handling for standalone MathCmd / CopyCmd. Latency = max(overhead_ns, _compute_ns(num_elements)): - overhead_ns: fixed per-invocation setup cost (from node attrs). - _compute_ns: SIMD cycle-based model (from vector_width + clock_freq). The larger of the two dominates (setup-bound vs compute-bound). + + CopyCmd (ADR-0063 §D3.1): vector-engine on-chip byte copy; cost + model = _compute_ns(prod(dst.shape)). """ - from kernbench.common.pe_commands import MathCmd + from kernbench.common.pe_commands import CopyCmd, MathCmd import math as _math cmd = pe_txn.command num_elements = 0 if isinstance(cmd, MathCmd) and cmd.out.shape: num_elements = _math.prod(cmd.out.shape) + elif isinstance(cmd, CopyCmd) and cmd.dst.shape: + num_elements = _math.prod(cmd.dst.shape) overhead_ns = float(self.node.attrs.get("overhead_ns", 0.0)) compute_ns = self._compute_ns(num_elements) diff --git a/src/kernbench/components/builtin/pe_scheduler.py b/src/kernbench/components/builtin/pe_scheduler.py index 87c73b9..8cd17ac 100644 --- a/src/kernbench/components/builtin/pe_scheduler.py +++ b/src/kernbench/components/builtin/pe_scheduler.py @@ -42,12 +42,16 @@ class PeSchedulerComponent(ComponentBase): def _ensure_dispatch_table(cls) -> None: if cls._CMD_DISPATCH: return - from kernbench.common.pe_commands import DmaReadCmd, DmaWriteCmd, GemmCmd, MathCmd + from kernbench.common.pe_commands import ( + CopyCmd, DmaReadCmd, DmaWriteCmd, GemmCmd, MathCmd, + ) cls._CMD_DISPATCH = { DmaReadCmd: "pe_dma", DmaWriteCmd: "pe_dma", GemmCmd: "pe_gemm", MathCmd: "pe_math", + # ADR-0063 §D3.1: tl.copy_to → vector engine. + CopyCmd: "pe_math", } def __init__(self, node: Node, ctx: ComponentContext | None = None) -> None: diff --git a/src/kernbench/sim_engine/data_executor.py b/src/kernbench/sim_engine/data_executor.py index a0a429c..5ecbb25 100644 --- a/src/kernbench/sim_engine/data_executor.py +++ b/src/kernbench/sim_engine/data_executor.py @@ -238,6 +238,12 @@ def _compute_math(op: str, inputs: list[np.ndarray], axis: int | None) -> np.nda x = inputs[0] + # ADR-0063 §D3.1: copy is the vector-engine byte move used by + # tl.copy_to to persist a scoped result to the persistent arena. + # In data mode the bytes flow through as-is — identity op. + if op == "copy": + return x + # Unary if op == "exp": return np.exp(x) diff --git a/src/kernbench/sim_engine/op_log.py b/src/kernbench/sim_engine/op_log.py index 1e8bb8f..922e9a2 100644 --- a/src/kernbench/sim_engine/op_log.py +++ b/src/kernbench/sim_engine/op_log.py @@ -186,7 +186,7 @@ class OpLogger: def _extract_op_info(msg: Any) -> tuple[str, str, dict[str, Any]]: """Extract op_kind, op_name, params from a data_op message.""" from kernbench.common.pe_commands import ( - DmaReadCmd, DmaWriteCmd, GemmCmd, MathCmd, CompositeCmd, + CompositeCmd, CopyCmd, DmaReadCmd, DmaWriteCmd, GemmCmd, MathCmd, ) if isinstance(msg, DmaReadCmd): return "memory", "dma_read", { @@ -237,6 +237,16 @@ def _extract_op_info(msg: Any) -> tuple[str, str, dict[str, Any]]: "dtype": msg.out.dtype, "axis": msg.axis, } + if isinstance(msg, CopyCmd): + return "math", "copy", { + "src_addr": msg.src.addr, + "src_space": getattr(msg.src, "space", "tcm"), + "dst_addr": msg.dst.addr, + "dst_space": getattr(msg.dst, "space", "tcm"), + "shape": msg.src.shape, + "dtype": msg.src.dtype, + "nbytes": msg.nbytes, + } if isinstance(msg, CompositeCmd): params: dict[str, Any] = { "op": msg.op, diff --git a/src/kernbench/triton_emu/kernel_runner.py b/src/kernbench/triton_emu/kernel_runner.py index ca25ea1..6b703ad 100644 --- a/src/kernbench/triton_emu/kernel_runner.py +++ b/src/kernbench/triton_emu/kernel_runner.py @@ -89,6 +89,7 @@ class KernelRunner: 4. Dispatches each command through SimPy components 5. Returns results to the kernel """ + from kernbench.common.cpu_issue_cost import DEFAULT_CPU_ISSUE_COST from kernbench.triton_emu.tl_context import TLContext self._parent = greenlet.getcurrent() @@ -102,6 +103,7 @@ class KernelRunner: runner=self, scratch_base=self._scratch_base, scratch_size=self._scratch_size, + issue_cost_table=DEFAULT_CPU_ISSUE_COST, ) self._tl = tl # exposed so switch_to_simpy can re-set on restore @@ -144,7 +146,9 @@ class KernelRunner: cmd = _switch_kernel() elif isinstance(cmd, DmaReadCmd): - # Dispatch DMA through SimPy components + # Legacy blocking path — retained as a fallback for any + # caller that bypasses the lazy ``tl.load`` surface. New + # lazy loads come in as ("load_issue", future) below. done_evt = env.event() pe_txn = PeInternalTxn( command=cmd, done=done_evt, pe_prefix=self._pe_prefix, @@ -245,6 +249,45 @@ class KernelRunner: } cmd = _switch_kernel(result) + elif isinstance(cmd, tuple) and len(cmd) == 2 and cmd[0] == "load_issue": + # ADR-0062 §D2: lazy tl.load. Post the DmaReadCmd, store the + # done event on the LoadFuture, switch back immediately — do + # NOT yield done_evt here. The auto-wait at first use + # ("load_await" below) is what eventually yields it. + future = cmd[1] + done_evt = env.event() + pe_txn = PeInternalTxn( + command=future.cmd, done=done_evt, pe_prefix=self._pe_prefix, + ) + yield self._out_ports[self._scheduler_id].put(pe_txn) + future.event = done_evt + cmd = _switch_kernel(None) + + elif isinstance(cmd, tuple) and len(cmd) == 2 and cmd[0] == "load_await": + # ADR-0062 §D2: auto-wait at first use. Yield on the future's + # DMA event if not yet triggered, then read data and attach + # it to the handle (frozen dataclass — mutate via object. + # __setattr__, the same controlled pattern Phase 1 + # blocking tl.load used). + handle = cmd[1] + future = handle.pending + if future is not None and not future.resolved: + if future.event is not None and not future.event.triggered: + yield future.event + data = None + if self._store is not None: + try: + data = self._store.read( + "hbm", future.cmd.src_addr, + shape=handle.shape, dtype=handle.dtype, + ) + except KeyError: + pass + future.data = data + future.resolved = True + object.__setattr__(handle, "data", data) + cmd = _switch_kernel(None) + elif isinstance(cmd, tuple) and len(cmd) == 2 and cmd[0] == "recv_async": # Non-blocking recv: post the IpcqRequest now, store the # event in the future, return None to kernel. diff --git a/src/kernbench/triton_emu/tl_context.py b/src/kernbench/triton_emu/tl_context.py index 67764b1..c38ad5c 100644 --- a/src/kernbench/triton_emu/tl_context.py +++ b/src/kernbench/triton_emu/tl_context.py @@ -22,6 +22,7 @@ from kernbench.common.pe_commands import ( EPILOGUE_OPS, CompletionHandle, CompositeCmd, + CopyCmd, DmaReadCmd, DmaWriteCmd, GemmCmd, @@ -42,13 +43,67 @@ _DTYPE_BYTES: dict[str, int] = { } +class LoadFuture: + """Lazy ``tl.load`` future (ADR-0062 §D2). + + Mirrors ``RecvFuture`` for IPCQ, generalised to HBM loads. Carries + the originating ``DmaReadCmd``, the SimPy completion event (set by + the runner once the DMA has been issued), and a resolved flag. + + Consumer ops auto-wait via ``TLContext._await_pending(handle)``, + which yields ``event`` if ``resolved`` is False, then reads the + DMA's bytes into ``handle.data`` and marks the future resolved. + """ + + __slots__ = ("cmd", "event", "resolved", "data") + + def __init__(self, cmd: DmaReadCmd) -> None: + self.cmd = cmd + self.event: object | None = None # simpy.Event set by runner + self.resolved: bool = False + self.data: object = None + + +class _ScratchScope: + """Context manager that recycles per-tile scratch (ADR-0063 D1). + + ``__enter__`` snapshots ``_scratch_cursor``; ``__exit__`` restores it, + so every handle allocated inside the ``with``-block has its address + freed for the next iteration. Persistent state (running ``(m, ℓ, O)``, + prefetch buffers) lives outside the scope per ADR-0063 D3. + """ + + def __init__(self, ctx: "TLContext") -> None: + self._ctx = ctx + self._save: int | None = None + + def __enter__(self) -> "_ScratchScope": + self._save = self._ctx._scratch_cursor + return self + + def __exit__(self, *exc_info: object) -> bool: + if self._save is not None: + self._ctx._scratch_cursor = self._save + return False + + class TLContext: """Fake Triton Language context. Args: pe_id: program instance index (returned by program_id). num_programs: total number of program instances. - dispatch_cycles: PE_CPU overhead per tl API call (auto-inserted). + dispatch_cycles: uniform PE_CPU overhead per tl API call. Used as + a fallback when ``issue_cost_table`` is None (ADR-0046 §D6 + back-compat). When ``issue_cost_table`` is provided, the + per-kind table value is used instead. + issue_cost_table: optional per-op-type CPU issue cost table + (ADR-0064 D1). When provided, each ``tl.*`` call charges the + table value keyed by op kind ("composite", "load", "store", + "dot", "math", "ipcq_send", "ipcq_recv", "copy_to"). Unknown + kinds fall back to ``dispatch_cycles``. Live PE_CPU paths + construct TLContext with ``DEFAULT_CPU_ISSUE_COST`` so the + hybrid's CPU-saturation lever is measurable. """ def __init__( @@ -61,12 +116,14 @@ class TLContext: num_cubes: int = 1, scratch_base: int = 0, scratch_size: int = 1 << 20, # 1 MiB per kernel invocation + issue_cost_table: dict[str, int] | None = None, ) -> None: self._pe_id = pe_id self._num_programs = num_programs self._cube_id = cube_id self._num_cubes = num_cubes self._dispatch_cycles = dispatch_cycles + self._issue_cost_table = issue_cost_table self._commands: list[PeCommand] = [] self._handle_counter = 0 self._completion_counter = 0 @@ -79,6 +136,22 @@ class TLContext: self._scratch_size = scratch_size self._scratch_cursor = 0 + def scratch_scope(self) -> _ScratchScope: + """Per-tile scratch recycling context manager (ADR-0063). + + Usage: + + with tl.scratch_scope(): + s = tl.dot(q, k_t) # per-tile temporaries — + p = tl.softmax(s) # their scratch is rewound + o_j = tl.dot(p, v) # on __exit__ + + Persistent state (running ``(m, ℓ, O)``, lazy-load prefetch + buffers) must be allocated **outside** the scope; only handles + allocated inside are recycled. + """ + return _ScratchScope(self) + def _scratch_alloc(self, nbytes: int) -> int: """Allocate a unique scratch address for an output TensorHandle. @@ -120,9 +193,23 @@ class TLContext: def _nbytes(self, shape: tuple[int, ...], dtype: str) -> int: return math.prod(shape) * self._dtype_bytes(dtype) - def _emit_dispatch_overhead(self) -> None: - if self._dispatch_cycles > 0: - self._emit(PeCpuOverheadCmd(cycles=self._dispatch_cycles)) + def _emit_dispatch_overhead(self, kind: str | None = None) -> None: + """Charge per-op-type CPU issue cost (ADR-0064 D1). + + When ``issue_cost_table`` was provided, look up the per-kind cost + and emit ``PeCpuOverheadCmd(cycles=N)`` if N > 0. Unknown kinds + fall back to the uniform ``dispatch_cycles`` for forward-compat + when a new ``tl.*`` op is added before the table is updated. + + When ``issue_cost_table`` is None, preserve the ADR-0046 §D6 + contract: emit ``PeCpuOverheadCmd(dispatch_cycles)`` if positive. + """ + if self._issue_cost_table is not None and kind is not None: + cycles = self._issue_cost_table.get(kind, self._dispatch_cycles) + else: + cycles = self._dispatch_cycles + if cycles > 0: + self._emit(PeCpuOverheadCmd(cycles=cycles)) def _make_handle( self, addr: int, shape: tuple[int, ...], dtype: str, @@ -177,37 +264,103 @@ class TLContext: def load( self, ptr: int, shape: tuple[int, ...], dtype: str = "f16", ) -> TensorHandle: - """Load tensor from HBM. Returns TensorHandle pointing at HBM[ptr]. + """Load tensor from HBM — **lazy** (ADR-0062 §D1/§D2). - In greenlet mode: returns TensorHandle with actual numpy data. - In command-list mode: returns TensorHandle with data=None. + Posts the ``DmaReadCmd`` to PE_DMA and returns immediately with a + ``TensorHandle`` whose ``pending`` field references a fresh + ``LoadFuture``. The actual DMA completion is awaited at the first + consuming op (``tl.dot``, MATH, ``tl.store``, ``tl.send``, + ``tl.copy_to``, ``tl.composite``) via ``_await_pending``. - The returned handle's ``space`` is "hbm" so subsequent ops (math, - send, store) using this handle as a source resolve via MemoryStore - at ``(hbm, ptr)`` — which is where the load's underlying data - actually lives in Phase 2 storage. + Command-list mode: emits the DmaReadCmd to ``self._commands``, + attaches a LoadFuture for structural compatibility (its event + stays None — no engine, no SimPy event to wait on). """ - self._emit_dispatch_overhead() - handle = self._make_handle( - addr=ptr, shape=shape, dtype=dtype, space="hbm", pinned=True, + self._emit_dispatch_overhead("load") + nbytes = self._nbytes(shape, dtype) + # LoadFuture is mutable; create it first, attach to the handle, + # then point its ``cmd`` at the DmaReadCmd that references the + # *final* handle. This guarantees ``handle.pending.cmd.handle is handle``. + future = LoadFuture.__new__(LoadFuture) + future.event = None + future.resolved = False + future.data = None + handle = TensorHandle( + id=self._next_handle_id(), + addr=ptr, shape=shape, dtype=dtype, nbytes=nbytes, + data=None, space="hbm", pinned=True, pending=future, ) - cmd = DmaReadCmd(handle=handle, src_addr=ptr, nbytes=handle.nbytes) - data = self._emit(cmd) - if data is not None: - # Greenlet mode: attach real data to handle (preserve space + pinned) - return TensorHandle( - id=handle.id, addr=handle.addr, shape=handle.shape, - dtype=handle.dtype, nbytes=handle.nbytes, data=data, - space=handle.space, pinned=handle.pinned, - ) + cmd = DmaReadCmd(handle=handle, src_addr=ptr, nbytes=nbytes) + future.cmd = cmd + if self._runner is not None: + # Lazy: runner posts the DmaReadCmd, sets future.event, then + # switches back immediately. No yield on completion here. + self._runner.switch_to_simpy(("load_issue", future)) + else: + self._commands.append(cmd) return handle + def _await_pending(self, *handles: TensorHandle | None) -> None: + """Auto-wait at first use (ADR-0062 §D2). + + For each handle carrying an unresolved ``LoadFuture``, yield on + the DMA completion event (greenlet → runner). Command-list mode + is a no-op. + """ + if self._runner is None: + return + for h in handles: + if h is None: + continue + pending = getattr(h, "pending", None) + if pending is None or pending.resolved: + continue + self._runner.switch_to_simpy(("load_await", h)) + def store(self, ptr: int, handle: TensorHandle) -> None: """Store tensor from TCM to HBM.""" - self._emit_dispatch_overhead() + self._await_pending(handle) + self._emit_dispatch_overhead("store") cmd = DmaWriteCmd(handle=handle, dst_addr=ptr, nbytes=handle.nbytes) self._emit(cmd) + def copy_to(self, dst: TensorHandle, src: TensorHandle) -> None: + """TCM-to-TCM byte copy (ADR-0063 §D3.1). + + Writes ``src``'s bytes into ``dst``'s address. Both handles must + live in TCM and have matching shape and dtype. Symmetric to + ``tl.store`` (which targets HBM) but stays on-chip so it doesn't + emit a DMA entry into op_log. + + Used inside ``tl.scratch_scope()`` to persist a scoped result — + typically an updated running ``(m, ℓ, O)`` — to an outside-scope + (persistent) handle so its bytes survive the scope's ``__exit__`` + cursor rewind (ADR-0063 §D3 two-arena pattern). + """ + if src.shape != dst.shape: + raise ValueError( + f"tl.copy_to: shape mismatch — src.shape={src.shape} " + f"vs dst.shape={dst.shape}" + ) + if src.dtype != dst.dtype: + raise ValueError( + f"tl.copy_to: dtype mismatch — src.dtype={src.dtype!r} " + f"vs dst.dtype={dst.dtype!r}" + ) + if dst.space != "tcm": + raise ValueError( + f"tl.copy_to: dst must be in TCM (got space={dst.space!r}); " + "writes to HBM go through tl.store" + ) + if src.space != "tcm": + raise ValueError( + f"tl.copy_to: src must be in TCM (got space={src.space!r}); " + "reads from HBM go through tl.load" + ) + self._await_pending(src) + self._emit_dispatch_overhead("copy_to") + self._emit(CopyCmd(src=src, dst=dst, nbytes=src.nbytes)) + # ── GEMM Engine (blocking) ──────────────────────────────────── def dot(self, a: TensorHandle, b: TensorHandle) -> TensorHandle: @@ -224,7 +377,8 @@ class TLContext: out_shape = (*a.shape[:-2], m, n) out_dtype = a.dtype out = self._make_compute_out(shape=out_shape, dtype=out_dtype) - self._emit_dispatch_overhead() + self._await_pending(a, b) + self._emit_dispatch_overhead("dot") self._emit(GemmCmd(a=a, b=b, out=out, m=m, k=k, n=n)) return out @@ -232,7 +386,8 @@ class TLContext: def _unary_math(self, op: str, x: TensorHandle) -> TensorHandle: out = self._make_compute_out(shape=x.shape, dtype=x.dtype) - self._emit_dispatch_overhead() + self._await_pending(x) + self._emit_dispatch_overhead("math") self._emit(MathCmd(op=op, inputs=(x,), out=out)) return out @@ -265,7 +420,8 @@ class TLContext: out_shape = list(x.shape) out_shape[axis] = 1 out = self._make_compute_out(shape=tuple(out_shape), dtype=x.dtype) - self._emit_dispatch_overhead() + self._await_pending(x) + self._emit_dispatch_overhead("math") self._emit(MathCmd(op=op, inputs=(x,), out=out, axis=axis)) return out @@ -284,7 +440,8 @@ class TLContext: self, op: str, a: TensorHandle, b: TensorHandle, ) -> TensorHandle: out = self._make_compute_out(shape=a.shape, dtype=a.dtype) - self._emit_dispatch_overhead() + self._await_pending(a, b) + self._emit_dispatch_overhead("math") self._emit(MathCmd(op=op, inputs=(a, b), out=out)) return out @@ -292,7 +449,8 @@ class TLContext: self, cond: TensorHandle, a: TensorHandle, b: TensorHandle, ) -> TensorHandle: out = self._make_compute_out(shape=a.shape, dtype=a.dtype) - self._emit_dispatch_overhead() + self._await_pending(cond, a, b) + self._emit_dispatch_overhead("math") self._emit(MathCmd(op="where", inputs=(cond, a, b), out=out)) return out @@ -309,7 +467,8 @@ class TLContext: ) -> TensorHandle: """Fused multiply-add: a * b + c (real Triton: tl.fma).""" out = self._make_compute_out(shape=a.shape, dtype=a.dtype) - self._emit_dispatch_overhead() + self._await_pending(a, b, c) + self._emit_dispatch_overhead("math") self._emit(MathCmd(op="fma", inputs=(a, b, c), out=out)) return out @@ -321,7 +480,8 @@ class TLContext: ) -> TensorHandle: """Clamp x to [min, max] (real Triton: tl.clamp).""" out = self._make_compute_out(shape=x.shape, dtype=x.dtype) - self._emit_dispatch_overhead() + self._await_pending(x, min, max) + self._emit_dispatch_overhead("math") self._emit(MathCmd(op="clamp", inputs=(x, min, max), out=out)) return out @@ -333,7 +493,8 @@ class TLContext: canonical (x - max) → exp → sum → div sequence. """ out = self._make_compute_out(shape=x.shape, dtype=x.dtype) - self._emit_dispatch_overhead() + self._await_pending(x) + self._emit_dispatch_overhead("math") self._emit(MathCmd(op="softmax", inputs=(x,), out=out, axis=axis)) return out @@ -427,13 +588,16 @@ class TLContext: space = getattr(src, "space", space) if src_addr is None or nbytes is None or shape is None: raise ValueError("tl.send: provide either a TensorHandle or src_addr/nbytes/shape") + # ADR-0062: if the source is a lazy-loaded handle, await first so + # the data snapshot below sees the real bytes. + self._await_pending(src) # Carry the handle's .data snapshot (if available). When the source # is a recv slot, .data holds the numpy array that was read from # MemoryStore at recv-time. This prevents a Phase 1 race where a # later IPCQ inbound overwrites the slot before the outbound # PE_DMA reads it. handle_data = getattr(src, "data", None) if src is not None else None - self._emit_dispatch_overhead() + self._emit_dispatch_overhead("ipcq_send") cmd = IpcqSendCmd( direction=dir, src_addr=src_addr, src_space=space, @@ -467,7 +631,7 @@ class TLContext: arrived. In greenlet/runner mode, ``handle.data`` carries the actual ndarray; in command-list mode the handle is a placeholder. """ - self._emit_dispatch_overhead() + self._emit_dispatch_overhead("ipcq_recv") if dst_addr is not None and dst_space is not None: cmd = IpcqRecvCmd( direction=dir, @@ -518,7 +682,7 @@ class TLContext: they receive. This API is segregated from ``tl.recv`` so the diagnostic flag can never accidentally be set in real workloads. """ - self._emit_dispatch_overhead() + self._emit_dispatch_overhead("ipcq_recv") cmd = IpcqRecvCmd( direction=dir, shape=shape, dtype=dtype, @@ -547,7 +711,7 @@ class TLContext: dtype: str = "f16", ) -> "RecvFuture": """Non-blocking recv. Returns a future to pass into ``tl.wait``.""" - self._emit_dispatch_overhead() + self._emit_dispatch_overhead("ipcq_recv") cmd = IpcqRecvCmd( direction=dir, shape=shape, dtype=dtype, @@ -582,6 +746,9 @@ class TLContext: Returns CompletionHandle for use with wait(). """ + # ADR-0062: composite operand DMA paths still need their inputs + # to be resolved before the composite reads them via PE_SCHEDULER. + self._await_pending(a, b) # Compute output size based on op if op == "gemm" and b is not None: m, k = a.shape[-2], a.shape[-1] @@ -609,7 +776,7 @@ class TLContext: ops_tuple = (head_spec, *epi_specs) completion = CompletionHandle(id=self._next_completion_id()) - self._emit_dispatch_overhead() + self._emit_dispatch_overhead("composite") self._emit(CompositeCmd( completion=completion, op=op, a=a, b=b, out_addr=out_ptr, out_nbytes=out_nbytes, diff --git a/tests/attention/test_attention_8kv_groups_diag.py b/tests/attention/test_attention_8kv_groups_diag.py deleted file mode 100644 index 46b78d8..0000000 --- a/tests/attention/test_attention_8kv_groups_diag.py +++ /dev/null @@ -1,286 +0,0 @@ -"""Diagnostic harness for the Llama-70B "1 Q-head per cube" target. - -Per the GQA Llama-70B sharding study at -``llm_paper_review/notes/GQA_MHA_sharding/scripts/_gen_llama70b_1M_4cases.py``, -the 1 Q-head/cube baseline uses 64 cubes (4 SIPs × 16 cubes/SIP) organized -into 8 KV-groups of 8 cubes each. Each KV-group occupies a ``2×4`` -sub-mesh within a SIP's ``4×4`` cube grid and runs the C2 2D row-then-col -AllReduce-mlo (ADR-0059 extension). This harness probes the gap between -validation and headline in three incrementally-larger steps: - - step_1_single_kv_group_at_full_breadth - ONE multi_user_decode launch on a 2×4 sub-mesh (8 cubes) of the - 4-SIP topology, via the 2D mesh-mlo kernel. Verifies the per-KV-group - 2D AllReduce works at full breadth. Smallest dim possible - (S_q=1, S_kv=16, h=1, d_head=64) to keep wall time bounded. - - step_2_four_kv_groups_one_per_sip - Four sequential multi_user_decode launches, each targeting a - different SIP. Verifies that per-SIP isolation works (each SIP holds - its own 2×4 KV-group; the SFR install only writes intra-SIP - E/W + N/S edges so the 4 groups don't see each other). - - step_3_eight_kv_groups_two_per_sip - The actual study target: 8 KV-groups, two per SIP (cubes 0..7 vs - cubes 8..15 within each SIP). Expected to FAIL with current infra — - DPPolicy doesn't take a cube offset and target_device is SIP-level, - so back-to-back launches both land on cubes 0..7 of their target SIP. - Captures what's needed to lift the 4-group cap to 8. - -Each step prints what it observed; the test asserts only the documented -expected outcomes so we can land it, watch CI, and iterate. Steps that -are *expected* to fail (step 3) are marked xfail with a precise reason. -""" -from __future__ import annotations - -import traceback -from pathlib import Path - -import pytest - -from kernbench.benches._attention_mesh_mlo_2d import attention_mesh_mlo_2d_kernel -from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config -from kernbench.ccl.sfr_config import configure_sfr_intercube_multisip -from kernbench.policy.placement.dp import DPPolicy -from kernbench.runtime_api.bench_runner import run_bench -from kernbench.runtime_api.types import DeviceSelector, resolve_device -from kernbench.sim_engine.engine import GraphEngine -from kernbench.topology.builder import resolve_topology - -TOPOLOGY_4SIP = ( - Path(__file__).resolve().parents[2] / "topologies" / "llama70b_4sip.yaml" -) -TOPOLOGY_DEFAULT = Path(__file__).resolve().parents[2] / "topology.yaml" - -S_Q_DECODE = 1 -S_KV_PER_RANK = 16 -H_Q = 1 -H_KV = 1 -D_HEAD = 64 -# 2×4 sub-mesh per KV-group (study: 8 cubes per KV-group at Q/cube=1). -MESH_ROWS = 2 -MESH_COLS = 4 -N_CUBES_PER_KV_GROUP = MESH_ROWS * MESH_COLS -DTYPE = "f16" - - -def _ccl_cfg(): - return resolve_algorithm_config( - load_ccl_config(), name="lrab_hierarchical_allreduce", - ) - - -def _engine_factory(t, d): - return GraphEngine(getattr(t, "topology_obj", t), enable_data=True) - - -def _make_one_kv_group_bench(mesh_rows: int, mesh_cols: int): - """Return a bench_fn that runs ONE multi_user_decode kernel on a - ``mesh_rows × mesh_cols`` sub-mesh.""" - n_cubes = mesh_rows * mesh_cols - - def _bench_fn(ctx): - configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg()) - dp_full = DPPolicy(cube="replicate", pe="replicate", - num_cubes=n_cubes, num_pes=8) - dp_kv = DPPolicy(cube="row_wise", pe="replicate", - num_cubes=n_cubes, num_pes=8) - q = ctx.zeros((S_Q_DECODE, H_Q * D_HEAD), - dtype=DTYPE, dp=dp_full, name="q") - k = ctx.zeros((S_KV_PER_RANK * n_cubes, H_KV * D_HEAD), - dtype=DTYPE, dp=dp_kv, name="k") - v = ctx.zeros((S_KV_PER_RANK * n_cubes, H_KV * D_HEAD), - dtype=DTYPE, dp=dp_kv, name="v") - o = ctx.empty((S_Q_DECODE, H_Q * D_HEAD), - dtype=DTYPE, dp=dp_full, name="o") - ctx.launch( - f"single_kv_group_{mesh_rows}x{mesh_cols}", - attention_mesh_mlo_2d_kernel, - q, k, v, o, - S_Q_DECODE, S_KV_PER_RANK, H_Q, H_KV, D_HEAD, - mesh_rows, mesh_cols, - 1, # rank_axis=1 → cube-level ring - 0, # cube_start=0 — single sub-mesh launch - _auto_dim_remap=False, - ) - - return _bench_fn - - -def _run_one_kv_group(topology_path: Path, mesh_rows: int, mesh_cols: int, - target_device=None): - topo = resolve_topology(str(topology_path)) - captured: dict = {"engine": None} - - def factory(t, d): - eng = _engine_factory(t, d) - captured["engine"] = eng - return eng - - exc = None - result = None - try: - result = run_bench( - topology=topo, - bench_fn=_make_one_kv_group_bench(mesh_rows, mesh_cols), - device=target_device or resolve_device(None), - engine_factory=factory, - ) - except BaseException as e: # noqa: BLE001 - exc = e - return exc, result, captured["engine"] - - -# ── Step 1 — single KV-group at the study's full breadth ────────── - - -def test_step_1_single_kv_group_at_full_breadth(): - """One multi_user_decode launch on a 2×4 sub-mesh, 4-SIP topology. - - Uses the C2 2D row-then-col AllReduce-mlo kernel: stage 1 reduces - across cols (E/W) within each row, stage 2 reduces across rows (N/S). - Expected to PASS — N/S edges are wired by - ``configure_sfr_intercube_multisip`` and the 2D fan-out avoids the - row-boundary IpcqInvalidDirection that the 1D kernel hit at cube 4. - """ - if not TOPOLOGY_4SIP.exists(): - pytest.skip(f"4-SIP topology missing: {TOPOLOGY_4SIP}") - exc, result, engine = _run_one_kv_group( - TOPOLOGY_4SIP, mesh_rows=MESH_ROWS, mesh_cols=MESH_COLS, - ) - if exc is not None: - oplog_len = len(getattr(engine, "op_log", []) or []) if engine else 0 - print(f"\nstep_1 FAIL — op_log records before crash: {oplog_len}") - traceback.print_exception(type(exc), exc, exc.__traceback__) - raise AssertionError(f"step_1 failed: {exc}") from exc - assert result is not None and result.completion.ok, ( - f"step_1: completion not ok — {result.completion if result else None}" - ) - - -# ── Step 2 — 4 KV-groups, one per SIP, sequential launches ──────── - - -def _make_multi_sip_bench_fn(sip_groups: list[tuple[int, str, int]]): - """One bench_fn that does one 2×4 multi_user_decode launch per item. - - Each ``(sip, tag, cube_start)`` tuple becomes one launch: - - ``ctx.ahbm.set_device(sip)`` switches allocations to that SIP - (mirrors ``milestone_1h_ccl.py:283-292``). - - ``cube_start`` selects which 8-cube sub-mesh within the SIP: - ``0`` → cubes 0..7 (rows 0..1), ``8`` → cubes 8..15 (rows 2..3). - - ``tag`` disambiguates tensor names so launches in the same - run_bench don't collide on the allocator namespace. - """ - n_cubes = MESH_ROWS * MESH_COLS - - def _bench_fn(ctx): - configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg()) - for sip, tag, cube_start in sip_groups: - ctx.ahbm.set_device(sip) - dp_full = DPPolicy(cube="replicate", pe="replicate", - num_cubes=n_cubes, num_pes=8, - cube_start=cube_start) - dp_kv = DPPolicy(cube="row_wise", pe="replicate", - num_cubes=n_cubes, num_pes=8, - cube_start=cube_start) - q = ctx.zeros((S_Q_DECODE, H_Q * D_HEAD), - dtype=DTYPE, dp=dp_full, name=f"q_{tag}") - k = ctx.zeros((S_KV_PER_RANK * n_cubes, H_KV * D_HEAD), - dtype=DTYPE, dp=dp_kv, name=f"k_{tag}") - v = ctx.zeros((S_KV_PER_RANK * n_cubes, H_KV * D_HEAD), - dtype=DTYPE, dp=dp_kv, name=f"v_{tag}") - o = ctx.empty((S_Q_DECODE, H_Q * D_HEAD), - dtype=DTYPE, dp=dp_full, name=f"o_{tag}") - ctx.launch( - f"kv_group_{tag}", attention_mesh_mlo_2d_kernel, - q, k, v, o, - S_Q_DECODE, S_KV_PER_RANK, H_Q, H_KV, D_HEAD, - MESH_ROWS, MESH_COLS, - 1, # rank_axis=1 → cube-level ring - cube_start, # converts physical id → launch-local rank - _auto_dim_remap=False, - ) - - return _bench_fn - - -def _run_multi_sip(sip_groups: list[tuple[int, str, int]]): - """Run a single run_bench call covering all (sip, tag) groups.""" - topo = resolve_topology(str(TOPOLOGY_4SIP)) - captured: dict = {"engine": None} - - def factory(t, d): - eng = _engine_factory(t, d) - captured["engine"] = eng - return eng - - exc = None - result = None - try: - result = run_bench( - topology=topo, - bench_fn=_make_multi_sip_bench_fn(sip_groups), - device=resolve_device(None), # "all" SIPs in scope - engine_factory=factory, - ) - except BaseException as e: # noqa: BLE001 - exc = e - return exc, result, captured["engine"] - - -def test_step_2_four_kv_groups_one_per_sip(): - """Four multi_user_decode launches, one per SIP, in ONE run_bench call. - - Uses the CCL milestone pattern (``milestone_1h_ccl.py:283-292``): - ``target_device="all"`` scopes the runtime to every SIP; then - ``ctx.ahbm.set_device(sip)`` before each ``ctx.zeros``/``launch`` - switches which SIP the next allocation+launch lands on. This is the - canonical sequential per-SIP pattern in the codebase — four separate - ``run_bench`` calls with ``DeviceSelector("sip:N")`` is a misuse. - - Expected to PASS — the SFR install draws intra-SIP edges only, so the - 4 KV-groups can't see each other. - """ - if not TOPOLOGY_4SIP.exists(): - pytest.skip(f"4-SIP topology missing: {TOPOLOGY_4SIP}") - sip_groups = [(sip, f"sip{sip}", 0) for sip in range(4)] - exc, result, engine = _run_multi_sip(sip_groups) - if exc is not None: - oplog_len = len(getattr(engine, "op_log", []) or []) if engine else 0 - print(f"\nstep_2 FAIL — op_log records before crash: {oplog_len}") - traceback.print_exception(type(exc), exc, exc.__traceback__) - raise AssertionError(f"step_2 failed: {exc}") from exc - assert result is not None and result.completion.ok, ( - f"step_2: completion not ok — {result.completion if result else None}" - ) - - -# ── Step 3 — 8 KV-groups, two per SIP (study target) ────────────── - - -def test_step_3_eight_kv_groups_two_per_sip(): - """Two launches per SIP × 4 SIPs = 8 KV-groups total in one run_bench. - - The headline target: 64 cubes serving 8 KV-groups, two disjoint 2×4 - sub-meshes per SIP. ``cube_start=0`` puts the first KV-group on - cubes 0..7 (rows 0..1); ``cube_start=8`` puts the second on cubes - 8..15 (rows 2..3). This is the use case ``DPPolicy.cube_start`` was - added to enable. - - Expected to PASS — the SFR install draws intra-SIP edges only, so - the 8 KV-groups can't see each other; ``cube_start`` ensures the - two halves of each SIP land on disjoint cubes. - """ - if not TOPOLOGY_4SIP.exists(): - pytest.skip(f"4-SIP topology missing: {TOPOLOGY_4SIP}") - sip_groups = [ - (sip, f"sip{sip}_half{half}", half * N_CUBES_PER_KV_GROUP) - for sip in range(4) for half in (0, 1) - ] - exc, result, engine = _run_multi_sip(sip_groups) - if exc is not None: - raise AssertionError(f"step_3 failed: {exc}") from exc - assert result is not None and result.completion.ok, ( - f"step_3: completion not ok — {result.completion if result else None}" - ) diff --git a/tests/attention/test_attention_mesh_panels_diag.py b/tests/attention/test_attention_mesh_panels_diag.py deleted file mode 100644 index 706071d..0000000 --- a/tests/attention/test_attention_mesh_panels_diag.py +++ /dev/null @@ -1,198 +0,0 @@ -"""End-to-end engine drives for the four GQA Llama-70B panels (sub-cycle 4c step 2). - -Mirrors the existing single_user_decode diag harness across all four panels -of the milestone-gqa-llama70b sweep (ADR-0057): - - single_user_prefill ring-K/V kernel, intracube PE ring (8 PEs / 1 cube) - single_user_decode allreduce-mlo kernel, intracube PE ring - multi_user_prefill ring-K/V kernel, intercube multisip (4 cubes) - multi_user_decode allreduce-mlo kernel, intercube multisip - -Each test runs the panel through ``run_bench`` with ``enable_data=True`` -and asserts ``result.completion.ok``. Failures dump the engine's op_log -tail and the exception, mirroring the decode-diag harness format. - -Validation-scale config matches ADR-0057 D4: - S_q_prefill=16, S_kv_per_rank=16, h_q=h_kv=1, d_head=64 - n_ranks_single_user=8, n_ranks_multi_user=4 -""" -from __future__ import annotations - -import traceback -from pathlib import Path - -import pytest - -from kernbench.benches._attention_mesh_kv import attention_mesh_kv_kernel -from kernbench.benches._attention_mesh_mlo import attention_mesh_mlo_kernel -from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config -from kernbench.ccl.sfr_config import ( - configure_sfr_intercube_multisip, - configure_sfr_intracube_pe_ring, -) -from kernbench.policy.placement.dp import DPPolicy -from kernbench.runtime_api.bench_runner import run_bench -from kernbench.runtime_api.types import resolve_device -from kernbench.sim_engine.engine import GraphEngine -from kernbench.topology.builder import resolve_topology - -TOPOLOGY_PATH = Path(__file__).resolve().parents[2] / "topology.yaml" - -S_Q_PREFILL = 16 -S_Q_DECODE = 1 -S_KV_PER_RANK = 16 -H_Q = 1 -H_KV = 1 -D_HEAD = 64 -N_RANKS_SINGLE_USER = 8 -N_RANKS_MULTI_USER = 4 -DTYPE = "f16" - - -# ── Helpers ────────────────────────────────────────────────────── - - -def _engine_factory(t, d): - return GraphEngine(getattr(t, "topology_obj", t), enable_data=True) - - -def _run_panel(bench_fn): - """Drive a panel through run_bench; return (exc, result, engine).""" - topo = resolve_topology(str(TOPOLOGY_PATH)) - captured: dict = {"engine": None} - - def factory(t, d): - eng = _engine_factory(t, d) - captured["engine"] = eng - return eng - - exc = None - result = None - try: - result = run_bench( - topology=topo, bench_fn=bench_fn, - device=resolve_device(None), engine_factory=factory, - ) - except BaseException as e: # noqa: BLE001 - exc = e - return exc, result, captured["engine"] - - -def _assert_ok(name: str, exc, result, engine) -> None: - if exc is not None: - oplog_len = len(getattr(engine, "op_log", []) or []) if engine else 0 - print(f"\n========== {name} FAIL ==========") - print(f"op_log records before crash: {oplog_len}") - print(f"{type(exc).__name__}: {exc}") - traceback.print_exception(type(exc), exc, exc.__traceback__) - raise AssertionError( - f"{name} failed at runtime: {exc}" - ) from exc - assert result is not None, f"{name}: no result" - assert result.completion.ok, f"{name}: completion not ok — {result.completion}" - - -# ── Panel bench fns ────────────────────────────────────────────── - - -def _bench_fn_single_user_prefill(ctx): - configure_sfr_intracube_pe_ring( - ctx.engine, ctx.spec, - resolve_algorithm_config(load_ccl_config(), name="lrab_hierarchical_allreduce"), - ) - n = N_RANKS_SINGLE_USER - dp_full = DPPolicy(cube="replicate", pe="replicate", num_cubes=1, num_pes=n) - dp_kv = DPPolicy(cube="replicate", pe="row_wise", num_cubes=1, num_pes=n) - q = ctx.zeros((S_Q_PREFILL, H_Q * D_HEAD), dtype=DTYPE, dp=dp_full, name="q") - k = ctx.zeros((S_KV_PER_RANK * n, H_KV * D_HEAD), dtype=DTYPE, dp=dp_kv, name="k") - v = ctx.zeros((S_KV_PER_RANK * n, H_KV * D_HEAD), dtype=DTYPE, dp=dp_kv, name="v") - o = ctx.empty((S_Q_PREFILL, H_Q * D_HEAD), dtype=DTYPE, dp=dp_full, name="o") - ctx.launch( - "single_user_prefill_mesh", attention_mesh_kv_kernel, - q, k, v, o, - S_Q_PREFILL, S_KV_PER_RANK, H_Q, H_KV, D_HEAD, n, - ) - - -def _bench_fn_single_user_decode(ctx): - configure_sfr_intracube_pe_ring( - ctx.engine, ctx.spec, - resolve_algorithm_config(load_ccl_config(), name="lrab_hierarchical_allreduce"), - ) - n = N_RANKS_SINGLE_USER - dp_full = DPPolicy(cube="replicate", pe="replicate", num_cubes=1, num_pes=n) - dp_kv = DPPolicy(cube="replicate", pe="row_wise", num_cubes=1, num_pes=n) - q = ctx.zeros((S_Q_DECODE, H_Q * D_HEAD), dtype=DTYPE, dp=dp_full, name="q") - k = ctx.zeros((S_KV_PER_RANK * n, H_KV * D_HEAD), dtype=DTYPE, dp=dp_kv, name="k") - v = ctx.zeros((S_KV_PER_RANK * n, H_KV * D_HEAD), dtype=DTYPE, dp=dp_kv, name="v") - o = ctx.empty((S_Q_DECODE, H_Q * D_HEAD), dtype=DTYPE, dp=dp_full, name="o") - ctx.launch( - "single_user_decode_mesh", attention_mesh_mlo_kernel, - q, k, v, o, - S_Q_DECODE, S_KV_PER_RANK, H_Q, H_KV, D_HEAD, n, - ) - - -def _bench_fn_multi_user_prefill(ctx): - configure_sfr_intercube_multisip( - ctx.engine, ctx.spec, - resolve_algorithm_config(load_ccl_config(), name="lrab_hierarchical_allreduce"), - ) - n = N_RANKS_MULTI_USER - dp_full = DPPolicy(cube="replicate", pe="replicate", num_cubes=n, num_pes=8) - dp_kv = DPPolicy(cube="row_wise", pe="replicate", num_cubes=n, num_pes=8) - q = ctx.zeros((S_Q_PREFILL, H_Q * D_HEAD), dtype=DTYPE, dp=dp_full, name="q") - k = ctx.zeros((S_KV_PER_RANK * n, H_KV * D_HEAD), dtype=DTYPE, dp=dp_kv, name="k") - v = ctx.zeros((S_KV_PER_RANK * n, H_KV * D_HEAD), dtype=DTYPE, dp=dp_kv, name="v") - o = ctx.empty((S_Q_PREFILL, H_Q * D_HEAD), dtype=DTYPE, dp=dp_full, name="o") - ctx.launch( - "multi_user_prefill_mesh", attention_mesh_kv_kernel, - q, k, v, o, - S_Q_PREFILL, S_KV_PER_RANK, H_Q, H_KV, D_HEAD, n, - 1, # rank_axis=1 → ring at cube level (ADR-0059 multi_user) - _auto_dim_remap=False, - ) - - -def _bench_fn_multi_user_decode(ctx): - configure_sfr_intercube_multisip( - ctx.engine, ctx.spec, - resolve_algorithm_config(load_ccl_config(), name="lrab_hierarchical_allreduce"), - ) - n = N_RANKS_MULTI_USER - dp_full = DPPolicy(cube="replicate", pe="replicate", num_cubes=n, num_pes=8) - dp_kv = DPPolicy(cube="row_wise", pe="replicate", num_cubes=n, num_pes=8) - q = ctx.zeros((S_Q_DECODE, H_Q * D_HEAD), dtype=DTYPE, dp=dp_full, name="q") - k = ctx.zeros((S_KV_PER_RANK * n, H_KV * D_HEAD), dtype=DTYPE, dp=dp_kv, name="k") - v = ctx.zeros((S_KV_PER_RANK * n, H_KV * D_HEAD), dtype=DTYPE, dp=dp_kv, name="v") - o = ctx.empty((S_Q_DECODE, H_Q * D_HEAD), dtype=DTYPE, dp=dp_full, name="o") - ctx.launch( - "multi_user_decode_mesh", attention_mesh_mlo_kernel, - q, k, v, o, - S_Q_DECODE, S_KV_PER_RANK, H_Q, H_KV, D_HEAD, n, - 1, # rank_axis=1 → ring at cube level (ADR-0059 multi_user) - _auto_dim_remap=False, - ) - - -# ── Tests ──────────────────────────────────────────────────────── - - -def test_single_user_prefill_through_engine(): - exc, result, engine = _run_panel(_bench_fn_single_user_prefill) - _assert_ok("single_user_prefill", exc, result, engine) - - -def test_single_user_decode_through_engine(): - exc, result, engine = _run_panel(_bench_fn_single_user_decode) - _assert_ok("single_user_decode", exc, result, engine) - - -def test_multi_user_prefill_through_engine(): - exc, result, engine = _run_panel(_bench_fn_multi_user_prefill) - _assert_ok("multi_user_prefill", exc, result, engine) - - -def test_multi_user_decode_through_engine(): - exc, result, engine = _run_panel(_bench_fn_multi_user_decode) - _assert_ok("multi_user_decode", exc, result, engine) diff --git a/tests/attention/test_gqa_decode.py b/tests/attention/test_gqa_decode.py new file mode 100644 index 0000000..1db00db --- /dev/null +++ b/tests/attention/test_gqa_decode.py @@ -0,0 +1,143 @@ +"""Phase 1 spec test for P1a GQA decode kernel (real GQA via M-fold). + +P1a is the first phase of the DDD-0060 plan, split out of the original +P1 (the composite-hybrid swap is P1b, deferred until the tl.composite +output-handle question is decided). P1a is the *correctness* unlock: +the kernel processes ONE KV head at a time and folds the G query heads +into the matmul M (row) dimension so a single Q·Kᵀ serves all G heads +sharing one K (ADR-0060 §5.2). This lifts the baseline's +``h_q == h_kv == 1`` cap pinned at +``tests/attention/test_milestone_gqa_llama70b.py:137-142``. + +P1a stays inside the existing ``tl`` API: the two attention GEMMs use the +blocking ``tl.dot`` so the chain ``Q·Kᵀ → softmax → P·V → store`` fits +without any composite-output chaining. The composite swap (P1b) will +revisit this once the API for feeding a composite's output into a +downstream MATH op is settled. + +Phase 1 (this commit): tests only — production code lands in Phase 2. +Tests fail at import in Phase 1 with ModuleNotFoundError; Phase 2 makes +them pass. +""" +from __future__ import annotations + +from pathlib import Path + +from kernbench.benches._gqa_decode import gqa_decode_kernel # noqa: F401 (Phase 2 deliverable) +from kernbench.policy.placement.dp import DPPolicy +from kernbench.runtime_api.bench_runner import run_bench +from kernbench.runtime_api.types import resolve_device +from kernbench.sim_engine.engine import GraphEngine +from kernbench.topology.builder import resolve_topology + +TOPOLOGY_DEFAULT = Path(__file__).resolve().parents[2] / "topology.yaml" + +# Decode shapes — P1a is one-shot, no tiling, single rank. P3 will tile. +T_Q = 1 +D_HEAD = 64 +S_KV = 16 +DTYPE = "f16" + + +def _engine_factory(t, d): + return GraphEngine(getattr(t, "topology_obj", t), enable_data=True) + + +def _run_decode_p1(*, h_q: int, h_kv: int): + """One-shot GQA decode on a single PE in a single CUBE (P=1, no SP). + + Tensor layout (natural K — same as the P2a SP tests; the kernel + reshapes K to (d_head, S_local) via byte-conserving load): + Q : (T_q, h_q · d_head) — natural Q layout; kernel reshapes to + (G·T_q, d_head) — byte-conserving and math-correct because + T_q=1 collapses axis ordering. + K : (S_kv, h_kv · d_head) — natural K layout; kernel loads as + (d_head, S_local) — reshape-as-transpose caveat (ADR-0060 §3 + / §B item 2), correct for zero inputs used here. + V : (S_kv, h_kv · d_head) — natural V layout. + O : (T_q, h_q · d_head) — same shape as Q. + """ + topo = resolve_topology(str(TOPOLOGY_DEFAULT)) + + def _bench_fn(ctx): + dp = DPPolicy(cube="replicate", pe="replicate", + num_cubes=1, num_pes=1) + q = ctx.zeros((T_Q, h_q * D_HEAD), + dtype=DTYPE, dp=dp, name=f"q_h{h_q}_kv{h_kv}") + k = ctx.zeros((S_KV, h_kv * D_HEAD), + dtype=DTYPE, dp=dp, name=f"k_h{h_q}_kv{h_kv}") + v = ctx.zeros((S_KV, h_kv * D_HEAD), + dtype=DTYPE, dp=dp, name=f"v_h{h_q}_kv{h_kv}") + o = ctx.empty((T_Q, h_q * D_HEAD), + dtype=DTYPE, dp=dp, name=f"o_h{h_q}_kv{h_kv}") + ctx.launch( + f"gqa_decode_p1_h{h_q}_kv{h_kv}", + gqa_decode_kernel, + q, k, v, o, + T_Q, S_KV, h_q, h_kv, D_HEAD, + 1, 1, # C=1, P=1 (no SP, degenerate) + _auto_dim_remap=False, + ) + + return run_bench( + topology=topo, + bench_fn=_bench_fn, + device=resolve_device(None), + engine_factory=_engine_factory, + ) + + +def _dma_read_count(op_log) -> int: + return sum(1 for r in op_log if r.op_name == "dma_read") + + +# ── Headline unlock: real GQA (h_q = G·h_kv) runs in data mode ────────── + + +def test_real_gqa_h_q_eight_h_kv_one_completes_in_data_mode(): + """ADR-0060 §A.1 headline unlock — the baseline raises + ``ValueError: Shape mismatch …`` in MemoryStore at h_q=8, h_kv=1 + because ``_view(K, (h_q·d, S_kv))`` only byte-conserves when h_q==h_kv. + M-fold processes one KV head at a time using only byte-conserving + reshapes, so this completes. + """ + result = _run_decode_p1(h_q=8, h_kv=1) + assert result.completion.ok, ( + f"real GQA (h_q=8, h_kv=1) decode failed: {result.completion}" + ) + + +# ── M-fold property: K/V loads do not scale with G ───────────────────── + + +def test_kv_dma_read_count_independent_of_g(): + """ADR-0060 TL;DR / §5.2: M-fold loads K and V once per KV head and + folds the G query heads into the GEMM M dim. The dma_read_count must + therefore be identical between (G=1, h_kv=1) and (G=8, h_kv=1) — both + issue exactly 3 reads (Q + K + V). This pins the GQA-reuse property + that the rest of the plan (composite streaming in P4, etc.) builds on. + """ + g1 = _run_decode_p1(h_q=1, h_kv=1) + g8 = _run_decode_p1(h_q=8, h_kv=1) + n_g1 = _dma_read_count(g1.engine.op_log) + n_g8 = _dma_read_count(g8.engine.op_log) + assert n_g1 == 3, f"G=1 dma_read_count must be 3 (Q+K+V); got {n_g1}" + assert n_g8 == 3, f"G=8 dma_read_count must be 3 (Q+K+V); got {n_g8}" + assert n_g1 == n_g8, ( + f"K/V dma_read_count must be independent of G; " + f"got G=1 -> {n_g1}, G=8 -> {n_g8}" + ) + + +# ── Backward-compat: degenerate G=1 still works ──────────────────────── + + +def test_degenerate_g_equals_one_still_works(): + """G=1 (h_q == h_kv == 1) is the baseline-compatible config. M-fold + degenerates to (T_q, d) = (1, 64) — the same shape the baseline + already exercises — so this proves no regression on that path. + """ + result = _run_decode_p1(h_q=1, h_kv=1) + assert result.completion.ok, ( + f"degenerate G=1 decode failed: {result.completion}" + ) diff --git a/tests/attention/test_gqa_decode_mc.py b/tests/attention/test_gqa_decode_mc.py new file mode 100644 index 0000000..9873677 --- /dev/null +++ b/tests/attention/test_gqa_decode_mc.py @@ -0,0 +1,182 @@ +"""Phase 1 spec test for P2b GQA decode multi-cube SP (both Level-2 + Level-1). + +P2b extends P2a to multiple CUBEs in one CUBE Group. The kernel uses the +canonical full SFR install (``configure_sfr_intercube_multisip``) which +provides disjoint direction namespaces: + + - ``intra_E / intra_W / intra_N / intra_S`` — PE↔PE within a CUBE + (logical 2×4 grid, no wrap) + - ``E / W / N / S`` — CUBE↔CUBE inter-CUBE + (mesh, no wrap) + +Reduce pattern (chain reduce-to-root, ADR-0060 §A.2 spirit, §4 chain +deviation noted): + + Level-2 (intra-CUBE, 2×4 grid): + row-then-col chain — each row reduces leftward along ``intra_W`` to + its col-0 PE, then PE 4 sends to PE 0 along ``intra_N``. 7 chain + steps per CUBE × 3 handles each = 21 ``ipcq_copy`` per CUBE. + + Level-1 (inter-CUBE): + only PE 0 of each CUBE participates. Chain leftward along ``W``. + (C-1) chain steps × 3 handles each. + + Final store at PE 0 of CUBE 0 only. + +Phase 1 (this commit): tests only — production code lands in Phase 2. +Phase 2 also updates ``test_gqa_decode.py`` (add C=1) and +``test_gqa_decode_sp.py`` (switch SFR + add C=1). +""" +from __future__ import annotations + +from pathlib import Path + +from kernbench.benches._gqa_decode import gqa_decode_kernel # noqa: F401 +from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config +from kernbench.ccl.sfr_config import configure_sfr_intercube_multisip +from kernbench.policy.placement.dp import DPPolicy +from kernbench.runtime_api.bench_runner import run_bench +from kernbench.runtime_api.types import resolve_device +from kernbench.sim_engine.engine import GraphEngine +from kernbench.topology.builder import resolve_topology + +TOPOLOGY_DEFAULT = Path(__file__).resolve().parents[2] / "topology.yaml" + +T_Q = 1 +D_HEAD = 64 +DTYPE = "f16" + + +def _ccl_cfg(): + return resolve_algorithm_config( + load_ccl_config(), name="lrab_hierarchical_allreduce", + ) + + +def _engine_factory(t, d): + return GraphEngine(getattr(t, "topology_obj", t), enable_data=True) + + +def _run_decode_mc(*, h_q: int, h_kv: int, C: int, P: int, S_kv: int): + """Multi-CUBE SP decode: C cubes × P PEs each share the work.""" + topo = resolve_topology(str(TOPOLOGY_DEFAULT)) + + def _bench_fn(ctx): + configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg()) + dp_full = DPPolicy(cube="replicate", pe="replicate", + num_cubes=C, num_pes=P) + dp_kv = DPPolicy(cube="row_wise", pe="row_wise", + num_cubes=C, num_pes=P) + # Total KV split across C×P ranks; each rank sees S_kv/(C·P) rows. + q = ctx.zeros((T_Q, h_q * D_HEAD), + dtype=DTYPE, dp=dp_full, + name=f"q_h{h_q}_kv{h_kv}_c{C}_p{P}") + k = ctx.zeros((S_kv, h_kv * D_HEAD), + dtype=DTYPE, dp=dp_kv, + name=f"k_h{h_q}_kv{h_kv}_c{C}_p{P}") + v = ctx.zeros((S_kv, h_kv * D_HEAD), + dtype=DTYPE, dp=dp_kv, + name=f"v_h{h_q}_kv{h_kv}_c{C}_p{P}") + o = ctx.empty((T_Q, h_q * D_HEAD), + dtype=DTYPE, dp=dp_full, + name=f"o_h{h_q}_kv{h_kv}_c{C}_p{P}") + ctx.launch( + f"gqa_decode_mc_h{h_q}_kv{h_kv}_c{C}_p{P}", + gqa_decode_kernel, + q, k, v, o, + T_Q, S_kv, h_q, h_kv, D_HEAD, + C, P, + _auto_dim_remap=False, + ) + + return run_bench( + topology=topo, + bench_fn=_bench_fn, + device=resolve_device(None), + engine_factory=_engine_factory, + ) + + +def _count(op_log, name: str) -> int: + return sum(1 for r in op_log if r.op_name == name) + + +# ── Degenerate C=1 P=1 ──────────────────────────────────────────────── + + +def test_mc_c_one_p_one_degenerate(): + """C=1, P=1: single rank, no SP. No IPCQ traffic; one dma_write.""" + result = _run_decode_mc(h_q=1, h_kv=1, C=1, P=1, S_kv=16) + assert result.completion.ok, ( + f"C=1 P=1 degenerate failed: {result.completion}" + ) + assert _count(result.engine.op_log, "ipcq_copy") == 0 + assert _count(result.engine.op_log, "dma_write") == 1 + + +# ── Intra-CUBE only (single CUBE, P=8 on 2×4 grid) ──────────────────── + + +def test_mc_c_one_p_eight_intracube_grid(): + """C=1, P=8: intra-cube row-then-col chain on the 2×4 grid. + 7 chain steps × 3 handles (m, ℓ, O) = 21 ipcq_copy.""" + result = _run_decode_mc(h_q=1, h_kv=1, C=1, P=8, S_kv=64) + assert result.completion.ok, ( + f"C=1 P=8 intra-cube failed: {result.completion}" + ) + assert _count(result.engine.op_log, "dma_write") == 1 + n_copy = _count(result.engine.op_log, "ipcq_copy") + assert n_copy == 21, ( + f"C=1 P=8: expected 21 ipcq_copy (7 chain × 3 handles); got {n_copy}" + ) + + +# ── Multi-CUBE root-only write ──────────────────────────────────────── + + +def test_mc_c_two_p_eight_root_only_writes_o(): + """C=2, P=8: 16 ranks total. Only PE 0 of CUBE 0 writes O.""" + result = _run_decode_mc(h_q=1, h_kv=1, C=2, P=8, S_kv=128) + assert result.completion.ok, ( + f"C=2 P=8 multi-cube failed: {result.completion}" + ) + n_writes = _count(result.engine.op_log, "dma_write") + assert n_writes == 1, ( + f"root-only write must hold for C=2 P=8 (16 ranks); got {n_writes}" + ) + + +# ── Multi-CUBE total IPCQ chain count ───────────────────────────────── + + +def test_mc_c_two_p_eight_total_ipcq_count(): + """C=2, P=8: 21 intra-cube ipcq_copy per CUBE × 2 CUBEs + 3 inter-cube + chain ipcq_copy (C-1=1 step × 3 handles) = 45 total.""" + result = _run_decode_mc(h_q=1, h_kv=1, C=2, P=8, S_kv=128) + assert result.completion.ok, ( + f"C=2 P=8 multi-cube failed: {result.completion}" + ) + n_copy = _count(result.engine.op_log, "ipcq_copy") + expected = 21 * 2 + (2 - 1) * 3 + assert n_copy == expected, ( + f"C=2 P=8: expected {expected} ipcq_copy (21 intra × 2 CUBEs + 3 " + f"inter); got {n_copy}" + ) + + +# ── Real GQA × multi-CUBE SP combined ───────────────────────────────── + + +def test_mc_real_gqa_c_two_p_eight(): + """Headline: real GQA (h_q = G·h_kv with G=8) AND multi-CUBE SP + (C=2, P=8) together — the case the original baseline can express + neither part of.""" + result = _run_decode_mc(h_q=8, h_kv=1, C=2, P=8, S_kv=128) + assert result.completion.ok, ( + f"real GQA + multi-CUBE SP combined failed: {result.completion}" + ) + n_writes = _count(result.engine.op_log, "dma_write") + assert n_writes == 1, ( + f"root-only write must hold under M-fold + multi-CUBE; " + f"got {n_writes}" + ) diff --git a/tests/attention/test_gqa_decode_sp.py b/tests/attention/test_gqa_decode_sp.py new file mode 100644 index 0000000..25a321c --- /dev/null +++ b/tests/attention/test_gqa_decode_sp.py @@ -0,0 +1,153 @@ +"""Phase 1 spec test for P2a GQA decode SP (chain reduce-to-root, Level-2 only). + +P2a is the first half of DDD-0060 P2: the kernel becomes multi-PE within +one CUBE and reduces to root (PE 0) using a chain over the 1D intra-cube +ring (W direction). This **replaces the baseline's bidirectional O(N) +fan-out** where every rank ends with O — ADR-0060 §A.2's headline. + +Deviation from DDD-0060 §7 P2 gate: the gate text asks for +``⌈log₂ P⌉`` reduce rounds. The intra-cube SFR install +(``configure_sfr_intracube_pe_ring``) wires only a 1D E/W ring, so a +true tree would require either multi-hop forwarding or a new SFR install +(future ADR). P2a uses a **chain reduce-to-root**: ``P-1`` rounds along +W. The architectural property the ADR cares about +(root-only output vs every-rank-has-O) is preserved; the logarithmic +collective is deferred. + +P2b (deferred) covers Level-1 inter-CUBE center-mesh reduce (C>1). + +Phase 1 (this commit): tests only — production code lands in Phase 2. +""" +from __future__ import annotations + +from pathlib import Path + +from kernbench.benches._gqa_decode import gqa_decode_kernel # noqa: F401 +from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config +from kernbench.ccl.sfr_config import configure_sfr_intercube_multisip +from kernbench.policy.placement.dp import DPPolicy +from kernbench.runtime_api.bench_runner import run_bench +from kernbench.runtime_api.types import resolve_device +from kernbench.sim_engine.engine import GraphEngine +from kernbench.topology.builder import resolve_topology + +TOPOLOGY_DEFAULT = Path(__file__).resolve().parents[2] / "topology.yaml" + +T_Q = 1 +D_HEAD = 64 +DTYPE = "f16" + + +def _ccl_cfg(): + return resolve_algorithm_config( + load_ccl_config(), name="lrab_hierarchical_allreduce", + ) + + +def _engine_factory(t, d): + return GraphEngine(getattr(t, "topology_obj", t), enable_data=True) + + +def _run_decode_sp(*, h_q: int, h_kv: int, P: int, S_kv: int): + """Single-CUBE SP decode: P PEs share the work along the intra-cube ring.""" + topo = resolve_topology(str(TOPOLOGY_DEFAULT)) + + def _bench_fn(ctx): + configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg()) + dp_full = DPPolicy(cube="replicate", pe="replicate", + num_cubes=1, num_pes=P) + dp_kv = DPPolicy(cube="replicate", pe="row_wise", + num_cubes=1, num_pes=P) + q = ctx.zeros((T_Q, h_q * D_HEAD), + dtype=DTYPE, dp=dp_full, name=f"q_h{h_q}_kv{h_kv}_p{P}") + # KV: total S_kv split across P PEs along axis 0 (row_wise sharding). + # Each PE sees (S_kv/P, h_kv·D_HEAD). + k = ctx.zeros((S_kv, h_kv * D_HEAD), + dtype=DTYPE, dp=dp_kv, name=f"k_h{h_q}_kv{h_kv}_p{P}") + v = ctx.zeros((S_kv, h_kv * D_HEAD), + dtype=DTYPE, dp=dp_kv, name=f"v_h{h_q}_kv{h_kv}_p{P}") + o = ctx.empty((T_Q, h_q * D_HEAD), + dtype=DTYPE, dp=dp_full, name=f"o_h{h_q}_kv{h_kv}_p{P}") + ctx.launch( + f"gqa_decode_sp_h{h_q}_kv{h_kv}_p{P}", + gqa_decode_kernel, + q, k, v, o, + T_Q, S_kv, h_q, h_kv, D_HEAD, + 1, P, # C=1, P=P (single-CUBE SP) + _auto_dim_remap=False, + ) + + return run_bench( + topology=topo, + bench_fn=_bench_fn, + device=resolve_device(None), + engine_factory=_engine_factory, + ) + + +def _count(op_log, name: str) -> int: + return sum(1 for r in op_log if r.op_name == name) + + +# ── Root-only write ──────────────────────────────────────────────────── + + +def test_sp_chain_reduce_root_only_writes_o(): + """ADR-0060 §A.2: only the root rank (PE 0) writes O. Baseline today + has every rank write the full final O (bidirectional fan-out).""" + result = _run_decode_sp(h_q=1, h_kv=1, P=8, S_kv=64) + assert result.completion.ok, f"P=8 chain reduce failed: {result.completion}" + n_writes = _count(result.engine.op_log, "dma_write") + assert n_writes == 1, ( + f"reduce-to-root must produce exactly 1 dma_write (PE 0); " + f"got {n_writes}" + ) + + +# ── Chain step count ─────────────────────────────────────────────────── + + +def test_sp_chain_reduce_p_minus_one_ipcq_pairs(): + """Chain reduce-to-root has P-1 send→recv pairs along the W chain; + each pair logs one ``ipcq_copy`` (inbound DMA, per + ``milestone_gqa_llama70b._summarize_op_log``). Each chain step ships + the triplet (m, ℓ, O) → 3 handles per step → 7 steps × 3 = 21.""" + result = _run_decode_sp(h_q=1, h_kv=1, P=8, S_kv=64) + assert result.completion.ok, f"P=8 chain reduce failed: {result.completion}" + n_copy = _count(result.engine.op_log, "ipcq_copy") + expected = (8 - 1) * 3 + assert n_copy == expected, ( + f"chain reduce: expected {expected} ipcq_copy (P-1=7 steps × " + f"3 handles m/ℓ/O); got {n_copy}" + ) + + +# ── Real GQA × SP combined ───────────────────────────────────────────── + + +def test_sp_real_gqa_h_q_eight_h_kv_one_p_eight(): + """The combined unlock: real GQA (h_q=G·h_kv with G=8) AND SP + (P=8) together — neither expressible by the baseline.""" + result = _run_decode_sp(h_q=8, h_kv=1, P=8, S_kv=64) + assert result.completion.ok, ( + f"real GQA + SP combined run failed: {result.completion}" + ) + n_writes = _count(result.engine.op_log, "dma_write") + assert n_writes == 1, ( + f"root-only write must hold under M-fold too; got {n_writes}" + ) + + +# ── Degenerate P=1 ──────────────────────────────────────────────────── + + +def test_sp_p_one_degenerate_no_ipcq_traffic(): + """P=1: SP degenerates to a single rank. No IPCQ traffic; one dma_write.""" + result = _run_decode_sp(h_q=8, h_kv=1, P=1, S_kv=16) + assert result.completion.ok, f"P=1 degenerate failed: {result.completion}" + n_send = _count(result.engine.op_log, "ipcq_send") + n_recv = _count(result.engine.op_log, "ipcq_recv") + assert n_send == 0, f"P=1 must have no ipcq_send; got {n_send}" + assert n_recv == 0, f"P=1 must have no ipcq_recv; got {n_recv}" + n_writes = _count(result.engine.op_log, "dma_write") + assert n_writes == 1, f"P=1: one dma_write; got {n_writes}" diff --git a/tests/attention/test_gqa_long_context.py b/tests/attention/test_gqa_long_context.py new file mode 100644 index 0000000..291dd1e --- /dev/null +++ b/tests/attention/test_gqa_long_context.py @@ -0,0 +1,157 @@ +"""Phase 1 spec test for Phase C: long-context regression for the GQA +kernels (ADR-0060 §B "long/short context split" + ADR-0063 §A.2). + +The headline panel today caps prefill at S_kv=16 / decode at S_kv≤128 — +NOT because the algorithm fails, but because the bump allocator never +recycles. ADR-0063 §A.2 (test req 3) requires a sweep at an S that +would overflow 1 MiB without recycling to complete after scope +discipline lands. + +Scratch-budget estimate at C=4, P=8 (current 1 MiB pool): + decode: per-rank S_local = S_kv / 32; intermediates ≲ 500 KB at + S_kv=32K (fits today — used as regression guard). + prefill: per-rank S_local = S_kv / 4; ~16·S_local bytes per ring + step × 4 steps ≈ 64·S_local bytes. Overflows at + ~64 K tokens (64 × 16K > 1 MiB). + +Phase 1 (this commit): tests only — production code lands in Phase 2. +The prefill test fails today with a RuntimeError("TLContext scratch +overflow"). After Phase 2 (scratch_scope + copy_to discipline) it +completes. +""" +from __future__ import annotations + +from pathlib import Path + +from kernbench.benches._gqa_decode import gqa_decode_kernel # noqa: F401 +from kernbench.benches._gqa_prefill import gqa_prefill_kernel # noqa: F401 +from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config +from kernbench.ccl.sfr_config import ( + configure_sfr_intercube_multisip, + configure_sfr_intercube_ring, +) +from kernbench.policy.placement.dp import DPPolicy +from kernbench.runtime_api.bench_runner import run_bench +from kernbench.runtime_api.types import resolve_device +from kernbench.sim_engine.engine import GraphEngine +from kernbench.topology.builder import resolve_topology + +TOPOLOGY_DEFAULT = Path(__file__).resolve().parents[2] / "topology.yaml" + +D_HEAD = 64 +DTYPE = "f16" + + +def _ccl_cfg(): + return resolve_algorithm_config( + load_ccl_config(), name="lrab_hierarchical_allreduce", + ) + + +def _engine_factory(t, d): + return GraphEngine(getattr(t, "topology_obj", t), enable_data=True) + + +# ── Decode at moderate-long context (regression guard) ─────────────── + + +def test_decode_long_context_32k_completes(): + """Decode at S_kv=32K (C=1, P=8) — per-rank S_local=4K. Should + complete with current scratch usage (~few KB intermediates) and + must continue to complete after the rewrite. + + This is a regression guard: scratch discipline shouldn't break + decode's existing long-context capability. + """ + topo = resolve_topology(str(TOPOLOGY_DEFAULT)) + + def _bench_fn(ctx): + configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg()) + P = 8 + S_kv = 32_768 + dp_full = DPPolicy(cube="replicate", pe="replicate", + num_cubes=1, num_pes=P) + dp_kv = DPPolicy(cube="replicate", pe="row_wise", + num_cubes=1, num_pes=P) + ctx.zeros((1, 8 * D_HEAD), dtype=DTYPE, dp=dp_full, name="q_long_dec") + k = ctx.zeros((S_kv, D_HEAD), + dtype=DTYPE, dp=dp_kv, name="k_long_dec") + v = ctx.zeros((S_kv, D_HEAD), + dtype=DTYPE, dp=dp_kv, name="v_long_dec") + o = ctx.empty((1, 8 * D_HEAD), + dtype=DTYPE, dp=dp_full, name="o_long_dec") + q = ctx.zeros((1, 8 * D_HEAD), + dtype=DTYPE, dp=dp_full, name="q_long_dec_2") + ctx.launch( + "gqa_decode_long_32k", + gqa_decode_kernel, + q, k, v, o, + 1, S_kv, 8, 1, D_HEAD, + 1, P, + _auto_dim_remap=False, + ) + + result = run_bench( + topology=topo, bench_fn=_bench_fn, + device=resolve_device(None), engine_factory=_engine_factory, + ) + assert result.completion.ok, ( + f"decode at S_kv=32K must complete; got {result.completion}" + ) + + +# ── Prefill at long context overflows without scope discipline ─────── + + +def test_prefill_long_context_completes_after_scope_discipline(): + """ADR-0063 §A.2 Test Req 3: a sweep at an ``S`` that would overflow + 1 MiB without recycling must complete after scope discipline lands. + + With C=4 and S_kv chosen so per-rank S_local·8·4 > 1 MiB + (~16K tokens per rank ⇒ S_kv ≥ 64K), the current kernel overflows + the 1 MiB scratch pool. After Phase 2 (scratch_scope wraps each + ring step's intermediates; copy_to persists running state), it + completes. + + This is the headline test that proves the S-ceiling is gone for + prefill. + """ + topo = resolve_topology(str(TOPOLOGY_DEFAULT)) + + def _bench_fn(ctx): + configure_sfr_intercube_ring( + ctx.engine, ctx.spec, _ccl_cfg(), ring_size=4, + ) + T_q = 4 + S_kv = 65_536 # 64 K — per-rank 16 K, ~2 MB scratch w/o scope + C = 4 + dp_q = DPPolicy(cube="replicate", pe="replicate", + num_cubes=C, num_pes=1) + dp_kv = DPPolicy(cube="row_wise", pe="replicate", + num_cubes=C, num_pes=1) + dp_o = DPPolicy(cube="row_wise", pe="replicate", + num_cubes=C, num_pes=1) + q = ctx.zeros((T_q, D_HEAD), + dtype=DTYPE, dp=dp_q, name="q_long_pre") + k = ctx.zeros((S_kv, D_HEAD), + dtype=DTYPE, dp=dp_kv, name="k_long_pre") + v = ctx.zeros((S_kv, D_HEAD), + dtype=DTYPE, dp=dp_kv, name="v_long_pre") + o = ctx.empty((T_q * C, D_HEAD), + dtype=DTYPE, dp=dp_o, name="o_long_pre") + ctx.launch( + "gqa_prefill_long_64k", + gqa_prefill_kernel, + q, k, v, o, + T_q, S_kv, D_HEAD, C, + _auto_dim_remap=False, + ) + + result = run_bench( + topology=topo, bench_fn=_bench_fn, + device=resolve_device(None), engine_factory=_engine_factory, + ) + assert result.completion.ok, ( + f"prefill at S_kv=64K must complete after scope discipline lands; " + f"got {result.completion}" + ) diff --git a/tests/attention/test_gqa_prefill.py b/tests/attention/test_gqa_prefill.py new file mode 100644 index 0000000..e086b55 --- /dev/null +++ b/tests/attention/test_gqa_prefill.py @@ -0,0 +1,118 @@ +"""Phase 1 spec test for P6a GQA prefill kernel (head-parallel, C=1 baseline). + +P6a introduces ``_gqa_prefill.py`` with the head-parallel structure (one +Q head per CUBE, per-CUBE distributed output, no reduce). C=1 is the +degenerate case — no Ring KV, no IPCQ traffic. Validates kernel +structure and T_q > 1 attention. + +P6b (deferred) adds the Ring KV rotation for C > 1, which needs either +a new SFR install function (intra_* + wrapped E/W at CUBE level) or a +topology-specific config — separate design call. + +The prefill kernel differs from decode (P1a/P2a/P2b) in three ways +(ADR-0060 §5.5 / TL;DR): + 1. Q has T_q > 1 rows (not just decode's single timestep). + 2. Head-parallel placement: each CUBE owns ONE Q head — no M-fold. + 3. Each CUBE writes its own head's output — NO reduce. + +Phase 1 (this commit): tests only — production code lands in Phase 2. +""" +from __future__ import annotations + +from pathlib import Path + +from kernbench.benches._gqa_prefill import gqa_prefill_kernel # noqa: F401 +from kernbench.policy.placement.dp import DPPolicy +from kernbench.runtime_api.bench_runner import run_bench +from kernbench.runtime_api.types import resolve_device +from kernbench.sim_engine.engine import GraphEngine +from kernbench.topology.builder import resolve_topology + +TOPOLOGY_DEFAULT = Path(__file__).resolve().parents[2] / "topology.yaml" + +D_HEAD = 64 +DTYPE = "f16" + + +def _engine_factory(t, d): + return GraphEngine(getattr(t, "topology_obj", t), enable_data=True) + + +def _run_prefill(*, T_q: int, S_kv: int, C: int = 1): + """C=1 head-parallel prefill: single CUBE owns the one head + full KV.""" + topo = resolve_topology(str(TOPOLOGY_DEFAULT)) + + def _bench_fn(ctx): + dp = DPPolicy(cube="replicate", pe="replicate", + num_cubes=C, num_pes=1) + # Q: (T_q, d_head) — one head per CUBE (head-parallel; for C=1 + # only one head total). 2D layout matches what the kernel loads. + # P6b will use a 3D (h_q, T_q, d_head) Q with cube_row_wise + # sharding so each CUBE owns its head. + q = ctx.zeros((T_q, D_HEAD), dtype=DTYPE, dp=dp, + name=f"q_t{T_q}_c{C}") + # K, V: full local for C=1 (no ring). Kernel loads K as + # (d_head, S_kv) via byte-conserving reshape. + k = ctx.zeros((S_kv, D_HEAD), dtype=DTYPE, dp=dp, + name=f"k_t{T_q}_c{C}") + v = ctx.zeros((S_kv, D_HEAD), dtype=DTYPE, dp=dp, + name=f"v_t{T_q}_c{C}") + # O: (T_q, d_head) — per-CUBE distributed output. + o = ctx.empty((T_q, D_HEAD), dtype=DTYPE, dp=dp, + name=f"o_t{T_q}_c{C}") + ctx.launch( + f"gqa_prefill_p6a_t{T_q}_s{S_kv}_c{C}", + gqa_prefill_kernel, + q, k, v, o, + T_q, S_kv, D_HEAD, C, + _auto_dim_remap=False, + ) + + return run_bench( + topology=topo, bench_fn=_bench_fn, + device=resolve_device(None), + engine_factory=_engine_factory, + ) + + +def _count(op_log, name: str) -> int: + return sum(1 for r in op_log if r.op_name == name) + + +def test_prefill_c_one_t_q_one_completes(): + """C=1, T_q=1: smallest workload (decode-like).""" + result = _run_prefill(T_q=1, S_kv=16, C=1) + assert result.completion.ok, ( + f"prefill C=1 T_q=1 failed: {result.completion}" + ) + + +def test_prefill_c_one_t_q_four_completes(): + """C=1, T_q=4: real prefill (Q has multiple rows) — distinguishes + prefill from decode (T_q=1).""" + result = _run_prefill(T_q=4, S_kv=16, C=1) + assert result.completion.ok, ( + f"prefill C=1 T_q=4 failed: {result.completion}" + ) + + +def test_prefill_c_one_no_ipcq_traffic(): + """C=1: no ring step, no IPCQ traffic.""" + result = _run_prefill(T_q=4, S_kv=16, C=1) + assert result.completion.ok + n_copy = _count(result.engine.op_log, "ipcq_copy") + assert n_copy == 0, ( + f"C=1 must have no IPCQ traffic (no ring); got {n_copy}" + ) + + +def test_prefill_c_one_one_dma_write(): + """C=1, one head: exactly one dma_write (per-CUBE distributed output; + no reduce). For C > 1 in P6b this becomes dma_write_count == C.""" + result = _run_prefill(T_q=4, S_kv=16, C=1) + assert result.completion.ok + n_writes = _count(result.engine.op_log, "dma_write") + assert n_writes == 1, ( + f"C=1 prefill: expected 1 dma_write (one head per cube); " + f"got {n_writes}" + ) diff --git a/tests/attention/test_gqa_prefill_ring.py b/tests/attention/test_gqa_prefill_ring.py new file mode 100644 index 0000000..9594ee4 --- /dev/null +++ b/tests/attention/test_gqa_prefill_ring.py @@ -0,0 +1,161 @@ +"""Phase 1 spec test for P6b GQA prefill Ring KV (head-parallel, C>1). + +P6b adds the Ring KV rotation (ADR-0060 §5.5) to the prefill kernel. +Each CUBE owns one Q head + its KV slice; over C ring steps the KV +blocks rotate around the C CUBEs so every CUBE sees every block. The +running ``(m, ℓ, O)`` is folded inside the loop. No reduce — each CUBE +writes its own head's output. + +Requires a new SFR install ``configure_sfr_intercube_ring(ring_size=C)`` +that wires: + - ``intra_*`` : 2×4 PE grid within a CUBE (same as multisip) + - ``E/W`` : 1D ring of CUBEs 0..ring_size-1 WITH WRAP + (symmetric to ``configure_sfr_intracube_pe_ring``, + applied at CUBE level) + - ``global_*`` : SIP-level (same as multisip) + +The kernel ring body: + for step in range(1, C): + tl.send(dir="W", src=Kc) + tl.send(dir="W", src=Vc) + Kc = tl.recv(dir="E", ...) + Vc = tl.recv(dir="E", ...) + ... local partial + online-softmax merge into (m, ℓ, O) ... + +Per CUBE per step: 2 sends (K, V) → 2 ``ipcq_copy``. Across all CUBEs: +``(C-1) * 2 * C`` total ipcq_copy. + +Restriction in P6b first cut: ``C ∈ {1, 2, 4}`` (single row of the 4×4 +cube mesh). C=8 ring would span rows — follow-on. + +Phase 1 (this commit): tests only — production code lands in Phase 2. +""" +from __future__ import annotations + +from pathlib import Path + +from kernbench.benches._gqa_prefill import gqa_prefill_kernel # noqa: F401 +from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config +from kernbench.ccl.sfr_config import configure_sfr_intercube_ring # noqa: F401 (Phase 2) +from kernbench.policy.placement.dp import DPPolicy +from kernbench.runtime_api.bench_runner import run_bench +from kernbench.runtime_api.types import resolve_device +from kernbench.sim_engine.engine import GraphEngine +from kernbench.topology.builder import resolve_topology + +TOPOLOGY_DEFAULT = Path(__file__).resolve().parents[2] / "topology.yaml" + +D_HEAD = 64 +DTYPE = "f16" + + +def _ccl_cfg(): + return resolve_algorithm_config( + load_ccl_config(), name="lrab_hierarchical_allreduce", + ) + + +def _engine_factory(t, d): + return GraphEngine(getattr(t, "topology_obj", t), enable_data=True) + + +def _run_prefill_ring(*, T_q: int, S_kv: int, C: int): + """Head-parallel prefill with Ring KV across C CUBEs.""" + topo = resolve_topology(str(TOPOLOGY_DEFAULT)) + + def _bench_fn(ctx): + # P6b: new SFR install with cube-level ring wrap. + configure_sfr_intercube_ring( + ctx.engine, ctx.spec, _ccl_cfg(), ring_size=C, + ) + # Q replicated on every CUBE (zeros for testing; per-CUBE head + # indexing is implicit). KV sequence-sharded by CUBE. O + # distributed — each CUBE writes its slice of (T_q*C, d_head). + dp_q = DPPolicy(cube="replicate", pe="replicate", + num_cubes=C, num_pes=1) + dp_kv = DPPolicy(cube="row_wise", pe="replicate", + num_cubes=C, num_pes=1) + dp_o = DPPolicy(cube="row_wise", pe="replicate", + num_cubes=C, num_pes=1) + q = ctx.zeros((T_q, D_HEAD), dtype=DTYPE, dp=dp_q, + name=f"q_t{T_q}_c{C}_ring") + k = ctx.zeros((S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv, + name=f"k_t{T_q}_c{C}_ring") + v = ctx.zeros((S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv, + name=f"v_t{T_q}_c{C}_ring") + # O: (T_q * C, d_head), each CUBE writes (T_q, d_head) at its + # slice. dma_write_count = C (per-CUBE distributed output). + o = ctx.empty((T_q * C, D_HEAD), dtype=DTYPE, dp=dp_o, + name=f"o_t{T_q}_c{C}_ring") + ctx.launch( + f"gqa_prefill_ring_t{T_q}_s{S_kv}_c{C}", + gqa_prefill_kernel, + q, k, v, o, + T_q, S_kv, D_HEAD, C, + _auto_dim_remap=False, + ) + + return run_bench( + topology=topo, bench_fn=_bench_fn, + device=resolve_device(None), + engine_factory=_engine_factory, + ) + + +def _count(op_log, name: str) -> int: + return sum(1 for r in op_log if r.op_name == name) + + +# ── C=2 Ring KV ──────────────────────────────────────────────────────── + + +def test_prefill_ring_c_two_completes(): + """C=2: 1 ring step rotates KV between the 2 CUBEs.""" + result = _run_prefill_ring(T_q=4, S_kv=16, C=2) + assert result.completion.ok, ( + f"prefill ring C=2 failed: {result.completion}" + ) + + +def test_prefill_ring_c_two_distributed_output(): + """C=2: per-CUBE distributed output, no reduce → dma_write_count == 2.""" + result = _run_prefill_ring(T_q=4, S_kv=16, C=2) + assert result.completion.ok + n_writes = _count(result.engine.op_log, "dma_write") + assert n_writes == 2, ( + f"C=2 prefill: expected 2 dma_write (one per CUBE); got {n_writes}" + ) + + +def test_prefill_ring_c_two_ipcq_count(): + """C=2: 1 ring step × 2 handles (K, V) × 2 CUBEs = 4 ipcq_copy.""" + result = _run_prefill_ring(T_q=4, S_kv=16, C=2) + assert result.completion.ok + n_copy = _count(result.engine.op_log, "ipcq_copy") + expected = (2 - 1) * 2 * 2 + assert n_copy == expected, ( + f"C=2 ring: expected {expected} ipcq_copy " + f"((C-1)·2·C = 1·2·2); got {n_copy}" + ) + + +# ── C=4 Ring KV (combined assertions) ───────────────────────────────── + + +def test_prefill_ring_c_four_combined(): + """C=4: 3 ring steps rotate KV around 4 CUBEs. + Expected: completes; 4 dma_writes; (4-1)·2·4 = 24 ipcq_copy.""" + result = _run_prefill_ring(T_q=4, S_kv=32, C=4) + assert result.completion.ok, ( + f"prefill ring C=4 failed: {result.completion}" + ) + n_writes = _count(result.engine.op_log, "dma_write") + assert n_writes == 4, ( + f"C=4 prefill: expected 4 dma_write (one per CUBE); got {n_writes}" + ) + n_copy = _count(result.engine.op_log, "ipcq_copy") + expected = (4 - 1) * 2 * 4 + assert n_copy == expected, ( + f"C=4 ring: expected {expected} ipcq_copy " + f"((C-1)·2·C = 3·2·4); got {n_copy}" + ) diff --git a/tests/attention/test_gqa_scoped_writeback.py b/tests/attention/test_gqa_scoped_writeback.py new file mode 100644 index 0000000..0da33cc --- /dev/null +++ b/tests/attention/test_gqa_scoped_writeback.py @@ -0,0 +1,198 @@ +"""Phase 1 spec test for Phase C: scratch_scope + tl.copy_to discipline +in the GQA kernels (ADR-0060 §5.2 / §5.5 + ADR-0063 §D3 / §D3.1). + +ADR-0060 §5.2 (decode pseudocode line 75 / §5.5 (prefill pseudocode line +96) both wrap per-tile / per-ring-step intermediates in +``with tl.scratch_scope():`` and persist the merged running ``(m, ℓ, O)`` +to a persistent arena allocated outside the scope. The original ADR-0063 +§D3 specifies the two-arena pattern; §D3.1 specifies the +``tl.copy_to(dst, src)`` writeback primitive used to persist scoped +results. + +Currently neither kernel uses ``scratch_scope`` or ``copy_to``; their +chain-merge / ring-merge bodies allocate every intermediate from the +bump cursor and never recycle. Result: op_log has 0 ``copy`` entries. + +After Phase 2: each per-tile / per-step merge writes the new running +``(m, ℓ, O)`` via ``copy_to`` to the persistent arena. Per merge step +→ 3 copy ops (m, ℓ, O). + +Phase 1 (this commit): tests only — production code lands in Phase 2. +""" +from __future__ import annotations + +from pathlib import Path + +from kernbench.benches._gqa_decode import gqa_decode_kernel # noqa: F401 +from kernbench.benches._gqa_prefill import gqa_prefill_kernel # noqa: F401 +from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config +from kernbench.ccl.sfr_config import ( + configure_sfr_intercube_multisip, + configure_sfr_intercube_ring, +) +from kernbench.policy.placement.dp import DPPolicy +from kernbench.runtime_api.bench_runner import run_bench +from kernbench.runtime_api.types import resolve_device +from kernbench.sim_engine.engine import GraphEngine +from kernbench.topology.builder import resolve_topology + +TOPOLOGY_DEFAULT = Path(__file__).resolve().parents[2] / "topology.yaml" + +D_HEAD = 64 +DTYPE = "f16" + + +def _ccl_cfg(): + return resolve_algorithm_config( + load_ccl_config(), name="lrab_hierarchical_allreduce", + ) + + +def _engine_factory(t, d): + return GraphEngine(getattr(t, "topology_obj", t), enable_data=True) + + +def _count(op_log, name: str) -> int: + return sum(1 for r in op_log if r.op_name == name) + + +# ── Decode chain merges must use scratch_scope + copy_to ───────────── + + +def _run_decode_sp(*, h_q: int, h_kv: int, P: int, S_kv: int): + """Single-CUBE SP decode (C=1, P PEs along intra-cube chain).""" + topo = resolve_topology(str(TOPOLOGY_DEFAULT)) + + def _bench_fn(ctx): + configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg()) + dp_full = DPPolicy(cube="replicate", pe="replicate", + num_cubes=1, num_pes=P) + dp_kv = DPPolicy(cube="replicate", pe="row_wise", + num_cubes=1, num_pes=P) + q = ctx.zeros((1, h_q * D_HEAD), + dtype=DTYPE, dp=dp_full, name=f"q_sc_{P}") + k = ctx.zeros((S_kv, h_kv * D_HEAD), + dtype=DTYPE, dp=dp_kv, name=f"k_sc_{P}") + v = ctx.zeros((S_kv, h_kv * D_HEAD), + dtype=DTYPE, dp=dp_kv, name=f"v_sc_{P}") + o = ctx.empty((1, h_q * D_HEAD), + dtype=DTYPE, dp=dp_full, name=f"o_sc_{P}") + ctx.launch( + f"gqa_decode_scoped_{P}", + gqa_decode_kernel, + q, k, v, o, + 1, S_kv, h_q, h_kv, D_HEAD, + 1, P, + _auto_dim_remap=False, + ) + + return run_bench( + topology=topo, bench_fn=_bench_fn, + device=resolve_device(None), engine_factory=_engine_factory, + ) + + +def test_decode_chain_merges_emit_copy_to_writeback(): + """ADR-0060 §5.2 + ADR-0063 §D3.1: each chain-merge step must wrap + its intermediates in ``scratch_scope`` and persist the new running + ``(m, ℓ, O)`` via ``tl.copy_to``. + + For (C=1, P=8): 7 intra-cube chain merges × 3 handles (m, ℓ, O) per + merge ⇒ 21 ``copy`` entries. + + Currently 0 because the kernel never calls ``tl.copy_to``. + """ + result = _run_decode_sp(h_q=8, h_kv=1, P=8, S_kv=64) + assert result.completion.ok, f"decode SP failed: {result.completion}" + n_copy = _count(result.engine.op_log, "copy") + assert n_copy > 0, ( + f"decode kernel must emit copy_to writeback per merge step " + f"(ADR-0060 §5.2 + ADR-0063 §D3.1); got 0 ``copy`` entries" + ) + + +# ── Prefill Ring KV merges must use scratch_scope + copy_to ────────── + + +def _run_prefill_ring(*, T_q: int, S_kv: int, C: int): + """Head-parallel prefill with Ring KV across C CUBEs.""" + topo = resolve_topology(str(TOPOLOGY_DEFAULT)) + + def _bench_fn(ctx): + configure_sfr_intercube_ring( + ctx.engine, ctx.spec, _ccl_cfg(), ring_size=C, + ) + dp_q = DPPolicy(cube="replicate", pe="replicate", + num_cubes=C, num_pes=1) + dp_kv = DPPolicy(cube="row_wise" if C > 1 else "replicate", + pe="replicate", num_cubes=C, num_pes=1) + dp_o = DPPolicy(cube="row_wise" if C > 1 else "replicate", + pe="replicate", num_cubes=C, num_pes=1) + q = ctx.zeros((T_q, D_HEAD), + dtype=DTYPE, dp=dp_q, name=f"q_ring_{C}") + k = ctx.zeros((S_kv, D_HEAD), + dtype=DTYPE, dp=dp_kv, name=f"k_ring_{C}") + v = ctx.zeros((S_kv, D_HEAD), + dtype=DTYPE, dp=dp_kv, name=f"v_ring_{C}") + o = ctx.empty((T_q * C, D_HEAD), + dtype=DTYPE, dp=dp_o, name=f"o_ring_{C}") + ctx.launch( + f"gqa_prefill_scoped_{C}", + gqa_prefill_kernel, + q, k, v, o, + T_q, S_kv, D_HEAD, C, + _auto_dim_remap=False, + ) + + return run_bench( + topology=topo, bench_fn=_bench_fn, + device=resolve_device(None), engine_factory=_engine_factory, + ) + + +def test_prefill_ring_step_merges_emit_copy_to_writeback(): + """ADR-0060 §5.5 + ADR-0063 §D3.1: each Ring KV step's online-softmax + merge must wrap its intermediates in ``scratch_scope`` and persist + the new running ``(m, ℓ, O)`` via ``tl.copy_to``. + + For C=4: 3 ring-step merges (steps 1..C-1) × 3 handles (m, ℓ, O) per + merge ⇒ 9 ``copy`` entries per participating CUBE. Aggregated across + C CUBEs: ⇒ 36 ``copy`` entries. + + Currently 0 because the kernel never calls ``tl.copy_to``. + """ + result = _run_prefill_ring(T_q=4, S_kv=16, C=4) + assert result.completion.ok, f"prefill ring failed: {result.completion}" + n_copy = _count(result.engine.op_log, "copy") + assert n_copy > 0, ( + f"prefill ring kernel must emit copy_to writeback per merge step " + f"(ADR-0060 §5.5 + ADR-0063 §D3.1); got 0 ``copy`` entries" + ) + + +# ── Scoped kernels still produce the same op_log shape (regression) ── + + +def test_decode_scoped_still_has_root_only_write(): + """ADR-0060 §A.2 root-only output must hold under the rewrite: + adding scratch_scope + copy_to should not change the reduce + topology; only the per-PE scratch usage. Single PE 0 writes O.""" + result = _run_decode_sp(h_q=8, h_kv=1, P=8, S_kv=64) + assert result.completion.ok + n_writes = _count(result.engine.op_log, "dma_write") + assert n_writes == 1, ( + f"scoped decode must still produce 1 dma_write (root-only); " + f"got {n_writes}" + ) + + +def test_prefill_scoped_still_has_per_cube_distributed_output(): + """ADR-0060 §5.5 per-CUBE distributed output must hold under the + rewrite: scoped prefill still writes one O slice per CUBE.""" + result = _run_prefill_ring(T_q=4, S_kv=16, C=4) + assert result.completion.ok + n_writes = _count(result.engine.op_log, "dma_write") + assert n_writes == 4, ( + f"scoped prefill must write one O per CUBE (C=4 → 4 dma_write); " + f"got {n_writes}" + ) diff --git a/tests/attention/test_gqa_short_context.py b/tests/attention/test_gqa_short_context.py new file mode 100644 index 0000000..ca99277 --- /dev/null +++ b/tests/attention/test_gqa_short_context.py @@ -0,0 +1,260 @@ +"""Phase 1 spec test for Phase D: short-context GQA kernels +(ADR-0060 §B "Items from the long/short context split", item B.split.2). + +The long-context kernels (§5.2 decode / §5.5 prefill) shard each KV +head row-wise across all CUBEs. At short context (S_kv < 256K), that +shard is too thin to feed the engines and the cube-level collective +overhead dominates. The short-context kernels drop cube-SP entirely: +each CUBE owns ``kv_per_cube`` *whole* KV heads, no S_kv sharding across +CUBEs. + +Design (per AskUserQuestion answers in this session): + - PE-parallel heads: P PEs split into ``kv_per_cube`` groups, each + group does PE-SP across (P/kv_per_cube) PEs for one owned head. + - scratch_scope + tl.copy_to discipline mirrors the long kernels. + - One short kernel handles kv_per_cube ∈ {1, 2, 4, 8} via a parameter. + +Group layout on the 2×4 PE grid: + kv_per_cube=1, group=8 PEs (full 2×4): row chain + col bridge. + kv_per_cube=2, group=4 PEs (one row): row chain only. + kv_per_cube=4, group=2 PEs (adj cols): 1-step chain. + kv_per_cube=8, group=1 PE: no chain. + +After chain reduce-to-group-root, the group's root PE writes its +owned head's output. No inter-CUBE reduce. + +Phase 1 (this commit): tests only — production code lands in Phase 2. +All tests fail because the short kernels (``_gqa_decode_short.py`` and +``_gqa_prefill_short.py``) do not exist yet. +""" +from __future__ import annotations + +from pathlib import Path + +import pytest + +from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config +from kernbench.ccl.sfr_config import configure_sfr_intercube_multisip +from kernbench.policy.placement.dp import DPPolicy +from kernbench.runtime_api.bench_runner import run_bench +from kernbench.runtime_api.types import resolve_device +from kernbench.sim_engine.engine import GraphEngine +from kernbench.topology.builder import resolve_topology + +TOPOLOGY_DEFAULT = Path(__file__).resolve().parents[2] / "topology.yaml" + +D_HEAD = 64 +DTYPE = "f16" + + +def _ccl_cfg(): + return resolve_algorithm_config( + load_ccl_config(), name="lrab_hierarchical_allreduce", + ) + + +def _engine_factory(t, d): + return GraphEngine(getattr(t, "topology_obj", t), enable_data=True) + + +def _count(op_log, name: str) -> int: + return sum(1 for r in op_log if r.op_name == name) + + +# ── Decode short-context kernel ────────────────────────────────────── + + +def _run_decode_short(*, h_q: int, h_kv: int, kv_per_cube: int, + C: int, P: int, S_kv: int): + """Run the short-context decode kernel with PE-parallel heads. + + Layout (after design iteration — see Phase D failure-recovery): + Q: (T_q, h_q·D_HEAD) replicated; kernel reshapes byte-conservingly. + K, V: (h_kv·S_kv, D_HEAD) head-stacked, ``cube=row_wise, pe=row_wise`` + so each PE gets a contiguous (S_local, D_HEAD) chunk at its own + addressable shard (no partial reads needed). + O: replicated; each group root writes the full byte-conserving + (h_q·T_q, D_HEAD) result. + """ + from kernbench.benches._gqa_decode_short import gqa_decode_short_kernel # Phase 2 + + topo = resolve_topology(str(TOPOLOGY_DEFAULT)) + + def _bench_fn(ctx): + configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg()) + dp_full = DPPolicy(cube="replicate", pe="replicate", + num_cubes=C, num_pes=P) + # Head-stacked KV with row_wise sharding so each PE's chunk is + # contiguous and exactly (S_local, D_HEAD) addressable. + dp_kv = DPPolicy(cube="row_wise", pe="row_wise", + num_cubes=C, num_pes=P) + q = ctx.zeros((1, h_q * D_HEAD), + dtype=DTYPE, dp=dp_full, name=f"q_short_{kv_per_cube}_{C}") + k = ctx.zeros((h_kv * S_kv, D_HEAD), + dtype=DTYPE, dp=dp_kv, name=f"k_short_{kv_per_cube}_{C}") + v = ctx.zeros((h_kv * S_kv, D_HEAD), + dtype=DTYPE, dp=dp_kv, name=f"v_short_{kv_per_cube}_{C}") + o = ctx.empty((1, h_q * D_HEAD), + dtype=DTYPE, dp=dp_full, name=f"o_short_{kv_per_cube}_{C}") + ctx.launch( + f"gqa_decode_short_{kv_per_cube}_{C}", + gqa_decode_short_kernel, + q, k, v, o, + 1, S_kv, h_q, h_kv, D_HEAD, C, P, kv_per_cube, + _auto_dim_remap=False, + ) + + return run_bench( + topology=topo, bench_fn=_bench_fn, + device=resolve_device(None), engine_factory=_engine_factory, + ) + + +def test_short_decode_smoke_kv_per_cube_2_C_4(): + """ADR-0060 §B.split.2 headline: kv_per_cube=2, C=4 — each CUBE + owns 2 heads (8 heads / 4 CUBEs). PE-parallel heads splits the 8 + PEs into 2 groups of 4, each group does PE-SP for one owned head. + Smoke: kernel completes.""" + result = _run_decode_short( + h_q=8, h_kv=8, kv_per_cube=2, C=4, P=8, S_kv=64, + ) + assert result.completion.ok, ( + f"short decode kv_per_cube=2 must complete; got {result.completion}" + ) + + +def test_short_decode_smoke_kv_per_cube_4_C_2(): + """kv_per_cube=4, C=2 — half the CUBEs participate, each owns 4 + heads, 4 PE groups of 2 PEs each.""" + result = _run_decode_short( + h_q=8, h_kv=8, kv_per_cube=4, C=2, P=8, S_kv=64, + ) + assert result.completion.ok, ( + f"short decode kv_per_cube=4 must complete; got {result.completion}" + ) + + +def test_short_decode_smoke_kv_per_cube_8_C_1(): + """kv_per_cube=8, C=1 — all heads on one CUBE. 8 PE groups of 1 PE + each → no chain reduce, each PE writes its head's output.""" + result = _run_decode_short( + h_q=8, h_kv=8, kv_per_cube=8, C=1, P=8, S_kv=64, + ) + assert result.completion.ok, ( + f"short decode kv_per_cube=8 must complete; got {result.completion}" + ) + + +def test_short_decode_dma_write_count_equals_h_kv(): + """ADR-0060 §B.split.2: each owned head produces exactly one + output (no inter-CUBE reduce). Total dma_writes = h_kv across all + participating CUBEs and groups. + + For h_kv=8, kv_per_cube=2, C=4: each CUBE writes 2 outputs → + 4 × 2 = 8 dma_writes total. + """ + result = _run_decode_short( + h_q=8, h_kv=8, kv_per_cube=2, C=4, P=8, S_kv=64, + ) + assert result.completion.ok + n_writes = _count(result.engine.op_log, "dma_write") + assert n_writes == 8, ( + f"short decode: expected 8 dma_writes (h_kv); got {n_writes}" + ) + + +def test_short_decode_no_inter_cube_traffic(): + """ADR-0060 §B.split.2: each head is fully owned by one CUBE → no + inter-CUBE reduce. The kernel must not invoke CUBE-level E/W IPCQ. + + Today's long-context kernel at C=4 emits ~12 inter-CUBE ipcq_copy + via direction "E"/"W". The short kernel must emit zero of those, + keeping all IPCQ traffic on the ``intra_*`` directions. + """ + result = _run_decode_short( + h_q=8, h_kv=8, kv_per_cube=2, C=4, P=8, S_kv=64, + ) + assert result.completion.ok + # Count IPCQ traffic that targeted the CUBE-level "E"/"W" directions. + inter_cube_ipcq = sum( + 1 for r in result.engine.op_log + if r.op_name == "ipcq_copy" + and r.params.get("direction") in ("E", "W") + ) + assert inter_cube_ipcq == 0, ( + f"short decode must have no inter-CUBE E/W IPCQ; got {inter_cube_ipcq}" + ) + + +# ── Prefill short-context kernel ───────────────────────────────────── + + +def _run_prefill_short(*, h_kv: int, kv_per_cube: int, + C: int, P: int, T_q: int, S_kv: int): + """Run the short-context prefill kernel with PE-parallel heads. + + Layout: same head-stacked K/V scheme as decode short, with Q/O + replicated and byte-conserving reshape inside the kernel. + """ + from kernbench.benches._gqa_prefill_short import gqa_prefill_short_kernel # Phase 2 + + topo = resolve_topology(str(TOPOLOGY_DEFAULT)) + + def _bench_fn(ctx): + configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg()) + dp_q = DPPolicy(cube="replicate", pe="replicate", + num_cubes=C, num_pes=P) + dp_kv = DPPolicy(cube="row_wise", pe="row_wise", + num_cubes=C, num_pes=P) + dp_o = DPPolicy(cube="replicate", pe="replicate", + num_cubes=C, num_pes=P) + q = ctx.zeros((T_q, h_kv * D_HEAD), + dtype=DTYPE, dp=dp_q, name=f"q_pre_short_{kv_per_cube}_{C}") + k = ctx.zeros((h_kv * S_kv, D_HEAD), + dtype=DTYPE, dp=dp_kv, name=f"k_pre_short_{kv_per_cube}_{C}") + v = ctx.zeros((h_kv * S_kv, D_HEAD), + dtype=DTYPE, dp=dp_kv, name=f"v_pre_short_{kv_per_cube}_{C}") + o = ctx.empty((T_q, h_kv * D_HEAD), + dtype=DTYPE, dp=dp_o, name=f"o_pre_short_{kv_per_cube}_{C}") + ctx.launch( + f"gqa_prefill_short_{kv_per_cube}_{C}", + gqa_prefill_short_kernel, + q, k, v, o, + T_q, S_kv, h_kv, D_HEAD, C, P, kv_per_cube, + _auto_dim_remap=False, + ) + + return run_bench( + topology=topo, bench_fn=_bench_fn, + device=resolve_device(None), engine_factory=_engine_factory, + ) + + +def test_short_prefill_smoke_kv_per_cube_2_C_4(): + """Short prefill kv_per_cube=2, C=4. Smoke: completes.""" + result = _run_prefill_short( + h_kv=8, kv_per_cube=2, C=4, P=8, T_q=4, S_kv=64, + ) + assert result.completion.ok, ( + f"short prefill kv_per_cube=2 must complete; got {result.completion}" + ) + + +def test_short_prefill_no_ring_KV_traffic(): + """ADR-0060 §B.split.2: short prefill DOES NOT use Ring KV — each + CUBE owns its KV heads fully, no rotation. The kernel must not + emit any KV-rotation IPCQ traffic at the CUBE level. + """ + result = _run_prefill_short( + h_kv=8, kv_per_cube=2, C=4, P=8, T_q=4, S_kv=64, + ) + assert result.completion.ok + inter_cube_ipcq = sum( + 1 for r in result.engine.op_log + if r.op_name == "ipcq_copy" + and r.params.get("direction") in ("E", "W") + ) + assert inter_cube_ipcq == 0, ( + f"short prefill must have no Ring KV (no inter-CUBE E/W IPCQ); " + f"got {inter_cube_ipcq}" + ) diff --git a/tests/attention/test_mesh_kernels_rank_axis.py b/tests/attention/test_mesh_kernels_rank_axis.py deleted file mode 100644 index be4229b..0000000 --- a/tests/attention/test_mesh_kernels_rank_axis.py +++ /dev/null @@ -1,172 +0,0 @@ -"""Phase 1 spec test for ``rank_axis`` parameter on the two mesh kernels. - -ADR-0059's mesh kernels currently hard-code ``rank = tl.program_id(axis=0)``, -which only works for single_user_* panels (rank == pe_id within cube). -For multi_user_* panels the ring is at the cube level — rank should be -``cube_id`` (axis=1), and the 7 non-rank-leader PEs in each cube should -not run the ring (they only hold KV replicas). - -This test pins the desired ``rank_axis`` kwarg semantics: - - rank_axis = 0 (default, single_user) - rank = tl.program_id(axis=0). Every PE in the cube runs the ring. - Existing behavior — no change. - - rank_axis = 1 (multi_user) - if tl.program_id(axis=0) != 0: return. (7/8 PEs early-exit.) - rank = tl.program_id(axis=1). - -Phase 1 expectation: tests fail today (kernels don't accept the kwarg). -Phase 2 lands the parameter on both kernels; tests turn green and the -multi_user_* diag harness clears its first send. -""" -from __future__ import annotations - -from kernbench.common.ipcq_types import IpcqRecvCmd, IpcqSendCmd -from kernbench.common.pe_commands import GemmCmd -from kernbench.triton_emu.tl_context import TLContext, run_kernel - -from kernbench.benches._attention_mesh_kv import attention_mesh_kv_kernel -from kernbench.benches._attention_mesh_mlo import attention_mesh_mlo_kernel - -S_Q_PREFILL = 16 -S_Q_DECODE = 1 -S_KV_PER_RANK = 16 -H_Q = 1 -H_KV = 1 -D_HEAD = 64 -N_RANKS_MULTI = 4 -PES_PER_CUBE = 8 - -Q_PTR = 0x10000 -K_PTR = 0x20000 -V_PTR = 0x30000 -O_PTR = 0x40000 - - -def _tl(pe_id: int, cube_id: int, num_pes: int, num_cubes: int) -> TLContext: - return TLContext( - pe_id=pe_id, - num_programs=num_pes, - cube_id=cube_id, - num_cubes=num_cubes, - dispatch_cycles=0, - scratch_base=0x80000, - scratch_size=1 << 20, - ) - - -# ── Default rank_axis=0 backward-compat ────────────────────────── - - -def test_mlo_kernel_default_rank_axis_zero_emits_commands_on_all_pes(): - """rank_axis defaults to 0 → kernel uses pe_id as rank, runs on every - PE. Verify by running rank=3 (interior PE) in a single-cube 8-rank - setup and asserting at least one GEMM and at least one IPCQ send - are emitted (interior ranks send in both directions).""" - tl = _tl(pe_id=3, cube_id=0, num_pes=8, num_cubes=1) - run_kernel( - attention_mesh_mlo_kernel, tl, - Q_PTR, K_PTR, V_PTR, O_PTR, - S_Q_DECODE, S_KV_PER_RANK, H_Q, H_KV, D_HEAD, 8, - ) - assert any(isinstance(c, GemmCmd) for c in tl.commands), \ - "default rank_axis=0 must run the kernel (≥1 GEMM)" - assert any(isinstance(c, IpcqSendCmd) for c in tl.commands), \ - "interior rank must emit ≥1 IpcqSendCmd" - - -def test_kv_kernel_default_rank_axis_zero_emits_commands_on_all_pes(): - tl = _tl(pe_id=3, cube_id=0, num_pes=8, num_cubes=1) - run_kernel( - attention_mesh_kv_kernel, tl, - Q_PTR, K_PTR, V_PTR, O_PTR, - S_Q_PREFILL, S_KV_PER_RANK, H_Q, H_KV, D_HEAD, 8, - ) - assert any(isinstance(c, GemmCmd) for c in tl.commands) - assert any(isinstance(c, IpcqSendCmd) for c in tl.commands) - - -# ── rank_axis=1 multi_user semantics ───────────────────────────── - - -def test_mlo_kernel_rank_axis_one_gates_non_zero_pe_to_no_commands(): - """rank_axis=1 + pe_id != 0 → kernel must early-return; no GEMM, - no DMA, no IPCQ. The 7 non-rank-leader PEs in a multi_user cube - must stay completely silent so the cube-level SFR install isn't - asked to route sends from PEs that have no neighbors installed.""" - tl = _tl(pe_id=2, cube_id=1, num_pes=PES_PER_CUBE, num_cubes=N_RANKS_MULTI) - run_kernel( - attention_mesh_mlo_kernel, tl, - Q_PTR, K_PTR, V_PTR, O_PTR, - S_Q_DECODE, S_KV_PER_RANK, H_Q, H_KV, D_HEAD, N_RANKS_MULTI, - rank_axis=1, - ) - assert not any(isinstance(c, GemmCmd) for c in tl.commands), \ - "pe_id=2 with rank_axis=1 must not emit GEMMs" - assert not any(isinstance(c, IpcqSendCmd) for c in tl.commands), \ - "pe_id=2 with rank_axis=1 must not emit IpcqSendCmd" - assert not any(isinstance(c, IpcqRecvCmd) for c in tl.commands), \ - "pe_id=2 with rank_axis=1 must not emit IpcqRecvCmd" - - -def test_kv_kernel_rank_axis_one_gates_non_zero_pe_to_no_commands(): - tl = _tl(pe_id=2, cube_id=1, num_pes=PES_PER_CUBE, num_cubes=N_RANKS_MULTI) - run_kernel( - attention_mesh_kv_kernel, tl, - Q_PTR, K_PTR, V_PTR, O_PTR, - S_Q_PREFILL, S_KV_PER_RANK, H_Q, H_KV, D_HEAD, N_RANKS_MULTI, - rank_axis=1, - ) - assert not any(isinstance(c, GemmCmd) for c in tl.commands) - assert not any(isinstance(c, IpcqSendCmd) for c in tl.commands) - assert not any(isinstance(c, IpcqRecvCmd) for c in tl.commands) - - -def test_mlo_kernel_rank_axis_one_pe_zero_uses_cube_id_as_rank(): - """rank_axis=1 + pe_id == 0 → kernel runs the ring with rank=cube_id. - For cube_id=1 in a 4-cube ring, rank=1 is an interior rank: has_E=True - AND has_W=True → IPCQ sends emitted in both E and W directions. - """ - tl = _tl(pe_id=0, cube_id=1, num_pes=PES_PER_CUBE, num_cubes=N_RANKS_MULTI) - run_kernel( - attention_mesh_mlo_kernel, tl, - Q_PTR, K_PTR, V_PTR, O_PTR, - S_Q_DECODE, S_KV_PER_RANK, H_Q, H_KV, D_HEAD, N_RANKS_MULTI, - rank_axis=1, - ) - sends = [c for c in tl.commands if isinstance(c, IpcqSendCmd)] - assert any(s.direction == "E" for s in sends), \ - "cube_id=1 (interior) must emit ≥1 E-send" - assert any(s.direction == "W" for s in sends), \ - "cube_id=1 (interior) must emit ≥1 W-send" - - -def test_kv_kernel_rank_axis_one_pe_zero_uses_cube_id_as_rank(): - tl = _tl(pe_id=0, cube_id=1, num_pes=PES_PER_CUBE, num_cubes=N_RANKS_MULTI) - run_kernel( - attention_mesh_kv_kernel, tl, - Q_PTR, K_PTR, V_PTR, O_PTR, - S_Q_PREFILL, S_KV_PER_RANK, H_Q, H_KV, D_HEAD, N_RANKS_MULTI, - rank_axis=1, - ) - sends = [c for c in tl.commands if isinstance(c, IpcqSendCmd)] - assert any(s.direction == "E" for s in sends) - assert any(s.direction == "W" for s in sends) - - -def test_mlo_kernel_rank_axis_one_west_edge_cube_no_west_sends(): - """cube_id=0 (west edge) with rank_axis=1: rank=0, has_W=False → no - W-direction IPCQ sends. has_E=True → ≥1 E-direction send.""" - tl = _tl(pe_id=0, cube_id=0, num_pes=PES_PER_CUBE, num_cubes=N_RANKS_MULTI) - run_kernel( - attention_mesh_mlo_kernel, tl, - Q_PTR, K_PTR, V_PTR, O_PTR, - S_Q_DECODE, S_KV_PER_RANK, H_Q, H_KV, D_HEAD, N_RANKS_MULTI, - rank_axis=1, - ) - sends = [c for c in tl.commands if isinstance(c, IpcqSendCmd)] - assert any(s.direction == "E" for s in sends), \ - "west-edge cube_id=0 must still emit ≥1 E-send" - assert not any(s.direction == "W" for s in sends), \ - "west-edge cube_id=0 must NOT emit any W-send (no W neighbor)" diff --git a/tests/attention/test_mesh_mlo_2d_correctness.py b/tests/attention/test_mesh_mlo_2d_correctness.py deleted file mode 100644 index 803717e..0000000 --- a/tests/attention/test_mesh_mlo_2d_correctness.py +++ /dev/null @@ -1,142 +0,0 @@ -"""Phase 1 spec test for the 2D row-then-col AllReduce-mlo decode kernel. - -The 2D kernel decomposes a ``(mesh_rows × mesh_cols)`` cube sub-mesh into a -two-stage AllReduce-mlo: stage 1 reduces across columns within each row -(E/W edges), stage 2 reduces across rows within each column (N/S edges). -After both stages every cube holds the same final ``(m, ℓ, o)``. - -This module is Phase 1 of C2 (see CLAUDE.md change protocol): it pins -the kernel's interface and observable behavior. Production code for the -kernel lands in Phase 2; until then this file fails to import. - -Test shapes (run on default ``topology.yaml`` — 2 SIPs × 4×4 cube_mesh): - - 1×4 sub-mesh (4 cubes, row 0 only) - Degenerates to a row-only AllReduce — equivalent in step count to - the existing 1D kernel at n_ranks=4. Verifies the kernel reduces - correctly when mesh_rows=1 (stage 2 collapses to no-op). - - 2×4 sub-mesh (8 cubes, rows 0+1) - The 8-KV-group target. Verifies that cubes 4..7 use ``dir="N"`` - (not ``dir="W"``) to reach row 0 — surfacing the IpcqInvalidDirection - bug that the 1D kernel hit at cube 4 (rank 4, no W neighbor). - - 4×4 sub-mesh (16 cubes, full SIP) - Full-SIP scale. Verifies the algorithm fans out over (cols-1)=3 - row steps followed by (rows-1)=3 col steps. -""" -from __future__ import annotations - -from pathlib import Path - -from kernbench.benches._attention_mesh_mlo_2d import attention_mesh_mlo_2d_kernel -from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config -from kernbench.ccl.sfr_config import configure_sfr_intercube_multisip -from kernbench.policy.placement.dp import DPPolicy -from kernbench.runtime_api.bench_runner import run_bench -from kernbench.runtime_api.types import resolve_device -from kernbench.sim_engine.engine import GraphEngine -from kernbench.topology.builder import resolve_topology - -TOPOLOGY_DEFAULT = Path(__file__).resolve().parents[2] / "topology.yaml" - -S_Q = 1 -S_KV_PER_RANK = 16 -H_Q = 1 -H_KV = 1 -D_HEAD = 64 -DTYPE = "f16" - - -def _ccl_cfg(): - return resolve_algorithm_config( - load_ccl_config(), name="lrab_hierarchical_allreduce", - ) - - -def _engine_factory(t, d): - return GraphEngine(getattr(t, "topology_obj", t), enable_data=True) - - -def _run_2d(mesh_rows: int, mesh_cols: int, cube_start: int = 0): - """Build a bench_fn and run it on the default topology.""" - n_cubes = mesh_rows * mesh_cols - topo = resolve_topology(str(TOPOLOGY_DEFAULT)) - - def _bench_fn(ctx): - configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg()) - dp_full = DPPolicy(cube="replicate", pe="replicate", - num_cubes=n_cubes, num_pes=8, - cube_start=cube_start) - dp_kv = DPPolicy(cube="row_wise", pe="replicate", - num_cubes=n_cubes, num_pes=8, - cube_start=cube_start) - q = ctx.zeros((S_Q, H_Q * D_HEAD), - dtype=DTYPE, dp=dp_full, name="q") - k = ctx.zeros((S_KV_PER_RANK * n_cubes, H_KV * D_HEAD), - dtype=DTYPE, dp=dp_kv, name="k") - v = ctx.zeros((S_KV_PER_RANK * n_cubes, H_KV * D_HEAD), - dtype=DTYPE, dp=dp_kv, name="v") - o = ctx.empty((S_Q, H_Q * D_HEAD), - dtype=DTYPE, dp=dp_full, name="o") - ctx.launch( - f"mesh_mlo_2d_{mesh_rows}x{mesh_cols}_start{cube_start}", - attention_mesh_mlo_2d_kernel, - q, k, v, o, - S_Q, S_KV_PER_RANK, H_Q, H_KV, D_HEAD, - mesh_rows, mesh_cols, - 1, # rank_axis=1 → cube-level ring - cube_start, - _auto_dim_remap=False, - ) - - return run_bench( - topology=topo, - bench_fn=_bench_fn, - device=resolve_device(None), - engine_factory=_engine_factory, - ) - - -def test_2d_kernel_1x4_row_only(): - """1×4 sub-mesh: row-only AllReduce, stage 2 collapses to no-op.""" - result = _run_2d(mesh_rows=1, mesh_cols=4) - assert result.completion.ok, ( - f"1x4: completion not ok - {result.completion}" - ) - - -def test_2d_kernel_2x4_eight_cubes(): - """2×4 sub-mesh: the 8-KV-group target. - - Verifies that cube 4 (row 1, col 0) uses ``dir="N"`` to reach cube 0 - (row 0, col 0) for stage 2, not ``dir="W"`` — the 1D kernel hit - IpcqInvalidDirection here. - """ - result = _run_2d(mesh_rows=2, mesh_cols=4) - assert result.completion.ok, ( - f"2x4: completion not ok - {result.completion}" - ) - - -def test_2d_kernel_4x4_full_sip(): - """4×4 sub-mesh: full-SIP scale (16 cubes).""" - result = _run_2d(mesh_rows=4, mesh_cols=4) - assert result.completion.ok, ( - f"4x4: completion not ok - {result.completion}" - ) - - -def test_2d_kernel_2x4_at_cube_start_eight(): - """2×4 sub-mesh at cube_start=8 (cubes 8..15, rows 2..3). - - The second KV-group per SIP in the 8-KV-group Llama-70B headline. - Verifies the kernel converts ``program_id(axis=1)`` (physical cube - id) back to launch-local rank via ``cube_start`` — without that - subtraction, cube 8 would compute my_row=2 (out of sub-mesh bounds) - and deadlock waiting on cube 4 which isn't in the launch. - """ - result = _run_2d(mesh_rows=2, mesh_cols=4, cube_start=8) - assert result.completion.ok, ( - f"2x4 @ cube_start=8: completion not ok - {result.completion}" - ) diff --git a/tests/attention/test_milestone_gqa_headline.py b/tests/attention/test_milestone_gqa_headline.py new file mode 100644 index 0000000..77b89d8 --- /dev/null +++ b/tests/attention/test_milestone_gqa_headline.py @@ -0,0 +1,148 @@ +"""Phase 1 spec test for P7: headline milestone-gqa bench with real GQA. + +P7 wires the new ``_gqa_decode`` and ``_gqa_prefill`` kernels into a +new 4-panel milestone bench (independent from the existing +``milestone-gqa-llama70b`` which still covers the baseline kernels). +Real GQA (``h_q > h_kv`` with G=8) runs end-to-end through a +milestone-style sweep + sweep.json output. + +Restriction in P7 first cut: + - C ≤ 4 (single-row inter-CUBE ring SFR; ADR-0060 §B leaves + multi-row rings for follow-on) + - Single SIP (default ``topology.yaml`` 4×4 cube mesh; 4-SIP + headline left to follow-on) + - No figure renderers (defer to a separate cycle) + +Panels (4 total): + single_user_prefill_gqa: C=1, T_q=4, S_kv=16 (no ring) + multi_user_prefill_gqa : C=4, T_q=4, S_kv=16 (Ring KV, 3 steps) + single_user_decode_gqa : C=1, P=8, h_q=8, h_kv=1, S_kv=64 (M-fold + intra-cube chain) + multi_user_decode_gqa : C=4, P=8, h_q=8, h_kv=1, S_kv=128 (M-fold + 2-level chain) + +Phase 1 (this commit): tests only — bench module lands in Phase 2. +""" +from __future__ import annotations + +import json + +import kernbench.benches.milestone_gqa_headline as bench_mod # noqa: F401 (Phase 2) +from kernbench.benches.registry import resolve +from kernbench.runtime_api.bench_runner import run_bench +from kernbench.runtime_api.types import resolve_device +from kernbench.sim_engine.engine import GraphEngine +from kernbench.topology.builder import resolve_topology + +BENCH_NAME = "milestone-gqa-headline" + +PANELS = ( + "single_user_prefill_gqa", + "multi_user_prefill_gqa", + "single_user_decode_gqa", + "multi_user_decode_gqa", +) + + +def _run_validation(): + topo = resolve_topology("topology.yaml") + return run_bench( + topology=topo, + bench_fn=resolve(BENCH_NAME).run, + device=resolve_device(None), + engine_factory=lambda t, d: GraphEngine( + getattr(t, "topology_obj", t), enable_data=True, + ), + ) + + +def _sweep_json(monkeypatch) -> dict: + monkeypatch.setenv("GQA_HEADLINE_RUN", "1") + out = bench_mod._OUTPUT_DIR / "sweep.json" + if not out.exists(): + result = _run_validation() + assert result.completion.ok, result.completion + assert out.exists(), f"missing {out}" + return json.loads(out.read_text()) + + +# ── Registration ─────────────────────────────────────────────────────── + + +def test_bench_registered(): + spec = resolve(BENCH_NAME) + assert spec.name == BENCH_NAME + assert callable(spec.run) + assert spec.description.strip(), "description must be non-empty" + + +# ── Validation run completes ────────────────────────────────────────── + + +def test_validation_run_completes_ok(monkeypatch): + monkeypatch.setenv("GQA_HEADLINE_RUN", "1") + result = _run_validation() + assert result.completion.ok, ( + f"headline validation run failed: {result.completion}" + ) + + +# ── sweep.json shape ────────────────────────────────────────────────── + + +def test_sweep_json_has_four_panels(monkeypatch): + data = _sweep_json(monkeypatch) + assert set(data["panels"]) == set(PANELS), ( + f"panels mismatch: expected {set(PANELS)}, got {set(data['panels'])}" + ) + assert len(data["rows"]) == 4 + assert {r["panel"] for r in data["rows"]} == set(PANELS) + + +def test_decode_panels_use_real_gqa(monkeypatch): + """ADR-0060 §A.1 headline: decode panels must use h_q = G·h_kv with G>1.""" + data = _sweep_json(monkeypatch) + cfg = data["config"] + assert cfg["h_q_decode"] > cfg["h_kv_decode"], ( + f"decode must use real GQA (h_q > h_kv); got " + f"h_q={cfg['h_q_decode']}, h_kv={cfg['h_kv_decode']}" + ) + + +# ── Per-panel architectural assertions ──────────────────────────────── + + +def _row(rows, panel: str) -> dict: + for r in rows: + if r["panel"] == panel: + return r + raise AssertionError(f"missing row for panel {panel!r}") + + +def test_prefill_ring_panel_has_ipcq_traffic(monkeypatch): + """multi_user_prefill_gqa uses Ring KV → IPCQ traffic > 0.""" + data = _sweep_json(monkeypatch) + row = _row(data["rows"], "multi_user_prefill_gqa") + n_copy = row["op_log_summary"].get("ipcq_copy_count", 0) + assert n_copy > 0, ( + f"multi_user_prefill_gqa must have Ring KV traffic; got " + f"ipcq_copy_count={n_copy}" + ) + + +def test_decode_reduce_panel_writes_once(monkeypatch): + """multi_user_decode_gqa: chain reduce-to-root → exactly 1 dma_write.""" + data = _sweep_json(monkeypatch) + row = _row(data["rows"], "multi_user_decode_gqa") + n_writes = row["op_log_summary"]["dma_write_count"] + assert n_writes == 1, ( + f"multi_user_decode_gqa root-only write: expected 1; got {n_writes}" + ) + + +def test_prefill_panel_distributes_output(monkeypatch): + """multi_user_prefill_gqa: per-CUBE distributed output → dma_write_count == C.""" + data = _sweep_json(monkeypatch) + row = _row(data["rows"], "multi_user_prefill_gqa") + n_writes = row["op_log_summary"]["dma_write_count"] + assert n_writes == 4, ( + f"multi_user_prefill_gqa per-CUBE distributed: expected 4; got {n_writes}" + ) diff --git a/tests/attention/test_milestone_gqa_llama70b.py b/tests/attention/test_milestone_gqa_llama70b.py deleted file mode 100644 index 335b139..0000000 --- a/tests/attention/test_milestone_gqa_llama70b.py +++ /dev/null @@ -1,222 +0,0 @@ -"""Phase 1 spec test for ``milestone-gqa-llama70b`` bench (sub-cycle 4a, all 4 panels). - -ADR-0057 (Proposed) defines an eval bench that drives both attention kernels -(ADR-0055 ring-K/V, ADR-0056 allreduce-mlo) and emits per-panel op_log -summaries into ``src/kernbench/benches/1H_milestone_output/gqa/sweep.json``. - -v1 (sub-cycle 4a) covers ALL FOUR panels: - - Panel name in JSON / test Study label SFR install used - ───────────────────────────────────────────────────────────────────────────── - single_user_prefill TL configure_sfr_intracube_pe_ring - multi_user_prefill TR configure_sfr_intercube_multisip - single_user_decode BL configure_sfr_intracube_pe_ring - multi_user_decode BR configure_sfr_intercube_multisip - -single_user_* panels became runnable after sub-cycle 4-pre delivered the -new SFR install function (ADR-0058). - -In Phase 1 the bench module does not exist; pytest collection fails with -``ModuleNotFoundError``. Once Phase 2 lands the bench module, every -assertion below must pass. - -Assertions: - - Bench is registered as ``milestone-gqa-llama70b``. - - A validation run (``GQA_VALIDATION=1``) completes ok via run_bench. - - sweep.json conforms to ADR-0057 D7 (v1 schema). - - All four panel rows present with sane op_log summaries. - - Both decode panels have gemm_count = 2 × n_ranks (one-shot per rank). - - Both prefill panels have gemm_count = 2 × n_ranks² (per-step GEMMs). -""" -from __future__ import annotations - -import json - -from kernbench.benches.registry import resolve -from kernbench.runtime_api.bench_runner import run_bench -from kernbench.runtime_api.types import resolve_device -from kernbench.sim_engine.engine import GraphEngine -from kernbench.topology.builder import resolve_topology - -# Production module (Phase 2 deliverable; absent in Phase 1). -import kernbench.benches.milestone_gqa_llama70b as gqa_bench - - -BENCH_NAME = "milestone-gqa-llama70b" - -PANELS_V1 = ( - "single_user_prefill", - "multi_user_prefill", - "single_user_decode", - "multi_user_decode", -) -SINGLE_USER_PANELS = ("single_user_prefill", "single_user_decode") -MULTI_USER_PANELS = ("multi_user_prefill", "multi_user_decode") -PREFILL_PANELS = ("single_user_prefill", "multi_user_prefill") -DECODE_PANELS = ("single_user_decode", "multi_user_decode") - - -def _run_validation(): - """Drive the bench through run_bench at validation scale.""" - topo = resolve_topology("topology.yaml") - return run_bench( - topology=topo, - bench_fn=resolve(BENCH_NAME).run, - device=resolve_device(None), - engine_factory=lambda t, d: GraphEngine( - getattr(t, "topology_obj", t), enable_data=True, - ), - ) - - -# ── Registration ──────────────────────────────────────────────────────── - - -def test_bench_registered(): - spec = resolve(BENCH_NAME) - assert spec.name == BENCH_NAME - assert callable(spec.run) - assert spec.description.strip(), "description must be non-empty" - - -# ── Validation run end-to-end ──────────────────────────────────────────── - - -def test_validation_run_completes_ok(monkeypatch): - monkeypatch.setenv("GQA_VALIDATION", "1") - result = _run_validation() - assert result.completion.ok, ( - f"validation run failed: {result.completion}" - ) - - -# ── JSON shape (ADR-0057 D7 amended for 4 panels) ────────────────────── - - -def _sweep_json(monkeypatch) -> dict: - """Run the bench (if needed) and return the parsed sweep.json.""" - monkeypatch.setenv("GQA_VALIDATION", "1") - out = gqa_bench._OUTPUT_DIR / "sweep.json" - if not out.exists(): - result = _run_validation() - assert result.completion.ok, result.completion - assert out.exists(), f"missing {out}" - return json.loads(out.read_text()) - - -def test_sweep_json_has_v1_schema(monkeypatch): - data = _sweep_json(monkeypatch) - assert data["version"] == 1 - assert data["validation_scale"] is True - assert isinstance(data["panels"], list) - assert isinstance(data["config"], dict) - assert isinstance(data["rows"], list) - - -def test_sweep_json_panels_are_all_four(monkeypatch): - """v1 covers all four panels — single_user_{prefill,decode} + - multi_user_{prefill,decode}. Q/cube sweep deferred to 4b.""" - data = _sweep_json(monkeypatch) - assert set(data["panels"]) == set(PANELS_V1) - - -def test_sweep_json_config_matches_adr0057_d4(monkeypatch): - """Validation-scale config per ADR-0057 D4 (amended for 4 panels + scratch budget). - - S_q_prefill and S_kv_per_rank are deliberately small (16 each) so the - simulator's 1 MB per-PE TCM kernel scratch (topology.yaml - ``pe_tcm.kernel_scratch_mb: 1``) is not exhausted by the - bump-allocated handle outputs of softmax/exp/dot/sum chains over - n_ranks ring steps. Headline-scale runs in 4c will lift these into a - config-driven sweep. - """ - data = _sweep_json(monkeypatch) - cfg = data["config"] - assert cfg["S_q_prefill"] == 16 - assert cfg["S_kv_per_rank"] == 16 - # v1 uses h_q == h_kv == 1 to avoid ADR-0055 D3's GQA broadcast view - # (which is symbolic and does not survive MemoryStore's nbytes check - # under simulator data execution). Real GQA (h_q > h_kv) is deferred - # to sub-cycle 4c (headline scale). - assert cfg["h_q"] == 1 - assert cfg["h_kv"] == 1 - assert cfg["d_head"] == 64 - # single_user_* uses the 8 PEs in one cube as ring ranks. - assert cfg["n_ranks_single_user"] == 8 - # multi_user_* uses cube-level ring; validation uses 4 cubes. - assert cfg["n_ranks_multi_user"] == 4 - - -def test_sweep_json_has_one_row_per_panel(monkeypatch): - data = _sweep_json(monkeypatch) - assert len(data["rows"]) == 4 - panels_in_rows = {r["panel"] for r in data["rows"]} - assert panels_in_rows == set(PANELS_V1) - - -# ── Per-row op_log summary sanity (ADR-0057 D7) ───────────────────────── - - -def _row(rows: list, panel: str) -> dict: - matches = [r for r in rows if r["panel"] == panel] - assert len(matches) == 1, f"expected exactly one {panel} row; got {len(matches)}" - return matches[0] - - -def _assert_sane_summary(row: dict) -> None: - s = row["op_log_summary"] - panel = row["panel"] - assert s["gemm_count"] > 0, f"{panel} must run GEMMs" - assert s["ipcq_send_count"] > 0, f"{panel} must send (ring/allreduce phase)" - assert s["ipcq_recv_count"] > 0, f"{panel} must recv" - assert s["dma_read_count"] >= 3, f"{panel}: Q + K + V loads" - assert s["dma_write_count"] >= 1, f"{panel}: final O store" - - -def test_single_user_prefill_row_has_sane_op_log_summary(monkeypatch): - data = _sweep_json(monkeypatch) - _assert_sane_summary(_row(data["rows"], "single_user_prefill")) - - -def test_multi_user_prefill_row_has_sane_op_log_summary(monkeypatch): - data = _sweep_json(monkeypatch) - _assert_sane_summary(_row(data["rows"], "multi_user_prefill")) - - -def test_single_user_decode_row_has_sane_op_log_summary(monkeypatch): - data = _sweep_json(monkeypatch) - _assert_sane_summary(_row(data["rows"], "single_user_decode")) - - -def test_multi_user_decode_row_has_sane_op_log_summary(monkeypatch): - data = _sweep_json(monkeypatch) - _assert_sane_summary(_row(data["rows"], "multi_user_decode")) - - -# ── Architectural invariant: decode = one-shot per rank ───────────────── - - -def test_single_user_decode_gemm_count_is_exactly_2_per_rank(monkeypatch): - """ADR-0056 D3: decode kernel does ONE local partial-attention pass per - rank → exactly 2 GEMMs per rank (Q·K^T + S·V). With n_ranks ranks the - total = 2 × n_ranks. This distinguishes decode from prefill where each - ring step has 2 GEMMs and the total scales as 2 × n_ranks².""" - data = _sweep_json(monkeypatch) - row = _row(data["rows"], "single_user_decode") - n_ranks = row["n_ranks"] - assert row["op_log_summary"]["gemm_count"] == 2 * n_ranks, ( - f"single_user_decode gemm_count must be 2 × n_ranks = {2 * n_ranks}; " - f"got {row['op_log_summary']['gemm_count']}" - ) - - -def test_multi_user_decode_gemm_count_is_exactly_2_per_rank(monkeypatch): - """Same one-shot invariant as single_user_decode — the kernel is the - same; what differs is who the ranks are (cubes vs PEs).""" - data = _sweep_json(monkeypatch) - row = _row(data["rows"], "multi_user_decode") - n_ranks = row["n_ranks"] - assert row["op_log_summary"]["gemm_count"] == 2 * n_ranks, ( - f"multi_user_decode gemm_count must be 2 × n_ranks = {2 * n_ranks}; " - f"got {row['op_log_summary']['gemm_count']}" - ) diff --git a/tests/gqa/_gqa_plot_helpers.py b/tests/gqa/_gqa_plot_helpers.py deleted file mode 100644 index 19cf155..0000000 --- a/tests/gqa/_gqa_plot_helpers.py +++ /dev/null @@ -1,25 +0,0 @@ -"""Thin re-export shim for the GQA figure tests. - -Not a test module (no ``test_`` prefix → pytest does not collect it). - -Mirrors ``tests/gemm/_gemm_plot_helpers.py``. The renderer logic lives in -``kernbench.benches.milestone_gqa_llama70b`` (production single home, -ADR-0054). Defaults still target the bench's ``_OUTPUT_DIR``. -""" -from __future__ import annotations - -from kernbench.benches.milestone_gqa_llama70b import ( - _OUTPUT_DIR as GQA_PLOTS_DIR, - _SWEEP_JSON as GQA_SWEEP_JSON, - emit_all_gqa_plots, - emit_gqa_comparison, - emit_panel_op_log_summary, -) - -__all__ = [ - "GQA_PLOTS_DIR", - "GQA_SWEEP_JSON", - "emit_all_gqa_plots", - "emit_gqa_comparison", - "emit_panel_op_log_summary", -] diff --git a/tests/gqa/test_plot_gqa_figures.py b/tests/gqa/test_plot_gqa_figures.py deleted file mode 100644 index 7f1e9d2..0000000 --- a/tests/gqa/test_plot_gqa_figures.py +++ /dev/null @@ -1,109 +0,0 @@ -"""Phase 1 spec test for GQA figure renderers (sub-cycle 4c). - -ADR-0057 D3 sub-cycle 4c adds 6 figure renderers; this test pins the -5 of 6 that don't depend on sub-cycle 4b's Q/cube sweep: - - - 4 per-panel op_log_summary PNGs (one per panel of v1's sweep.json) - - 1 cross-panel ``gqa_comparison.png`` (4-panel grouped bars over the - 5 op_log_summary counts: gemm, ipcq_send, ipcq_recv, dma_read, dma_write) - -The 6th, ``gqa_scaling.png``, needs the Q/cube ∈ {1, 2, 4} sweep from -sub-cycle 4b and is deferred. - -Each test depends on the committed -``benches/1H_milestone_output/gqa/sweep.json`` (landed in commit -``e748a62``); they assert the renderer writes a non-empty PNG at the -expected path. - -Phase 1 expectation: tests fail at import (renderer functions don't -exist yet on the bench module). Phase 2 lands them and the tests -turn green. -""" -from __future__ import annotations - -from pathlib import Path - -import pytest - -from tests.gqa._gqa_plot_helpers import ( - GQA_PLOTS_DIR, - GQA_SWEEP_JSON, - emit_all_gqa_plots, - emit_gqa_comparison, - emit_panel_op_log_summary, -) - - -_PANELS = ( - "single_user_prefill", - "multi_user_prefill", - "single_user_decode", - "multi_user_decode", -) - - -@pytest.mark.skipif( - not GQA_SWEEP_JSON.exists(), - reason="gqa sweep.json absent; run milestone-gqa-llama70b first", -) -@pytest.mark.parametrize("panel", _PANELS) -def test_emit_panel_op_log_summary_writes_png_for_each_panel(panel): - out = emit_panel_op_log_summary(panel) - assert out is not None, f"{panel}: renderer returned None" - path = Path(out) - assert path.exists(), f"{panel}: expected PNG at {path}" - assert path.suffix == ".png", f"{panel}: not a PNG: {path}" - assert path.stat().st_size > 0, f"{panel}: empty PNG: {path}" - assert panel in path.stem, ( - f"{panel}: panel name not in filename {path.name}" - ) - - -@pytest.mark.skipif( - not GQA_SWEEP_JSON.exists(), - reason="gqa sweep.json absent; run milestone-gqa-llama70b first", -) -def test_emit_gqa_comparison_writes_png(): - out = emit_gqa_comparison() - assert out is not None - path = Path(out) - assert path.exists() - assert path.name == "gqa_comparison.png" - assert path.stat().st_size > 0 - - -@pytest.mark.skipif( - not GQA_SWEEP_JSON.exists(), - reason="gqa sweep.json absent; run milestone-gqa-llama70b first", -) -def test_emit_all_gqa_plots_writes_five_figures(): - """emit_all returns a list of 5 written PNG paths (deferring the - 6th gqa_scaling.png to after sub-cycle 4b lands the Q/cube sweep).""" - paths = emit_all_gqa_plots() - assert isinstance(paths, list) - # 4 per-panel + 1 comparison. - assert len(paths) == 5, f"expected 5 PNGs, got {len(paths)}: {paths}" - for p in paths: - assert Path(p).exists() and Path(p).stat().st_size > 0 - names = {Path(p).name for p in paths} - assert "gqa_comparison.png" in names - for panel in _PANELS: - assert any(panel in n for n in names), ( - f"no per-panel PNG for {panel} in {names}" - ) - - -def test_emit_all_gqa_plots_output_dir_matches_bench_output_dir(): - """The renderers must write under the bench's own _OUTPUT_DIR so - MILESTONE_FAST=1 reuse (and committed baselines) all point at the - same on-disk location.""" - # Stub assertion that fails until emit_all_gqa_plots exists with a - # default ``out_dir`` argument identical to GQA_PLOTS_DIR. - import inspect - - sig = inspect.signature(emit_all_gqa_plots) - assert "out_dir" in sig.parameters - default = sig.parameters["out_dir"].default - assert Path(default) == GQA_PLOTS_DIR, ( - f"default out_dir {default} != bench _OUTPUT_DIR {GQA_PLOTS_DIR}" - ) diff --git a/tests/test_kernel_runner.py b/tests/test_kernel_runner.py index 3cb2c64..f673719 100644 --- a/tests/test_kernel_runner.py +++ b/tests/test_kernel_runner.py @@ -28,7 +28,13 @@ def _mock_scheduler(env, inbox): def test_kernel_runner_basic_load(): - """Kernel with tl.load runs through greenlet without hanging.""" + """Kernel with tl.load runs through greenlet without hanging. + + Under ADR-0062 (lazy ``tl.load``) the handle's ``data`` is None + until a consumer op triggers auto-wait at first use. A no-op math + op (``tl.exp``) consumes ``a`` and forces resolution before we + inspect ``a.data``. + """ env = simpy.Environment() store = MemoryStore() data = np.ones((4, 4), dtype=np.float16) @@ -39,6 +45,7 @@ def test_kernel_runner_basic_load(): def kernel(a_ptr, tl): a = tl.load(a_ptr, (4, 4), "f16") + tl.exp(a) # consumer op → auto-wait at first use (ADR-0062 §D2) assert a.data is not None assert a.data.shape == (4, 4) @@ -50,7 +57,12 @@ def test_kernel_runner_basic_load(): def test_kernel_runner_load_returns_data(): - """tl.load returns actual numpy data from MemoryStore.""" + """tl.load returns actual numpy data from MemoryStore. + + Under ADR-0062 lazy semantics the data is attached at auto-wait + time (first consumer op), so we read ``a.data`` after a consumer + op fires the wait. + """ env = simpy.Environment() store = MemoryStore() data = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float16) @@ -63,6 +75,7 @@ def test_kernel_runner_load_returns_data(): def kernel(ptr, tl): a = tl.load(ptr, (2, 2), "f16") + tl.exp(a) # consumer op → auto-wait at first use (ADR-0062 §D2) results["data"] = a.data def run(): @@ -106,6 +119,11 @@ def test_kernel_runner_dynamic_branch(): def kernel(flag_ptr, tl): flag = tl.load(flag_ptr, (1,), "f32") + # ADR-0062: under lazy tl.load, dynamic-branching kernels must + # force resolution by consuming the handle (any tl.* op works). + # Without this, flag.data is None at branch time and control + # always takes the not-taken path. + tl.exp(flag) if flag.data is not None and flag.data[0] > 0.5: results["branch"] = "taken" else: diff --git a/tests/test_phase_e_cpu_issue_cost.py b/tests/test_phase_e_cpu_issue_cost.py new file mode 100644 index 0000000..023b9a9 --- /dev/null +++ b/tests/test_phase_e_cpu_issue_cost.py @@ -0,0 +1,382 @@ +"""Phase 1 spec tests for ADR-0064 (per-op-type CPU issue cost model). + +Phase E lands the cost-table machinery and turns it on by default so the +hybrid's CPU-saturation lever (ADR-0060 §1) becomes measurable instead of +modelled away. Today every ``tl.*`` call goes through +``_emit_dispatch_overhead()`` with a single uniform ``dispatch_cycles`` +scalar that is hard-coded to 0 on the live PE_CPU paths +(``pe_cpu.py:_execute_legacy`` and ``kernel_runner.py:run``). + +These tests assume the post-Phase-2 surface: + +- ``kernbench.common.cpu_issue_cost`` exports ``OpKind``, + ``DEFAULT_CPU_ISSUE_COST`` (the ADR-0064 D1 table) and + ``get_issue_cost(kind, table=None) -> int``. +- ``TLContext.__init__`` accepts ``issue_cost_table: dict[str, int] | None``. + When provided, ``_emit_dispatch_overhead(kind)`` looks up the per-kind + cost; when absent, it falls back to the uniform ``dispatch_cycles`` + (back-compat with the existing ADR-0046 §D6 contract). +- Live PE_CPU paths (greenlet + legacy replay) construct TLContext with + ``issue_cost_table=DEFAULT_CPU_ISSUE_COST``. + +Phase 1 (this commit): tests only. All tests FAIL until Phase 2. +""" +from __future__ import annotations + +from pathlib import Path + +import pytest + +from kernbench.common.pe_commands import ( + CompositeCmd, + DmaReadCmd, + DmaWriteCmd, + GemmCmd, + MathCmd, + PeCpuOverheadCmd, +) +from kernbench.policy.address.phyaddr import PhysAddr +from kernbench.runtime_api.kernel import KernelLaunchMsg, KernelRef +from kernbench.sim_engine.engine import GraphEngine +from kernbench.sim_engine.transaction import Transaction +from kernbench.topology.builder import load_topology +from kernbench.triton_emu.registry import clear_registry, register_kernel +from kernbench.triton_emu.tl_context import TLContext, run_kernel + +TOPOLOGY_PATH = Path(__file__).parent.parent / "topology.yaml" + + +def _engine(): + return GraphEngine(load_topology(TOPOLOGY_PATH)) + + +def _hbm_pa(sip: int = 0, cube: int = 0, pe_id: int = 0) -> int: + slice_bytes = 48 * (1 << 30) // 8 + pa = PhysAddr.pe_hbm_addr( + sip_id=sip, die_id=cube, pe_id=pe_id, + pe_local_hbm_offset=0x1000, slice_size_bytes=slice_bytes, + ) + return pa.encode() + + +# ── T1: default table shape ────────────────────────────────────── + + +def test_default_cost_table_has_expected_keys(): + """ADR-0064 D1 table — 8 keys with composite ≫ primitive ratio.""" + from kernbench.common.cpu_issue_cost import DEFAULT_CPU_ISSUE_COST + + expected_keys = { + "composite", "load", "store", "dot", "math", + "ipcq_send", "ipcq_recv", "copy_to", + } + assert set(DEFAULT_CPU_ISSUE_COST.keys()) == expected_keys, ( + f"DEFAULT_CPU_ISSUE_COST keys must match ADR-0064 D1; " + f"got {set(DEFAULT_CPU_ISSUE_COST.keys())}" + ) + # Composite is the lever — ratio against primitives must be ≥ 4×. + assert DEFAULT_CPU_ISSUE_COST["composite"] == 40 + for primitive in ("load", "store", "dot", "math", + "ipcq_send", "ipcq_recv", "copy_to"): + assert DEFAULT_CPU_ISSUE_COST[primitive] == 5, ( + f"primitive {primitive!r} default cost must be 5 ns" + ) + + +# ── T2: get_issue_cost lookup ──────────────────────────────────── + + +def test_get_issue_cost_lookup(): + """Helper returns table value; unknown kind returns 0 (no charge).""" + from kernbench.common.cpu_issue_cost import ( + DEFAULT_CPU_ISSUE_COST, + get_issue_cost, + ) + + assert get_issue_cost("composite") == 40 + assert get_issue_cost("load") == 5 + assert get_issue_cost("unknown_kind") == 0 + # Custom table override + custom = {"composite": 100, "load": 1} + assert get_issue_cost("composite", table=custom) == 100 + assert get_issue_cost("load", table=custom) == 1 + assert get_issue_cost("store", table=custom) == 0 + + +# ── T3: TLContext consumes a passed cost table ─────────────────── + + +def test_tlcontext_accepts_issue_cost_table(): + """TLContext(issue_cost_table=...) → tl.load emits the per-kind cycles.""" + tl = TLContext( + pe_id=0, num_programs=1, + dispatch_cycles=0, + issue_cost_table={"load": 7, "store": 3, "dot": 2, "math": 2}, + ) + tl.load(0x1000, shape=(4, 4), dtype="f16") + overheads = [c for c in tl.commands if isinstance(c, PeCpuOverheadCmd)] + assert len(overheads) == 1, ( + f"expected exactly one PeCpuOverheadCmd before the DmaReadCmd; " + f"got {len(overheads)} (cmds={[type(c).__name__ for c in tl.commands]})" + ) + assert overheads[0].cycles == 7, ( + f"load issue cost from table must be 7; got {overheads[0].cycles}" + ) + + +# ── T4: composite ≫ primitive issue cost differential ──────────── + + +def test_tlcontext_composite_vs_primitive_charges_differ(): + """Composite kernel charges once (40); primitive sequence charges per-op.""" + from kernbench.common.cpu_issue_cost import DEFAULT_CPU_ISSUE_COST + + # Composite kernel: 1 load + 1 composite = 5 + 40 = 45 cycles of issue cost. + tl_comp = TLContext( + pe_id=0, num_programs=1, + dispatch_cycles=0, + issue_cost_table=DEFAULT_CPU_ISSUE_COST, + ) + a = tl_comp.load(0x1000, shape=(8, 16), dtype="f16") + b_ref = tl_comp.ref(0x2000, shape=(16, 8), dtype="f16") + tl_comp.composite("gemm", a, b_ref, out_ptr=0x3000) + comp_cycles = sum( + c.cycles for c in tl_comp.commands if isinstance(c, PeCpuOverheadCmd) + ) + + # Primitive kernel: 2 loads + 1 dot + 1 math (exp) = 5 + 5 + 5 + 5 = 20 cycles. + tl_prim = TLContext( + pe_id=0, num_programs=1, + dispatch_cycles=0, + issue_cost_table=DEFAULT_CPU_ISSUE_COST, + ) + a = tl_prim.load(0x1000, shape=(8, 16), dtype="f16") + b = tl_prim.load(0x2000, shape=(16, 8), dtype="f16") + c_out = tl_prim.dot(a, b) + tl_prim.exp(c_out) + prim_cycles = sum( + c.cycles for c in tl_prim.commands if isinstance(c, PeCpuOverheadCmd) + ) + + # ADR-0064 ratio: composite issue cost >> primitive issue cost per op. + # For these specific kernels: composite=45 (5+40), primitive=20 (4×5). + assert comp_cycles == 45, f"composite kernel: expected 45, got {comp_cycles}" + assert prim_cycles == 20, f"primitive kernel: expected 20, got {prim_cycles}" + # Headline assertion: a single composite charges more than all 3 + # post-load primitives combined (40 > 3×5) — the ADR-0060 §1 lever. + assert comp_cycles - 5 > 3 * 5, ( + f"composite issue charge {comp_cycles - 5} must exceed " + f"3× primitive issue charge {3 * 5} (ADR-0064 ratio)" + ) + + +# ── T5: cost table is purely additive (Q2 — no double-count) ───── + + +def test_cost_table_is_additive_only(): + """Q2 invariant: the cost table is additive on PE_CPU, NOT folded into + DMA/GEMM/MATH command shapes. Two TLContexts running the same kernel + with different cost tables must produce identical non-overhead command + sequences (same DmaReadCmd, GemmCmd, MathCmd, addrs, shapes, dtypes). + Only the PeCpuOverheadCmd ``cycles`` field is allowed to differ. + """ + from kernbench.common.cpu_issue_cost import DEFAULT_CPU_ISSUE_COST + + def kernel(tl): + a = tl.load(0x1000, shape=(8, 16), dtype="f16") + b = tl.load(0x2000, shape=(16, 8), dtype="f16") + c = tl.dot(a, b) + d = tl.exp(c) + tl.store(0x3000, d) + + tl_zero = TLContext( + pe_id=0, num_programs=1, + dispatch_cycles=0, + issue_cost_table={}, # empty table → every kind = 0 cost + ) + run_kernel(kernel, tl_zero) + + tl_default = TLContext( + pe_id=0, num_programs=1, + dispatch_cycles=0, + issue_cost_table=DEFAULT_CPU_ISSUE_COST, + ) + run_kernel(kernel, tl_default) + + # Strip PeCpuOverheadCmd from both streams; what remains must match. + non_overhead_zero = [ + c for c in tl_zero.commands if not isinstance(c, PeCpuOverheadCmd) + ] + non_overhead_default = [ + c for c in tl_default.commands if not isinstance(c, PeCpuOverheadCmd) + ] + assert len(non_overhead_zero) == len(non_overhead_default), ( + f"non-overhead command count differs: " + f"{len(non_overhead_zero)} vs {len(non_overhead_default)}" + ) + for a_cmd, b_cmd in zip(non_overhead_zero, non_overhead_default): + assert type(a_cmd) is type(b_cmd), ( + f"non-overhead command type changed under cost table: " + f"{type(a_cmd).__name__} vs {type(b_cmd).__name__}" + ) + # Total overhead under zero-table must be 0; under default must be > 0. + cycles_zero = sum( + c.cycles for c in tl_zero.commands if isinstance(c, PeCpuOverheadCmd) + ) + cycles_default = sum( + c.cycles for c in tl_default.commands if isinstance(c, PeCpuOverheadCmd) + ) + assert cycles_zero == 0, f"empty table must add 0 cycles; got {cycles_zero}" + assert cycles_default > 0, ( + f"default table must add > 0 cycles; got {cycles_default}" + ) + + +# ── T6: greenlet vs legacy replay use same cost table (review #6) ─ + + +def test_greenlet_and_legacy_path_parity(): + """Both PE_CPU execution paths read the same cost table. + + The greenlet path (kernel_runner.py:run) and the legacy replay path + (pe_cpu.py:_execute_legacy) must construct TLContext with the same + default cost table so the same kernel produces identical PE_CPU + overhead cycles via either route. This is ADR-0064 review item #6. + """ + from kernbench.common.cpu_issue_cost import DEFAULT_CPU_ISSUE_COST + + # Build the command list for a representative kernel via TLContext + # using the default table — this is what both live paths should see. + def kernel(tl): + a = tl.load(0x1000, shape=(4, 4), dtype="f16") + b = tl.load(0x2000, shape=(4, 4), dtype="f16") + c = tl.dot(a, b) + tl.store(0x3000, c) + + tl1 = TLContext( + pe_id=0, num_programs=1, + dispatch_cycles=0, + issue_cost_table=DEFAULT_CPU_ISSUE_COST, + ) + run_kernel(kernel, tl1) + cycles1 = sum( + c.cycles for c in tl1.commands if isinstance(c, PeCpuOverheadCmd) + ) + + tl2 = TLContext( + pe_id=0, num_programs=1, + dispatch_cycles=0, + issue_cost_table=DEFAULT_CPU_ISSUE_COST, + ) + run_kernel(kernel, tl2) + cycles2 = sum( + c.cycles for c in tl2.commands if isinstance(c, PeCpuOverheadCmd) + ) + + assert cycles1 == cycles2, ( + f"same kernel via same default table must produce identical " + f"overhead cycles; got {cycles1} vs {cycles2}" + ) + # Concrete expected for this kernel: 2×load(5) + 1×dot(5) + 1×store(5) = 20. + assert cycles1 == 20, ( + f"expected 20 cycles total (2 load + 1 dot + 1 store at 5 ns); " + f"got {cycles1}" + ) + + +# ── T7: back-compat — no table → uniform dispatch_cycles ───────── + + +def test_back_compat_no_table_uses_dispatch_cycles(): + """ADR-0046 §D6 contract preserved when no issue_cost_table is passed. + + Existing call sites doing ``TLContext(dispatch_cycles=0)`` must + continue to emit zero overhead. Existing call sites doing + ``TLContext(dispatch_cycles=1)`` must continue to emit uniform 1. + """ + # Zero path (most existing tests use this). + tl_zero = TLContext(pe_id=0, num_programs=1, dispatch_cycles=0) + tl_zero.load(0x1000, shape=(4, 4), dtype="f16") + overheads = [c for c in tl_zero.commands if isinstance(c, PeCpuOverheadCmd)] + assert overheads == [], ( + f"dispatch_cycles=0 with no table must emit no overhead; " + f"got {[c.cycles for c in overheads]}" + ) + + # Uniform path (test_dispatch_overhead_inserted relies on this). + tl_one = TLContext(pe_id=0, num_programs=1, dispatch_cycles=1) + tl_one.load(0x1000, shape=(4, 4), dtype="f16") + overheads = [c for c in tl_one.commands if isinstance(c, PeCpuOverheadCmd)] + assert overheads == [PeCpuOverheadCmd(cycles=1)], ( + f"dispatch_cycles=1 with no table must emit cycles=1; " + f"got {[c.cycles for c in overheads]}" + ) + + +# ── T8: live PE_CPU constructs TLContext with the default table ── + + +def test_live_pe_cpu_uses_default_table(monkeypatch): + """End-to-end: live PE_CPU greenlet path constructs TLContext with + ``issue_cost_table=DEFAULT_CPU_ISSUE_COST``. This is the wiring assertion + for ADR-0064 D4 "active by default" + review item #6 (path parity). + + We patch ``TLContext.__init__`` to record its kwargs and run a + single-load kernel through the live PE_CPU. The recorded + ``issue_cost_table`` must equal ``DEFAULT_CPU_ISSUE_COST``. + """ + from kernbench.common.cpu_issue_cost import DEFAULT_CPU_ISSUE_COST + from kernbench.triton_emu import tl_context as _tlc + + captured: list[dict] = [] + real_init = _tlc.TLContext.__init__ + + def spy_init(self, *args, **kwargs): + captured.append(dict(kwargs)) + real_init(self, *args, **kwargs) + + monkeypatch.setattr(_tlc.TLContext, "__init__", spy_init) + + clear_registry() + hbm_pa = _hbm_pa(sip=0, cube=0, pe_id=0) + + def single_load_kernel(tl): + tl.load(hbm_pa, shape=(4, 4), dtype="f16") + + register_kernel("test_active_default_table", single_load_kernel) + + engine = _engine() + pe_cpu_id = "sip0.cube0.pe0.pe_cpu" + done = engine._env.event() + txn = Transaction( + request=KernelLaunchMsg( + correlation_id="t", request_id="r", + kernel_ref=KernelRef(name="test_active_default_table", kind="builtin"), + args=(), + ), + path=[pe_cpu_id], step=0, nbytes=0, done=done, + ) + + def inject(): + yield engine._components[pe_cpu_id]._inbox.put(txn) + yield done + + engine._env.process(inject()) + engine._env.run() + clear_registry() + + # Live PE_CPU constructs TLContext on either the greenlet path + # (kernel_runner.py:run) or the legacy replay path + # (pe_cpu.py:_execute_legacy) — whichever is chosen depends on whether + # MemoryStore is wired. Both must pass DEFAULT_CPU_ISSUE_COST. + pe_ctx_calls = [c for c in captured if c.get("pe_id") == 0 + and "issue_cost_table" in c] + assert len(pe_ctx_calls) >= 1, ( + f"expected at least one TLContext constructed by a live PE_CPU " + f"path with an issue_cost_table kwarg; got {len(pe_ctx_calls)} " + f"(all captures: {captured})" + ) + last = pe_ctx_calls[-1] + assert last["issue_cost_table"] == DEFAULT_CPU_ISSUE_COST, ( + f"live PE_CPU must use DEFAULT_CPU_ISSUE_COST; got {last['issue_cost_table']}" + ) diff --git a/tests/test_tl_copy_to.py b/tests/test_tl_copy_to.py new file mode 100644 index 0000000..659a120 --- /dev/null +++ b/tests/test_tl_copy_to.py @@ -0,0 +1,166 @@ +"""Phase 1 spec test for ``tl.copy_to`` (ADR-0063 §D3.1). + +ADR-0063 §D3.1: ``tl.copy_to(dst, src)`` is a TCM-to-TCM byte copy that +lets a kernel write a scoped result's bytes to a persistent (outside- +scope) address. Required for the two-arena flash pattern in §D3. + +Emit-time validation (ADR-0063 §D3.1 "Mechanics" / "Emit-time validation"): + - ``dst.shape == src.shape`` + - ``dst.dtype == src.dtype`` + - ``dst.space == "tcm"`` (TCM-only — HBM goes through ``tl.store``) + - ``src.space == "tcm"`` (same reason) + +op_log shape (ADR-0063 §D3.1 "Mechanics"): + - ``op_kind="math"`` (runs on the vector engine) + - ``op_name="copy"`` + +Phase 1 (this commit): tests only — production code lands in Phase 2. +The tests are written against the ``CopyCmd`` dataclass and the +``TLContext.copy_to`` method that ADR-0063 §D3.1 declares; both are +absent today, so every test in this file should currently fail with +``AttributeError`` / ``ImportError``. +""" +from __future__ import annotations + +import pytest + +from kernbench.triton_emu.tl_context import TLContext + + +def _ctx() -> TLContext: + """TLContext with a real scratch base so _make_compute_out allocates.""" + return TLContext( + pe_id=0, num_programs=1, dispatch_cycles=0, + scratch_base=1 << 24, scratch_size=1 << 20, + ) + + +# ── 1. Method exists on the tl surface ─────────────────────────────── + + +def test_copy_to_method_exists(): + """ADR-0063 §D3.1: TLContext must expose ``copy_to(dst, src)``.""" + ctx = _ctx() + assert hasattr(ctx, "copy_to"), ( + "TLContext must expose copy_to(dst, src) per ADR-0063 §D3.1" + ) + + +# ── 2. Emits a CopyCmd with the documented fields ──────────────────── + + +def test_copy_to_emits_copy_command(): + """ADR-0063 §D3.1 Mechanics: a new ``CopyCmd(src, dst, nbytes)`` + command is recorded; data_op=True (so it appears in op_log).""" + from kernbench.common.pe_commands import CopyCmd # added in Phase 2 + + ctx = _ctx() + src = ctx._make_compute_out(shape=(8, 64), dtype="f16") + dst = ctx._make_compute_out(shape=(8, 64), dtype="f16") + ctx.copy_to(dst, src) + + copy_cmds = [c for c in ctx.commands if isinstance(c, CopyCmd)] + assert len(copy_cmds) == 1, ( + f"exactly one CopyCmd expected; got {len(copy_cmds)}" + ) + cmd = copy_cmds[0] + assert cmd.src is src + assert cmd.dst is dst + assert cmd.nbytes == 8 * 64 * 2 + assert cmd.data_op is True + + +# ── 3. Shape mismatch is rejected at emit time ─────────────────────── + + +def test_copy_to_shape_mismatch_rejected(): + """ADR-0063 §D3.1: ``dst.shape == src.shape`` enforced at emit time.""" + ctx = _ctx() + src = ctx._make_compute_out(shape=(8, 64), dtype="f16") + dst = ctx._make_compute_out(shape=(8, 128), dtype="f16") # mismatched + with pytest.raises((ValueError, AssertionError)) as excinfo: + ctx.copy_to(dst, src) + assert "shape" in str(excinfo.value).lower() + + +# ── 4. Dtype mismatch is rejected at emit time ─────────────────────── + + +def test_copy_to_dtype_mismatch_rejected(): + """ADR-0063 §D3.1: ``dst.dtype == src.dtype`` enforced at emit time.""" + ctx = _ctx() + src = ctx._make_compute_out(shape=(8, 64), dtype="f16") + dst = ctx._make_compute_out(shape=(8, 64), dtype="f32") # mismatched + with pytest.raises((ValueError, AssertionError)) as excinfo: + ctx.copy_to(dst, src) + assert "dtype" in str(excinfo.value).lower() + + +# ── 5. Non-TCM dst is rejected at emit time ────────────────────────── + + +def test_copy_to_dst_must_be_tcm(): + """ADR-0063 §D3.1: ``dst.space == 'tcm'`` enforced at emit time. + + Writing to HBM goes through ``tl.store``, not ``tl.copy_to`` — + keeping copy_to TCM-only avoids polluting op_log with DMA entries + that the bump-allocator scope mechanism doesn't model. + """ + ctx = _ctx() + src = ctx._make_compute_out(shape=(8, 64), dtype="f16") + # Fabricate an HBM-resident dst (e.g. a load handle). + dst_hbm = ctx.load(0x10_000, shape=(8, 64), dtype="f16") + assert dst_hbm.space == "hbm", "fixture sanity check" + with pytest.raises((ValueError, AssertionError)) as excinfo: + ctx.copy_to(dst_hbm, src) + msg = str(excinfo.value).lower() + assert "tcm" in msg or "space" in msg or "hbm" in msg + + +# ── 6. Non-TCM src is rejected at emit time ────────────────────────── + + +def test_copy_to_src_must_be_tcm(): + """ADR-0063 §D3.1: src.space == 'tcm' (symmetric to dst). + + Reading from HBM goes through ``tl.load``, not ``tl.copy_to``. + """ + ctx = _ctx() + src_hbm = ctx.load(0x10_000, shape=(8, 64), dtype="f16") + assert src_hbm.space == "hbm" + dst = ctx._make_compute_out(shape=(8, 64), dtype="f16") + with pytest.raises((ValueError, AssertionError)) as excinfo: + ctx.copy_to(dst, src_hbm) + msg = str(excinfo.value).lower() + assert "tcm" in msg or "space" in msg or "hbm" in msg + + +# ── 7. op_log routes CopyCmd as ("math", "copy", ...) ──────────────── + + +def test_copy_to_op_log_classifies_as_math_copy(): + """ADR-0063 §D3.1 Mechanics: op_kind='math', op_name='copy'. + + The vector engine handles the copy; op_log classification reflects + that (engine-class accounting in bench summaries can sum over + op_kind='math' to include copy time alongside softmax/exp). + """ + from kernbench.common.pe_commands import CopyCmd # added in Phase 2 + from kernbench.sim_engine.op_log import _extract_op_info + + ctx = _ctx() + src = ctx._make_compute_out(shape=(8, 64), dtype="f16") + dst = ctx._make_compute_out(shape=(8, 64), dtype="f16") + ctx.copy_to(dst, src) + + cmd = next(c for c in ctx.commands if isinstance(c, CopyCmd)) + op_kind, op_name, params = _extract_op_info(cmd) + assert op_kind == "math", ( + f"copy runs on the vector engine: op_kind='math'; got {op_kind!r}" + ) + assert op_name == "copy", ( + f"op_name must be 'copy' per ADR-0063 §D3.1; got {op_name!r}" + ) + # Params must let DataExecutor replay (read src, write dst). + assert params.get("dst_addr") == dst.addr + assert params.get("dst_space", "tcm") == "tcm" diff --git a/tests/test_tl_load_lazy.py b/tests/test_tl_load_lazy.py new file mode 100644 index 0000000..c27b266 --- /dev/null +++ b/tests/test_tl_load_lazy.py @@ -0,0 +1,167 @@ +"""Phase 1 spec test for lazy ``tl.load`` (ADR-0062). + +ADR-0062 §D1/§D2: ``tl.load`` is non-blocking. It issues a ``DmaReadCmd`` +and returns immediately with a ``TensorHandle`` carrying a pending +``LoadFuture``. The runtime auto-inserts a wait on that future at the +first consuming op (``tl.dot``, MATH ops, ``tl.store``, ``tl.send``, +``tl.copy_to``, ``tl.composite`` operands). Symmetric with the existing +``recv_async`` / ``RecvFuture`` pattern, generalised to HBM loads. + +Phase 1 (this commit): tests only — production code lands in Phase 2. +All tests currently fail because: + - ``TensorHandle`` has no ``pending`` field + - ``LoadFuture`` class doesn't exist +""" +from __future__ import annotations + +from kernbench.common.pe_commands import DmaReadCmd +from kernbench.triton_emu.tl_context import TLContext + + +def _ctx() -> TLContext: + """TLContext with a real scratch base so _make_compute_out allocates.""" + return TLContext( + pe_id=0, num_programs=1, dispatch_cycles=0, + scratch_base=1 << 24, scratch_size=1 << 20, + ) + + +# ── 1. tl.load handle carries a `pending` field ────────────────────── + + +def test_tl_load_handle_carries_pending_field(): + """ADR-0062 §D2: tl.load handle has a ``pending`` attribute that is + not None — it references the LoadFuture the consumer will await.""" + ctx = _ctx() + handle = ctx.load(0x1000, shape=(8, 64), dtype="f16") + assert hasattr(handle, "pending"), ( + "TensorHandle must have a `pending` field per ADR-0062 §D2" + ) + assert handle.pending is not None, ( + "tl.load handle must carry a non-None pending (LoadFuture)" + ) + + +# ── 2. Constant / math-output handles have pending=None ────────────── + + +def test_constant_handle_pending_is_none(): + """Non-loaded handles (tl.zeros, tl.full, _make_compute_out, tl.arange) + have ``pending=None`` — there is no DMA to wait on. Consumer ops can + safely skip auto-wait for these inputs.""" + ctx = _ctx() + z = ctx.zeros((8, 64), dtype="f16") + assert getattr(z, "pending", "MISSING") is None, ( + "tl.zeros handle must have pending=None" + ) + f = ctx.full((8, 64), value=0.0, dtype="f16") + assert getattr(f, "pending", "MISSING") is None, ( + "tl.full handle must have pending=None" + ) + a = ctx.arange(0, 8, dtype="i32") + assert getattr(a, "pending", "MISSING") is None, ( + "tl.arange handle must have pending=None" + ) + + +def test_compute_out_pending_is_none(): + """Scratch-allocated output handles (via _make_compute_out) have + pending=None. Math ops produce them; no DMA → no auto-wait needed + when they're consumed downstream.""" + ctx = _ctx() + out = ctx._make_compute_out(shape=(8, 64), dtype="f16") + assert getattr(out, "pending", "MISSING") is None, ( + "compute-out handle must have pending=None" + ) + + +# ── 3. LoadFuture class exists ─────────────────────────────────────── + + +def test_loadfuture_class_exists(): + """ADR-0062 §D2: LoadFuture class wraps the DMA done event and a + resolved-flag, analogous to RecvFuture for IPCQ. Lives in + tl_context.py to keep all greenlet bridge types in one module.""" + from kernbench.triton_emu.tl_context import LoadFuture # added Phase 2 + assert LoadFuture is not None + + +# ── 4. Distinct loads produce distinct LoadFutures ─────────────────── + + +def test_distinct_loads_have_distinct_pending(): + """ADR-0062 §D2: two tl.load calls (different addresses) produce two + independent LoadFuture objects. Each consumer must wait on the + right one — sharing the same future would over-serialise.""" + ctx = _ctx() + h1 = ctx.load(0x1000, shape=(8,), dtype="f16") + h2 = ctx.load(0x2000, shape=(8,), dtype="f16") + assert h1.pending is not None and h2.pending is not None + assert h1.pending is not h2.pending, ( + "distinct loads must have distinct LoadFuture objects" + ) + + +# ── 5. LoadFuture starts unresolved ────────────────────────────────── + + +def test_loadfuture_starts_unresolved(): + """ADR-0062 §D2: a fresh LoadFuture has resolved=False — the consumer + op must await it before reading. Once awaited (in greenlet mode by + auto-wait at first use), the SimPy event triggers and a subsequent + consumer can skip the yield.""" + ctx = _ctx() + h = ctx.load(0x1000, shape=(8,), dtype="f16") + assert h.pending.resolved is False, ( + "LoadFuture must start unresolved; got resolved=True at issue time" + ) + + +# ── 6. op_log: dma_read_count unchanged from blocking version ──────── + + +def test_op_log_dma_read_count_unchanged(): + """ADR-0062 §D2 (last paragraph) / Test Req 4: each tl.load still + emits exactly one DmaReadCmd. The asynchrony is a scheduling + property, not a new op kind — dma_read_count and existing op_log + consumers see no change.""" + ctx = _ctx() + ctx.load(0x1000, shape=(8, 64), dtype="f16") + ctx.load(0x2000, shape=(8, 64), dtype="f16") + ctx.load(0x3000, shape=(8, 64), dtype="f16") + dma_cmds = [c for c in ctx.commands if isinstance(c, DmaReadCmd)] + assert len(dma_cmds) == 3, ( + f"lazy tl.load must still emit one DmaReadCmd per call; " + f"got {len(dma_cmds)}" + ) + + +# ── 7. Handle metadata fields preserved across the lazy change ─────── + + +def test_tl_load_handle_metadata_unchanged(): + """ADR-0062 §D1: the tl surface is unchanged. The handle's addr, + shape, dtype, nbytes, space, pinned fields are identical to the + blocking version — only `pending` is new.""" + ctx = _ctx() + h = ctx.load(0x1000, shape=(8, 64), dtype="f16") + assert h.addr == 0x1000 + assert h.shape == (8, 64) + assert h.dtype == "f16" + assert h.nbytes == 8 * 64 * 2 + assert h.space == "hbm" + assert h.pinned is True + + +# ── 8. LoadFuture carries the DmaReadCmd it wraps ──────────────────── + + +def test_loadfuture_carries_cmd(): + """ADR-0062 §D2: LoadFuture references the DmaReadCmd it wraps so + the runtime can resolve handle → command → completion event at + auto-wait time. Mirrors RecvFuture.cmd.""" + ctx = _ctx() + h = ctx.load(0x1000, shape=(8, 64), dtype="f16") + assert hasattr(h.pending, "cmd"), "LoadFuture.cmd must exist" + assert isinstance(h.pending.cmd, DmaReadCmd) + assert h.pending.cmd.handle is h diff --git a/tests/test_tl_scratch_scope.py b/tests/test_tl_scratch_scope.py new file mode 100644 index 0000000..510adee --- /dev/null +++ b/tests/test_tl_scratch_scope.py @@ -0,0 +1,213 @@ +"""Phase 1 spec test for ``tl.scratch_scope`` (ADR-0063, P3a). + +ADR-0063 D1: ``tl.scratch_scope()`` is a context manager whose ``__enter__`` +snapshots ``_scratch_cursor`` and whose ``__exit__`` restores it, freeing +every handle allocated inside the ``with``-block. This pins the bump +allocator behaviour without changing any command emission. + +Phase 1 (this commit): tests only — production code lands in Phase 2. +""" +from __future__ import annotations + +from kernbench.triton_emu.tl_context import TLContext + + +# Configured TLContext with a real (non-zero) scratch base — required for +# _make_compute_out to allocate, see tl_context.py:_scratch_alloc. +def _ctx() -> TLContext: + return TLContext( + pe_id=0, num_programs=1, + dispatch_cycles=0, + scratch_base=1 << 24, # 16 MiB base + scratch_size=1 << 20, # 1 MiB pool + ) + + +# ── 1. Method exists and returns a usable ctx-mgr ───────────────────── + + +def test_scratch_scope_method_exists(): + ctx = _ctx() + assert hasattr(ctx, "scratch_scope"), ( + "TLContext must expose scratch_scope() per ADR-0063 D1" + ) + with ctx.scratch_scope() as scope: + assert scope is not None + + +# ── 2. Recycling within a loop bounds peak cursor ──────────────────── + + +def test_scratch_scope_recycles_within_loop(): + """ADR-0063 D1 / Test Req 1: a loop inside scratch_scope keeps peak + _scratch_cursor bounded by one iteration's footprint.""" + ctx = _ctx() + # One iteration footprint: allocate three small handles via + # _make_compute_out. + one_iter_peaks = [] + for _ in range(10): + with ctx.scratch_scope(): + ctx._make_compute_out(shape=(64,), dtype="f16") + ctx._make_compute_out(shape=(64,), dtype="f16") + ctx._make_compute_out(shape=(64,), dtype="f16") + one_iter_peaks.append(ctx._scratch_cursor) + # After the loop the cursor must be at its post-loop start (0 here), + # not 10× one iteration. + assert ctx._scratch_cursor == 0, ( + f"cursor must reset on each scope exit; final cursor = " + f"{ctx._scratch_cursor}" + ) + # Every iteration sees the same peak — proves recycling. + assert len(set(one_iter_peaks)) == 1, ( + f"every iteration must reach the same peak (recycled); " + f"got distinct peaks {one_iter_peaks}" + ) + + +# ── 3. Allocations outside the scope are not recycled ──────────────── + + +def test_scratch_scope_preserves_outside_allocations(): + """ADR-0063 D3: handles allocated outside the scope (persistent + arena) keep their addresses; only inside-scope allocations are + rewound.""" + ctx = _ctx() + # Persistent allocation BEFORE the scope. + persistent = ctx._make_compute_out(shape=(64,), dtype="f16") + cursor_after_persistent = ctx._scratch_cursor + with ctx.scratch_scope(): + ctx._make_compute_out(shape=(64,), dtype="f16") + ctx._make_compute_out(shape=(64,), dtype="f16") + # On exit the cursor must rewind to the save-point (which equals + # cursor_after_persistent), NOT to 0 — the persistent allocation is + # preserved. + assert ctx._scratch_cursor == cursor_after_persistent, ( + f"cursor must rewind to scope-enter point ({cursor_after_persistent}), " + f"not 0; got {ctx._scratch_cursor}" + ) + # A new allocation after the scope must not collide with the + # persistent one. + fresh = ctx._make_compute_out(shape=(64,), dtype="f16") + assert fresh.addr != persistent.addr, ( + "post-scope allocation must not reuse the persistent address" + ) + + +# ── 4. Nested scopes rewind to their own save-points ───────────────── + + +def test_scratch_scope_nesting(): + """ADR-0063 D4: nested scopes rewind to the matching enter-point.""" + ctx = _ctx() + with ctx.scratch_scope(): + ctx._make_compute_out(shape=(64,), dtype="f16") + outer_save = ctx._scratch_cursor + with ctx.scratch_scope(): + ctx._make_compute_out(shape=(64,), dtype="f16") + ctx._make_compute_out(shape=(64,), dtype="f16") + inner_peak = ctx._scratch_cursor + # Inner exit: rewind to outer_save. + assert ctx._scratch_cursor == outer_save, ( + f"inner exit must rewind to {outer_save}; " + f"got {ctx._scratch_cursor}" + ) + assert inner_peak > outer_save, ( + "sanity: inner scope must have allocated something" + ) + # Outer exit: rewind all the way to 0. + assert ctx._scratch_cursor == 0, ( + f"outer exit must rewind to 0; got {ctx._scratch_cursor}" + ) + + +# ── 5. Control: without scope, allocations pile up ─────────────────── + + +def test_without_scope_grows_unbounded(): + """Sanity check that the bump allocator without scratch_scope grows + linearly — the failure mode ADR-0063 §A.2 cites for the S=16 cap.""" + ctx = _ctx() + for _ in range(10): + ctx._make_compute_out(shape=(64,), dtype="f16") + ctx._make_compute_out(shape=(64,), dtype="f16") + ctx._make_compute_out(shape=(64,), dtype="f16") + # 30 allocations × aligned(64 * 2 = 128 → 128) B = 3840 B + assert ctx._scratch_cursor >= 30 * 128, ( + f"without scope, cursor must grow with allocation count; got " + f"{ctx._scratch_cursor}" + ) + + +# ── 6. Overflow avoided by recycling ───────────────────────────────── + + +def test_overflow_avoided_with_scope(): + """ADR-0063 Test Req 3: a loop that would overflow 1 MiB without + scope completes when wrapped in scratch_scope.""" + # 1 MiB pool / one alloc of 64 KiB = 16 max without scope. + # 100 iterations of 64 KiB without recycling = 6.4 MiB → overflow. + # With scratch_scope: peak ≈ 64 KiB per iter, well under 1 MiB. + ctx = _ctx() + big_shape = (32 * 1024,) # 32 K elements × 2 B = 64 KiB per handle + for _ in range(100): + with ctx.scratch_scope(): + ctx._make_compute_out(shape=big_shape, dtype="f16") + # Completes without raising — the recycle keeps us under the cap. + assert ctx._scratch_cursor == 0 + + +# ── 7. scratch_scope + copy_to integration (ADR-0063 §D3.1) ────────── + + +def test_scoped_loop_with_copy_to_keeps_cursor_bounded(): + """ADR-0063 §D3.1: the two-arena flash pattern. + + Persistent ``running`` allocated OUTSIDE the scope; per-iteration + work allocates inside; the last act before scope exit is + ``tl.copy_to(running, new_running)`` to persist state. After the + loop the cursor must equal the post-persistent-allocation value + (running survived) and never have grown beyond that + one + iteration's footprint. + """ + ctx = _ctx() + # Persistent arena: one running-state handle outside any scope. + running = ctx._make_compute_out(shape=(8, 64), dtype="f16") + cursor_after_persist = ctx._scratch_cursor + assert running.addr != 0 + + n_iters = 50 + for _ in range(n_iters): + with ctx.scratch_scope(): + # Simulate per-tile work: a few scoped allocations. + ctx._make_compute_out(shape=(8, 64), dtype="f16") + ctx._make_compute_out(shape=(8, 64), dtype="f16") + new_running = ctx._make_compute_out(shape=(8, 64), dtype="f16") + # Persist back to the outside-scope handle so its bytes + # survive __exit__. + ctx.copy_to(running, new_running) + # After scope exit cursor returns to the persistent footprint — + # scoped allocations are recycled, running survives. + assert ctx._scratch_cursor == cursor_after_persist, ( + f"after scope exit cursor must rewind to " + f"{cursor_after_persist}; got {ctx._scratch_cursor}" + ) + + assert ctx._scratch_cursor == cursor_after_persist + + +def test_persistent_handle_survives_after_copy_to_in_scope(): + """ADR-0063 §D2 / §D3: a handle allocated outside a scratch_scope + must keep its address even after a copy_to(...) targeting it from + inside the scope and after that scope exits.""" + ctx = _ctx() + persistent = ctx._make_compute_out(shape=(64,), dtype="f16") + persistent_addr_before = persistent.addr + + with ctx.scratch_scope(): + scoped = ctx._make_compute_out(shape=(64,), dtype="f16") + ctx.copy_to(persistent, scoped) + # Out of scope; persistent's address is unchanged. + assert persistent.addr == persistent_addr_before + # A fresh allocation after the scope must not collide with persistent. + fresh = ctx._make_compute_out(shape=(64,), dtype="f16") + assert fresh.addr != persistent.addr diff --git a/topologies/llama70b_4sip.yaml b/topologies/llama70b_4sip.yaml deleted file mode 100644 index ad37d23..0000000 --- a/topologies/llama70b_4sip.yaml +++ /dev/null @@ -1,152 +0,0 @@ -# Llama-70B GQA 4-SIP topology — sub-cycle 3 (ADR-0055 + ADR-0056 context). -# -# Identical to repo-root topology.yaml except for system.sips.count: 2 → 4, -# matching the GQA Llama-70B sharding study's TL/TR baseline at 1 Q-head -# per cube (h_q=64, 8 cubes per KV-group × 8 KV-groups = 64 cubes = 4 SIPs). -# See: -# llm_paper_review/notes/GQA_MHA_sharding/scripts/_gen_llama70b_1M_4cases.py -# lines 17-19 (model dims) and 287-302 (Q/cube scaling configs). -# -# Opt-in only — the milestone-gqa-llama70b bench (sub-cycle 4) points at -# this file via env var. The repo-root topology.yaml is unchanged and -# continues to drive milestone-1h-ccl / -gemm at their original 2-SIP -# scale (CLAUDE.md "Surgical Changes"). - -system: - ns_per_mm: 0.01 # wire propagation delay: 10 ps/mm (on-chip silicon) - - sips: - count: 4 - topology: ring_1d - - components: - switch: { kind: switch, impl: builtin.switch, attrs: { overhead_ns: 5.0 } } - - links: - io_ep_to_switch: - kind: pcie - bw_gbs_per_ep: 768.0 - distance_mm: 20.0 - -sip: - cube_mesh: { w: 4, h: 4 } - - iochiplet: - components: - pcie_ep: { kind: pcie_ep, impl: builtin.pcie_ep, attrs: { overhead_ns: 5.0 } } - io_cpu: { kind: io_cpu, impl: builtin.io_cpu, attrs: { overhead_ns: 10.0 } } - io_noc: { kind: io_noc, impl: builtin.forwarding, attrs: { overhead_ns: 0.0 } } - links: - pcie_ep_to_noc_bw_gbs: 256.0 - pcie_ep_to_noc_mm: 1.0 - io_cpu_to_noc_bw_gbs: 256.0 - io_cpu_to_noc_mm: 0.5 - ucie: - overhead_ns: 8.0 - n_connections: 4 - per_connection_bw_gbs: 128.0 # 4 × 128 = 512 GB/s = PHY BW - noc_to_ucie_mm: 0.5 - instances: - - id: io0 - place: { side: N, offset_norm: 0.5 } - ucie: { phy_bw_gbs: 512.0, phys: [P0, P1, P2, P3] } - cube_ports: - - { cube: {xy: [0,0]}, cube_side: N, phy: P0, distance_mm: 2.0 } - - { cube: {xy: [1,0]}, cube_side: N, phy: P1, distance_mm: 2.0 } - - { cube: {xy: [2,0]}, cube_side: N, phy: P2, distance_mm: 2.0 } - - { cube: {xy: [3,0]}, cube_side: N, phy: P3, distance_mm: 2.0 } - - links: - inter_cube_mesh: - bw_gbs_per_ucie_phy: 512.0 - distance_mm_across_seam: 1.0 - routing: { algo: xy } - -cube: - geometry: - cube_mm: { w: 17.0, h: 14.0 } - hbm_mm: { w: 9.0, h: 5.0 } - ucie_mm: { size: 2.0 } - - pe_layout: - corners: [NW, NE, SW, SE] # N corners → top PE rows; S corners → bottom PE rows - pe_per_corner: 2 # total PEs per cube: 4 * 2 = 8 - - pe_template: - components: - pe_cpu: { kind: pe_cpu, impl: builtin.pe_cpu, attrs: { overhead_ns: 2.0 } } - pe_scheduler: { kind: pe_scheduler, impl: builtin.pe_scheduler, attrs: { overhead_ns: 1.0 } } - pe_dma: { kind: pe_dma, impl: builtin.pe_dma, attrs: { rd_engines: 1, wr_engines: 1 } } - pe_gemm: { kind: pe_gemm, impl: builtin.pe_gemm, attrs: { overhead_ns: 0.0, shared_resource: accel_slot, peak_tflops_f16: 8.0 } } - pe_math: { kind: pe_math, impl: builtin.pe_math, attrs: { overhead_ns: 0.0, shared_resource: accel_slot } } - pe_fetch_store: { kind: pe_fetch_store, impl: builtin.pe_fetch_store, attrs: { overhead_ns: 0.0 } } - pe_mmu: { kind: pe_mmu, impl: builtin.pe_mmu, attrs: { tlb_overhead_ns: 0.5, page_size: 4096 } } - pe_tcm: { kind: pe_tcm, impl: builtin.pe_tcm, attrs: { size_mb: 16, read_bw_gbs: 512.0, write_bw_gbs: 512.0, kernel_scratch_mb: 1 } } - pe_ipcq: { kind: pe_ipcq, impl: builtin.pe_ipcq, attrs: { overhead_ns: 0.0 } } - links: - pe_cpu_to_scheduler_mm: 0.5 - scheduler_to_dma_mm: 0.5 - scheduler_to_gemm_mm: 0.5 - scheduler_to_math_mm: 0.5 - scheduler_to_fetch_store_mm: 0.5 - dma_to_tcm_bw_gbs: 512.0 - dma_to_tcm_mm: 0.5 - dma_to_fetch_store_mm: 0.0 # DMA → fetch_store chaining (ADR-0014 D6) - fetch_store_to_tcm_bw_gbs: 512.0 - fetch_store_to_tcm_mm: 0.0 - fetch_store_to_gemm_mm: 0.0 # fetch → GEMM chaining (ADR-0014 D6) - fetch_store_to_math_mm: 0.0 # fetch → MATH chaining (ADR-0014 D6) - gemm_to_fetch_store_mm: 0.0 # GEMM → store chaining (ADR-0014 D6) - gemm_to_math_mm: 0.0 # GEMM → MATH epilogue chaining (ADR-0014 D6) - math_to_fetch_store_mm: 0.0 # MATH → store chaining (ADR-0014 D6) - fetch_store_to_dma_mm: 0.0 # store → DMA writeback chaining (ADR-0014 D6) - gemm_to_tcm_bw_gbs: 512.0 - gemm_to_tcm_mm: 0.5 - math_to_tcm_bw_gbs: 512.0 - math_to_tcm_mm: 0.5 - cpu_to_ipcq_mm: 0.5 # PE_CPU → PE_IPCQ (ADR-0023) - ipcq_to_dma_mm: 0.0 # PE_IPCQ → PE_DMA token forwarding (ADR-0023) - dma_to_ipcq_mm: 0.0 # PE_DMA → PE_IPCQ metadata arrival (ADR-0023) - - memory_map: - hbm_total_gb_per_cube: 48 - hbm_slices_per_cube: 8 - hbm_total_bw_gbs: 1024.0 - hbm_mapping_mode: n_to_one # one_to_one | n_to_one (ADR-0017 D8) - hbm_pseudo_channels: 64 # total pseudo channels per cube - hbm_channels_per_pe: 8 # = pseudo_channels / pes_per_cube - hbm_channel_bw_gbs: 32.0 # per-channel bandwidth (GB/s) - - components: - noc_router: { kind: noc_router, impl: builtin.forwarding, attrs: { overhead_ns: 2.0 } } - m_cpu: { kind: m_cpu, impl: builtin.m_cpu, attrs: { overhead_ns: 5.0 } } - hbm_ctrl: { kind: hbm_ctrl, impl: builtin.hbm_ctrl, attrs: { capacity: 1, efficiency: 1.0, num_pcs: 8, burst_bytes: 256, switch_penalty_ns: 0.0 } } - sram: { kind: sram, impl: builtin.sram, attrs: { size_mb: 32, overhead_ns: 2.0 } } - - # Physical placement of non-PE components (mm coordinates) - placement: - m_cpu: { pos_mm: [7.5, 3.0] } # top center, below UCIe-N - sram: { pos_mm: [1.5, 9.0] } # left side, below HBM zone - - ucie: - decompose: true - ports: [N, S, E, W] - overhead_ns: 8.0 - n_connections: 4 # independent NOC↔UCIe connections per port - per_connection_bw_gbs: 128.0 # BW per connection; 4 × 128 = 512 GB/s = UCIe PHY BW - - links: - # Router mesh links (ADR-0017 D5) - router_link_bw_gbs: 256.0 # inter-router XY mesh link BW - router_overhead_ns: 2.0 # per-router switching overhead - pe_to_router_bw_gbs: 256.0 # PE_DMA ↔ router (= N × channel_bw) - hbm_to_router_bw_gbs: 256.0 # HBM_CTRL ↔ router (= N × channel_bw) - sram_to_router_bw_gbs: 128.0 # SRAM ↔ router - m_cpu_to_router_mm: 0.0 # M_CPU ↔ router distance - pe_dma_to_noc_bw_gbs: 256.0 # PE → router BW (= HBM slice BW, no bottleneck) - noc_to_pe_cpu_mm: 0.0 # router → PE_CPU distance (command path) - -visualization: - emit_views: [system, sip, cube] - sip_ids: [0] - cubes: [0, 9, 15]