gqa(milestone-1h-gqa): add umbrella bench + first sweep outputs
Single @bench entry that drives the prefill + decode long-context
4-cases sweeps in one invocation, mirroring milestone-1h-gemm. The
individual gqa_helpers/long_ctx/gqa_{prefill,decode}_long_ctx_4cases
modules are now pure helpers (no @bench), reached only via this
umbrella.
Env var contract:
GQA_1H_RUN=1 (required gate)
GQA_1H_SWEEPS=prefill,decode (default: both; pick subset to skip)
GQA_1H_TOPOLOGY=topology.yaml (override)
Output layout:
benches/1H_milestone_output/gqa/long_ctx/
sweep_prefill.json
sweep_decode.json
gqa_prefill_long_ctx_4cases_{latency,traffic,memory,parallelism}.png
gqa_decode_long_ctx_4cases_{latency,traffic,memory,parallelism}.png
The decode bench config drops S_kv from 131_072 to 8_192 so the umbrella
finishes in minutes — the comparative-story bars are the same shape at
8K. The 131K production headline number is recoverable by overriding
the dispatch; the helper docstring notes this.
Outputs (2 sweep JSONs + 8 PNGs) ship in this commit alongside the
umbrella that produced them, following the milestone-1h-gemm pattern
where derived artifacts live with the bench that emits them.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,78 @@
|
||||
"""milestone-1h-gqa: umbrella GQA bench for the 1H code-sign milestone.
|
||||
|
||||
Single ``@bench`` entry that drives all GQA panels by delegating to
|
||||
internal helpers under ``kernbench.benches.gqa_helpers``. Mirrors the
|
||||
``milestone_1h_gemm`` umbrella pattern.
|
||||
|
||||
Currently exercises (long-context only — short-context panels are
|
||||
future work):
|
||||
- 4-cases prefill comparative study (gqa_helpers.long_ctx.gqa_prefill_long_ctx_4cases)
|
||||
- 4-cases decode comparative study (gqa_helpers.long_ctx.gqa_decode_long_ctx_4cases)
|
||||
|
||||
Each sub-sweep writes its own ``sweep_{prefill,decode}.json`` into the
|
||||
shared output dir ``benches/1H_milestone_output/gqa/gqa_long_ctx/``.
|
||||
Selection via the env var ``GQA_1H_SWEEPS=prefill,decode`` (default
|
||||
runs both). Toggle individual sweeps with ``GQA_1H_SWEEPS=prefill``
|
||||
or ``GQA_1H_SWEEPS=decode``.
|
||||
|
||||
Gated by ``GQA_1H_RUN=1`` to keep CI fast.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
|
||||
from kernbench.benches.gqa_helpers.long_ctx.gqa_decode_long_ctx_4cases import (
|
||||
run_sweep as _run_decode_sweep,
|
||||
)
|
||||
from kernbench.benches.gqa_helpers.long_ctx.gqa_prefill_long_ctx_4cases import (
|
||||
run_sweep as _run_prefill_sweep,
|
||||
)
|
||||
from kernbench.benches.registry import bench
|
||||
from kernbench.policy.placement.dp import DPPolicy
|
||||
|
||||
|
||||
@bench(
|
||||
name="milestone-1h-gqa",
|
||||
description=(
|
||||
"Umbrella GQA milestone — drives long-context prefill + decode "
|
||||
"4-cases sweeps on the LLaMA-3.1-70B single-KV-head group "
|
||||
"(8 cubes × 8 PEs). Gated by GQA_1H_RUN=1."
|
||||
),
|
||||
)
|
||||
def run(torch) -> None:
|
||||
"""Drive selected GQA sub-sweeps; each writes its own sweep_*.json.
|
||||
|
||||
Env vars:
|
||||
GQA_1H_RUN=1 (required gate)
|
||||
GQA_1H_TOPOLOGY=topology.yaml (override topology path)
|
||||
GQA_1H_SWEEPS=prefill,decode (default: both; comma-separated)
|
||||
"""
|
||||
if not os.environ.get("GQA_1H_RUN"):
|
||||
raise RuntimeError("milestone-1h-gqa needs GQA_1H_RUN=1.")
|
||||
|
||||
topology = os.environ.get("GQA_1H_TOPOLOGY", "topology.yaml")
|
||||
requested = os.environ.get("GQA_1H_SWEEPS", "prefill,decode")
|
||||
sweeps = [s.strip() for s in requested.split(",") if s.strip()]
|
||||
|
||||
runners = {
|
||||
"prefill": _run_prefill_sweep,
|
||||
"decode": _run_decode_sweep,
|
||||
}
|
||||
unknown = [s for s in sweeps if s not in runners]
|
||||
if unknown:
|
||||
raise RuntimeError(
|
||||
f"GQA_1H_SWEEPS contains unknown sweep(s) {unknown}; "
|
||||
f"valid: {sorted(runners)}"
|
||||
)
|
||||
|
||||
for s in sweeps:
|
||||
runners[s](topology)
|
||||
|
||||
# Sentinel tensor (ADR-0045 D4 / ADR-0054 D2 carve-out) — the
|
||||
# sub-sweeps each spin up their own GraphEngine via ``run_bench``,
|
||||
# so this outer @bench needs to submit at least one request.
|
||||
torch.zeros(
|
||||
(1, 1), dtype="f16",
|
||||
dp=DPPolicy(cube="row_wise", pe="replicate", num_cubes=1, num_pes=1),
|
||||
name="milestone_1h_gqa_sentinel",
|
||||
)
|
||||
Reference in New Issue
Block a user