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,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}"
|
||||
)
|
||||
Reference in New Issue
Block a user