gqa: S_kv tile sweep for decode_long/decode_short/prefill_short (ADR-0063 §A.2)
Replace each kernel's local one-shot partial with a Tile-0 bootstrap + a ``scratch_scope``-wrapped merge loop. Per-rank scratch is now bounded by ``TILE_S_KV = 1024`` regardless of ``S_local``, lifting the long-context S ceiling. Backward compatible at ``S_local <= TILE_S_KV`` (loop body is empty; op_log structurally identical to today). Framework: extend ``memory_store.read`` to support partial reads at offsets within a stored region. Models real Triton ``tl.load(ptr+offset, shape=...)`` — needed because each tile loads ``k_ptr + tile_start*row_bytes`` for its sub-slice of the per-rank KV region. prefill_long is intentionally left untiled. Its ring loop carries full-slice ``Kc``/``Vc`` for ``tl.send`` between CUBEs, so tile-sweeping step 0 wouldn't shrink the kernel's actual scratch footprint. Lifting prefill_long's ceiling requires a tile-granular ring rewrite — separate phase. Docstring + inline comments cleaned to reflect the current shape. Docstrings across all 4 GQA kernels: drop P1a/P2a/P2b/P6a/P6b lineage paragraphs and historical deviation lists; describe each kernel by what it does, not how it got here. Add explicit ``# Local attention`` / ``# Communication`` section headers. Tests: new ``tests/attention/test_gqa_tile_sweep.py`` with 5 tests covering single-tile op_log stability, multi-tile ``copy_to`` emission across all 3 refactored kernels, and the headline ceiling-lift (decode_long at S_kv=256K). Full focused regression green: 85/85 across ``tests/attention/`` + Phase E + TL discipline tests. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,290 @@
|
||||
"""Phase 1 spec test for P3b: S_kv tile sweep in the GQA kernels
|
||||
(ADR-0063 §A.2 + ADR-0060 §B "long/short context split").
|
||||
|
||||
The local one-shot partial in three kernels — ``_gqa_decode_long.py``,
|
||||
``_gqa_decode_short.py``, ``_gqa_prefill_short.py`` — loads its rank's
|
||||
entire ``(d_head, S_local)`` / ``(S_local, d_head)`` KV slice in one
|
||||
shot, so per-rank scratch grows linearly with ``S_local``. The
|
||||
``_attention_*`` baselines hit a TCM ceiling once ``S_local`` exceeds
|
||||
the 1 MiB pool divided by the per-rank intermediate footprint.
|
||||
|
||||
P3b replaces the one-shot partial with a tile sweep:
|
||||
- Module-level ``TILE_S_KV = 1024`` per kernel file.
|
||||
- Tile 0 establishes the persistent ``(m_local, l_local, O_local)``.
|
||||
- Tiles 1..n_tiles-1 wrap their intermediates (K_T tile, V tile,
|
||||
scores, centered, exp_scores, partials) in ``tl.scratch_scope()``
|
||||
and persist the merged ``(m, ℓ, O)`` to the persistent handles via
|
||||
``tl.copy_to`` (ADR-0063 §D3 / §D3.1).
|
||||
|
||||
When ``S_local ≤ TILE_S_KV`` the loop body never runs and the op_log
|
||||
is structurally identical to today's one-shot path — every existing
|
||||
validation-scale test continues to pass.
|
||||
|
||||
``_gqa_prefill_long.py`` is **out of scope** for P3b. Its inner step
|
||||
is IPCQ partial-recv (Ring KV rotation), not an HBM load; tiling it
|
||||
intersects the ring-step structure and is a separate phase.
|
||||
|
||||
Phase 1 (this commit): tests only — production code lands in Phase 2.
|
||||
The four multi-tile tests fail today because the kernels never call
|
||||
``copy_to`` outside the chain-reduce merges (so isolating to a
|
||||
chain-free config gives ``copy_to == 0`` today). The 128K test fails
|
||||
today with a ``TLContext`` scratch overflow.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from kernbench.benches._gqa_decode_long import gqa_decode_long_kernel # noqa: F401
|
||||
from kernbench.benches._gqa_decode_short import gqa_decode_short_kernel # noqa: F401
|
||||
from kernbench.benches._gqa_prefill_short import gqa_prefill_short_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"
|
||||
|
||||
|
||||
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)
|
||||
|
||||
|
||||
# ── decode_long runner (chain-free configs isolate tile-sweep behavior) ──
|
||||
|
||||
|
||||
def _run_decode_long(*, C: int, P: int, S_kv: int, h_q: int = 1, h_kv: int = 1):
|
||||
"""Run the long-context decode kernel.
|
||||
|
||||
For tile-sweep isolation, callers pass C=1, P=1 — this leaves both
|
||||
chain-reduce levels inactive so any ``copy_to`` in op_log comes
|
||||
solely from the tile-merge body.
|
||||
"""
|
||||
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" if P > 1 else "replicate",
|
||||
num_cubes=C, num_pes=P)
|
||||
q = ctx.zeros((1, h_q * D_HEAD),
|
||||
dtype=DTYPE, dp=dp_full, name=f"q_tl_c{C}_p{P}_s{S_kv}")
|
||||
k = ctx.zeros((S_kv, h_kv * D_HEAD),
|
||||
dtype=DTYPE, dp=dp_kv, name=f"k_tl_c{C}_p{P}_s{S_kv}")
|
||||
v = ctx.zeros((S_kv, h_kv * D_HEAD),
|
||||
dtype=DTYPE, dp=dp_kv, name=f"v_tl_c{C}_p{P}_s{S_kv}")
|
||||
o = ctx.empty((1, h_q * D_HEAD),
|
||||
dtype=DTYPE, dp=dp_full, name=f"o_tl_c{C}_p{P}_s{S_kv}")
|
||||
ctx.launch(
|
||||
f"gqa_decode_long_tile_c{C}_p{P}_s{S_kv}",
|
||||
gqa_decode_long_kernel,
|
||||
q, k, v, o,
|
||||
1, S_kv, h_q, h_kv, D_HEAD, C, P,
|
||||
_auto_dim_remap=False,
|
||||
)
|
||||
|
||||
return run_bench(
|
||||
topology=topo, bench_fn=_bench_fn,
|
||||
device=resolve_device(None), engine_factory=_engine_factory,
|
||||
)
|
||||
|
||||
|
||||
# ── decode_short / prefill_short runners ─────────────────────────────
|
||||
|
||||
|
||||
def _run_decode_short(*, kv_per_cube: int, C: int, P: int, S_kv: int,
|
||||
h_q: int = 8, h_kv: int = 8):
|
||||
"""For tile-sweep isolation: kv_per_cube=P → group_size=1 (no chain)."""
|
||||
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)
|
||||
q = ctx.zeros((1, h_q * D_HEAD),
|
||||
dtype=DTYPE, dp=dp_full, name=f"q_dsh_s{S_kv}")
|
||||
k = ctx.zeros((h_kv * S_kv, D_HEAD),
|
||||
dtype=DTYPE, dp=dp_kv, name=f"k_dsh_s{S_kv}")
|
||||
v = ctx.zeros((h_kv * S_kv, D_HEAD),
|
||||
dtype=DTYPE, dp=dp_kv, name=f"v_dsh_s{S_kv}")
|
||||
o = ctx.empty((1, h_q * D_HEAD),
|
||||
dtype=DTYPE, dp=dp_full, name=f"o_dsh_s{S_kv}")
|
||||
ctx.launch(
|
||||
f"gqa_decode_short_tile_s{S_kv}",
|
||||
gqa_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 _run_prefill_short(*, kv_per_cube: int, C: int, P: int,
|
||||
T_q: int, S_kv: int, h_kv: int = 8):
|
||||
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)
|
||||
q = ctx.zeros((T_q, h_kv * D_HEAD),
|
||||
dtype=DTYPE, dp=dp_full, name=f"q_psh_s{S_kv}")
|
||||
k = ctx.zeros((h_kv * S_kv, D_HEAD),
|
||||
dtype=DTYPE, dp=dp_kv, name=f"k_psh_s{S_kv}")
|
||||
v = ctx.zeros((h_kv * S_kv, D_HEAD),
|
||||
dtype=DTYPE, dp=dp_kv, name=f"v_psh_s{S_kv}")
|
||||
o = ctx.empty((T_q, h_kv * D_HEAD),
|
||||
dtype=DTYPE, dp=dp_full, name=f"o_psh_s{S_kv}")
|
||||
ctx.launch(
|
||||
f"gqa_prefill_short_tile_s{S_kv}",
|
||||
gqa_prefill_short_kernel,
|
||||
q, k, v, o,
|
||||
T_q, S_kv, 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,
|
||||
)
|
||||
|
||||
|
||||
# ── T1: single-tile path is op_log-stable (regression guard) ─────────
|
||||
|
||||
|
||||
def test_decode_long_single_tile_path_unchanged():
|
||||
"""ADR-0063 §A.2: when S_local <= TILE_S_KV the tile-sweep loop
|
||||
body never runs and op_log is structurally identical to today's
|
||||
one-shot path.
|
||||
|
||||
Passes today (one-shot path) AND after Phase 2 (n_tiles=1 skips
|
||||
the merge loop). The witness: with chain reduce inactive (C=1,
|
||||
P=1), there must be zero ``copy_to`` entries — neither today nor
|
||||
after Phase 2 should the kernel emit tile-merge writebacks here.
|
||||
"""
|
||||
result = _run_decode_long(C=1, P=1, S_kv=64)
|
||||
assert result.completion.ok, (
|
||||
f"single-tile decode_long must complete; got {result.completion}"
|
||||
)
|
||||
n_copy = _count(result.engine.op_log, "copy")
|
||||
assert n_copy == 0, (
|
||||
f"S_local <= TILE_S_KV must not emit tile-merge copy_to; got {n_copy}"
|
||||
)
|
||||
|
||||
|
||||
# ── T2: decode_long multi-tile emits tile-merge copy_to ──────────────
|
||||
|
||||
|
||||
def test_decode_long_multi_tile_emits_copy_to_merges():
|
||||
"""ADR-0063 §A.2 + §D3.1: when S_local > TILE_S_KV the kernel must
|
||||
wrap each subsequent tile's intermediates in ``scratch_scope`` and
|
||||
persist the merged ``(m, ℓ, O)`` via ``tl.copy_to``.
|
||||
|
||||
Chain-free config (C=1, P=1) isolates the tile-merge body — any
|
||||
``copy_to`` in op_log comes from the tile sweep, not chain reduce.
|
||||
|
||||
S_kv=2048, C=1, P=1 → S_local=2048 → 2 tiles → 1 merge × 3 handles
|
||||
(m, ℓ, O) ⇒ 3 ``copy`` entries.
|
||||
|
||||
Currently 0 because the kernel performs a one-shot partial.
|
||||
"""
|
||||
result = _run_decode_long(C=1, P=1, S_kv=2048)
|
||||
assert result.completion.ok, (
|
||||
f"multi-tile decode_long must complete; got {result.completion}"
|
||||
)
|
||||
n_copy = _count(result.engine.op_log, "copy")
|
||||
assert n_copy >= 3, (
|
||||
f"decode_long multi-tile must emit >= 3 copy_to entries "
|
||||
f"(1 merge × 3 handles); got {n_copy}"
|
||||
)
|
||||
|
||||
|
||||
# ── T3: decode_short multi-tile emits tile-merge copy_to ─────────────
|
||||
|
||||
|
||||
def test_decode_short_multi_tile_emits_copy_to_merges():
|
||||
"""Same property as T2 for the short decode kernel.
|
||||
|
||||
kv_per_cube=8, P=8, C=1 → group_size=1 (no chain reduce); S_kv=2048
|
||||
→ S_local=2048 → 2 tiles per PE → 3 ``copy`` entries per PE × 8 PEs
|
||||
= 24 total.
|
||||
|
||||
Currently 0 because the kernel performs a one-shot partial.
|
||||
"""
|
||||
result = _run_decode_short(kv_per_cube=8, C=1, P=8, S_kv=2048)
|
||||
assert result.completion.ok, (
|
||||
f"multi-tile decode_short must complete; got {result.completion}"
|
||||
)
|
||||
n_copy = _count(result.engine.op_log, "copy")
|
||||
assert n_copy >= 3 * 8, (
|
||||
f"decode_short multi-tile must emit >= 24 copy_to entries "
|
||||
f"(1 merge × 3 handles × 8 PEs); got {n_copy}"
|
||||
)
|
||||
|
||||
|
||||
# ── T4: prefill_short multi-tile emits tile-merge copy_to ────────────
|
||||
|
||||
|
||||
def test_prefill_short_multi_tile_emits_copy_to_merges():
|
||||
"""Same property as T3 for the short prefill kernel.
|
||||
|
||||
kv_per_cube=8, P=8, C=1, T_q=4, S_kv=2048 → group_size=1 (no chain),
|
||||
S_local=2048 → 2 tiles per PE → 3 ``copy`` entries per PE × 8 PEs
|
||||
= 24 total.
|
||||
|
||||
Currently 0 because the kernel performs a one-shot partial.
|
||||
"""
|
||||
result = _run_prefill_short(
|
||||
kv_per_cube=8, C=1, P=8, T_q=4, S_kv=2048,
|
||||
)
|
||||
assert result.completion.ok, (
|
||||
f"multi-tile prefill_short must complete; got {result.completion}"
|
||||
)
|
||||
n_copy = _count(result.engine.op_log, "copy")
|
||||
assert n_copy >= 3 * 8, (
|
||||
f"prefill_short multi-tile must emit >= 24 copy_to entries "
|
||||
f"(1 merge × 3 handles × 8 PEs); got {n_copy}"
|
||||
)
|
||||
|
||||
|
||||
# ── T5: decode_long at S_kv=256K completes (TCM ceiling lifted) ──────
|
||||
|
||||
|
||||
def test_decode_long_context_256k_completes():
|
||||
"""ADR-0063 §A.2 (test req 3): headline ceiling-lift. C=1 P=8
|
||||
decode at S_kv=256K → S_local=32K. Today the per-rank score stack
|
||||
(3 ×M·S_local·2 bytes, M=G·T_q=8) is ~1.5 MB → exceeds the 1 MiB
|
||||
scratch pool → TLContext scratch overflow. After Phase 2 the tile
|
||||
sweep bounds per-tile score stack at 3·M·TILE_S_KV·2 ≈ 48 KB and
|
||||
the run completes regardless of S_kv.
|
||||
"""
|
||||
result = _run_decode_long(C=1, P=8, S_kv=262144, h_q=8, h_kv=1)
|
||||
assert result.completion.ok, (
|
||||
f"decode_long at S_kv=256K must complete after tile sweep; "
|
||||
f"got {result.completion}"
|
||||
)
|
||||
Reference in New Issue
Block a user