gqa: ADR-0060/0062/0063/0064 unified GQA kernels + CPU cost model
Land the new GQA fused-attention kernels (ADR-0060) for prefill/decode
across long and short context, the TL discipline primitives they depend
on (ADR-0062 lazy load, ADR-0063 scratch_scope + copy_to), and the
per-op-type CPU issue cost model (ADR-0064). Remove the pre-ADR-0060
mesh-attention baseline now that the unified kernels supersede it.
ADR-0060 (long context)
- _gqa_decode.py: M-fold + 2-level chain reduce-to-root (Level-2
intra-CUBE row-then-col + Level-1 inter-CUBE) — root-only output.
- _gqa_prefill.py: head-parallel + Ring KV rotation around C CUBEs,
online-softmax merge per ring step, per-CUBE distributed output.
- Each merge stage wraps in scratch_scope() and persists running
(m, l, O) via copy_to() to lift the 1 MiB scratch ceiling.
ADR-0060 §B.split.2 (short context, kv_per_cube in {1,2,4,8})
- _gqa_decode_short.py / _gqa_prefill_short.py: no cube-SP; each CUBE
owns whole KV heads; PE-parallel heads with intra-group chain
reduce. Prefill has no Ring KV (each head fully resident).
ADR-0062 (lazy tl.load): future-bearing TensorHandle, auto-wait at
first consuming op (dot/MATH/store/send/copy_to/composite).
ADR-0063 (tl.scratch_scope + tl.copy_to): scoped per-tile arena with
copy_to writeback primitive for persistent running state.
ADR-0064 (CPU issue cost model)
- common/cpu_issue_cost.py: per-op-type table (composite=40 ns,
primitives=5 ns); ratios are load-bearing per D1.
- TLContext: issue_cost_table param; _emit_dispatch_overhead(kind)
consults table with dispatch_cycles fallback (ADR-0046 §D6
back-compat).
- Live PE_CPU paths (greenlet + legacy) construct TLContext with
DEFAULT_CPU_ISSUE_COST so saturation lever (ADR-0060 §1) is
measurable end-to-end.
P7 headline bench: milestone-gqa-headline writes per-panel
op_log_summary to 1H_milestone_output/gqa_headline/sweep.json. No
figure renderers yet (deferred).
Removals (pre-ADR-0060 baseline now superseded):
- benches: _attention_mesh_kv.py, _attention_mesh_mlo.py,
_attention_mesh_mlo_2d.py, milestone_gqa_llama70b.py
- tests: test_attention_*, test_mesh_*, test_milestone_gqa_llama70b
- topology: llama70b_4sip.yaml (only consumer was the deleted diag)
- artifacts: 1H_milestone_output/gqa/ (sweep.json + 5 PNGs)
- tests/gqa/ plot helper + test (broken on Windows Tcl/Tkinter)
- ADR-0060/0061 references to deleted file paths cleaned up
(EN + KO kept in sync).
Tests: 124/124 focused regression green (attention + Phase E + TL
discipline + triton_emu + pe_components). Full regression: 764 pass,
2 pre-existing test_bench_registry failures (stale EXPECTED_NAMES
across multiple benches, not introduced here).
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
@@ -237,3 +237,119 @@ def configure_sfr_intracube_pe_ring(
|
||||
algo_module=mock_module,
|
||||
rank_to_pe=pe_idx_to_pe,
|
||||
)
|
||||
|
||||
|
||||
# ── Inter-cube 1D ring (ADR-0060 §5.5 prefill Ring KV) ─────────────────
|
||||
|
||||
|
||||
def configure_sfr_intercube_ring(
|
||||
engine: Any,
|
||||
spec: dict,
|
||||
cfg: dict,
|
||||
*,
|
||||
ring_size: int | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Install intra-cube PE grid + a 1D CUBE-level ring with wrap.
|
||||
|
||||
Direction namespaces (disjoint, same as
|
||||
``configure_sfr_intercube_multisip``):
|
||||
|
||||
- ``intra_N/S/E/W`` : 2×4 PE grid within each cube (no wrap)
|
||||
- ``E/W`` : 1D ring of cubes 0..ring_size-1 WITH WRAP
|
||||
(symmetric to ``configure_sfr_intracube_pe_ring``
|
||||
at PE level — wrap applied at CUBE level here)
|
||||
- ``global_*`` : SIP topology (same as multisip)
|
||||
|
||||
N/S at CUBE level are intentionally NOT installed — use
|
||||
``configure_sfr_intercube_multisip`` for the full 4×4 cube mesh.
|
||||
|
||||
Args:
|
||||
ring_size: number of CUBEs in the ring (wrap applies to cubes
|
||||
0..ring_size-1). Defaults to the full cube_mesh count.
|
||||
Must be ≤ mesh_w (single row); multi-row rings span
|
||||
non-neighbour boundaries.
|
||||
"""
|
||||
cm = spec["sip"]["cube_mesh"]
|
||||
mesh_w = int(cm["w"])
|
||||
mesh_h = int(cm["h"])
|
||||
n_cubes = mesh_w * mesh_h
|
||||
sips_cfg = spec.get("system", {}).get("sips", {})
|
||||
n_sips = int(sips_cfg.get("count", 1))
|
||||
sip_topology = str(sips_cfg.get("topology", "ring_1d"))
|
||||
sip_w = sips_cfg.get("w")
|
||||
sip_h = sips_cfg.get("h")
|
||||
sip_w = int(sip_w) if sip_w is not None else None
|
||||
sip_h = int(sip_h) if sip_h is not None else None
|
||||
|
||||
if ring_size is None:
|
||||
ring_size = n_cubes
|
||||
if ring_size > mesh_w:
|
||||
raise ValueError(
|
||||
f"intercube_ring ring_size={ring_size} > mesh_w={mesh_w}; "
|
||||
"multi-row rings cross non-neighbour boundaries"
|
||||
)
|
||||
|
||||
if sip_topology not in _TOPO_BUILTINS:
|
||||
raise ValueError(
|
||||
f"Unknown sip topology '{sip_topology}'. "
|
||||
f"Available: {list(_TOPO_BUILTINS)}"
|
||||
)
|
||||
_sip_topo_fn_raw = _TOPO_BUILTINS[sip_topology]
|
||||
|
||||
def sip_topo_fn(rank: int, ws: int) -> dict:
|
||||
if sip_w is not None and sip_h is not None:
|
||||
try:
|
||||
return _sip_topo_fn_raw(rank, ws, w=sip_w, h=sip_h)
|
||||
except TypeError:
|
||||
pass
|
||||
return _sip_topo_fn_raw(rank, ws)
|
||||
|
||||
pes_per_cube = _PES_PER_CUBE
|
||||
world_size = n_sips * n_cubes * pes_per_cube
|
||||
pe_idx_to_pe: list[tuple[int, int, int]] = [
|
||||
(sip, cube, pe)
|
||||
for sip in range(n_sips)
|
||||
for cube in range(n_cubes)
|
||||
for pe in range(pes_per_cube)
|
||||
]
|
||||
|
||||
def _pe_idx(sip: int, cube: int, pe: int) -> int:
|
||||
return (sip * n_cubes + cube) * pes_per_cube + pe
|
||||
|
||||
def _neighbors(pe_idx: int, ws: int, _base: dict) -> dict[str, int]:
|
||||
tmp = pe_idx
|
||||
pe = tmp % pes_per_cube
|
||||
tmp //= pes_per_cube
|
||||
cube = tmp % n_cubes
|
||||
sip = tmp // n_cubes
|
||||
|
||||
nbrs: dict[str, int] = {}
|
||||
|
||||
# ── Intra-cube (intra_N/S/E/W) ──
|
||||
for d, peer_pe in _intra_cube_neighbors(pe).items():
|
||||
nbrs[d] = _pe_idx(sip, cube, peer_pe)
|
||||
|
||||
# ── Cube ring (E/W with wrap for cubes 0..ring_size-1) ──
|
||||
if cube < ring_size:
|
||||
nbrs["E"] = _pe_idx(sip, (cube + 1) % ring_size, pe)
|
||||
nbrs["W"] = _pe_idx(sip, (cube - 1) % ring_size, pe)
|
||||
|
||||
# ── Inter-SIP same-(cube, pe) (global_*) ──
|
||||
if n_sips > 1:
|
||||
sip_nbrs = sip_topo_fn(sip, n_sips)
|
||||
for d, peer_sip in sip_nbrs.items():
|
||||
nbrs[f"global_{d}"] = _pe_idx(peer_sip, cube, pe)
|
||||
|
||||
return nbrs
|
||||
|
||||
mock_module = types.SimpleNamespace(neighbors=_neighbors)
|
||||
|
||||
cfg_copy = dict(cfg)
|
||||
cfg_copy["world_size"] = world_size
|
||||
cfg_copy["topology"] = "none"
|
||||
|
||||
return install_ipcq(
|
||||
engine, spec, cfg_copy,
|
||||
algo_module=mock_module,
|
||||
rank_to_pe=pe_idx_to_pe,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user