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

add three gqa attention kernel variants for short context length

Browse Source
This commit is contained in:
eusonice
2026-06-24 14:46:55 -07:00
parent a98e93110b
commit 1971ac731c
7 changed files with 575 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/attention/test_gqa_short_context.py
+129 -6
View File
@@ -64,7 +64,7 @@ def _count(op_log, name: str) -> int:
def _run_prefill(*, kv_per_cube: int, C: int, T_q: int, S_kv: int,
h_q: int = 8):
"""Run the unified prefill kernel in the given mode."""
from kernbench.benches._gqa_attention_prefill_short import (
from kernbench.benches.gqa_helpers.short_ctx._gqa_attention_prefill_short import (
_validate_config as _validate_prefill_config,
gqa_attention_prefill_short_kernel,
)
@@ -182,7 +182,7 @@ def test_prefill_multitile_scaling(kv_per_cube, C):
def _run_decode(*, kv_per_cube: int, C: int, S_kv: int, h_q: int = 8):
"""Run the unified decode kernel in the given mode."""
from kernbench.benches._gqa_attention_decode_short import (
from kernbench.benches.gqa_helpers.short_ctx._gqa_attention_decode_short import (
_validate_config as _validate_decode_config,
gqa_attention_decode_short_kernel,
)
@@ -387,7 +387,7 @@ def test_prefill_per_cube_dma_balanced(kv_per_cube, C):
def _run_prefill_composite(*, kv_per_cube: int, C: int, T_q: int,
S_kv: int, h_q: int = 8):
"""Same helper as _run_prefill but for the composite-GEMM kernel."""
from kernbench.benches._gqa_attention_prefill_short_composite import (
from kernbench.benches.gqa_helpers.short_ctx._gqa_attention_prefill_short_composite import (
_validate_config as _validate_prefill_cmp_config,
gqa_attention_prefill_short_composite_kernel,
)
@@ -428,7 +428,7 @@ def _run_prefill_composite(*, kv_per_cube: int, C: int, T_q: int,
def _run_decode_composite(*, kv_per_cube: int, C: int, S_kv: int,
h_q: int = 8):
"""Same helper as _run_decode but for the composite-GEMM kernel."""
from kernbench.benches._gqa_attention_decode_short_composite import (
from kernbench.benches.gqa_helpers.short_ctx._gqa_attention_decode_short_composite import (
_validate_config as _validate_decode_cmp_config,
gqa_attention_decode_short_composite_kernel,
)
@@ -467,9 +467,92 @@ def _run_decode_composite(*, kv_per_cube: int, C: int, S_kv: int,
)
def _run_prefill_composite_fused(*, kv_per_cube: int, C: int, T_q: int,
S_kv: int, h_q: int = 8):
"""Variant (3): GEMM-only composite + softmax_merge prologue fusion."""
from kernbench.benches.gqa_helpers.short_ctx._gqa_attention_prefill_short_composite_fused import (
_validate_config as _validate_prefill_fuse_config,
gqa_attention_prefill_short_composite_fused_kernel,
)
_validate_prefill_fuse_config(kv_per_cube=kv_per_cube, T_q=T_q, P=P,
C=C, h_q=h_q, h_kv=H_KV, S_kv=S_kv)
n_tiles = S_kv // TILE_S_KV
Q_ROWS = kv_per_cube * T_q
Q_COLS = (h_q * D_HEAD) // kv_per_cube
topo = resolve_topology(str(TOPOLOGY_DEFAULT))
def _bench_fn(ctx):
configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
dp_q = DPPolicy(cube="column_wise", pe="replicate", num_cubes=C, num_pes=P)
dp_kv = DPPolicy(cube="row_wise", pe="replicate", num_cubes=C, num_pes=P)
dp_o = DPPolicy(cube="column_wise", pe="replicate", num_cubes=C, num_pes=P)
q = ctx.zeros((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_q,
name=f"q_pre_fuse_kv{kv_per_cube}")
k = ctx.zeros((H_KV * n_tiles * D_HEAD, TILE_S_KV),
dtype=DTYPE, dp=dp_kv, name=f"k_pre_fuse_kv{kv_per_cube}")
v = ctx.zeros((H_KV * S_kv, D_HEAD),
dtype=DTYPE, dp=dp_kv, name=f"v_pre_fuse_kv{kv_per_cube}")
o = ctx.empty((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_o,
name=f"o_pre_fuse_kv{kv_per_cube}")
ctx.launch(
f"gqa_prefill_fuse_kv{kv_per_cube}",
gqa_attention_prefill_short_composite_fused_kernel,
q, k, v, o,
T_q, S_kv, h_q, H_KV, D_HEAD, C, P, kv_per_cube,
_auto_dim_remap=False,
)
return run_bench(
topology=topo, bench_fn=_bench_fn,
device=resolve_device(None), engine_factory=_engine_factory,
)
def _run_decode_composite_fused(*, kv_per_cube: int, C: int, S_kv: int,
h_q: int = 8):
"""Variant (3): GEMM-only composite + softmax_merge prologue fusion."""
from kernbench.benches.gqa_helpers.short_ctx._gqa_attention_decode_short_composite_fused import (
_validate_config as _validate_decode_fuse_config,
gqa_attention_decode_short_composite_fused_kernel,
)
T_q = 1
_validate_decode_fuse_config(kv_per_cube=kv_per_cube, T_q=T_q, P=P,
C=C, h_q=h_q, h_kv=H_KV, S_kv=S_kv)
Q_ROWS = kv_per_cube * T_q
Q_COLS = (h_q * D_HEAD) // kv_per_cube
k_rows = (H_KV * S_kv * D_HEAD) // TILE_S_KV
topo = resolve_topology(str(TOPOLOGY_DEFAULT))
def _bench_fn(ctx):
configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
dp_q = DPPolicy(cube="column_wise", pe="replicate", num_cubes=C, num_pes=P)
dp_kv = DPPolicy(cube="row_wise", pe="row_wise", num_cubes=C, num_pes=P)
dp_o = DPPolicy(cube="column_wise", pe="replicate", num_cubes=C, num_pes=P)
q = ctx.zeros((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_q,
name=f"q_dec_fuse_kv{kv_per_cube}")
k = ctx.zeros((k_rows, TILE_S_KV), dtype=DTYPE, dp=dp_kv,
name=f"k_dec_fuse_kv{kv_per_cube}")
v = ctx.zeros((H_KV * S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv,
name=f"v_dec_fuse_kv{kv_per_cube}")
o = ctx.empty((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_o,
name=f"o_dec_fuse_kv{kv_per_cube}")
ctx.launch(
f"gqa_decode_fuse_kv{kv_per_cube}",
gqa_attention_decode_short_composite_fused_kernel,
q, k, v, o,
T_q, S_kv, h_q, H_KV, D_HEAD, C, P, kv_per_cube,
_auto_dim_remap=False,
)
return run_bench(
topology=topo, bench_fn=_bench_fn,
device=resolve_device(None), engine_factory=_engine_factory,
)
@pytest.mark.parametrize("kv_per_cube,C", MODES)
def test_prefill_composite_smoke(kv_per_cube, C):
"""Composite-GEMM prefill completes in all 4 modes."""
"""Variant (2) GEMM-only composite prefill — all 4 modes."""
r = _run_prefill_composite(kv_per_cube=kv_per_cube, C=C, T_q=8, S_kv=1024)
assert r.completion.ok, (
f"prefill-composite kv_per_cube={kv_per_cube}: {r.completion}"
@@ -478,8 +561,48 @@ def test_prefill_composite_smoke(kv_per_cube, C):
@pytest.mark.parametrize("kv_per_cube,C", MODES)
def test_decode_composite_smoke(kv_per_cube, C):
"""Composite-GEMM decode completes in all 4 modes."""
"""Variant (2) GEMM-only composite decode — all 4 modes."""
r = _run_decode_composite(kv_per_cube=kv_per_cube, C=C, S_kv=8192)
assert r.completion.ok, (
f"decode-composite kv_per_cube={kv_per_cube}: {r.completion}"
)
@pytest.mark.parametrize("kv_per_cube,C", MODES)
def test_decode_composite_fused_smoke(kv_per_cube, C):
"""Variant (3) composite + softmax_merge fused decode — all 4 modes."""
r = _run_decode_composite_fused(kv_per_cube=kv_per_cube, C=C, S_kv=8192)
assert r.completion.ok, (
f"decode-composite-fused kv_per_cube={kv_per_cube}: {r.completion}"
)
@pytest.mark.parametrize("kv_per_cube,C", MODES)
def test_decode_composite_fused_multitile(kv_per_cube, C):
"""Fused decode with n_tiles_per_pe == 2 in EVERY mode, so the tile-1+
fused composite loop actually runs.
The smoke test at S_kv=8192 gives A1 (group_size=8) n_tiles_per_pe=1,
which never enters the fused path. S_kv = group_size·2·TILE_S_KV forces
2 tiles per PE for all modes (mirrors test_decode_multitile_per_pe)."""
group_size = P // kv_per_cube
S_kv = group_size * 2 * TILE_S_KV
r = _run_decode_composite_fused(kv_per_cube=kv_per_cube, C=C, S_kv=S_kv)
assert r.completion.ok, (
f"decode-composite-fused-multitile kv={kv_per_cube} S_kv={S_kv}: "
f"{r.completion}"
)
@pytest.mark.parametrize("kv_per_cube,C", MODES)
def test_prefill_composite_fused_smoke(kv_per_cube, C):
"""Variant (3) composite + softmax_merge fused prefill — all 4 modes.
S_kv=2048 (n_tiles=2) so the tile-1+ fused composite loop actually
runs; S_kv=1024 gives n_tiles=1 and never enters the fused path.
Multi-cube (A1/A2/A4) works now that recv'd K/V slots are pinned."""
r = _run_prefill_composite_fused(kv_per_cube=kv_per_cube, C=C,
T_q=8, S_kv=2048)
assert r.completion.ok, (
f"prefill-composite-fused kv_per_cube={kv_per_cube}: {r.completion}"
)