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,236 @@
|
||||
"""milestone-gqa-headline bench: real GQA + 2-level SP + Ring KV.
|
||||
|
||||
Wires the new ``_gqa_decode`` and ``_gqa_prefill`` kernels through 4
|
||||
panels with real GQA (h_q = G·h_kv, G > 1 on the decode side), writing
|
||||
per-panel ``op_log_summary`` into ``sweep.json``. Independent from the
|
||||
existing ``milestone-gqa-llama70b`` validation-scale bench (which stays
|
||||
on the legacy baseline kernels).
|
||||
|
||||
Restrictions (P7 first cut):
|
||||
- C ≤ 4 (single-row inter-CUBE ring SFR; multi-row deferred)
|
||||
- Single SIP (default ``topology.yaml`` 4×4 cube mesh; 4-SIP
|
||||
headline deferred)
|
||||
- No figure renderers (defer to a separate cycle)
|
||||
|
||||
Panels:
|
||||
single_user_prefill_gqa : prefill C=1, T_q=4, S_kv=16
|
||||
multi_user_prefill_gqa : prefill C=4 Ring KV, T_q=4, S_kv=16
|
||||
single_user_decode_gqa : decode C=1, P=8, h_q=8, h_kv=1, S_kv=64
|
||||
(M-fold + intra-cube row-then-col chain)
|
||||
multi_user_decode_gqa : decode C=4, P=8, h_q=8, h_kv=1, S_kv=128
|
||||
(M-fold + 2-level chain reduce-to-root)
|
||||
|
||||
Gated by ``GQA_HEADLINE_RUN=1``.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
from kernbench.benches._gqa_decode import gqa_decode_kernel
|
||||
from kernbench.benches._gqa_prefill import gqa_prefill_kernel
|
||||
from kernbench.benches.registry import bench
|
||||
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
|
||||
|
||||
_OUTPUT_DIR = Path(__file__).resolve().parent / "1H_milestone_output" / "gqa_headline"
|
||||
_SWEEP_JSON = _OUTPUT_DIR / "sweep.json"
|
||||
|
||||
# ── Panel configs ────────────────────────────────────────────────────
|
||||
|
||||
_DTYPE = "f16"
|
||||
_D_HEAD = 64
|
||||
_T_Q_PREFILL = 4
|
||||
_T_Q_DECODE = 1
|
||||
_S_KV_PREFILL = 16
|
||||
_H_Q_DECODE = 8 # real GQA: G = H_Q_DECODE / H_KV_DECODE = 8
|
||||
_H_KV_DECODE = 1
|
||||
|
||||
_PANELS = (
|
||||
"single_user_prefill_gqa",
|
||||
"multi_user_prefill_gqa",
|
||||
"single_user_decode_gqa",
|
||||
"multi_user_decode_gqa",
|
||||
)
|
||||
|
||||
# Each entry: (kind, panel-specific params)
|
||||
_PANEL_DISPATCH: dict[str, tuple[str, dict]] = {
|
||||
"single_user_prefill_gqa": ("prefill", {"C": 1, "S_kv": _S_KV_PREFILL}),
|
||||
"multi_user_prefill_gqa": ("prefill", {"C": 4, "S_kv": _S_KV_PREFILL}),
|
||||
"single_user_decode_gqa": ("decode", {"C": 1, "P": 8, "S_kv": 64}),
|
||||
"multi_user_decode_gqa": ("decode", {"C": 4, "P": 8, "S_kv": 128}),
|
||||
}
|
||||
|
||||
|
||||
def _ccl_cfg():
|
||||
return resolve_algorithm_config(
|
||||
load_ccl_config(), name="lrab_hierarchical_allreduce",
|
||||
)
|
||||
|
||||
|
||||
# ── Per-kind launch helpers ──────────────────────────────────────────
|
||||
|
||||
|
||||
def _run_prefill_panel(ctx, *, panel: str, C: int, S_kv: int) -> None:
|
||||
if C > 1:
|
||||
configure_sfr_intercube_ring(
|
||||
ctx.engine, ctx.spec, _ccl_cfg(), ring_size=C,
|
||||
)
|
||||
dp_q = DPPolicy(cube="replicate", pe="replicate",
|
||||
num_cubes=C, num_pes=1)
|
||||
dp_kv = DPPolicy(cube="row_wise" if C > 1 else "replicate",
|
||||
pe="replicate", num_cubes=C, num_pes=1)
|
||||
dp_o = DPPolicy(cube="row_wise" if C > 1 else "replicate",
|
||||
pe="replicate", num_cubes=C, num_pes=1)
|
||||
q = ctx.zeros((_T_Q_PREFILL, _D_HEAD),
|
||||
dtype=_DTYPE, dp=dp_q, name=f"{panel}_q")
|
||||
k = ctx.zeros((S_kv, _D_HEAD),
|
||||
dtype=_DTYPE, dp=dp_kv, name=f"{panel}_k")
|
||||
v = ctx.zeros((S_kv, _D_HEAD),
|
||||
dtype=_DTYPE, dp=dp_kv, name=f"{panel}_v")
|
||||
o = ctx.empty((_T_Q_PREFILL * C, _D_HEAD),
|
||||
dtype=_DTYPE, dp=dp_o, name=f"{panel}_o")
|
||||
ctx.launch(
|
||||
panel, gqa_prefill_kernel,
|
||||
q, k, v, o,
|
||||
_T_Q_PREFILL, S_kv, _D_HEAD, C,
|
||||
_auto_dim_remap=False,
|
||||
)
|
||||
|
||||
|
||||
def _run_decode_panel(ctx, *, panel: str, C: int, P: int, S_kv: int) -> None:
|
||||
configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
|
||||
dp_full = DPPolicy(cube="replicate", pe="replicate",
|
||||
num_cubes=C, num_pes=P)
|
||||
dp_kv = DPPolicy(cube="row_wise" if C > 1 else "replicate",
|
||||
pe="row_wise", num_cubes=C, num_pes=P)
|
||||
q = ctx.zeros((_T_Q_DECODE, _H_Q_DECODE * _D_HEAD),
|
||||
dtype=_DTYPE, dp=dp_full, name=f"{panel}_q")
|
||||
k = ctx.zeros((S_kv, _H_KV_DECODE * _D_HEAD),
|
||||
dtype=_DTYPE, dp=dp_kv, name=f"{panel}_k")
|
||||
v = ctx.zeros((S_kv, _H_KV_DECODE * _D_HEAD),
|
||||
dtype=_DTYPE, dp=dp_kv, name=f"{panel}_v")
|
||||
o = ctx.empty((_T_Q_DECODE, _H_Q_DECODE * _D_HEAD),
|
||||
dtype=_DTYPE, dp=dp_full, name=f"{panel}_o")
|
||||
ctx.launch(
|
||||
panel, gqa_decode_kernel,
|
||||
q, k, v, o,
|
||||
_T_Q_DECODE, S_kv, _H_Q_DECODE, _H_KV_DECODE, _D_HEAD, C, P,
|
||||
_auto_dim_remap=False,
|
||||
)
|
||||
|
||||
|
||||
def _make_bench_fn(panel: str):
|
||||
kind, params = _PANEL_DISPATCH[panel]
|
||||
|
||||
def _bench_fn(ctx):
|
||||
if kind == "prefill":
|
||||
_run_prefill_panel(ctx, panel=panel, **params)
|
||||
else:
|
||||
_run_decode_panel(ctx, panel=panel, **params)
|
||||
|
||||
return _bench_fn
|
||||
|
||||
|
||||
# ── Op-log summary ──────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _summarize_op_log(op_log) -> dict[str, int]:
|
||||
"""Per-panel op_log counts (gemm, ipcq_copy, dma_read, dma_write)."""
|
||||
gemm_count = 0
|
||||
ipcq_copy_count = 0
|
||||
dma_read_count = 0
|
||||
dma_write_count = 0
|
||||
for r in op_log:
|
||||
if r.op_kind == "gemm":
|
||||
gemm_count += 1
|
||||
elif r.op_name == "dma_read":
|
||||
dma_read_count += 1
|
||||
elif r.op_name == "dma_write":
|
||||
dma_write_count += 1
|
||||
elif r.op_name == "ipcq_copy":
|
||||
ipcq_copy_count += 1
|
||||
return {
|
||||
"gemm_count": gemm_count,
|
||||
"ipcq_copy_count": ipcq_copy_count,
|
||||
"dma_read_count": dma_read_count,
|
||||
"dma_write_count": dma_write_count,
|
||||
}
|
||||
|
||||
|
||||
def _run_panel(panel: str, topology: str) -> dict:
|
||||
"""Run one panel in a fresh engine; return its row dict."""
|
||||
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
|
||||
|
||||
topo = resolve_topology(topology)
|
||||
result = run_bench(
|
||||
topology=topo, bench_fn=_make_bench_fn(panel),
|
||||
device=resolve_device(None),
|
||||
engine_factory=lambda t, d: GraphEngine(
|
||||
getattr(t, "topology_obj", t), enable_data=True,
|
||||
),
|
||||
)
|
||||
if not result.completion.ok:
|
||||
raise RuntimeError(
|
||||
f"milestone-gqa-headline panel {panel!r} failed: "
|
||||
f"{result.completion}"
|
||||
)
|
||||
kind, params = _PANEL_DISPATCH[panel]
|
||||
return {
|
||||
"panel": panel,
|
||||
"kind": kind,
|
||||
**params,
|
||||
"op_log_summary": _summarize_op_log(result.engine.op_log),
|
||||
}
|
||||
|
||||
|
||||
# ── Bench entry ──────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@bench(
|
||||
name="milestone-gqa-headline",
|
||||
description="Headline GQA milestone — real GQA h_q=8/h_kv=1 + 2-level SP (decode) + Ring KV (prefill).",
|
||||
)
|
||||
def run(torch) -> None:
|
||||
"""Drive 4 headline panels through the new GQA kernels; write sweep.json.
|
||||
|
||||
Gated by GQA_HEADLINE_RUN=1.
|
||||
"""
|
||||
_OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
|
||||
if not os.environ.get("GQA_HEADLINE_RUN"):
|
||||
raise RuntimeError(
|
||||
"milestone-gqa-headline needs GQA_HEADLINE_RUN=1."
|
||||
)
|
||||
|
||||
topology = os.environ.get("GQA_HEADLINE_TOPOLOGY", "topology.yaml")
|
||||
rows = [_run_panel(panel, topology) for panel in _PANELS]
|
||||
sweep = {
|
||||
"version": 1,
|
||||
"panels": list(_PANELS),
|
||||
"config": {
|
||||
"T_q_prefill": _T_Q_PREFILL,
|
||||
"T_q_decode": _T_Q_DECODE,
|
||||
"S_kv_prefill": _S_KV_PREFILL,
|
||||
"h_q_decode": _H_Q_DECODE,
|
||||
"h_kv_decode": _H_KV_DECODE,
|
||||
"d_head": _D_HEAD,
|
||||
},
|
||||
"rows": rows,
|
||||
}
|
||||
_SWEEP_JSON.write_text(json.dumps(sweep, indent=2))
|
||||
print(f" milestone-gqa-headline: {len(rows)} rows -> {_SWEEP_JSON}")
|
||||
|
||||
# Sentinel tensor (ADR-0045 D4 / ADR-0054 D2 carve-out).
|
||||
torch.zeros(
|
||||
(1, 1), dtype="f16",
|
||||
dp=DPPolicy(cube="row_wise", pe="replicate", num_cubes=1, num_pes=1),
|
||||
name="milestone_gqa_headline_sentinel",
|
||||
)
|
||||
Reference in New Issue
Block a user