diff --git a/src/kernbench/benches/_gqa_decode.py b/src/kernbench/benches/_gqa_decode_long.py similarity index 99% rename from src/kernbench/benches/_gqa_decode.py rename to src/kernbench/benches/_gqa_decode_long.py index b83aa86..483a8cd 100644 --- a/src/kernbench/benches/_gqa_decode.py +++ b/src/kernbench/benches/_gqa_decode_long.py @@ -67,7 +67,7 @@ def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl): return m_new, l_new, O_new -def gqa_decode_kernel( +def gqa_decode_long_kernel( q_ptr: int, k_ptr: int, v_ptr: int, diff --git a/src/kernbench/benches/_gqa_prefill.py b/src/kernbench/benches/_gqa_prefill_long.py similarity index 99% rename from src/kernbench/benches/_gqa_prefill.py rename to src/kernbench/benches/_gqa_prefill_long.py index 4b8c0aa..ea4beca 100644 --- a/src/kernbench/benches/_gqa_prefill.py +++ b/src/kernbench/benches/_gqa_prefill_long.py @@ -25,7 +25,7 @@ Deviations from ADR-0060 §5.5 — deferred to later phases: from __future__ import annotations -def gqa_prefill_kernel( +def gqa_prefill_long_kernel( q_ptr: int, k_ptr: int, v_ptr: int, diff --git a/src/kernbench/benches/milestone_gqa_headline.py b/src/kernbench/benches/milestone_gqa_headline.py index eb3946a..461deee 100644 --- a/src/kernbench/benches/milestone_gqa_headline.py +++ b/src/kernbench/benches/milestone_gqa_headline.py @@ -28,8 +28,8 @@ import json import os from pathlib import Path -from kernbench.benches._gqa_decode import gqa_decode_kernel -from kernbench.benches._gqa_prefill import gqa_prefill_kernel +from kernbench.benches._gqa_decode_long import gqa_decode_long_kernel +from kernbench.benches._gqa_prefill_long import gqa_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 ( @@ -96,7 +96,7 @@ def _run_prefill_panel(ctx, *, panel: str, C: int, S_kv: int) -> None: o = ctx.empty((_T_Q_PREFILL * C, _D_HEAD), dtype=_DTYPE, dp=dp_o, name=f"{panel}_o") ctx.launch( - panel, gqa_prefill_kernel, + panel, gqa_prefill_long_kernel, q, k, v, o, _T_Q_PREFILL, S_kv, _D_HEAD, C, _auto_dim_remap=False, @@ -118,7 +118,7 @@ def _run_decode_panel(ctx, *, panel: str, C: int, P: int, S_kv: int) -> None: o = ctx.empty((_T_Q_DECODE, _H_Q_DECODE * _D_HEAD), dtype=_DTYPE, dp=dp_full, name=f"{panel}_o") ctx.launch( - panel, gqa_decode_kernel, + panel, gqa_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, diff --git a/tests/attention/test_gqa_decode.py b/tests/attention/test_gqa_decode.py index 1db00db..5bbcf51 100644 --- a/tests/attention/test_gqa_decode.py +++ b/tests/attention/test_gqa_decode.py @@ -23,7 +23,7 @@ from __future__ import annotations from pathlib import Path -from kernbench.benches._gqa_decode import gqa_decode_kernel # noqa: F401 (Phase 2 deliverable) +from kernbench.benches._gqa_decode_long import gqa_decode_long_kernel # noqa: F401 (Phase 2 deliverable) from kernbench.policy.placement.dp import DPPolicy from kernbench.runtime_api.bench_runner import run_bench from kernbench.runtime_api.types import resolve_device @@ -72,7 +72,7 @@ def _run_decode_p1(*, h_q: int, h_kv: int): dtype=DTYPE, dp=dp, name=f"o_h{h_q}_kv{h_kv}") ctx.launch( f"gqa_decode_p1_h{h_q}_kv{h_kv}", - gqa_decode_kernel, + gqa_decode_long_kernel, q, k, v, o, T_Q, S_KV, h_q, h_kv, D_HEAD, 1, 1, # C=1, P=1 (no SP, degenerate) diff --git a/tests/attention/test_gqa_decode_mc.py b/tests/attention/test_gqa_decode_mc.py index 9873677..c2d16f6 100644 --- a/tests/attention/test_gqa_decode_mc.py +++ b/tests/attention/test_gqa_decode_mc.py @@ -31,7 +31,7 @@ from __future__ import annotations from pathlib import Path -from kernbench.benches._gqa_decode import gqa_decode_kernel # noqa: F401 +from kernbench.benches._gqa_decode_long import gqa_decode_long_kernel # noqa: F401 from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config from kernbench.ccl.sfr_config import configure_sfr_intercube_multisip from kernbench.policy.placement.dp import DPPolicy @@ -82,7 +82,7 @@ def _run_decode_mc(*, h_q: int, h_kv: int, C: int, P: int, S_kv: int): name=f"o_h{h_q}_kv{h_kv}_c{C}_p{P}") ctx.launch( f"gqa_decode_mc_h{h_q}_kv{h_kv}_c{C}_p{P}", - gqa_decode_kernel, + gqa_decode_long_kernel, q, k, v, o, T_Q, S_kv, h_q, h_kv, D_HEAD, C, P, diff --git a/tests/attention/test_gqa_decode_sp.py b/tests/attention/test_gqa_decode_sp.py index 25a321c..21e4e26 100644 --- a/tests/attention/test_gqa_decode_sp.py +++ b/tests/attention/test_gqa_decode_sp.py @@ -22,7 +22,7 @@ from __future__ import annotations from pathlib import Path -from kernbench.benches._gqa_decode import gqa_decode_kernel # noqa: F401 +from kernbench.benches._gqa_decode_long import gqa_decode_long_kernel # noqa: F401 from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config from kernbench.ccl.sfr_config import configure_sfr_intercube_multisip from kernbench.policy.placement.dp import DPPolicy @@ -70,7 +70,7 @@ def _run_decode_sp(*, h_q: int, h_kv: int, P: int, S_kv: int): dtype=DTYPE, dp=dp_full, name=f"o_h{h_q}_kv{h_kv}_p{P}") ctx.launch( f"gqa_decode_sp_h{h_q}_kv{h_kv}_p{P}", - gqa_decode_kernel, + gqa_decode_long_kernel, q, k, v, o, T_Q, S_kv, h_q, h_kv, D_HEAD, 1, P, # C=1, P=P (single-CUBE SP) diff --git a/tests/attention/test_gqa_long_context.py b/tests/attention/test_gqa_long_context.py index 291dd1e..f3baf6e 100644 --- a/tests/attention/test_gqa_long_context.py +++ b/tests/attention/test_gqa_long_context.py @@ -23,8 +23,8 @@ from __future__ import annotations from pathlib import Path -from kernbench.benches._gqa_decode import gqa_decode_kernel # noqa: F401 -from kernbench.benches._gqa_prefill import gqa_prefill_kernel # noqa: F401 +from kernbench.benches._gqa_decode_long import gqa_decode_long_kernel # noqa: F401 +from kernbench.benches._gqa_prefill_long import gqa_prefill_long_kernel # noqa: F401 from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config from kernbench.ccl.sfr_config import ( configure_sfr_intercube_multisip, @@ -84,7 +84,7 @@ def test_decode_long_context_32k_completes(): dtype=DTYPE, dp=dp_full, name="q_long_dec_2") ctx.launch( "gqa_decode_long_32k", - gqa_decode_kernel, + gqa_decode_long_kernel, q, k, v, o, 1, S_kv, 8, 1, D_HEAD, 1, P, @@ -141,7 +141,7 @@ def test_prefill_long_context_completes_after_scope_discipline(): dtype=DTYPE, dp=dp_o, name="o_long_pre") ctx.launch( "gqa_prefill_long_64k", - gqa_prefill_kernel, + gqa_prefill_long_kernel, q, k, v, o, T_q, S_kv, D_HEAD, C, _auto_dim_remap=False, diff --git a/tests/attention/test_gqa_prefill.py b/tests/attention/test_gqa_prefill.py index e086b55..ac9ca2e 100644 --- a/tests/attention/test_gqa_prefill.py +++ b/tests/attention/test_gqa_prefill.py @@ -1,6 +1,6 @@ """Phase 1 spec test for P6a GQA prefill kernel (head-parallel, C=1 baseline). -P6a introduces ``_gqa_prefill.py`` with the head-parallel structure (one +P6a introduces ``_gqa_prefill_long.py`` with the head-parallel structure (one Q head per CUBE, per-CUBE distributed output, no reduce). C=1 is the degenerate case — no Ring KV, no IPCQ traffic. Validates kernel structure and T_q > 1 attention. @@ -21,7 +21,7 @@ from __future__ import annotations from pathlib import Path -from kernbench.benches._gqa_prefill import gqa_prefill_kernel # noqa: F401 +from kernbench.benches._gqa_prefill_long import gqa_prefill_long_kernel # noqa: F401 from kernbench.policy.placement.dp import DPPolicy from kernbench.runtime_api.bench_runner import run_bench from kernbench.runtime_api.types import resolve_device @@ -62,7 +62,7 @@ def _run_prefill(*, T_q: int, S_kv: int, C: int = 1): name=f"o_t{T_q}_c{C}") ctx.launch( f"gqa_prefill_p6a_t{T_q}_s{S_kv}_c{C}", - gqa_prefill_kernel, + gqa_prefill_long_kernel, q, k, v, o, T_q, S_kv, D_HEAD, C, _auto_dim_remap=False, diff --git a/tests/attention/test_gqa_prefill_ring.py b/tests/attention/test_gqa_prefill_ring.py index 9594ee4..16be9c6 100644 --- a/tests/attention/test_gqa_prefill_ring.py +++ b/tests/attention/test_gqa_prefill_ring.py @@ -34,7 +34,7 @@ from __future__ import annotations from pathlib import Path -from kernbench.benches._gqa_prefill import gqa_prefill_kernel # noqa: F401 +from kernbench.benches._gqa_prefill_long import gqa_prefill_long_kernel # noqa: F401 from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config from kernbench.ccl.sfr_config import configure_sfr_intercube_ring # noqa: F401 (Phase 2) from kernbench.policy.placement.dp import DPPolicy @@ -89,7 +89,7 @@ def _run_prefill_ring(*, T_q: int, S_kv: int, C: int): name=f"o_t{T_q}_c{C}_ring") ctx.launch( f"gqa_prefill_ring_t{T_q}_s{S_kv}_c{C}", - gqa_prefill_kernel, + gqa_prefill_long_kernel, q, k, v, o, T_q, S_kv, D_HEAD, C, _auto_dim_remap=False, diff --git a/tests/attention/test_gqa_scoped_writeback.py b/tests/attention/test_gqa_scoped_writeback.py index 0da33cc..825cf2f 100644 --- a/tests/attention/test_gqa_scoped_writeback.py +++ b/tests/attention/test_gqa_scoped_writeback.py @@ -23,8 +23,8 @@ from __future__ import annotations from pathlib import Path -from kernbench.benches._gqa_decode import gqa_decode_kernel # noqa: F401 -from kernbench.benches._gqa_prefill import gqa_prefill_kernel # noqa: F401 +from kernbench.benches._gqa_decode_long import gqa_decode_long_kernel # noqa: F401 +from kernbench.benches._gqa_prefill_long import gqa_prefill_long_kernel # noqa: F401 from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config from kernbench.ccl.sfr_config import ( configure_sfr_intercube_multisip, @@ -79,7 +79,7 @@ def _run_decode_sp(*, h_q: int, h_kv: int, P: int, S_kv: int): dtype=DTYPE, dp=dp_full, name=f"o_sc_{P}") ctx.launch( f"gqa_decode_scoped_{P}", - gqa_decode_kernel, + gqa_decode_long_kernel, q, k, v, o, 1, S_kv, h_q, h_kv, D_HEAD, 1, P, @@ -138,7 +138,7 @@ def _run_prefill_ring(*, T_q: int, S_kv: int, C: int): dtype=DTYPE, dp=dp_o, name=f"o_ring_{C}") ctx.launch( f"gqa_prefill_scoped_{C}", - gqa_prefill_kernel, + gqa_prefill_long_kernel, q, k, v, o, T_q, S_kv, D_HEAD, C, _auto_dim_remap=False,