ywkang
24c705419e
Complete the composite-command study across both attention regimes and write up the result, resolving when the composite command helps latency. Decode (memory-bound, T_q=1, M=8): the command form is latency-neutral — the kernel is bound by KV streaming and the MAC array is near-idle, so the composite's only benefit here is host-issue offload (O(n_tiles) -> O(1) PE_CPU commands). Figure: gqa_decode_long_ctx_composite (latency flat, command count saturates). Prefill (compute-bound, large M=G*T_q tile-filling): add three command-form variants of a single-rank FlashAttention prefill kernel (_gqa_prefill_compute_bound: primitive / composite / composite_extended). Here the composite's scheduler-internal per-tile DMA<->compute pipeline keeps the MAC array fed while the primitive's blocking tl.dot starves it, so composite wins on MAC utilization (68% flat -> 83%) and wall-clock (19% faster at ctx=1024), with the margin growing with context (deeper P.V reduction = more tiles to pipeline) — the compute-bound mirror of the GEMM result in section 3. Figure: gqa_prefill_compute_bound (latency + MAC util). Sweep wired as GQA_1H_SWEEPS=prefill_cb; tests cover structure, e2e, and composite<primitive in the compute-bound regime. The synthesis: composite has two benefits set by roofline position — host-issue offload (always) and MAC-array feeding (compute-bound only). Decode exercises the first, prefill both. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
88 lines
3.2 KiB
Python
88 lines
3.2 KiB
Python
"""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)
|
||
- Case-6 composite-command study (gqa_helpers.long_ctx.gqa_decode_long_ctx_composite)
|
||
|
||
Each sub-sweep writes its own ``sweep_*.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 prefill+decode). The composite study is opt-in (it sweeps the
|
||
data-mode engine over several S_kv points): ``GQA_1H_SWEEPS=composite``.
|
||
|
||
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_decode_long_ctx_composite import (
|
||
run_sweep as _run_composite_sweep,
|
||
)
|
||
from kernbench.benches.gqa_helpers.long_ctx.gqa_prefill_compute_bound import (
|
||
run_sweep as _run_prefill_cb_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,
|
||
"composite": _run_composite_sweep,
|
||
"prefill_cb": _run_prefill_cb_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",
|
||
)
|