diff --git a/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_decode_long_ctx_4cases_latency.png b/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_decode_long_ctx_4cases_latency.png new file mode 100644 index 0000000..69a179d Binary files /dev/null and b/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_decode_long_ctx_4cases_latency.png differ diff --git a/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_decode_long_ctx_4cases_memory.png b/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_decode_long_ctx_4cases_memory.png new file mode 100644 index 0000000..52889b9 Binary files /dev/null and b/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_decode_long_ctx_4cases_memory.png differ diff --git a/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_decode_long_ctx_4cases_parallelism.png b/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_decode_long_ctx_4cases_parallelism.png new file mode 100644 index 0000000..5e47d35 Binary files /dev/null and b/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_decode_long_ctx_4cases_parallelism.png differ diff --git a/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_decode_long_ctx_4cases_traffic.png b/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_decode_long_ctx_4cases_traffic.png new file mode 100644 index 0000000..12fafe2 Binary files /dev/null and b/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_decode_long_ctx_4cases_traffic.png differ diff --git a/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_prefill_long_ctx_4cases_latency.png b/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_prefill_long_ctx_4cases_latency.png new file mode 100644 index 0000000..a675a1f Binary files /dev/null and b/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_prefill_long_ctx_4cases_latency.png differ diff --git a/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_prefill_long_ctx_4cases_memory.png b/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_prefill_long_ctx_4cases_memory.png new file mode 100644 index 0000000..3a24c87 Binary files /dev/null and b/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_prefill_long_ctx_4cases_memory.png differ diff --git a/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_prefill_long_ctx_4cases_parallelism.png b/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_prefill_long_ctx_4cases_parallelism.png new file mode 100644 index 0000000..fce10f2 Binary files /dev/null and b/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_prefill_long_ctx_4cases_parallelism.png differ diff --git a/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_prefill_long_ctx_4cases_traffic.png b/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_prefill_long_ctx_4cases_traffic.png new file mode 100644 index 0000000..a5532e5 Binary files /dev/null and b/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_prefill_long_ctx_4cases_traffic.png differ diff --git a/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/sweep_decode.json b/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/sweep_decode.json new file mode 100644 index 0000000..0de23f4 --- /dev/null +++ b/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/sweep_decode.json @@ -0,0 +1,115 @@ +{ + "version": 1, + "panels": [ + "single_kv_group_decode_long_ctx_gqa_cube_sp_pe_sp", + "single_kv_group_decode_long_ctx_gqa_cube_repl_pe_tp", + "single_kv_group_decode_long_ctx_gqa_cube_repl_pe_sp", + "single_kv_group_decode_long_ctx_gqa_cube_sp_pe_tp" + ], + "rows": [ + { + "panel": "single_kv_group_decode_long_ctx_gqa_cube_sp_pe_sp", + "kind": "decode_long_ctx_cube_sp_pe_sp", + "C": 8, + "P": 8, + "T_q": 1, + "S_kv": 8192, + "d_head": 128, + "h_q": 8, + "h_kv": 1, + "op_log_summary": { + "gemm_count": 128, + "ipcq_copy_count": 189, + "dma_read_count": 192, + "dma_write_count": 1 + }, + "latency_ns": 34030.00300000079, + "engine_occupancy_ns": { + "pe_gemm": 4194.303999999538, + "pe_math": 2282.0, + "pe_dma": 1138201.8199999991, + "pe_fetch_store": 0, + "pe_ipcq": 0, + "pe_cpu": 0 + } + }, + { + "panel": "single_kv_group_decode_long_ctx_gqa_cube_repl_pe_tp", + "kind": "decode_long_ctx_cube_repl_pe_tp", + "C": 8, + "P": 8, + "T_q": 1, + "S_kv": 8192, + "d_head": 128, + "h_q": 8, + "h_kv": 1, + "op_log_summary": { + "gemm_count": 16, + "ipcq_copy_count": 0, + "dma_read_count": 17, + "dma_write_count": 1 + }, + "latency_ns": 25428.38200085424, + "engine_occupancy_ns": { + "pe_gemm": 4194.303999997675, + "pe_math": 714.0, + "pe_dma": 18217.080000881106, + "pe_fetch_store": 0, + "pe_ipcq": 0, + "pe_cpu": 0 + } + }, + { + "panel": "single_kv_group_decode_long_ctx_gqa_cube_repl_pe_sp", + "kind": "decode_long_ctx_cube_repl_pe_sp", + "C": 8, + "P": 8, + "T_q": 1, + "S_kv": 8192, + "d_head": 128, + "h_q": 8, + "h_kv": 1, + "op_log_summary": { + "gemm_count": 128, + "ipcq_copy_count": 168, + "dma_read_count": 192, + "dma_write_count": 1 + }, + "latency_ns": 20152.09400000231, + "engine_occupancy_ns": { + "pe_gemm": 33554.431999996305, + "pe_math": 5684.0, + "pe_dma": 853218.7100000716, + "pe_fetch_store": 0, + "pe_ipcq": 0, + "pe_cpu": 0 + } + }, + { + "panel": "single_kv_group_decode_long_ctx_gqa_cube_sp_pe_tp", + "kind": "decode_long_ctx_cube_sp_pe_tp", + "C": 8, + "P": 8, + "T_q": 1, + "S_kv": 8192, + "d_head": 128, + "h_q": 8, + "h_kv": 1, + "op_log_summary": { + "gemm_count": 16, + "ipcq_copy_count": 21, + "dma_read_count": 24, + "dma_write_count": 1 + }, + "latency_ns": 26987.291500001098, + "engine_occupancy_ns": { + "pe_gemm": 4194.303999999538, + "pe_math": 714.0, + "pe_dma": 126624.99000000354, + "pe_fetch_store": 0, + "pe_ipcq": 0, + "pe_cpu": 0 + } + } + ] +} \ No newline at end of file diff --git a/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/sweep_prefill.json b/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/sweep_prefill.json new file mode 100644 index 0000000..70989fb --- /dev/null +++ b/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/sweep_prefill.json @@ -0,0 +1,116 @@ +{ + "version": 1, + "panels": [ + "single_kv_group_prefill_long_ctx_gqa_cube_sp_pe_sp", + "single_kv_group_prefill_long_ctx_gqa_cube_repl_pe_tp", + "single_kv_group_prefill_long_ctx_gqa_cube_repl_pe_sp", + "single_kv_group_prefill_long_ctx_gqa_cube_sp_pe_tp" + ], + "rows": [ + { + "panel": "single_kv_group_prefill_long_ctx_gqa_cube_sp_pe_sp", + "kind": "prefill_long_ctx_cube_sp_pe_sp", + "C": 8, + "P": 8, + "T_q": 512, + "S_kv": 512, + "d_head": 128, + "h_q": 8, + "h_kv": 1, + "op_log_summary": { + "gemm_count": 1024, + "ipcq_copy_count": 1064, + "dma_read_count": 192, + "dma_write_count": 8 + }, + "latency_ns": 86702.76799999712, + "engine_occupancy_ns": { + "pe_gemm": 134217.72799998685, + "pe_math": 546656.0, + "pe_dma": 2388650.3709999938, + "pe_fetch_store": 0, + "pe_ipcq": 0, + "pe_cpu": 0 + } + }, + { + "panel": "single_kv_group_prefill_long_ctx_gqa_cube_repl_pe_tp", + "kind": "prefill_long_ctx_cube_repl_pe_tp", + "C": 8, + "P": 8, + "T_q": 512, + "S_kv": 512, + "d_head": 128, + "h_q": 8, + "h_kv": 1, + "op_log_summary": { + "gemm_count": 1024, + "ipcq_copy_count": 0, + "dma_read_count": 1088, + "dma_write_count": 64 + }, + "latency_ns": 93564.33200003332, + "engine_occupancy_ns": { + "pe_gemm": 134217.72799998525, + "pe_math": 81279.99999999997, + "pe_dma": 177951.73600022023, + "pe_fetch_store": 0, + "pe_ipcq": 0, + "pe_cpu": 0 + } + }, + { + "panel": "single_kv_group_prefill_long_ctx_gqa_cube_repl_pe_sp", + "kind": "prefill_long_ctx_cube_repl_pe_sp", + "C": 8, + "P": 8, + "T_q": 512, + "S_kv": 512, + "d_head": 128, + "h_q": 8, + "h_kv": 1, + "op_log_summary": { + "gemm_count": 8192, + "ipcq_copy_count": 10752, + "dma_read_count": 12288, + "dma_write_count": 64 + }, + "latency_ns": 190579.57100032456, + "engine_occupancy_ns": { + "pe_gemm": 1073741.8240003586, + "pe_math": 635904.0, + "pe_dma": 1756854.05751659, + "pe_fetch_store": 0, + "pe_ipcq": 0, + "pe_cpu": 0 + } + }, + { + "panel": "single_kv_group_prefill_long_ctx_gqa_cube_sp_pe_tp", + "kind": "prefill_long_ctx_cube_sp_pe_tp", + "C": 8, + "P": 8, + "T_q": 512, + "S_kv": 512, + "d_head": 128, + "h_q": 8, + "h_kv": 1, + "op_log_summary": { + "gemm_count": 1024, + "ipcq_copy_count": 896, + "dma_read_count": 192, + "dma_write_count": 64 + }, + "latency_ns": 38876.32300000038, + "engine_occupancy_ns": { + "pe_gemm": 134217.7280000001, + "pe_math": 81280.0, + "pe_dma": 916284.0330000015, + "pe_fetch_store": 0, + "pe_ipcq": 0, + "pe_cpu": 0 + } + } + ], + "failures": [] +} \ No newline at end of file diff --git a/src/kernbench/benches/milestone_1h_gqa.py b/src/kernbench/benches/milestone_1h_gqa.py new file mode 100644 index 0000000..451cade --- /dev/null +++ b/src/kernbench/benches/milestone_1h_gqa.py @@ -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", + )