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,161 @@
|
||||
"""Phase 1 spec test for P6b GQA prefill Ring KV (head-parallel, C>1).
|
||||
|
||||
P6b adds the Ring KV rotation (ADR-0060 §5.5) to the prefill kernel.
|
||||
Each CUBE owns one Q head + its KV slice; over C ring steps the KV
|
||||
blocks rotate around the C CUBEs so every CUBE sees every block. The
|
||||
running ``(m, ℓ, O)`` is folded inside the loop. No reduce — each CUBE
|
||||
writes its own head's output.
|
||||
|
||||
Requires a new SFR install ``configure_sfr_intercube_ring(ring_size=C)``
|
||||
that wires:
|
||||
- ``intra_*`` : 2×4 PE grid within a CUBE (same as multisip)
|
||||
- ``E/W`` : 1D ring of CUBEs 0..ring_size-1 WITH WRAP
|
||||
(symmetric to ``configure_sfr_intracube_pe_ring``,
|
||||
applied at CUBE level)
|
||||
- ``global_*`` : SIP-level (same as multisip)
|
||||
|
||||
The kernel ring body:
|
||||
for step in range(1, C):
|
||||
tl.send(dir="W", src=Kc)
|
||||
tl.send(dir="W", src=Vc)
|
||||
Kc = tl.recv(dir="E", ...)
|
||||
Vc = tl.recv(dir="E", ...)
|
||||
... local partial + online-softmax merge into (m, ℓ, O) ...
|
||||
|
||||
Per CUBE per step: 2 sends (K, V) → 2 ``ipcq_copy``. Across all CUBEs:
|
||||
``(C-1) * 2 * C`` total ipcq_copy.
|
||||
|
||||
Restriction in P6b first cut: ``C ∈ {1, 2, 4}`` (single row of the 4×4
|
||||
cube mesh). C=8 ring would span rows — follow-on.
|
||||
|
||||
Phase 1 (this commit): tests only — production code lands in Phase 2.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
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_ring # noqa: F401 (Phase 2)
|
||||
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 _run_prefill_ring(*, T_q: int, S_kv: int, C: int):
|
||||
"""Head-parallel prefill with Ring KV across C CUBEs."""
|
||||
topo = resolve_topology(str(TOPOLOGY_DEFAULT))
|
||||
|
||||
def _bench_fn(ctx):
|
||||
# P6b: new SFR install with cube-level ring wrap.
|
||||
configure_sfr_intercube_ring(
|
||||
ctx.engine, ctx.spec, _ccl_cfg(), ring_size=C,
|
||||
)
|
||||
# Q replicated on every CUBE (zeros for testing; per-CUBE head
|
||||
# indexing is implicit). KV sequence-sharded by CUBE. O
|
||||
# distributed — each CUBE writes its slice of (T_q*C, d_head).
|
||||
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=f"q_t{T_q}_c{C}_ring")
|
||||
k = ctx.zeros((S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv,
|
||||
name=f"k_t{T_q}_c{C}_ring")
|
||||
v = ctx.zeros((S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv,
|
||||
name=f"v_t{T_q}_c{C}_ring")
|
||||
# O: (T_q * C, d_head), each CUBE writes (T_q, d_head) at its
|
||||
# slice. dma_write_count = C (per-CUBE distributed output).
|
||||
o = ctx.empty((T_q * C, D_HEAD), dtype=DTYPE, dp=dp_o,
|
||||
name=f"o_t{T_q}_c{C}_ring")
|
||||
ctx.launch(
|
||||
f"gqa_prefill_ring_t{T_q}_s{S_kv}_c{C}",
|
||||
gqa_prefill_kernel,
|
||||
q, k, v, o,
|
||||
T_q, S_kv, D_HEAD, C,
|
||||
_auto_dim_remap=False,
|
||||
)
|
||||
|
||||
return run_bench(
|
||||
topology=topo, bench_fn=_bench_fn,
|
||||
device=resolve_device(None),
|
||||
engine_factory=_engine_factory,
|
||||
)
|
||||
|
||||
|
||||
def _count(op_log, name: str) -> int:
|
||||
return sum(1 for r in op_log if r.op_name == name)
|
||||
|
||||
|
||||
# ── C=2 Ring KV ────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_prefill_ring_c_two_completes():
|
||||
"""C=2: 1 ring step rotates KV between the 2 CUBEs."""
|
||||
result = _run_prefill_ring(T_q=4, S_kv=16, C=2)
|
||||
assert result.completion.ok, (
|
||||
f"prefill ring C=2 failed: {result.completion}"
|
||||
)
|
||||
|
||||
|
||||
def test_prefill_ring_c_two_distributed_output():
|
||||
"""C=2: per-CUBE distributed output, no reduce → dma_write_count == 2."""
|
||||
result = _run_prefill_ring(T_q=4, S_kv=16, C=2)
|
||||
assert result.completion.ok
|
||||
n_writes = _count(result.engine.op_log, "dma_write")
|
||||
assert n_writes == 2, (
|
||||
f"C=2 prefill: expected 2 dma_write (one per CUBE); got {n_writes}"
|
||||
)
|
||||
|
||||
|
||||
def test_prefill_ring_c_two_ipcq_count():
|
||||
"""C=2: 1 ring step × 2 handles (K, V) × 2 CUBEs = 4 ipcq_copy."""
|
||||
result = _run_prefill_ring(T_q=4, S_kv=16, C=2)
|
||||
assert result.completion.ok
|
||||
n_copy = _count(result.engine.op_log, "ipcq_copy")
|
||||
expected = (2 - 1) * 2 * 2
|
||||
assert n_copy == expected, (
|
||||
f"C=2 ring: expected {expected} ipcq_copy "
|
||||
f"((C-1)·2·C = 1·2·2); got {n_copy}"
|
||||
)
|
||||
|
||||
|
||||
# ── C=4 Ring KV (combined assertions) ─────────────────────────────────
|
||||
|
||||
|
||||
def test_prefill_ring_c_four_combined():
|
||||
"""C=4: 3 ring steps rotate KV around 4 CUBEs.
|
||||
Expected: completes; 4 dma_writes; (4-1)·2·4 = 24 ipcq_copy."""
|
||||
result = _run_prefill_ring(T_q=4, S_kv=32, C=4)
|
||||
assert result.completion.ok, (
|
||||
f"prefill ring C=4 failed: {result.completion}"
|
||||
)
|
||||
n_writes = _count(result.engine.op_log, "dma_write")
|
||||
assert n_writes == 4, (
|
||||
f"C=4 prefill: expected 4 dma_write (one per CUBE); got {n_writes}"
|
||||
)
|
||||
n_copy = _count(result.engine.op_log, "ipcq_copy")
|
||||
expected = (4 - 1) * 2 * 4
|
||||
assert n_copy == expected, (
|
||||
f"C=4 ring: expected {expected} ipcq_copy "
|
||||
f"((C-1)·2·C = 3·2·4); got {n_copy}"
|
||||
)
|
||||
Reference in New Issue
Block a user