ywkang/kernbench2
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

gqa: rename long-context kernels to *_long for symmetry with *_short

Browse Source 
Make the file + function naming symmetric:
  _gqa_decode.py        -> _gqa_decode_long.py
  _gqa_prefill.py       -> _gqa_prefill_long.py
  gqa_decode_kernel     -> gqa_decode_long_kernel
  gqa_prefill_kernel    -> gqa_prefill_long_kernel

Mirrors the existing _gqa_{decode,prefill}_short.py naming. Updates the
two imports + two call sites in milestone_gqa_headline.py and the 9
attention tests that import the kernels.

Tests: 72/72 focused regression green (tests/attention/ + Phase E + TL
discipline). milestone-gqa-headline bench passes its 7 panel/schema
assertions.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
Mukesh Garg
2026-06-10 10:05:54 -07:00
parent d282144339
commit 5a76ed4f6a
10 changed files with 25 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/kernbench/benches/_gqa_decode.py → src/kernbench/benches/_gqa_decode_long.py
+1 -1
View File
@@ -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,
src/kernbench/benches/_gqa_prefill.py → src/kernbench/benches/_gqa_prefill_long.py
+1 -1
View File
@@ -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,
src/kernbench/benches/milestone_gqa_headline.py
+4 -4
View File
@@ -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,
tests/attention/test_gqa_decode.py
+2 -2
View File
@@ -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)
tests/attention/test_gqa_decode_mc.py
+2 -2
View File
@@ -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,
tests/attention/test_gqa_decode_sp.py
+2 -2
View File
@@ -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)
tests/attention/test_gqa_long_context.py
+4 -4
View File
@@ -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,
tests/attention/test_gqa_prefill.py
+3 -3
View File
@@ -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,
tests/attention/test_gqa_prefill_ring.py
+2 -2
View File
@@ -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,
tests/attention/test_gqa_scoped_writeback.py
+4 -4
View File
@@ -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,