test short context length attention kernel
This commit is contained in:
eusonice
@@ -1,31 +1,17 @@
|
||||
"""Phase 1 spec test for Phase D: short-context GQA kernels
|
||||
(ADR-0060 §B "Items from the long/short context split", item B.split.2).
|
||||
"""Short-context GQA kernel tests — unified A1/A2/A4/B (ADR-0070).
|
||||
|
||||
The long-context kernels (§5.2 decode / §5.5 prefill) shard each KV
|
||||
head row-wise across all CUBEs. At short context (S_kv < 256K), that
|
||||
shard is too thin to feed the engines and the cube-level collective
|
||||
overhead dominates. The short-context kernels drop cube-SP entirely:
|
||||
each CUBE owns ``kv_per_cube`` *whole* KV heads, no S_kv sharding across
|
||||
CUBEs.
|
||||
Both ``_gqa_attention_prefill_short.py`` and
|
||||
``_gqa_attention_decode_short.py`` are single unified kernels selected
|
||||
at launch via ``kv_per_cube ∈ {1, 2, 4, 8}``:
|
||||
|
||||
Design (per AskUserQuestion answers in this session):
|
||||
- PE-parallel heads: P PEs split into ``kv_per_cube`` groups, each
|
||||
group does PE-SP across (P/kv_per_cube) PEs for one owned head.
|
||||
- scratch_scope + tl.copy_to discipline mirrors the long kernels.
|
||||
- One short kernel handles kv_per_cube ∈ {1, 2, 4, 8} via a parameter.
|
||||
Mode kv_per_cube C group_size
|
||||
---- ----------- ------- ----------
|
||||
A1 1 h_kv P (=8)
|
||||
A2 2 h_kv/2 P/2 (=4)
|
||||
A4 4 h_kv/4 P/4 (=2)
|
||||
B 8 1 1
|
||||
|
||||
Group layout on the 2×4 PE grid:
|
||||
kv_per_cube=1, group=8 PEs (full 2×4): row chain + col bridge.
|
||||
kv_per_cube=2, group=4 PEs (one row): row chain only.
|
||||
kv_per_cube=4, group=2 PEs (adj cols): 1-step chain.
|
||||
kv_per_cube=8, group=1 PE: no chain.
|
||||
|
||||
After chain reduce-to-group-root, the group's root PE writes its
|
||||
owned head's output. No inter-CUBE reduce.
|
||||
|
||||
Phase 1 (this commit): tests only — production code lands in Phase 2.
|
||||
All tests fail because the short kernels (``_gqa_attention_decode_short.py`` and
|
||||
``_gqa_attention_prefill_short.py``) do not exist yet.
|
||||
Tests are mode-parametrized over the four (kv_per_cube, C) tuples.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -45,6 +31,17 @@ TOPOLOGY_DEFAULT = Path(__file__).resolve().parents[2] / "topology.yaml"
|
||||
|
||||
D_HEAD = 64
|
||||
DTYPE = "f16"
|
||||
TILE_S_KV = 1024
|
||||
P = 8
|
||||
H_KV = 8
|
||||
|
||||
# (kv_per_cube, C) for the four modes.
|
||||
MODES = [
|
||||
pytest.param(1, 8, id="A1"),
|
||||
pytest.param(2, 4, id="A2"),
|
||||
pytest.param(4, 2, id="A4"),
|
||||
pytest.param(8, 1, id="B"),
|
||||
]
|
||||
|
||||
|
||||
def _ccl_cfg():
|
||||
@@ -61,166 +58,41 @@ def _count(op_log, name: str) -> int:
|
||||
return sum(1 for r in op_log if r.op_name == name)
|
||||
|
||||
|
||||
# ── Decode short-context kernel ──────────────────────────────────────
|
||||
# ── Prefill ──────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _run_decode_short(*, h_q: int, h_kv: int, kv_per_cube: int,
|
||||
C: int, P: int, S_kv: int):
|
||||
"""Run the short-context decode kernel with PE-parallel heads.
|
||||
|
||||
Layout (after design iteration — see Phase D failure-recovery):
|
||||
Q: (T_q, h_q·D_HEAD) replicated; kernel reshapes byte-conservingly.
|
||||
K, V: (h_kv·S_kv, D_HEAD) head-stacked, ``cube=row_wise, pe=row_wise``
|
||||
so each PE gets a contiguous (S_local, D_HEAD) chunk at its own
|
||||
addressable shard (no partial reads needed).
|
||||
O: replicated; each group root writes the full byte-conserving
|
||||
(h_q·T_q, D_HEAD) result.
|
||||
"""
|
||||
from kernbench.benches.gqa_helpers.short_ctx._gqa_attention_decode_short import gqa_attention_decode_short_kernel # Phase 2
|
||||
|
||||
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 (
|
||||
_validate_config as _validate_prefill_config,
|
||||
gqa_attention_prefill_short_kernel,
|
||||
)
|
||||
_validate_prefill_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_full = DPPolicy(cube="replicate", pe="replicate",
|
||||
num_cubes=C, num_pes=P)
|
||||
# Head-stacked KV with row_wise sharding so each PE's chunk is
|
||||
# contiguous and exactly (S_local, D_HEAD) addressable.
|
||||
dp_kv = DPPolicy(cube="row_wise", pe="row_wise",
|
||||
num_cubes=C, num_pes=P)
|
||||
q = ctx.zeros((1, h_q * D_HEAD),
|
||||
dtype=DTYPE, dp=dp_full, name=f"q_short_{kv_per_cube}_{C}")
|
||||
k = ctx.zeros((h_kv * S_kv, D_HEAD),
|
||||
dtype=DTYPE, dp=dp_kv, name=f"k_short_{kv_per_cube}_{C}")
|
||||
v = ctx.zeros((h_kv * S_kv, D_HEAD),
|
||||
dtype=DTYPE, dp=dp_kv, name=f"v_short_{kv_per_cube}_{C}")
|
||||
o = ctx.empty((1, h_q * D_HEAD),
|
||||
dtype=DTYPE, dp=dp_full, name=f"o_short_{kv_per_cube}_{C}")
|
||||
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_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_kv{kv_per_cube}")
|
||||
v = ctx.zeros((H_KV * S_kv, D_HEAD),
|
||||
dtype=DTYPE, dp=dp_kv, name=f"v_pre_kv{kv_per_cube}")
|
||||
o = ctx.empty((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_o,
|
||||
name=f"o_pre_kv{kv_per_cube}")
|
||||
ctx.launch(
|
||||
f"gqa_decode_short_{kv_per_cube}_{C}",
|
||||
gqa_attention_decode_short_kernel,
|
||||
q, k, v, o,
|
||||
1, 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 test_short_decode_smoke_kv_per_cube_2_C_4():
|
||||
"""ADR-0060 §B.split.2 headline: kv_per_cube=2, C=4 — each CUBE
|
||||
owns 2 heads (8 heads / 4 CUBEs). PE-parallel heads splits the 8
|
||||
PEs into 2 groups of 4, each group does PE-SP for one owned head.
|
||||
Smoke: kernel completes."""
|
||||
result = _run_decode_short(
|
||||
h_q=8, h_kv=8, kv_per_cube=2, C=4, P=8, S_kv=64,
|
||||
)
|
||||
assert result.completion.ok, (
|
||||
f"short decode kv_per_cube=2 must complete; got {result.completion}"
|
||||
)
|
||||
|
||||
|
||||
def test_short_decode_smoke_kv_per_cube_4_C_2():
|
||||
"""kv_per_cube=4, C=2 — half the CUBEs participate, each owns 4
|
||||
heads, 4 PE groups of 2 PEs each."""
|
||||
result = _run_decode_short(
|
||||
h_q=8, h_kv=8, kv_per_cube=4, C=2, P=8, S_kv=64,
|
||||
)
|
||||
assert result.completion.ok, (
|
||||
f"short decode kv_per_cube=4 must complete; got {result.completion}"
|
||||
)
|
||||
|
||||
|
||||
def test_short_decode_smoke_kv_per_cube_8_C_1():
|
||||
"""kv_per_cube=8, C=1 — all heads on one CUBE. 8 PE groups of 1 PE
|
||||
each → no chain reduce, each PE writes its head's output."""
|
||||
result = _run_decode_short(
|
||||
h_q=8, h_kv=8, kv_per_cube=8, C=1, P=8, S_kv=64,
|
||||
)
|
||||
assert result.completion.ok, (
|
||||
f"short decode kv_per_cube=8 must complete; got {result.completion}"
|
||||
)
|
||||
|
||||
|
||||
def test_short_decode_dma_write_count_equals_h_kv():
|
||||
"""ADR-0060 §B.split.2: each owned head produces exactly one
|
||||
output (no inter-CUBE reduce). Total dma_writes = h_kv across all
|
||||
participating CUBEs and groups.
|
||||
|
||||
For h_kv=8, kv_per_cube=2, C=4: each CUBE writes 2 outputs →
|
||||
4 × 2 = 8 dma_writes total.
|
||||
"""
|
||||
result = _run_decode_short(
|
||||
h_q=8, h_kv=8, kv_per_cube=2, C=4, P=8, S_kv=64,
|
||||
)
|
||||
assert result.completion.ok
|
||||
n_writes = _count(result.engine.op_log, "dma_write")
|
||||
assert n_writes == 8, (
|
||||
f"short decode: expected 8 dma_writes (h_kv); got {n_writes}"
|
||||
)
|
||||
|
||||
|
||||
def test_short_decode_no_inter_cube_traffic():
|
||||
"""ADR-0060 §B.split.2: each head is fully owned by one CUBE → no
|
||||
inter-CUBE reduce. The kernel must not invoke CUBE-level E/W IPCQ.
|
||||
|
||||
Today's long-context kernel at C=4 emits ~12 inter-CUBE ipcq_copy
|
||||
via direction "E"/"W". The short kernel must emit zero of those,
|
||||
keeping all IPCQ traffic on the ``intra_*`` directions.
|
||||
"""
|
||||
result = _run_decode_short(
|
||||
h_q=8, h_kv=8, kv_per_cube=2, C=4, P=8, S_kv=64,
|
||||
)
|
||||
assert result.completion.ok
|
||||
# Count IPCQ traffic that targeted the CUBE-level "E"/"W" directions.
|
||||
inter_cube_ipcq = sum(
|
||||
1 for r in result.engine.op_log
|
||||
if r.op_name == "ipcq_copy"
|
||||
and r.params.get("direction") in ("E", "W")
|
||||
)
|
||||
assert inter_cube_ipcq == 0, (
|
||||
f"short decode must have no inter-CUBE E/W IPCQ; got {inter_cube_ipcq}"
|
||||
)
|
||||
|
||||
|
||||
# ── Prefill short-context kernel ─────────────────────────────────────
|
||||
|
||||
|
||||
def _run_prefill_short(*, h_kv: int, kv_per_cube: int,
|
||||
C: int, P: int, T_q: int, S_kv: int):
|
||||
"""Run the short-context prefill kernel with PE-parallel heads.
|
||||
|
||||
Layout: same head-stacked K/V scheme as decode short, with Q/O
|
||||
replicated and byte-conserving reshape inside the kernel.
|
||||
"""
|
||||
from kernbench.benches.gqa_helpers.short_ctx._gqa_attention_prefill_short import gqa_attention_prefill_short_kernel # Phase 2
|
||||
|
||||
topo = resolve_topology(str(TOPOLOGY_DEFAULT))
|
||||
|
||||
def _bench_fn(ctx):
|
||||
configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
|
||||
dp_q = DPPolicy(cube="replicate", 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="replicate", pe="replicate",
|
||||
num_cubes=C, num_pes=P)
|
||||
q = ctx.zeros((T_q, h_kv * D_HEAD),
|
||||
dtype=DTYPE, dp=dp_q, name=f"q_pre_short_{kv_per_cube}_{C}")
|
||||
k = ctx.zeros((h_kv * S_kv, D_HEAD),
|
||||
dtype=DTYPE, dp=dp_kv, name=f"k_pre_short_{kv_per_cube}_{C}")
|
||||
v = ctx.zeros((h_kv * S_kv, D_HEAD),
|
||||
dtype=DTYPE, dp=dp_kv, name=f"v_pre_short_{kv_per_cube}_{C}")
|
||||
o = ctx.empty((T_q, h_kv * D_HEAD),
|
||||
dtype=DTYPE, dp=dp_o, name=f"o_pre_short_{kv_per_cube}_{C}")
|
||||
ctx.launch(
|
||||
f"gqa_prefill_short_{kv_per_cube}_{C}",
|
||||
f"gqa_prefill_kv{kv_per_cube}",
|
||||
gqa_attention_prefill_short_kernel,
|
||||
q, k, v, o,
|
||||
T_q, S_kv, h_kv, D_HEAD, C, P, kv_per_cube,
|
||||
T_q, S_kv, h_q, H_KV, D_HEAD, C, P, kv_per_cube,
|
||||
_auto_dim_remap=False,
|
||||
)
|
||||
|
||||
@@ -230,31 +102,384 @@ def _run_prefill_short(*, h_kv: int, kv_per_cube: int,
|
||||
)
|
||||
|
||||
|
||||
def test_short_prefill_smoke_kv_per_cube_2_C_4():
|
||||
"""Short prefill kv_per_cube=2, C=4. Smoke: completes."""
|
||||
result = _run_prefill_short(
|
||||
h_kv=8, kv_per_cube=2, C=4, P=8, T_q=4, S_kv=64,
|
||||
)
|
||||
assert result.completion.ok, (
|
||||
f"short prefill kv_per_cube=2 must complete; got {result.completion}"
|
||||
@pytest.mark.parametrize("kv_per_cube,C", MODES)
|
||||
def test_prefill_smoke(kv_per_cube, C):
|
||||
"""All four prefill modes complete on a single-tile mini config."""
|
||||
r = _run_prefill(kv_per_cube=kv_per_cube, C=C, T_q=8, S_kv=1024)
|
||||
assert r.completion.ok, (
|
||||
f"prefill kv_per_cube={kv_per_cube}: {r.completion}"
|
||||
)
|
||||
|
||||
|
||||
def test_short_prefill_no_ring_KV_traffic():
|
||||
"""ADR-0060 §B.split.2: short prefill DOES NOT use Ring KV — each
|
||||
CUBE owns its KV heads fully, no rotation. The kernel must not
|
||||
emit any KV-rotation IPCQ traffic at the CUBE level.
|
||||
@pytest.mark.parametrize("kv_per_cube,C", MODES)
|
||||
def test_prefill_qtile_split_dma_writes(kv_per_cube, C):
|
||||
"""Every PE in every active group stores its q-tile output:
|
||||
dma_writes = group_size · kv_per_cube · C."""
|
||||
r = _run_prefill(kv_per_cube=kv_per_cube, C=C, T_q=8, S_kv=1024)
|
||||
assert r.completion.ok
|
||||
group_size = P // kv_per_cube
|
||||
expected = group_size * kv_per_cube * C
|
||||
n_writes = _count(r.engine.op_log, "dma_write")
|
||||
assert n_writes == expected, (
|
||||
f"prefill kv_per_cube={kv_per_cube}: expected {expected} "
|
||||
f"dma_writes; got {n_writes}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kv_per_cube,C", MODES)
|
||||
def test_prefill_kv_hbm_read_collapsed(kv_per_cube, C):
|
||||
"""IPCQ KV broadcast collapses HBM K/V reads to one per group per tile:
|
||||
dma_reads = (P · C) Q-reads + (kv_per_cube · 2 · n_tiles · C) KV-reads.
|
||||
"""
|
||||
result = _run_prefill_short(
|
||||
h_kv=8, kv_per_cube=2, C=4, P=8, T_q=4, S_kv=64,
|
||||
r = _run_prefill(kv_per_cube=kv_per_cube, C=C, T_q=8, S_kv=1024)
|
||||
assert r.completion.ok
|
||||
n_tiles = 1024 // TILE_S_KV
|
||||
expected = (P * C) + (kv_per_cube * 2 * n_tiles * C)
|
||||
n_reads = _count(r.engine.op_log, "dma_read")
|
||||
assert n_reads == expected, (
|
||||
f"prefill kv_per_cube={kv_per_cube}: expected {expected} dma_reads "
|
||||
f"(Q + IPCQ-broadcasted KV); got {n_reads}"
|
||||
)
|
||||
assert result.completion.ok
|
||||
|
||||
|
||||
def test_prefill_no_ring_KV_traffic():
|
||||
"""Short prefill never uses Ring KV — no inter-CUBE E/W IPCQ."""
|
||||
r = _run_prefill(kv_per_cube=2, C=4, T_q=8, S_kv=1024)
|
||||
assert r.completion.ok
|
||||
inter_cube_ipcq = sum(
|
||||
1 for r in result.engine.op_log
|
||||
if r.op_name == "ipcq_copy"
|
||||
and r.params.get("direction") in ("E", "W")
|
||||
1 for rec in r.engine.op_log
|
||||
if rec.op_name == "ipcq_copy"
|
||||
and rec.params.get("direction") in ("E", "W")
|
||||
)
|
||||
assert inter_cube_ipcq == 0, (
|
||||
f"short prefill must have no Ring KV (no inter-CUBE E/W IPCQ); "
|
||||
f"short prefill must have no inter-CUBE E/W IPCQ; "
|
||||
f"got {inter_cube_ipcq}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kv_per_cube,C", MODES)
|
||||
def test_prefill_multitile_scaling(kv_per_cube, C):
|
||||
"""Prefill dma_read scales linearly with n_tiles in every mode.
|
||||
|
||||
Per-mode formula (IPCQ broadcast: cube's group-root PE reads K/V):
|
||||
dma_reads = P·C + 2·kv_per_cube·C·n_tiles
|
||||
= (Q: 1 per PE) + (K + V: 1 per group per tile)
|
||||
"""
|
||||
for S_kv in (1024, 2048, 4096):
|
||||
r = _run_prefill(kv_per_cube=kv_per_cube, C=C, T_q=8, S_kv=S_kv)
|
||||
assert r.completion.ok, f"kv={kv_per_cube} S_kv={S_kv}"
|
||||
n_tiles = S_kv // TILE_S_KV
|
||||
expected = (P * C) + (2 * kv_per_cube * C * n_tiles)
|
||||
n_reads = _count(r.engine.op_log, "dma_read")
|
||||
assert n_reads == expected, (
|
||||
f"prefill kv={kv_per_cube} n_tiles={n_tiles}: expected "
|
||||
f"{expected} dma_reads; got {n_reads}"
|
||||
)
|
||||
|
||||
|
||||
# ── Decode ───────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
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 (
|
||||
_validate_config as _validate_decode_config,
|
||||
gqa_attention_decode_short_kernel,
|
||||
)
|
||||
T_q = 1
|
||||
_validate_decode_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 total elements = h_kv · S_kv · d_head; mode-invariant deploy.
|
||||
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_kv{kv_per_cube}")
|
||||
k = ctx.zeros((k_rows, TILE_S_KV), dtype=DTYPE, dp=dp_kv,
|
||||
name=f"k_dec_kv{kv_per_cube}")
|
||||
v = ctx.zeros((H_KV * S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv,
|
||||
name=f"v_dec_kv{kv_per_cube}")
|
||||
o = ctx.empty((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_o,
|
||||
name=f"o_dec_kv{kv_per_cube}")
|
||||
ctx.launch(
|
||||
f"gqa_decode_kv{kv_per_cube}",
|
||||
gqa_attention_decode_short_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_decode_smoke(kv_per_cube, C):
|
||||
"""All four decode modes complete at S_kv = 8K (min multi-tile-aligned
|
||||
config that satisfies every mode's group-size constraint)."""
|
||||
r = _run_decode(kv_per_cube=kv_per_cube, C=C, S_kv=8192)
|
||||
assert r.completion.ok, (
|
||||
f"decode kv_per_cube={kv_per_cube}: {r.completion}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kv_per_cube,C", MODES)
|
||||
def test_decode_chain_reduce_store_count(kv_per_cube, C):
|
||||
"""Decode chain-reduce: only group root (PE 0 per group) stores —
|
||||
dma_writes = kv_per_cube · C (one per group root × cubes)."""
|
||||
r = _run_decode(kv_per_cube=kv_per_cube, C=C, S_kv=8192)
|
||||
assert r.completion.ok
|
||||
expected = kv_per_cube * C
|
||||
n_writes = _count(r.engine.op_log, "dma_write")
|
||||
assert n_writes == expected, (
|
||||
f"decode kv_per_cube={kv_per_cube}: expected {expected} dma_writes "
|
||||
f"(1 per group root); got {n_writes}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kv_per_cube,C", MODES)
|
||||
def test_decode_multitile_per_pe(kv_per_cube, C):
|
||||
"""Decode multi-tile sweep path exercised in every mode.
|
||||
|
||||
S_kv is chosen so each PE owns 2 tiles (n_tiles_per_pe == 2),
|
||||
forcing the post-tile-0 sweep loop to run.
|
||||
"""
|
||||
group_size = P // kv_per_cube
|
||||
S_kv = group_size * 2 * TILE_S_KV
|
||||
r = _run_decode(kv_per_cube=kv_per_cube, C=C, S_kv=S_kv)
|
||||
assert r.completion.ok, (
|
||||
f"decode kv={kv_per_cube} S_kv={S_kv}: {r.completion}"
|
||||
)
|
||||
|
||||
|
||||
# ── ADR-0011 D-VA1 contract regression tests ───────────────────────
|
||||
|
||||
|
||||
def _per_cube_dma_busy(op_log):
|
||||
"""Sum of dma_read t_end-t_start grouped by cube index."""
|
||||
from collections import defaultdict
|
||||
busy = defaultdict(float)
|
||||
for rec in op_log:
|
||||
if rec.op_name != "dma_read":
|
||||
continue
|
||||
cid = rec.component_id or ""
|
||||
for part in cid.split("."):
|
||||
if part.startswith("cube"):
|
||||
try:
|
||||
busy[int(part[4:])] += rec.t_end - rec.t_start
|
||||
except ValueError:
|
||||
pass
|
||||
break
|
||||
return dict(busy)
|
||||
|
||||
|
||||
def _per_cube_large_read_addrs(op_log, *, min_nbytes: int = 1024):
|
||||
"""Group dma_read src_addrs (>= min_nbytes) by issuing cube."""
|
||||
from collections import defaultdict
|
||||
per_cube: dict[int, set[int]] = defaultdict(set)
|
||||
for rec in op_log:
|
||||
if rec.op_name != "dma_read":
|
||||
continue
|
||||
if rec.params.get("nbytes", 0) < min_nbytes:
|
||||
continue
|
||||
cid = rec.component_id or ""
|
||||
for part in cid.split("."):
|
||||
if part.startswith("cube"):
|
||||
try:
|
||||
per_cube[int(part[4:])].add(rec.params["src_addr"])
|
||||
except ValueError:
|
||||
pass
|
||||
break
|
||||
return per_cube
|
||||
|
||||
|
||||
def _assert_cubes_disjoint(per_cube_addrs, label):
|
||||
seen = sorted(per_cube_addrs.items())
|
||||
for i, (ci, ai) in enumerate(seen):
|
||||
for cj, aj in seen[i + 1:]:
|
||||
assert ai.isdisjoint(aj), (
|
||||
f"{label}: cube{ci} and cube{cj} share {len(ai & aj)} "
|
||||
f"dma_read src_addrs; cubes must target disjoint HBM regions."
|
||||
)
|
||||
|
||||
|
||||
def _assert_cube_balanced(busy, label, max_ratio=1.1):
|
||||
if len(busy) <= 1:
|
||||
return # single-cube modes auto-pass
|
||||
vals = list(busy.values())
|
||||
ratio = max(vals) / min(vals) if min(vals) > 0 else 0
|
||||
assert ratio < max_ratio, (
|
||||
f"{label}: per-cube DMA_READ unbalanced: max/min = {ratio:.2f}× "
|
||||
f"(expected < {max_ratio}×). Per-cube busy (μs): "
|
||||
f"{[f'cube{c}={v/1000:.1f}' for c,v in sorted(busy.items())]}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kv_per_cube,C", MODES)
|
||||
def test_decode_per_cube_disjoint_src_addrs(kv_per_cube, C):
|
||||
"""Decode: cubes must read from disjoint HBM regions (per ADR-0011
|
||||
D-VA1, kernel computes cube/PE shard offset from program_id).
|
||||
|
||||
Regression guard: pre-fix all 64 PEs read offset 0 → all cubes
|
||||
share the same src_addrs.
|
||||
"""
|
||||
r = _run_decode(kv_per_cube=kv_per_cube, C=C, S_kv=8192)
|
||||
assert r.completion.ok
|
||||
per_cube = _per_cube_large_read_addrs(r.engine.op_log)
|
||||
assert len(per_cube) == C, (
|
||||
f"decode kv={kv_per_cube}: expected {C} cubes; got {len(per_cube)}"
|
||||
)
|
||||
_assert_cubes_disjoint(per_cube, f"decode kv={kv_per_cube}")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kv_per_cube,C", MODES)
|
||||
def test_decode_per_cube_dma_balanced(kv_per_cube, C):
|
||||
"""Decode: per-cube DMA_READ busy must be ~equal (cubes operate on
|
||||
independent local HBM). max/min < 1.1× for multi-cube modes.
|
||||
|
||||
Regression guard: pre-fix 11.5× ratio (cross-cube traffic to cube0).
|
||||
"""
|
||||
r = _run_decode(kv_per_cube=kv_per_cube, C=C, S_kv=8192)
|
||||
assert r.completion.ok
|
||||
busy = _per_cube_dma_busy(r.engine.op_log)
|
||||
assert len(busy) == C, (
|
||||
f"decode kv={kv_per_cube}: expected {C} cubes; got {sorted(busy.keys())}"
|
||||
)
|
||||
_assert_cube_balanced(busy, f"decode kv={kv_per_cube}")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kv_per_cube,C", MODES)
|
||||
def test_prefill_per_cube_disjoint_src_addrs(kv_per_cube, C):
|
||||
"""Prefill: cubes must read from disjoint HBM regions."""
|
||||
r = _run_prefill(kv_per_cube=kv_per_cube, C=C, T_q=8, S_kv=1024)
|
||||
assert r.completion.ok
|
||||
per_cube = _per_cube_large_read_addrs(r.engine.op_log)
|
||||
assert len(per_cube) == C, (
|
||||
f"prefill kv={kv_per_cube}: expected {C} cubes; got {len(per_cube)}"
|
||||
)
|
||||
_assert_cubes_disjoint(per_cube, f"prefill kv={kv_per_cube}")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kv_per_cube,C", MODES)
|
||||
def test_prefill_per_cube_dma_balanced(kv_per_cube, C):
|
||||
"""Prefill: per-cube DMA_READ busy must be ~equal (cube-parallel)."""
|
||||
r = _run_prefill(kv_per_cube=kv_per_cube, C=C, T_q=8, S_kv=1024)
|
||||
assert r.completion.ok
|
||||
busy = _per_cube_dma_busy(r.engine.op_log)
|
||||
assert len(busy) == C, (
|
||||
f"prefill kv={kv_per_cube}: expected {C} cubes; got {sorted(busy.keys())}"
|
||||
)
|
||||
_assert_cube_balanced(busy, f"prefill kv={kv_per_cube}")
|
||||
|
||||
|
||||
# ── Composite (second-level) variant smoke tests ────────────────────
|
||||
|
||||
|
||||
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 (
|
||||
_validate_config as _validate_prefill_cmp_config,
|
||||
gqa_attention_prefill_short_composite_kernel,
|
||||
)
|
||||
_validate_prefill_cmp_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_cmp_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_cmp_kv{kv_per_cube}")
|
||||
v = ctx.zeros((H_KV * S_kv, D_HEAD),
|
||||
dtype=DTYPE, dp=dp_kv, name=f"v_pre_cmp_kv{kv_per_cube}")
|
||||
o = ctx.empty((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_o,
|
||||
name=f"o_pre_cmp_kv{kv_per_cube}")
|
||||
ctx.launch(
|
||||
f"gqa_prefill_cmp_kv{kv_per_cube}",
|
||||
gqa_attention_prefill_short_composite_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(*, 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 (
|
||||
_validate_config as _validate_decode_cmp_config,
|
||||
gqa_attention_decode_short_composite_kernel,
|
||||
)
|
||||
T_q = 1
|
||||
_validate_decode_cmp_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_cmp_kv{kv_per_cube}")
|
||||
k = ctx.zeros((k_rows, TILE_S_KV), dtype=DTYPE, dp=dp_kv,
|
||||
name=f"k_dec_cmp_kv{kv_per_cube}")
|
||||
v = ctx.zeros((H_KV * S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv,
|
||||
name=f"v_dec_cmp_kv{kv_per_cube}")
|
||||
o = ctx.empty((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_o,
|
||||
name=f"o_dec_cmp_kv{kv_per_cube}")
|
||||
ctx.launch(
|
||||
f"gqa_decode_cmp_kv{kv_per_cube}",
|
||||
gqa_attention_decode_short_composite_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."""
|
||||
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}"
|
||||
)
|
||||
|
||||
|
||||
@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."""
|
||||
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}"
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user