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