gqa: rename long-context kernels to *_long for symmetry with *_short
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:
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user