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,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}"
|
||||
)
|
||||
Reference in New Issue
Block a user