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