Intercube allreduce: pe0 cube-mesh reduce + multi-SIP ring/torus/mesh

New intercube allreduce kernel replacing the old flat ring algorithms.
Reduces across the 4x4 cube mesh within each SIP (pe0-only, same-lane),
then inter-SIP exchange on root cube, then broadcast back. Supports
ring_1d, torus_2d, and mesh_2d_no_wrap SIP topologies driven by
topology.yaml. Integrated with dist.init_process_group / dist.all_reduce.

New files:
- src/kernbench/ccl/algorithms/intercube_allreduce.py (kernel)
- src/kernbench/ccl/sfr_config.py (configure_sfr_intercube_multisip)
- tests/test_allreduce_multidevice.py (config-driven, 3 topologies)
- tests/test_distributed_intercube_allreduce.py (full distributed path)
- tests/test_intercube_sfr_config.py (SFR wiring verification)

Modified:
- distributed.py: AhbmCCLBackend uses configure_sfr_intercube_multisip
- topologies.py: added torus_2d, mesh_2d_no_wrap
- install.py: global_E/W/N/S in _OPPOSITE_DIR
- topology.yaml: added system.sips.topology
- ccl.yaml: single intercube_allreduce algorithm
- benches/ccl_allreduce.py: row_wise cube-mesh tensor layout

Removed old flat-ring algorithms and their tests.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-04-16 17:33:42 -07:00
parent cfc2d74ec4
commit 1d8b9401e5
30 changed files with 876 additions and 2892 deletions
+104
View File
@@ -0,0 +1,104 @@
"""SFR configuration for intercube + inter-SIP IPCQ wiring.
Provides ``configure_sfr_intercube_multisip`` which programs PE_IPCQ
neighbor tables for:
1. Intercube within each SIP — pe0 of every cube connects to pe0 of
its N/S/E/W mesh neighbors (no wrap-around).
2. Inter-SIP on ALL cubes — pe0 of cube_c on sip_A connects to pe0 of
cube_c on each peer SIP, using ``global_E``/``global_W`` (ring) or
``global_N``/``global_S``/``global_E``/``global_W`` (mesh/torus)
direction labels. Wiring all cubes allows the kernel to
dynamically elect the root cube at runtime.
SIP-level topology is read from ``topology.yaml`` →
``system.sips.topology`` (e.g. ``ring_1d``, ``mesh_2d``).
Intercube mesh dimensions come from ``sip.cube_mesh.w/h``.
Internally delegates to ``install_ipcq`` with a computed ``rank_to_pe``
(pe0-only) and a closure-captured ``neighbors()`` function.
"""
from __future__ import annotations
import types
from typing import Any
from kernbench.ccl.install import install_ipcq
from kernbench.ccl.topologies import _BUILTIN as _TOPO_BUILTINS
def configure_sfr_intercube_multisip(
engine: Any,
spec: dict,
cfg: dict,
) -> dict[str, Any]:
"""Wire IPCQ for intercube (pe0, mesh) + inter-SIP (pe0, all cubes).
Args:
engine: GraphEngine with ``_components``.
spec: topology spec dict (from topology.yaml).
cfg: merged algorithm config (from ``resolve_algorithm_config``).
Returns:
The install plan dict from ``install_ipcq``.
"""
cm = spec["sip"]["cube_mesh"]
mesh_w = int(cm["w"])
mesh_h = int(cm["h"])
n_cubes = mesh_w * mesh_h
n_sips = int(spec.get("system", {}).get("sips", {}).get("count", 1))
sip_topology = str(
spec.get("system", {}).get("sips", {}).get("topology", "ring_1d")
)
if sip_topology not in _TOPO_BUILTINS:
raise ValueError(
f"Unknown sip topology '{sip_topology}'. "
f"Available: {list(_TOPO_BUILTINS)}"
)
sip_topo_fn = _TOPO_BUILTINS[sip_topology]
world_size = n_sips * n_cubes
pe_idx_to_pe: list[tuple[int, int, int]] = [
(sip, cube, 0)
for sip in range(n_sips)
for cube in range(n_cubes)
]
def _neighbors(pe_idx: int, ws: int, _base: dict) -> dict[str, int]:
sip = pe_idx // n_cubes
cube = pe_idx % n_cubes
row = cube // mesh_w
col = cube % mesh_w
nbrs: dict[str, int] = {}
# Intercube within SIP (mesh, no wrap-around)
if col < mesh_w - 1:
nbrs["E"] = sip * n_cubes + (row * mesh_w + col + 1)
if col > 0:
nbrs["W"] = sip * n_cubes + (row * mesh_w + col - 1)
if row < mesh_h - 1:
nbrs["S"] = sip * n_cubes + ((row + 1) * mesh_w + col)
if row > 0:
nbrs["N"] = sip * n_cubes + ((row - 1) * mesh_w + col)
# Inter-SIP on ALL cubes
if n_sips > 1:
sip_nbrs = sip_topo_fn(sip, n_sips)
for d, peer_sip in sip_nbrs.items():
nbrs[f"global_{d}"] = peer_sip * n_cubes + cube
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,
)