gqa: single-KV-group LLaMA-3.1-70B prefill milestone (Increments 1-5)
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:
@@ -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,
|
||||
)
|
||||
Reference in New Issue
Block a user