gqa: reorganize benches into gqa_helpers/ subpackage; drop legacy headline
Splits the GQA helpers into a dedicated subpackage to make room for the
prefill 4-cases study (next commit) and a single umbrella bench
(milestone-1h-gqa, after that).
Layout:
benches/gqa_helpers/
long_ctx/ — decode 4-cases kernels + sweep runner
short_ctx/ — prefill/decode short-context kernels
shared/ — _gqa_panel_helpers + decode_opt2 (context-agnostic)
The registry audit now skips subpackages so gqa_helpers/ (without a
leading underscore) doesn't get audited for @bench decorators.
Also drops the legacy milestone-gqa-headline bench, its
_gqa_attention_prefill_long kernel, 6 dependent prefill tests, the
paper_gqa_latency.py report harness, and the 3 stale headline-derived
PNGs the §6 wire-up referenced (paper will re-pull from the new
1H_milestone_output/gqa/long_ctx/ once §6 is updated).
The _ccl_cfg and _summarize_op_log helpers used to live in the
headline bench; extracted them to gqa_helpers/shared/_gqa_panel_helpers.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
@@ -1,31 +0,0 @@
|
||||
{
|
||||
"version": 1,
|
||||
"panels": [
|
||||
"single_kv_group_prefill_gqa_c8_p8"
|
||||
],
|
||||
"config": {
|
||||
"T_q_prefill": 4,
|
||||
"T_q_decode": 1,
|
||||
"S_kv_prefill": 16,
|
||||
"h_q_decode": 8,
|
||||
"h_kv_decode": 1,
|
||||
"d_head": 64
|
||||
},
|
||||
"rows": [
|
||||
{
|
||||
"panel": "single_kv_group_prefill_gqa_c8_p8",
|
||||
"kind": "prefill",
|
||||
"C": 8,
|
||||
"P": 8,
|
||||
"T_q": 1024,
|
||||
"S_kv": 1024,
|
||||
"d_head": 128,
|
||||
"op_log_summary": {
|
||||
"gemm_count": 1024,
|
||||
"ipcq_copy_count": 896,
|
||||
"dma_read_count": 192,
|
||||
"dma_write_count": 64
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -1,190 +0,0 @@
|
||||
"""GQA fused-attention prefill kernel — long context (ADR-0060 §5.5).
|
||||
|
||||
Head-parallel: each CUBE owns one Q head and one KV slice. Over C ring
|
||||
steps the KV blocks rotate around the C CUBEs (W-send / E-recv) so every
|
||||
CUBE sees every block; the online-softmax merge folds each step into the
|
||||
running ``(m, ℓ, O)``. No inter-CUBE reduce — each CUBE writes its own
|
||||
head's output.
|
||||
|
||||
The ring is **tile-granular** (ADR-0060 §5.5 + ADR-0063 §A.2): each
|
||||
ring step transmits ``n_tiles`` tiles of size ``TILE_S_KV`` rather than
|
||||
one full ``S_local`` slice. The kernel loop is nested over
|
||||
``(ring_step, tile_idx)``, with the bootstrap at ``(k=0, t=0)``
|
||||
establishing the persistent ``(m, ℓ, O)`` and every subsequent
|
||||
iteration folding its tile in inside a ``tl.scratch_scope``. Per-rank
|
||||
persistent scratch is ``(m, ℓ, O)`` only (~1 KB); per-tile in-scope
|
||||
scratch is bounded by ``TILE_S_KV`` regardless of ``S_local``.
|
||||
|
||||
Topology / SFR:
|
||||
- Requires ``configure_sfr_intercube_ring`` — either ``ring_size=C``
|
||||
(1D row, ``C ≤ mesh_w``) or ``submesh_shape=(rows, cols)`` (snake
|
||||
Hamiltonian cycle through a rectangular sub-mesh, for ``C > mesh_w``).
|
||||
- ``P == 1`` (default): only PE 0 of each CUBE participates;
|
||||
intra-CUBE PE parallelism is disabled.
|
||||
- ``P > 1``: all P PEs of each CUBE participate via query-axis
|
||||
split (ADR-0060 §5.5 last bullet + §B-item-3) — each PE owns
|
||||
``T_q/P`` disjoint query rows. Output rows are disjoint per PE,
|
||||
so no intra-CUBE reduce is needed. Each PE drives its own
|
||||
same-lane ring via independent IPCQ channels (P parallel rings).
|
||||
Requires ``T_q % P == 0``.
|
||||
|
||||
Layout caveats:
|
||||
- GEMMs use ``tl.dot`` (no composite epilogue / ``softmax_scale``).
|
||||
- K loaded as ``(d_head, tile_s)`` via byte-conserving reshape of
|
||||
the deployed ``(tile_s, d_head)`` shard (ADR-0060 §3
|
||||
reshape-not-transpose caveat).
|
||||
- No causal masking / step-skip; blocking ``tl.recv``.
|
||||
- ``TILE_S_KV`` is assumed to divide ``S_local`` evenly; last-tile
|
||||
padding is not modelled.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
TILE_S_KV = 1024 # ADR-0063 §A.2 S_kv-axis tile sweep (per-tile width).
|
||||
|
||||
|
||||
def _merge_running(m, l, O, m_step, l_step, O_step, *, tl):
|
||||
"""Online-softmax merge of two partial ``(m, ℓ, O)`` triples."""
|
||||
m_new = tl.maximum(m, m_step)
|
||||
scale_old = tl.exp(m - m_new)
|
||||
scale_step = tl.exp(m_step - m_new)
|
||||
l_new = l * scale_old + l_step * scale_step
|
||||
O_new = O * scale_old + O_step * scale_step
|
||||
return m_new, l_new, O_new
|
||||
|
||||
|
||||
def gqa_attention_prefill_long_kernel(
|
||||
q_ptr: int,
|
||||
k_ptr: int,
|
||||
v_ptr: int,
|
||||
o_ptr: int,
|
||||
T_q: int,
|
||||
S_kv: int,
|
||||
d_head: int,
|
||||
C: int,
|
||||
P: int = 1,
|
||||
*,
|
||||
tl,
|
||||
) -> None:
|
||||
"""Head-parallel prefill attention with tile-granular Ring KV (ADR-0060 §5.5).
|
||||
|
||||
Tensor layout (per-rank shapes):
|
||||
Q : (T_q_local, d_head) one head per CUBE; replicated within
|
||||
a CUBE for P=1, or sharded ``pe="row_wise"`` for P>1 so each
|
||||
PE owns ``T_q_local = T_q // P`` query rows.
|
||||
K : (S_kv, d_head) sharded cube_row_wise → each CUBE owns
|
||||
(S_local, d_head). Tiles loaded as (d_head, tile_s) via
|
||||
byte-conserving reshape.
|
||||
V : (S_kv, d_head) sharded cube_row_wise → each CUBE owns
|
||||
(S_local, d_head). Tiles loaded as (tile_s, d_head).
|
||||
O : (T_q * C, d_head) sharded cube_row_wise → each CUBE
|
||||
writes its own ``T_q`` rows; for P>1 each PE writes a
|
||||
disjoint ``(T_q_local, d_head)`` slice (no intra-CUBE
|
||||
reduce — disjoint output rows). NO inter-CUBE reduce.
|
||||
|
||||
Algorithm: nested loop over (ring_step k, tile_idx t). At k=0 each
|
||||
rank loads its own block's tiles from HBM; at k > 0 it receives
|
||||
tiles from its E neighbour. Tiles are forwarded W to the next
|
||||
ring step. Each tile's partial is folded into the running
|
||||
(m, ℓ, O) via online-softmax merge.
|
||||
"""
|
||||
pe_id = tl.program_id(axis=0)
|
||||
if P == 1:
|
||||
# Head-parallel only: PE 0 of each CUBE participates.
|
||||
if pe_id != 0:
|
||||
return
|
||||
T_q_local = T_q
|
||||
else:
|
||||
# Intra-CUBE PE-SP (ADR-0060 §5.5 last bullet + §B-item-3):
|
||||
# query-axis split across P PEs of each CUBE. Output rows are
|
||||
# disjoint per PE ⇒ no intra-CUBE reduce. Each PE drives its
|
||||
# own same-lane ring via independent IPCQ channels.
|
||||
if pe_id >= P:
|
||||
return
|
||||
if T_q % P != 0:
|
||||
raise ValueError(
|
||||
f"T_q={T_q} must be divisible by P={P} when P > 1; "
|
||||
"use P=1 for the PE-0-only fallback path."
|
||||
)
|
||||
T_q_local = T_q // P
|
||||
|
||||
S_local = S_kv // C
|
||||
n_tiles = (S_local + TILE_S_KV - 1) // TILE_S_KV
|
||||
KV_ROW_BYTES = d_head * 2 # f16
|
||||
Q = tl.load(q_ptr, shape=(T_q_local, d_head), dtype="f16")
|
||||
|
||||
# ── Bootstrap: (t=0, k=0) — load my own tile 0, establish (m, ℓ, O) ──
|
||||
#
|
||||
# Cannot be folded into the (t, k) loop (kernbench-only limitation):
|
||||
# - persistent (m, ℓ, O) must live OUTSIDE ``tl.scratch_scope``,
|
||||
# otherwise scope teardown discards them before the next tile's
|
||||
# merge can read them;
|
||||
# - kernbench has no scratch-backed initializer — ``tl.zeros`` /
|
||||
# ``tl.full`` return addr=0 handles with no backing storage, so
|
||||
# they cannot be overwritten via ``tl.copy_to`` to seed (-inf, 0, 0);
|
||||
# - the ring-step ordering also requires k=0 to send its loaded
|
||||
# tile W before any k>0 iteration can recv from E, so the very
|
||||
# first (t=0, k=0) step has to run outside the scoped loop body
|
||||
# where the send is conditional on ``C > 1``.
|
||||
# Triton port: limitation does not apply (SSA tensors stay live
|
||||
# across iterations); the bootstrap collapses into the loop.
|
||||
tile_s = min(TILE_S_KV, S_local)
|
||||
K_t = tl.load(k_ptr, shape=(d_head, tile_s), dtype="f16")
|
||||
V_t = tl.load(v_ptr, shape=(tile_s, d_head), dtype="f16")
|
||||
if C > 1:
|
||||
tl.send(dir="W", src=K_t)
|
||||
tl.send(dir="W", src=V_t)
|
||||
scores = tl.dot(Q, K_t)
|
||||
m = tl.max(scores, axis=-1)
|
||||
exp_scores = tl.exp(scores - m)
|
||||
l = tl.sum(exp_scores, axis=-1)
|
||||
O = tl.dot(exp_scores, V_t)
|
||||
|
||||
# ── Nested loop: outer tile, inner ring step ──
|
||||
# Each tile propagates all the way through the ring before the next
|
||||
# tile starts; IPCQ in-flight depth stays at 1 per direction.
|
||||
#
|
||||
# Per outer ``t``:
|
||||
# k=0 : load my own tile ``t`` from HBM
|
||||
# k=1..C-1 : receive tile ``t`` of a peer's block from E
|
||||
# (sent by my E neighbour during their previous k step)
|
||||
# k<C-1 : forward the tile to W
|
||||
#
|
||||
# The ``(t=0, k=0)`` iteration is the bootstrap above; the loop skips
|
||||
# it via ``k_start``. Every other iteration uses ``scratch_scope`` +
|
||||
# ``copy_to`` to recycle per-tile intermediates.
|
||||
#
|
||||
# Triton port: drop ``with tl.scratch_scope():`` and replace each
|
||||
# ``copy_to`` with a Python rebind (e.g. ``m = m_new``).
|
||||
for t in range(n_tiles):
|
||||
tile_start = t * TILE_S_KV
|
||||
tile_s = min(TILE_S_KV, S_local - tile_start)
|
||||
k_start = 1 if t == 0 else 0
|
||||
for k in range(k_start, C):
|
||||
with tl.scratch_scope():
|
||||
if k == 0:
|
||||
K_t = tl.load(k_ptr + tile_start * KV_ROW_BYTES,
|
||||
shape=(d_head, tile_s), dtype="f16")
|
||||
V_t = tl.load(v_ptr + tile_start * KV_ROW_BYTES,
|
||||
shape=(tile_s, d_head), dtype="f16")
|
||||
else:
|
||||
K_t = tl.recv(dir="E", shape=(d_head, tile_s), dtype="f16")
|
||||
V_t = tl.recv(dir="E", shape=(tile_s, d_head), dtype="f16")
|
||||
if k < C - 1:
|
||||
tl.send(dir="W", src=K_t)
|
||||
tl.send(dir="W", src=V_t)
|
||||
scores = tl.dot(Q, K_t)
|
||||
m_step = tl.max(scores, axis=-1)
|
||||
exp_scores = tl.exp(scores - m_step)
|
||||
l_step = tl.sum(exp_scores, axis=-1)
|
||||
O_step = tl.dot(exp_scores, V_t)
|
||||
m_new, l_new, O_new = _merge_running(
|
||||
m, l, O, m_step, l_step, O_step, tl=tl,
|
||||
)
|
||||
tl.copy_to(m, m_new)
|
||||
tl.copy_to(l, l_new)
|
||||
tl.copy_to(O, O_new)
|
||||
|
||||
# ── Final normalise + store (each CUBE writes its own head) ──
|
||||
O_final = O / l
|
||||
tl.store(o_ptr, O_final)
|
||||
@@ -0,0 +1,9 @@
|
||||
"""GQA attention kernels, panel runners, and helpers (internal subpackage).
|
||||
|
||||
The bench-package auto-discovery in ``kernbench.benches.registry``
|
||||
skips subpackages, so this folder is not audited for ``@bench``
|
||||
decorators. The only public-facing GQA bench is
|
||||
``kernbench.benches.milestone_1h_gqa`` (umbrella that drives all
|
||||
panels by importing this subpackage's ``gqa_decode_long_ctx_4cases``
|
||||
and ``gqa_prefill_long_ctx_4cases`` modules).
|
||||
"""
|
||||
@@ -0,0 +1 @@
|
||||
"""GQA long-context kernels + 4-cases comparative-study runners."""
|
||||
+19
-39
@@ -26,32 +26,32 @@ import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
from kernbench.benches._gqa_attention_decode_long_ctx_cube_repl_pe_sp import (
|
||||
from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_repl_pe_sp import (
|
||||
gqa_attention_decode_long_ctx_cube_repl_pe_sp_kernel,
|
||||
)
|
||||
from kernbench.benches._gqa_attention_decode_long_ctx_cube_repl_pe_tp import (
|
||||
from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_repl_pe_tp import (
|
||||
gqa_attention_decode_long_ctx_cube_repl_pe_tp_kernel,
|
||||
)
|
||||
from kernbench.benches._gqa_attention_decode_long_ctx_cube_sp_pe_sp import (
|
||||
from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_sp import (
|
||||
gqa_attention_decode_long_ctx_cube_sp_pe_sp_kernel,
|
||||
)
|
||||
from kernbench.benches._gqa_attention_decode_long_ctx_cube_sp_pe_tp import (
|
||||
from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_tp import (
|
||||
gqa_attention_decode_long_ctx_cube_sp_pe_tp_kernel,
|
||||
)
|
||||
from kernbench.benches.milestone_gqa_headline import (
|
||||
from kernbench.benches.gqa_helpers.shared._gqa_panel_helpers import (
|
||||
_ccl_cfg,
|
||||
_summarize_op_log,
|
||||
)
|
||||
from kernbench.benches.registry import bench
|
||||
from kernbench.ccl.sfr_config import configure_sfr_intercube_multisip
|
||||
from kernbench.policy.placement.dp import DPPolicy
|
||||
|
||||
# File is at benches/gqa_helpers/long_ctx/ — go up 2 parents to reach
|
||||
# benches/, then into 1H_milestone_output/gqa/long_ctx/.
|
||||
_OUTPUT_DIR = (
|
||||
Path(__file__).resolve().parent
|
||||
/ "1H_milestone_output"
|
||||
/ "gqa_decode_long_ctx_4cases"
|
||||
Path(__file__).resolve().parents[2]
|
||||
/ "1H_milestone_output" / "gqa" / "long_ctx"
|
||||
)
|
||||
_SWEEP_JSON = _OUTPUT_DIR / "sweep.json"
|
||||
_SWEEP_JSON = _OUTPUT_DIR / "sweep_decode.json"
|
||||
|
||||
|
||||
# ── Panel registry ───────────────────────────────────────────────────
|
||||
@@ -75,7 +75,7 @@ _PANEL_DISPATCH: dict[str, tuple[str, dict]] = {
|
||||
# The sub_w=4/sub_h=2 lrab center-root geometry (root cube 6) is
|
||||
# baked into the Case-4 wrapper kernel.
|
||||
"C": 8, "P": 8,
|
||||
"T_q": 1, "S_kv": 131_072,
|
||||
"T_q": 1, "S_kv": 8_192,
|
||||
"d_head": 128, "h_q": 8, "h_kv": 1,
|
||||
}),
|
||||
"single_kv_group_decode_long_ctx_gqa_cube_repl_pe_tp": ("decode_long_ctx_cube_repl_pe_tp", {
|
||||
@@ -83,7 +83,7 @@ _PANEL_DISPATCH: dict[str, tuple[str, dict]] = {
|
||||
# on batch. For B=1 only one rank works (slide-11 PE-TP waste).
|
||||
# No inter-rank communication.
|
||||
"C": 8, "P": 8,
|
||||
"T_q": 1, "S_kv": 131_072,
|
||||
"T_q": 1, "S_kv": 8_192,
|
||||
"d_head": 128, "h_q": 8, "h_kv": 1,
|
||||
}),
|
||||
"single_kv_group_decode_long_ctx_gqa_cube_repl_pe_sp": ("decode_long_ctx_cube_repl_pe_sp", {
|
||||
@@ -91,7 +91,7 @@ _PANEL_DISPATCH: dict[str, tuple[str, dict]] = {
|
||||
# within each cube. Intra-CUBE 8-way reduce; no inter-CUBE comm
|
||||
# (every cube ends with full answer; designated writer = cube 0).
|
||||
"C": 8, "P": 8,
|
||||
"T_q": 1, "S_kv": 131_072,
|
||||
"T_q": 1, "S_kv": 8_192,
|
||||
"d_head": 128, "h_q": 8, "h_kv": 1,
|
||||
}),
|
||||
"single_kv_group_decode_long_ctx_gqa_cube_sp_pe_tp": ("decode_long_ctx_cube_sp_pe_tp", {
|
||||
@@ -99,7 +99,7 @@ _PANEL_DISPATCH: dict[str, tuple[str, dict]] = {
|
||||
# PEs TP on batch — at B=1 only PE 0 of each cube works (PE-TP
|
||||
# waste). Inter-CUBE lrab AR (root = cube 6); no intra-CUBE comm.
|
||||
"C": 8, "P": 8,
|
||||
"T_q": 1, "S_kv": 131_072,
|
||||
"T_q": 1, "S_kv": 8_192,
|
||||
"d_head": 128, "h_q": 8, "h_kv": 1,
|
||||
}),
|
||||
}
|
||||
@@ -316,31 +316,12 @@ def _run_panel(panel: str, topology: str) -> dict:
|
||||
}
|
||||
|
||||
|
||||
# ── Bench entry ──────────────────────────────────────────────────────
|
||||
# ── Sweep entry (called by the umbrella ``milestone_1h_gqa``) ────────
|
||||
|
||||
|
||||
@bench(
|
||||
name="milestone-gqa-decode-long-ctx-4cases",
|
||||
description=(
|
||||
"Long-context decode 4-cases comparative study on the "
|
||||
"LLaMA-3.1-70B single-KV-head group (8 cubes × 8 PEs)."
|
||||
),
|
||||
)
|
||||
def run(torch) -> None:
|
||||
"""Drive the registered decode case panels; write sweep.json.
|
||||
|
||||
Gated by GQA_DECODE_LONG_CTX_4CASES_RUN=1.
|
||||
"""
|
||||
def run_sweep(topology: str = "topology.yaml") -> int:
|
||||
"""Drive all decode case panels; write sweep.json. Returns row count."""
|
||||
_OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
|
||||
if not os.environ.get("GQA_DECODE_LONG_CTX_4CASES_RUN"):
|
||||
raise RuntimeError(
|
||||
"milestone-gqa-decode-long-ctx-4cases needs "
|
||||
"GQA_DECODE_LONG_CTX_4CASES_RUN=1."
|
||||
)
|
||||
|
||||
topology = os.environ.get(
|
||||
"GQA_DECODE_LONG_CTX_4CASES_TOPOLOGY", "topology.yaml",
|
||||
)
|
||||
rows = [_run_panel(panel, topology) for panel in _PANELS]
|
||||
sweep = {
|
||||
"version": 1,
|
||||
@@ -348,6 +329,5 @@ def run(torch) -> None:
|
||||
"rows": rows,
|
||||
}
|
||||
_SWEEP_JSON.write_text(json.dumps(sweep, indent=2))
|
||||
print(
|
||||
f" milestone-gqa-decode-long-ctx-4cases: {len(rows)} rows -> {_SWEEP_JSON}"
|
||||
)
|
||||
print(f" gqa-decode-long-ctx-4cases: {len(rows)} rows -> {_SWEEP_JSON}")
|
||||
return len(rows)
|
||||
@@ -0,0 +1,5 @@
|
||||
"""GQA helpers shared across long-ctx and short-ctx variants.
|
||||
|
||||
Includes the panel-metrics helpers (``_gqa_panel_helpers``) and the
|
||||
context-agnostic decode-opt2 kernel.
|
||||
"""
|
||||
@@ -0,0 +1,40 @@
|
||||
"""Shared helpers for the GQA 4-cases milestone benches.
|
||||
|
||||
Used by both ``milestone_gqa_decode_long_ctx_4cases`` and
|
||||
``milestone_gqa_prefill_long_ctx_4cases``. Underscore-prefixed so the
|
||||
benches package auto-discovery (which skips ``_*.py``) doesn't try to
|
||||
audit it for a ``@bench`` decorator.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config
|
||||
|
||||
|
||||
def _ccl_cfg() -> dict:
|
||||
"""Resolve the lrab hierarchical AllReduce CCL config."""
|
||||
return resolve_algorithm_config(
|
||||
load_ccl_config(), name="lrab_hierarchical_allreduce",
|
||||
)
|
||||
|
||||
|
||||
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,
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
"""GQA short-context kernels."""
|
||||
@@ -1,231 +0,0 @@
|
||||
"""milestone-gqa-headline bench: real GQA prefill — single-KV-group target.
|
||||
|
||||
Drives the LLaMA-3.1-70B single-KV-head-group prefill panel through
|
||||
``_gqa_attention_prefill_long`` (C=8 snake Ring KV + intra-CUBE PE-SP,
|
||||
all 64 ranks), writing ``op_log_summary`` into ``sweep.json``.
|
||||
Independent from the existing ``milestone-gqa-llama70b`` validation-scale
|
||||
bench (which stays on the legacy baseline kernels).
|
||||
|
||||
Decode panels live in ``milestone_gqa_decode_long_ctx_4cases.py`` (the
|
||||
4-cases long-context decode comparative study).
|
||||
|
||||
Restrictions:
|
||||
- Single SIP (default ``topology.yaml`` 4×4 cube mesh; 4-SIP
|
||||
headline deferred per ADR-0060 §B-item-1)
|
||||
- T_q = S_kv = 1024 (scratch-limited; 32K headline awaits Q-axis
|
||||
kernel tiling)
|
||||
- No figure renderers (defer to a separate cycle)
|
||||
|
||||
Panels:
|
||||
single_kv_group_prefill_gqa_c8_p8: prefill C=8 snake Ring KV +
|
||||
intra-CUBE PE-SP, T_q=S_kv=1K,
|
||||
d_head=128 — the LLaMA-3.1-70B
|
||||
single-KV-group prefill target.
|
||||
|
||||
Gated by ``GQA_HEADLINE_RUN=1``.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
from kernbench.benches._gqa_attention_prefill_long import gqa_attention_prefill_long_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_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
|
||||
_S_KV_PREFILL = 16
|
||||
|
||||
_PANELS = (
|
||||
"single_kv_group_prefill_gqa_c8_p8",
|
||||
)
|
||||
|
||||
# Each entry: (kind, panel-specific params)
|
||||
_PANEL_DISPATCH: dict[str, tuple[str, dict]] = {
|
||||
"single_kv_group_prefill_gqa_c8_p8": ("prefill", {
|
||||
"C": 8, "P": 8,
|
||||
# T_q = S_kv = 1024 (one-shot long-context prefill, scratch-limited).
|
||||
# The bootstrap section of the prefill kernel leaves K_t/V_t/scores/
|
||||
# exp_scores persistent (outside tl.scratch_scope), inflating the
|
||||
# baseline. At T_q=2048 the peak hits ~1.05 MB vs 1.0 MB budget; at
|
||||
# 1024 it fits comfortably. The true LLaMA 32K headline awaits a
|
||||
# future increment to (a) add Q-axis tiling and (b) move the
|
||||
# bootstrap into scratch discipline.
|
||||
"T_q": 1_024, "S_kv": 1_024,
|
||||
"d_head": 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,
|
||||
P: int = 1,
|
||||
T_q: int = _T_Q_PREFILL,
|
||||
d_head: int = _D_HEAD,
|
||||
) -> None:
|
||||
if C > 1:
|
||||
mesh_w = int(ctx.spec["sip"]["cube_mesh"]["w"])
|
||||
if C <= mesh_w:
|
||||
configure_sfr_intercube_ring(
|
||||
ctx.engine, ctx.spec, _ccl_cfg(), ring_size=C,
|
||||
)
|
||||
else:
|
||||
if C % mesh_w != 0:
|
||||
raise ValueError(
|
||||
f"C={C} > mesh_w={mesh_w} requires C divisible by mesh_w"
|
||||
)
|
||||
configure_sfr_intercube_ring(
|
||||
ctx.engine, ctx.spec, _ccl_cfg(),
|
||||
submesh_shape=(C // mesh_w, mesh_w),
|
||||
)
|
||||
# Q and O switch to pe="row_wise" when intra-CUBE PE-SP is active
|
||||
# (ADR-0060 §5.5 + §B-item-3: disjoint query-row split across PEs).
|
||||
q_pe = "row_wise" if P > 1 else "replicate"
|
||||
o_pe = "row_wise" if P > 1 else "replicate"
|
||||
dp_q = DPPolicy(cube="replicate", pe=q_pe,
|
||||
num_cubes=C, num_pes=P)
|
||||
dp_kv = DPPolicy(cube="row_wise" if C > 1 else "replicate",
|
||||
pe="replicate", num_cubes=C, num_pes=P)
|
||||
dp_o = DPPolicy(cube="row_wise" if C > 1 else "replicate",
|
||||
pe=o_pe, num_cubes=C, num_pes=P)
|
||||
q = ctx.zeros((T_q, 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 * C, d_head),
|
||||
dtype=_DTYPE, dp=dp_o, name=f"{panel}_o")
|
||||
ctx.launch(
|
||||
panel, gqa_attention_prefill_long_kernel,
|
||||
q, k, v, o,
|
||||
T_q, S_kv, d_head, C, P,
|
||||
_auto_dim_remap=False,
|
||||
)
|
||||
|
||||
|
||||
def _make_bench_fn(panel: str):
|
||||
kind, params = _PANEL_DISPATCH[panel]
|
||||
assert kind == "prefill", (
|
||||
f"milestone-gqa-headline only registers prefill panels; got {kind!r}"
|
||||
)
|
||||
|
||||
def _bench_fn(ctx):
|
||||
_run_prefill_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 prefill milestone — LLaMA-3.1-70B single-KV-group target (C=8, P=8).",
|
||||
)
|
||||
def run(torch) -> None:
|
||||
"""Drive the headline prefill panel; 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,
|
||||
"S_kv_prefill": _S_KV_PREFILL,
|
||||
"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",
|
||||
)
|
||||
@@ -99,7 +99,11 @@ def _audit_modules(imported: list[str], registered: set[str]) -> None:
|
||||
def _eager_import_and_audit(pkg_path: list[str], pkg_name: str) -> None:
|
||||
imported: list[str] = []
|
||||
for m in iter_modules(pkg_path):
|
||||
if m.name == "registry" or m.name.startswith("_"):
|
||||
# Skip helper-name modules, the registry itself, and any
|
||||
# subpackages (audit policy applies to ``.py`` bench files
|
||||
# only — subpackages are explicitly imported by the benches
|
||||
# that need them).
|
||||
if m.ispkg or m.name == "registry" or m.name.startswith("_"):
|
||||
continue
|
||||
mod = import_module(f"{pkg_name}.{m.name}")
|
||||
imported.append(mod.__name__)
|
||||
|
||||
Reference in New Issue
Block a user