mukesh
5b4d9cb597
The headline T_q=S_kv=32K target overflows the 1 MB per-PE scratch pool: at T_q_local=4K and tile_s=1024 the scores matrix alone is 8 MB. The prefill kernel's bootstrap section also leaves K_t/V_t/scores/ exp_scores persistent (outside tl.scratch_scope), inflating baseline. Scale-down to T_q=S_kv=1K (T_q_local=128, fits comfortably) preserves the C=8 + P=8 architecture demonstration; the true LLaMA 32K headline awaits a future increment to add Q-axis tiling and tighten bootstrap scratch discipline. Verified end-to-end: kernbench run --bench milestone-gqa-headline now produces sweep.json with all 5 panels. The new panel shows ipcq_copy=896 (matches (C-1)·n_tiles·2·C·P = 7·1·2·8·8) and dma_write=64 (one per PE, head-parallel + intra-CUBE PE-SP). Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
277 lines
10 KiB
Python
277 lines
10 KiB
Python
"""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:
|
||
- Decode side capped at C ≤ 4 (1D chain reduce; 2D mesh wiring on
|
||
the decode panels deferred to the 4-cases decode comparative study)
|
||
- Single SIP (default ``topology.yaml`` 4×4 cube mesh; 4-SIP
|
||
headline deferred per ADR-0060 §B-item-1)
|
||
- 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_kv_group_prefill_gqa_c8_p8: prefill C=8 snake Ring KV +
|
||
intra-CUBE PE-SP (all 64 ranks),
|
||
T_q=S_kv=1K, d_head=128 — the
|
||
LLaMA-3.1-70B single-KV-group target
|
||
(scratch-limited; 32K headline awaits
|
||
Q-axis kernel tiling)
|
||
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_attention_decode_long import gqa_attention_decode_long_kernel
|
||
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_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_kv_group_prefill_gqa_c8_p8",
|
||
"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_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,
|
||
}),
|
||
"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,
|
||
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 _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_attention_decode_long_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",
|
||
)
|