gqa: ADR-0060/0062/0063/0064 unified GQA kernels + CPU cost model
Land the new GQA fused-attention kernels (ADR-0060) for prefill/decode
across long and short context, the TL discipline primitives they depend
on (ADR-0062 lazy load, ADR-0063 scratch_scope + copy_to), and the
per-op-type CPU issue cost model (ADR-0064). Remove the pre-ADR-0060
mesh-attention baseline now that the unified kernels supersede it.
ADR-0060 (long context)
- _gqa_decode.py: M-fold + 2-level chain reduce-to-root (Level-2
intra-CUBE row-then-col + Level-1 inter-CUBE) — root-only output.
- _gqa_prefill.py: head-parallel + Ring KV rotation around C CUBEs,
online-softmax merge per ring step, per-CUBE distributed output.
- Each merge stage wraps in scratch_scope() and persists running
(m, l, O) via copy_to() to lift the 1 MiB scratch ceiling.
ADR-0060 §B.split.2 (short context, kv_per_cube in {1,2,4,8})
- _gqa_decode_short.py / _gqa_prefill_short.py: no cube-SP; each CUBE
owns whole KV heads; PE-parallel heads with intra-group chain
reduce. Prefill has no Ring KV (each head fully resident).
ADR-0062 (lazy tl.load): future-bearing TensorHandle, auto-wait at
first consuming op (dot/MATH/store/send/copy_to/composite).
ADR-0063 (tl.scratch_scope + tl.copy_to): scoped per-tile arena with
copy_to writeback primitive for persistent running state.
ADR-0064 (CPU issue cost model)
- common/cpu_issue_cost.py: per-op-type table (composite=40 ns,
primitives=5 ns); ratios are load-bearing per D1.
- TLContext: issue_cost_table param; _emit_dispatch_overhead(kind)
consults table with dispatch_cycles fallback (ADR-0046 §D6
back-compat).
- Live PE_CPU paths (greenlet + legacy) construct TLContext with
DEFAULT_CPU_ISSUE_COST so saturation lever (ADR-0060 §1) is
measurable end-to-end.
P7 headline bench: milestone-gqa-headline writes per-panel
op_log_summary to 1H_milestone_output/gqa_headline/sweep.json. No
figure renderers yet (deferred).
Removals (pre-ADR-0060 baseline now superseded):
- benches: _attention_mesh_kv.py, _attention_mesh_mlo.py,
_attention_mesh_mlo_2d.py, milestone_gqa_llama70b.py
- tests: test_attention_*, test_mesh_*, test_milestone_gqa_llama70b
- topology: llama70b_4sip.yaml (only consumer was the deleted diag)
- artifacts: 1H_milestone_output/gqa/ (sweep.json + 5 PNGs)
- tests/gqa/ plot helper + test (broken on Windows Tcl/Tkinter)
- ADR-0060/0061 references to deleted file paths cleaned up
(EN + KO kept in sync).
Tests: 124/124 focused regression green (attention + Phase E + TL
discipline + triton_emu + pe_components). Full regression: 764 pass,
2 pre-existing test_bench_registry failures (stale EXPECTED_NAMES
across multiple benches, not introduced here).
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
@@ -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 커널;
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user