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:
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user