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

gqa: single-KV-group LLaMA-3.1-70B prefill milestone (Increments 1-5)

Browse Source 
End-to-end wires C=8 P=8 d_head=128 prefill with snake-ring inter-CUBE
SFR, intra-CUBE PE-SP (all 64 ranks active), and the milestone bench
panel. Decode kernel gains lrab-adapted center-root reduce for the
2×4 sub-mesh per ADR-0060 §4.2.

Increment 1 — SFR multi-row snake
  src/kernbench/ccl/sfr_config.py: configure_sfr_intercube_ring gains
  submesh_shape / submesh_origin kwargs; installs a Hamiltonian snake
  ring through a rectangular sub-mesh (every hop is 1-hop physical
  neighbour). Backward-compat: 1D-row behaviour preserved when
  submesh_shape is None.
  tests/test_intercube_snake_ring.py (12 tests)

Increment 2 — Decode lrab-adapted center-root reduce
  src/kernbench/benches/_gqa_attention_decode_long.py: new sub_w param
  (default 0 = existing 1D-chain). sub_w >= 2 selects the ADR-0060
  §4.2 prescribed lrab-adapted Phase 1+2 reduce (bidirectional row +
  bidirectional col converge to the center cube), with log-sum-exp
  _merge_running replacing the plain + of lrab.
  tests/attention/test_gqa_decode_long_2d_reduce.py (4 tests)

Increment 3 — Prefill kernel at C=8 (no production change)
  Verified by inspection that the existing prefill_long kernel +
  Increment 1's snake SFR already work at C=8 without any kernel
  edit. The kernel speaks logical W/E; the snake routes it.
  tests/attention/test_gqa_prefill_long_c8_snake.py (3 tests)

Increment 4 — Intra-CUBE PE-SP in prefill (all 64 ranks)
  src/kernbench/benches/_gqa_attention_prefill_long.py: new P param
  (default 1 = existing PE-0-only). P > 1 splits T_q query-axis-wise
  across the P PEs of each CUBE; output rows are disjoint per PE so
  no intra-CUBE reduce is needed; each PE drives its own same-lane
  ring (P parallel rings).
  tests/attention/test_gqa_prefill_long_pe_sp.py (5 tests)

Increment 5 — LLaMA-scale milestone bench panel
  src/kernbench/benches/milestone_gqa_headline.py: new panel
  single_kv_group_prefill_gqa_c8_p8 (C=8, P=8, T_q=S_kv=32K,
  d_head=128). _run_prefill_panel extended with P/T_q/d_head
  defaults; routes snake SFR when C > mesh_w.
  tests/attention/test_milestone_gqa_single_kv_group_prefill_panel.py (3 tests)

