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:
@@ -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
|
||||
Reference in New Issue
Block a user