Total: 4 production files modified, 5 new test files, 27 new tests.
Followed the Phase 1/2 protocol per CLAUDE.md throughout.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
Mukesh Garg
2026-06-15 13:42:31 -07:00
parent 9270d3435a
commit 9e1242039b
9 changed files with 1332 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/attention/test_gqa_decode_long_2d_reduce.py
+191
View File
@@ -0,0 +1,191 @@
"""Tests for lrab-adapted center-root inter-CUBE reduce in decode_long.
Verifies the ADR-0060 §4.2 prescribed Level-1 collective for the
single-KV-group LLaMA-3.1-70B target (C=8 over a 2×4 sub-mesh, root at
the geometric center CUBE).
Activation contract:
- ``sub_w == 0`` (default; omitted from launch args) → existing 1D
inter-CUBE chain that converges at CUBE 0. Byte-for-byte unchanged.
- ``sub_w >= 2`` with ``sub_h = C // sub_w >= 2`` → lrab-adapted
center-root mesh reduce. Root cube is
``(sub_h//2)*sub_w + (sub_w//2)``.
Degenerate (``sub_h < 2`` or ``sub_w < 2``) combinations are rejected at
launch time per ADR-0060 §4.2 + CLAUDE.md "Simplicity First": those
configurations belong to the 1D-chain code path, not lrab.
Phase 1 (this commit): tests only — production code lands in Phase 2.
T1, T2, T4 fail today (kernel signature does not accept ``sub_w``).
T3 passes today as the backward-compat anchor for the existing
1D-chain path.
"""
from __future__ import annotations
import re
from pathlib import Path
import pytest
from kernbench.benches._gqa_attention_decode_long import gqa_attention_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
from kernbench.runtime_api.bench_runner import run_bench
from kernbench.runtime_api.types import resolve_device
from kernbench.sim_engine.engine import GraphEngine
from kernbench.topology.builder import resolve_topology
TOPOLOGY_DEFAULT = Path(__file__).resolve().parents[2] / "topology.yaml"
D_HEAD = 64
DTYPE = "f16"
_CUBE_RE = re.compile(r"\bcube(\d+)\b")
def _ccl_cfg():
return resolve_algorithm_config(
load_ccl_config(), name="lrab_hierarchical_allreduce",
)
def _engine_factory(t, d):
return GraphEngine(getattr(t, "topology_obj", t), enable_data=True)
def _dma_write_cubes(op_log) -> list[int]:
"""Return the list of CUBE ids that emitted a ``dma_write``.
Decode's final-store path uses one ``tl.store`` from the root rank.
Multiple HBM-channel-level write records may correspond to the same
logical store; this just reports the source CUBE for each.
"""
cubes: list[int] = []
for r in op_log:
if r.op_name != "dma_write":
continue
m = _CUBE_RE.search(r.component_id)
if m is not None:
cubes.append(int(m.group(1)))
return cubes
def _run_decode_long(
*, C: int, P: int, S_kv: int, sub_w: int | None,
):
"""Drive a decode_long launch. ``sub_w=None`` means omit the arg
entirely (exercises the kernel's current default behaviour)."""
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)
dp_kv = DPPolicy(cube="row_wise" if C > 1 else "replicate",
pe="row_wise", num_cubes=C, num_pes=P)
T_q = 1
h_q = 8
h_kv = 1
ctx.zeros((T_q, h_q * D_HEAD),
dtype=DTYPE, dp=dp_full, name=f"q_dec_2d_c{C}")
k = ctx.zeros((S_kv, h_kv * D_HEAD),
dtype=DTYPE, dp=dp_kv, name=f"k_dec_2d_c{C}")
v = ctx.zeros((S_kv, h_kv * D_HEAD),
dtype=DTYPE, dp=dp_kv, name=f"v_dec_2d_c{C}")
o = ctx.empty((T_q, h_q * D_HEAD),
dtype=DTYPE, dp=dp_full, name=f"o_dec_2d_c{C}")
q = ctx.zeros((T_q, h_q * D_HEAD),
dtype=DTYPE, dp=dp_full, name=f"q_dec_2d_c{C}_2")
launch_args = [q, k, v, o, T_q, S_kv, h_q, h_kv, D_HEAD, C, P]
if sub_w is not None:
launch_args.append(sub_w)
ctx.launch(
f"gqa_decode_long_2d_c{C}_sw{sub_w}",
gqa_attention_decode_long_kernel,
*launch_args,
_auto_dim_remap=False,
)
return run_bench(
topology=topo, bench_fn=_bench_fn,
device=resolve_device(None), engine_factory=_engine_factory,
)
# ── T1: C=8 lrab-adapted center-root completes ───────────────────────
def test_decode_2d_sub_w4_sub_h2_completes():
"""ADR-0060 §4.2 prescribed center-root mesh reduce at the
milestone target: C=8 over a 2×4 sub-mesh (sub_w=4, sub_h=2),
P=8 PEs per CUBE. With zero inputs, output is trivially zero;
the test guards that the kernel reaches completion without
deadlock or scratch overflow.
"""
result = _run_decode_long(C=8, P=8, S_kv=2048, sub_w=4)
assert result.completion.ok, (
f"decode at C=8, sub_w=4 (lrab-adapted) must complete; "
f"got {result.completion}"
)
# ── T2: root lives at the geometric center CUBE (cube 6) ─────────────
def test_decode_2d_sub_w4_sub_h2_root_at_center_cube_6():
"""For sub_w=4, sub_h=2: root_col=2, root_row=1, root_cube=6.
Only the root CUBE issues the final ``tl.store`` — every other
rank short-circuits. The dma_write records must come exclusively
from CUBE 6.
"""
result = _run_decode_long(C=8, P=8, S_kv=2048, sub_w=4)
assert result.completion.ok
cubes = _dma_write_cubes(result.engine.op_log)
assert cubes, "expected at least one dma_write for the final O store"
distinct = set(cubes)
assert distinct == {6}, (
f"final dma_write must come exclusively from CUBE 6 (sub_w=4, "
f"sub_h=2 root); got cubes={sorted(distinct)}"
)
# ── T3: backward-compat — sub_w omitted → existing 1D chain at C=4 ───
def test_decode_2d_backward_compat_sub_w0_default():
"""When ``sub_w`` is omitted from the launch, the kernel must
behave identically to today: 1D inter-CUBE chain along W,
converging at CUBE 0. This serves as a regression anchor — Phase 2
must not change the existing C=4 panel behaviour.
Passes today as well as after Phase 2.
"""
result = _run_decode_long(C=4, P=8, S_kv=2048, sub_w=None)
assert result.completion.ok, (
f"decode at C=4 with default sub_w (1D chain) must complete; "
f"got {result.completion}"
)
cubes = _dma_write_cubes(result.engine.op_log)
assert cubes, "expected at least one dma_write for the final O store"
distinct = set(cubes)
assert distinct == {0}, (
f"1D-chain root must be CUBE 0; got cubes={sorted(distinct)}"
)
# ── T4: degenerate sub_w configurations are rejected ─────────────────
def test_decode_2d_invalid_sub_w_rejected_when_sub_h_lt_2():
"""ADR-0060 §4.2 + CLAUDE.md "Simplicity First": only sub-meshes
with both sub_w >= 2 and sub_h >= 2 use lrab. C=4 with sub_w=4
would give sub_h=1 (degenerate; Phase 2 of lrab becomes a no-op).
Callers must use the 1D-chain path (sub_w=0/omitted) for that case.
The kernel must reject the degenerate combination with a clear
error rather than silently producing a 1D-chain result at the
wrong root location.
"""
with pytest.raises((ValueError, AssertionError), match=r"sub_[wh]"):
_run_decode_long(C=4, P=8, S_kv=2048, sub_w=4)
tests/attention/test_gqa_prefill_long_c8_snake.py
+154
View File
@@ -0,0 +1,154 @@
"""Tests for prefill_long Ring KV at C=8 over a 2×4 snake sub-mesh.
ADR-0060 §5.5 design target for the single-KV-group LLaMA-3.1-70B
configuration: ``C = G = 8`` (one Q head per CUBE), Ring KV rotates
the 8 KV slices around all 8 CUBEs.
Increment 1 wired ``configure_sfr_intercube_ring(submesh_shape=(2, 4))``
to map the kernel's logical E/W to a snake/serpentine 1-hop path
through the top 2×4 sub-mesh of the 4×4 SIP CUBE mesh. The kernel
itself uses only logical ``dir="W"`` / ``dir="E"`` — the snake is
transparent at kernel level.
This file verifies the assembly works end-to-end at C=8:
T1 kernel completes (no scratch overflow, no deadlock)
T2 per-CUBE head-parallel output (8 dma_writes, one per CUBE)
T3 tile-granular ring traffic at C=8 follows the same
``(C-1)·n_tiles·2·C`` formula as the existing tile-ring tests
"""
from __future__ import annotations
from pathlib import Path
from kernbench.benches._gqa_attention_prefill_long import gqa_attention_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
from kernbench.policy.placement.dp import DPPolicy
from kernbench.runtime_api.bench_runner import run_bench
from kernbench.runtime_api.types import resolve_device
from kernbench.sim_engine.engine import GraphEngine
from kernbench.topology.builder import resolve_topology
TOPOLOGY_DEFAULT = Path(__file__).resolve().parents[2] / "topology.yaml"
D_HEAD = 64
DTYPE = "f16"
def _ccl_cfg():
return resolve_algorithm_config(
load_ccl_config(), name="lrab_hierarchical_allreduce",
)
def _engine_factory(t, d):
return GraphEngine(getattr(t, "topology_obj", t), enable_data=True)
def _count(op_log, name: str) -> int:
return sum(1 for r in op_log if r.op_name == name)
def _run_prefill_c8_snake(*, T_q: int, S_kv: int):
"""Head-parallel prefill at C=8 with snake-mapped Ring KV over the
top 2×4 sub-mesh of the 4×4 SIP CUBE mesh."""
topo = resolve_topology(str(TOPOLOGY_DEFAULT))
C = 8
def _bench_fn(ctx):
configure_sfr_intercube_ring(
ctx.engine, ctx.spec, _ccl_cfg(),
submesh_shape=(2, 4),
)
dp_q = DPPolicy(cube="replicate", pe="replicate",
num_cubes=C, num_pes=1)
dp_kv = DPPolicy(cube="row_wise", pe="replicate",
num_cubes=C, num_pes=1)
dp_o = DPPolicy(cube="row_wise", pe="replicate",
num_cubes=C, num_pes=1)
q = ctx.zeros((T_q, D_HEAD), dtype=DTYPE, dp=dp_q,
name=f"q_c8_snake_t{T_q}_s{S_kv}")
k = ctx.zeros((S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv,
name=f"k_c8_snake_t{T_q}_s{S_kv}")
v = ctx.zeros((S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv,
name=f"v_c8_snake_t{T_q}_s{S_kv}")
o = ctx.empty((T_q * C, D_HEAD), dtype=DTYPE, dp=dp_o,
name=f"o_c8_snake_t{T_q}_s{S_kv}")
ctx.launch(
f"gqa_prefill_long_c8_snake_t{T_q}_s{S_kv}",
gqa_attention_prefill_long_kernel,
q, k, v, o,
T_q, S_kv, D_HEAD, C,
_auto_dim_remap=False,
)
return run_bench(
topology=topo, bench_fn=_bench_fn,
device=resolve_device(None),
engine_factory=_engine_factory,
)
# ── T1: kernel completes end-to-end at C=8 ───────────────────────────
def test_prefill_long_c8_snake_completes():
"""ADR-0060 §5.5 design target: prefill Ring KV at C=G=8 over the
2×4 snake sub-mesh. With zero inputs, output is trivially zero;
this test guards that the kernel reaches completion without
deadlock, scratch overflow, or routing failure.
Configuration: T_q=4, S_kv=8192 → S_local=1024, n_tiles=1
(TILE_S_KV=1024). Bounded scratch.
"""
result = _run_prefill_c8_snake(T_q=4, S_kv=8192)
assert result.completion.ok, (
f"prefill at C=8 with snake ring must complete; "
f"got {result.completion}"
)
# ── T2: 8 dma_writes (head-parallel, one head per CUBE) ──────────────
def test_prefill_long_c8_snake_distributed_output_count():
"""Head-parallel: each CUBE owns one Q head and writes its own
head's output (T_q, d_head) rows. No inter-CUBE reduce on the
output side — so ``dma_write_count == C == 8``.
This is the C=8 analogue of
``test_prefill_long_tile_ring_dma_write_count`` at C=4.
"""
result = _run_prefill_c8_snake(T_q=4, S_kv=8192)
assert result.completion.ok
n_writes = _count(result.engine.op_log, "dma_write")
assert n_writes == 8, (
f"head-parallel C=8: expected 8 dma_writes (one per CUBE); "
f"got {n_writes}"
)
# ── T3: tile-granular ring ipcq_copy count scales as (C-1)·n_tiles·2·C
def test_prefill_long_c8_snake_tile_ipcq_count():
"""ADR-0060 §5.5.1 tile-granular ring: per-CUBE send count is
``2·n_tiles·(C-1)`` (K + V per tile per ring step). Aggregated
across C CUBEs, total ipcq_copy = ``(C-1)·n_tiles·2·C``.
Configuration: T_q=4, S_kv=8192, C=8 → S_local=1024, n_tiles=1
(TILE_S_KV=1024). Expected total = ``7·1·2·8 = 112``.
Verifies that the snake-mapped 1-hop physical links carry the
same logical ring traffic as the existing C=4 1D-row ring tests.
"""
C = 8
n_tiles = 1 # S_local=1024 / TILE_S_KV=1024
result = _run_prefill_c8_snake(T_q=4, S_kv=8192)
assert result.completion.ok
n_copy = _count(result.engine.op_log, "ipcq_copy")
expected = (C - 1) * n_tiles * 2 * C
assert n_copy == expected, (
f"tile-granular ring at C=8: expected {expected} ipcq_copy "
f"((C-1)·n_tiles·2·C = {C - 1}·{n_tiles}·2·{C}); got {n_copy}"
)
tests/attention/test_gqa_prefill_long_pe_sp.py
+233
View File
@@ -0,0 +1,233 @@
"""Tests for intra-CUBE PE-SP in prefill_long (query-axis split).
ADR-0060 §5.5 last bullet + §B-item-3: split the head's query rows
``[T_q, d_head]`` across the ``P`` PEs of each CUBE so all P PEs work
in parallel. Output rows are disjoint across PEs ⇒ no intra-CUBE
reduce needed.
Same-lane SFR wiring: PE ``i`` in CUBE A has its own E/W ring link to
PE ``i`` in CUBE B (the snake's prev/next). All P PEs of a CUBE see
the same K, V (HBM-resident, ``pe="replicate"``) ⇒ P parallel rings
run in lockstep, each PE rotating its own K/V copies via its own IPCQ
channels.
Activation contract:
- ``P == 1`` (default; omitted from launch args) → existing
PE-0-only behavior. Byte-for-byte unchanged.
- ``P > 1`` → all P PEs active; each handles ``T_q // P`` query
rows. Requires ``T_q % P == 0`` (degenerate T_q < P is rejected
— caller must use ``P=1`` for that workload, ADR-0060 §B-item-3).
Phase 1: tests only — production code lands in Phase 2.
T1, T2, T3, T5 fail today (TypeError: kernel signature has no P).
T4 passes today as the backward-compat anchor.
"""
from __future__ import annotations
from pathlib import Path
import pytest
from kernbench.benches._gqa_attention_prefill_long import gqa_attention_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
from kernbench.policy.placement.dp import DPPolicy
from kernbench.runtime_api.bench_runner import run_bench
from kernbench.runtime_api.types import resolve_device
from kernbench.sim_engine.engine import GraphEngine
from kernbench.topology.builder import resolve_topology
TOPOLOGY_DEFAULT = Path(__file__).resolve().parents[2] / "topology.yaml"
D_HEAD = 64
DTYPE = "f16"
def _ccl_cfg():
return resolve_algorithm_config(
load_ccl_config(), name="lrab_hierarchical_allreduce",
)
def _engine_factory(t, d):
return GraphEngine(getattr(t, "topology_obj", t), enable_data=True)
def _count(op_log, name: str) -> int:
return sum(1 for r in op_log if r.op_name == name)
def _run_prefill_pe_sp(
*, T_q: int, S_kv: int, C: int, P: int | None,
snake: bool,
):
"""Drive a prefill_long launch with optional intra-CUBE PE-SP.
``P=None`` → omit P from the launch args (exercises the kernel's
default behaviour).
``snake=True`` → install the 2×4 snake ring SFR (Increment 1);
else install the 1D-row ring at ``ring_size=C``.
"""
topo = resolve_topology(str(TOPOLOGY_DEFAULT))
def _bench_fn(ctx):
if snake:
configure_sfr_intercube_ring(
ctx.engine, ctx.spec, _ccl_cfg(),
submesh_shape=(2, 4),
)
else:
configure_sfr_intercube_ring(
ctx.engine, ctx.spec, _ccl_cfg(),
ring_size=C,
)
num_pes = P if (P is not None and P > 1) else 1
# When PE-SP is active, Q and O are split row-wise across PEs;
# K, V remain replicated within a CUBE.
q_pe = "row_wise" if num_pes > 1 else "replicate"
o_pe = "row_wise" if num_pes > 1 else "replicate"
dp_q = DPPolicy(cube="replicate", pe=q_pe,
num_cubes=C, num_pes=num_pes)
dp_kv = DPPolicy(cube="row_wise" if C > 1 else "replicate",
pe="replicate", num_cubes=C, num_pes=num_pes)
dp_o = DPPolicy(cube="row_wise" if C > 1 else "replicate",
pe=o_pe, num_cubes=C, num_pes=num_pes)
suffix = f"c{C}_p{P}_t{T_q}_s{S_kv}"
q = ctx.zeros((T_q, D_HEAD), dtype=DTYPE, dp=dp_q,
name=f"q_pe_sp_{suffix}")
k = ctx.zeros((S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv,
name=f"k_pe_sp_{suffix}")
v = ctx.zeros((S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv,
name=f"v_pe_sp_{suffix}")
o = ctx.empty((T_q * C, D_HEAD), dtype=DTYPE, dp=dp_o,
name=f"o_pe_sp_{suffix}")
launch_args = [q, k, v, o, T_q, S_kv, D_HEAD, C]
if P is not None:
launch_args.append(P)
ctx.launch(
f"gqa_prefill_long_pe_sp_{suffix}",
gqa_attention_prefill_long_kernel,
*launch_args,
_auto_dim_remap=False,
)
return run_bench(
topology=topo, bench_fn=_bench_fn,
device=resolve_device(None),
engine_factory=_engine_factory,
)
# ── T1: C=8 P=8 completes end-to-end ─────────────────────────────────
def test_prefill_c8_p8_pe_sp_completes():
"""ADR-0060 §5.5 + §B-item-3 design target: 64 ranks active
(8 CUBEs × 8 PEs), query-axis split across PEs, snake-mapped
Ring KV. Guards no scratch overflow, no deadlock across the
P parallel same-lane rings.
"""
result = _run_prefill_pe_sp(
T_q=8, S_kv=8192, C=8, P=8, snake=True,
)
assert result.completion.ok, (
f"prefill at C=8, P=8 (PE-SP) must complete; "
f"got {result.completion}"
)
# ── T2: 64 dma_writes (one per PE, disjoint query rows) ──────────────
def test_prefill_c8_p8_pe_sp_64_dma_writes():
"""With query-axis split, each PE owns ``T_q/P = 1`` query row
and stores its own ``(1, d_head)`` slice. Across all 64 ranks
(8 CUBEs × 8 PEs), ``dma_write_count == 64``.
This is the structural signal that PE-SP is actually wired:
PE-0-only would give 8 dma_writes (one per CUBE).
"""
result = _run_prefill_pe_sp(
T_q=8, S_kv=8192, C=8, P=8, snake=True,
)
assert result.completion.ok
n_writes = _count(result.engine.op_log, "dma_write")
assert n_writes == 64, (
f"PE-SP at C=8, P=8: expected 64 dma_writes "
f"(8 cubes × 8 PEs, each writes its T_q/P=1 query row); "
f"got {n_writes}"
)
# ── T3: P parallel rings — ipcq_copy scales by P ─────────────────────
def test_prefill_c8_p8_pe_sp_ring_ipcq_count():
"""Per-PE same-lane rings: each PE runs its own ring traffic via
its own IPCQ channels. Total inter-CUBE ipcq_copy:
``(C-1) · n_tiles · 2 · C · P`` (= existing C=8 formula × P).
Configuration: T_q=8, S_kv=8192, C=8 → S_local=1024, n_tiles=1
(TILE_S_KV=1024). Expected = ``7·1·2·8·8`` = **896**.
"""
C = 8
P = 8
n_tiles = 1
result = _run_prefill_pe_sp(
T_q=8, S_kv=8192, C=C, P=P, snake=True,
)
assert result.completion.ok
n_copy = _count(result.engine.op_log, "ipcq_copy")
expected = (C - 1) * n_tiles * 2 * C * P
assert n_copy == expected, (
f"PE-SP parallel rings at C=8, P=8: expected {expected} "
f"ipcq_copy ((C-1)·n_tiles·2·C·P = "
f"{C - 1}·{n_tiles}·2·{C}·{P}); got {n_copy}"
)
# ── T4: backward-compat — P omitted → existing PE-0-only behaviour ───
def test_prefill_p_default_1_backward_compat():
"""When ``P`` is omitted from the launch args, the kernel must
behave identically to today: only PE 0 of each CUBE participates;
one head per CUBE; one dma_write per CUBE.
At C=4 this gives 4 dma_writes (matches the existing
``test_prefill_long_tile_ring_dma_write_count``). Phase 2 must
not regress this path.
Passes today AND after Phase 2.
"""
C = 4
result = _run_prefill_pe_sp(
T_q=4, S_kv=8192, C=C, P=None, snake=False,
)
assert result.completion.ok, (
f"prefill at C=4 with default P (PE-0-only) must complete; "
f"got {result.completion}"
)
n_writes = _count(result.engine.op_log, "dma_write")
assert n_writes == C, (
f"default P=1 (PE-0-only): expected {C} dma_writes "
f"(one per CUBE); got {n_writes}"
)
# ── T5: validation — T_q must be divisible by P when P > 1 ───────────
def test_prefill_pe_sp_rejects_non_divisible_t_q():
"""ADR-0060 §B-item-3 fallback (KV-block split + intra-CUBE
reduce for T_q < P) is deferred. Callers must request ``P=1`` for
workloads where T_q < P or T_q % P != 0.
C=4, P=8, T_q=4 violates ``T_q % P == 0`` (and also T_q < P).
The kernel must raise ValueError with a clear error message
rather than silently producing wrong results.
"""
with pytest.raises((ValueError, AssertionError), match=r"T_q"):
_run_prefill_pe_sp(
T_q=4, S_kv=8192, C=4, P=8, snake=False,
)
tests/attention/test_milestone_gqa_single_kv_group_prefill_panel.py
+146
View File
@@ -0,0 +1,146 @@
"""Tests for the single-KV-group prefill panel in milestone_gqa_headline.
The single-KV-group (LLaMA-3.1-70B target) prefill panel wires together
all Increment 1–4 work in a single bench panel:
- C=8 head-parallel + Ring KV via snake-mapped 2×4 sub-mesh (Inc 1, 3)
- P=8 intra-CUBE PE-SP (Inc 4)
- d_head=128 (LLaMA-3.1-70B), one-shot prefill T_q = S_kv = 32K
This file verifies the bench-config wiring:
T1 the new ``single_kv_group_prefill_gqa_c8_p8`` panel is registered
with the expected dims in ``_PANELS`` + ``_PANEL_DISPATCH``.
T2 ``_run_prefill_panel`` accepts the new ``P``, ``T_q``, ``d_head``
kwargs and drives the prefill kernel to completion at the
milestone-target ``(C, P) = (8, 8)``. Uses smaller T_q/S_kv to
keep test time bounded — full 32K runs come from
``kernbench run milestone-gqa-headline``.
T3 Existing prefill panels still work when ``_run_prefill_panel``
is called without the new kwargs (backward compat anchor).
Phase 1: tests only — production code lands in Phase 2.
T1 fails today (panel not registered).
T2 fails today (helper doesn't accept the new kwargs).
T3 passes today as the backward-compat anchor.
"""
from __future__ import annotations
from pathlib import Path
from kernbench.benches.milestone_gqa_headline import ( # noqa: F401
_PANEL_DISPATCH,
_PANELS,
_run_prefill_panel,
)
from kernbench.runtime_api.bench_runner import run_bench
from kernbench.runtime_api.types import resolve_device
from kernbench.sim_engine.engine import GraphEngine
from kernbench.topology.builder import resolve_topology
TOPOLOGY_DEFAULT = Path(__file__).resolve().parents[2] / "topology.yaml"
def _engine_factory(t, d):
return GraphEngine(getattr(t, "topology_obj", t), enable_data=True)
# ── T1: new panel registered ─────────────────────────────────────────
def test_single_kv_group_prefill_panel_registered():
"""The ``single_kv_group_prefill_gqa_c8_p8`` panel must be in ``_PANELS`` and
its dispatch entry must have the expected LLaMA-3.1-70B params.
Headline config:
C = 8 (head-parallel, snake sub-mesh)
P = 8 (intra-CUBE PE-SP)
T_q = 32_768 (one-shot long-context prefill)
S_kv = 32_768
d_head = 128 (LLaMA-3.1-70B)
"""
panel_name = "single_kv_group_prefill_gqa_c8_p8"
assert panel_name in _PANELS, (
f"{panel_name!r} not in _PANELS; got {_PANELS}"
)
assert panel_name in _PANEL_DISPATCH, (
f"{panel_name!r} not in _PANEL_DISPATCH"
)
kind, params = _PANEL_DISPATCH[panel_name]
assert kind == "prefill", f"kind={kind!r}, expected 'prefill'"
assert params.get("C") == 8, f"C={params.get('C')}, expected 8"
assert params.get("P") == 8, f"P={params.get('P')}, expected 8"
assert params.get("T_q") == 32_768, (
f"T_q={params.get('T_q')}, expected 32_768"
)
assert params.get("S_kv") == 32_768, (
f"S_kv={params.get('S_kv')}, expected 32_768"
)
assert params.get("d_head") == 128, (
f"d_head={params.get('d_head')}, expected 128"
)
# ── T2: helper drives the C=8, P=8 prefill kernel to completion ──────
def test_single_kv_group_prefill_panel_runner_smoke():
"""``_run_prefill_panel`` must accept the new ``P``, ``T_q``,
``d_head`` kwargs and successfully launch the prefill kernel at
``(C, P) = (8, 8)`` with the snake-mapped 2×4 SFR.
Uses **smaller** T_q/S_kv than the headline panel so the test
completes quickly. The headline 32K dims are exercised via
``kernbench run milestone-gqa-headline``, not pytest.
"""
topo = resolve_topology(str(TOPOLOGY_DEFAULT))
def _bench_fn(ctx):
_run_prefill_panel(
ctx,
panel="single_kv_group_prefill_gqa_c8_p8",
C=8, P=8,
T_q=8, S_kv=8192, # smoke dims, not the 32K headline
d_head=128,
)
result = run_bench(
topology=topo, bench_fn=_bench_fn,
device=resolve_device(None),
engine_factory=_engine_factory,
)
assert result.completion.ok, (
f"single_kv_group prefill panel smoke at C=8 P=8 must complete; "
f"got {result.completion}"
)
# ── T3: existing panel calls still work (backward compat) ────────────
def test_existing_prefill_panel_runner_backward_compat():
"""Existing callers of ``_run_prefill_panel`` (no ``P``/``T_q``/
``d_head`` overrides) must continue to work after the signature
extension. Exercises the ``multi_user_prefill_gqa`` panel dims
from ``_PANEL_DISPATCH``.
Passes today AND after Phase 2.
"""
panel_name = "multi_user_prefill_gqa"
_kind, params = _PANEL_DISPATCH[panel_name]
topo = resolve_topology(str(TOPOLOGY_DEFAULT))
def _bench_fn(ctx):
_run_prefill_panel(
ctx,
panel=panel_name,
C=params["C"], S_kv=params["S_kv"],
)
result = run_bench(
topology=topo, bench_fn=_bench_fn,
device=resolve_device(None),
engine_factory=_engine_factory,
)
assert result.completion.ok, (
f"existing prefill panel {panel_name!r} must still run with "
f"the original kwargs; got {result.completion}"
)
tests/test_intercube_snake_ring.py
+273
View File
@@ -0,0 +1,273 @@
"""Tests for snake/serpentine ring extension of configure_sfr_intercube_ring.
Verifies the multi-row snake ring SFR wiring used for ADR-0060 §5.5
prefill Ring KV at C=8 on a 2×4 sub-mesh of the 4×4 CUBE mesh.
The snake path for ``submesh_shape=(2, 4), origin=(0, 0)`` is::
(0,0)→(0,1)→(0,2)→(0,3)→(1,3)→(1,2)→(1,1)→(1,0)→wrap to (0,0)
In cube indices on a mesh_w=4 grid: ``[0, 1, 2, 3, 7, 6, 5, 4]``.
Every consecutive pair (including the wrap) is a 1-hop CUBE NOC neighbour.
The kernel sees a 1D logical E/W ring; the snake is invisible to it.
"""
from __future__ import annotations
from pathlib import Path
import pytest
from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config
from kernbench.ccl.sfr_config import configure_sfr_intercube_ring
from kernbench.sim_engine.engine import GraphEngine
from kernbench.topology.builder import resolve_topology
TOPOLOGY_PATH = Path(__file__).parent.parent / "topology.yaml"
N_CUBES = 16 # 4×4 SIP CUBE mesh
PES_PER_CUBE = 8
MESH_W = 4
# Canonical snake path for (2,4) sub-mesh at origin (0,0).
SNAKE_2X4_O00 = [0, 1, 2, 3, 7, 6, 5, 4]
# Canonical snake path for (2,4) sub-mesh at origin (1,0) (rows 1-2).
SNAKE_2X4_O10 = [4, 5, 6, 7, 11, 10, 9, 8]
def _engine_and_spec():
topo = resolve_topology(str(TOPOLOGY_PATH))
engine = GraphEngine(topo.topology_obj, enable_data=True)
return engine, topo.topology_obj.spec
def _merged_cfg():
cfg = load_ccl_config()
return resolve_algorithm_config(cfg, name="lrab_hierarchical_allreduce")
def _qp(engine, sip: int, cube: int, pe: int):
return engine._components[f"sip{sip}.cube{cube}.pe{pe}.pe_ipcq"].queue_pairs
def _row_col(cube: int) -> tuple[int, int]:
return cube // MESH_W, cube % MESH_W
def _l1(c1: int, c2: int) -> int:
r1, col1 = _row_col(c1)
r2, col2 = _row_col(c2)
return abs(r1 - r2) + abs(col1 - col2)
class TestSnakeRing2x4Origin00:
"""Snake ring through the top 2×4 sub-mesh of the 4×4 SIP."""
def test_snake_ring_2x4_world_size(self):
"""All PEs on all CUBEs of all SIPs still enumerated (snake only
restricts which CUBEs have ring links, not the world)."""
engine, spec = _engine_and_spec()
cfg = _merged_cfg()
plan = configure_sfr_intercube_ring(
engine, spec, cfg, submesh_shape=(2, 4),
)
n_sips = int(spec["system"]["sips"]["count"])
assert plan["world_size"] == n_sips * N_CUBES * PES_PER_CUBE
assert len(plan["rank_to_pe"]) == plan["world_size"]
def test_snake_ring_2x4_path_corners_E(self):
"""The four 'interesting' E hops on the snake path:
cube0.E → cube1 (start of row 0)
cube3.E → cube7 (row-bridge: top-right corner down)
cube7.E → cube6 (row 1, going leftward)
cube4.E → cube0 (wrap: bottom-left to top-left)
"""
engine, spec = _engine_and_spec()
cfg = _merged_cfg()
configure_sfr_intercube_ring(
engine, spec, cfg, submesh_shape=(2, 4),
)
assert _qp(engine, 0, 0, 0)["E"]["peer"].cube == 1
assert _qp(engine, 0, 3, 0)["E"]["peer"].cube == 7
assert _qp(engine, 0, 7, 0)["E"]["peer"].cube == 6
assert _qp(engine, 0, 4, 0)["E"]["peer"].cube == 0
def test_snake_ring_2x4_path_corners_W(self):
"""Symmetric W hops (W is reverse of E along the snake):
cube0.W → cube4 (wrap reverse)
cube7.W → cube3 (row-bridge reverse: up)
cube6.W → cube7 (row 1 reverse: rightward)
cube4.W → cube5 (row-1 leftmost, W goes right within row 1)
Note: under the snake path [0,1,2,3,7,6,5,4], W(cube_i) is the
predecessor on that path. So W(4)=5 (previous on path).
"""
engine, spec = _engine_and_spec()
cfg = _merged_cfg()
configure_sfr_intercube_ring(
engine, spec, cfg, submesh_shape=(2, 4),
)
assert _qp(engine, 0, 0, 0)["W"]["peer"].cube == 4
assert _qp(engine, 0, 7, 0)["W"]["peer"].cube == 3
assert _qp(engine, 0, 6, 0)["W"]["peer"].cube == 7
assert _qp(engine, 0, 4, 0)["W"]["peer"].cube == 5
def test_snake_ring_2x4_interior(self):
"""Interior hops along each row (not at row-bridge or wrap)."""
engine, spec = _engine_and_spec()
cfg = _merged_cfg()
configure_sfr_intercube_ring(
engine, spec, cfg, submesh_shape=(2, 4),
)
# Row 0 interior: cube1 between cube0 and cube2.
assert _qp(engine, 0, 1, 0)["E"]["peer"].cube == 2
assert _qp(engine, 0, 1, 0)["W"]["peer"].cube == 0
# Row 1 interior: cube5 between cube6 and cube4 along the snake.
assert _qp(engine, 0, 5, 0)["E"]["peer"].cube == 4
assert _qp(engine, 0, 5, 0)["W"]["peer"].cube == 6
def test_snake_ring_2x4_all_hops_are_1_hop(self):
"""Every E hop on the snake is L1-distance 1 in the cube mesh.
Defensive — holds by construction (snake = Hamiltonian cycle
on the 2×4 grid graph), but explicit guards against future
path-builder regressions.
"""
engine, spec = _engine_and_spec()
cfg = _merged_cfg()
configure_sfr_intercube_ring(
engine, spec, cfg, submesh_shape=(2, 4),
)
for cube in SNAKE_2X4_O00:
qp = _qp(engine, 0, cube, 0)
e_peer = qp["E"]["peer"].cube
w_peer = qp["W"]["peer"].cube
assert _l1(cube, e_peer) == 1, (
f"snake E hop cube{cube}→cube{e_peer} is not 1-hop"
)
assert _l1(cube, w_peer) == 1, (
f"snake W hop cube{cube}→cube{w_peer} is not 1-hop"
)
def test_snake_ring_2x4_off_path_cubes_have_no_ring_links(self):
"""CUBEs not on the snake path (rows 2-3, i.e. cubes 8..15)
get no E/W ring entries — consistent with the current
``if cube < ring_size`` behaviour."""
engine, spec = _engine_and_spec()
cfg = _merged_cfg()
configure_sfr_intercube_ring(
engine, spec, cfg, submesh_shape=(2, 4),
)
for cube in range(8, 16):
qp = _qp(engine, 0, cube, 0)
assert "E" not in qp, (
f"sip0.cube{cube}.pe0 has E but is off the snake path"
)
assert "W" not in qp
def test_snake_ring_2x4_intra_namespace_unchanged(self):
"""Snake doesn't disturb the intra-cube 2×4 PE grid wiring
(intra_N/S/E/W) — it only changes cube-level E/W."""
from kernbench.ccl.sfr_config import _intra_cube_neighbors
engine, spec = _engine_and_spec()
cfg = _merged_cfg()
configure_sfr_intercube_ring(
engine, spec, cfg, submesh_shape=(2, 4),
)
# Spot-check a few cubes (including off-path) and all 8 PEs.
for cube in (0, 3, 4, 7, 10, 15):
for pe in range(PES_PER_CUBE):
qp = _qp(engine, 0, cube, pe)
expected = _intra_cube_neighbors(pe)
for d, expected_pe in expected.items():
assert d in qp, f"cube{cube}.pe{pe} missing {d}"
assert qp[d]["peer"].pe == expected_pe
assert qp[d]["peer"].cube == cube
class TestSnakeRing2x4OriginShift:
"""Snake ring through rows 1-2 (origin=(1, 0)) — proves the
origin parameter shifts the sub-mesh."""
def test_snake_ring_origin_shift_path(self):
"""For origin=(1, 0), snake path is [4,5,6,7,11,10,9,8]:
row 1 left→right: cube4 → cube5 → cube6 → cube7
↓
row 2 right→left: cube11 → cube10 → cube9 → cube8
↺ wraps to cube4
"""
engine, spec = _engine_and_spec()
cfg = _merged_cfg()
configure_sfr_intercube_ring(
engine, spec, cfg, submesh_shape=(2, 4), submesh_origin=(1, 0),
)
# Path corners.
assert _qp(engine, 0, 4, 0)["E"]["peer"].cube == 5
assert _qp(engine, 0, 7, 0)["E"]["peer"].cube == 11
assert _qp(engine, 0, 11, 0)["E"]["peer"].cube == 10
assert _qp(engine, 0, 8, 0)["E"]["peer"].cube == 4 # wrap
def test_snake_ring_origin_shift_off_path_no_links(self):
"""For origin=(1, 0), cubes 0..3 (row 0) and 12..15 (row 3)
get no ring links."""
engine, spec = _engine_and_spec()
cfg = _merged_cfg()
configure_sfr_intercube_ring(
engine, spec, cfg, submesh_shape=(2, 4), submesh_origin=(1, 0),
)
for cube in list(range(0, 4)) + list(range(12, 16)):
qp = _qp(engine, 0, cube, 0)
assert "E" not in qp
assert "W" not in qp
class TestSnakeRingBackwardCompat:
"""The existing 1D-row API (no submesh_shape) must behave identically
to before the snake extension lands."""
def test_snake_ring_backward_compat_1d_default(self):
"""ring_size=4 (single-row) without submesh_shape — identical
behaviour to today's 1D-row ring: cube0..3 with E/W wrap;
cubes 4..15 have no ring links."""
engine, spec = _engine_and_spec()
cfg = _merged_cfg()
configure_sfr_intercube_ring(
engine, spec, cfg, ring_size=4,
)
# Row-0 1D ring with wrap.
assert _qp(engine, 0, 0, 0)["E"]["peer"].cube == 1
assert _qp(engine, 0, 0, 0)["W"]["peer"].cube == 3 # wrap
assert _qp(engine, 0, 3, 0)["E"]["peer"].cube == 0 # wrap
# Off-row cubes get no ring links.
for cube in range(4, 16):
qp = _qp(engine, 0, cube, 0)
assert "E" not in qp
assert "W" not in qp
class TestSnakeRingValidation:
"""Input validation of the submesh_shape / ring_size combinations."""
def test_snake_ring_invalid_submesh_overflow(self):
"""submesh_shape that doesn't fit the cube mesh is rejected."""
engine, spec = _engine_and_spec()
cfg = _merged_cfg()
with pytest.raises(ValueError, match=r"sub.?mesh"):
configure_sfr_intercube_ring(
engine, spec, cfg, submesh_shape=(2, 5), # col overflow
)
def test_snake_ring_ring_size_mismatch(self):
"""submesh_shape and ring_size disagree → ValueError."""
engine, spec = _engine_and_spec()
cfg = _merged_cfg()
with pytest.raises(ValueError, match=r"ring_size"):
configure_sfr_intercube_ring(
engine, spec, cfg,
submesh_shape=(2, 4), ring_size=4,
)