kernbench2/docs/adr/ADR-0050-par-ccl-algorithm-module-contract.md

ADR-0050: CCL Algorithm Module Contract — ccl/algorithms/*.py

Status

Accepted (2026-05-22).

Pins down the interface, kernel signature, and addition workflow that a module under src/kernbench/ccl/algorithms/ must satisfy in order to be used as a collective algorithm by the AHBM CCL backend (ADR-0047). ADR-0047 D3 states only that "the algorithm module must expose kernel, kernel_args, optionally TOPO_NAME_TO_KIND"; the contract an algorithm-module author needs to follow has had no ADR-level coverage. This ADR pairs with ADR-0045's bench-module contract.

First action

An algorithm module is imported at two moments:

1. AHBM backend entry: when user code calls dist.init_process_group(backend="ahbm"), AhbmCCLBackend.__init__ runs self._algo_module = importlib.import_module(self._merged["module"]). At module level, the following occur first:

   • Topology-kind integer constants like SIP_TOPO_RING/TORUS/MESH are bound in the module namespace.
   • The TOPO_NAME_TO_KIND dict is bound; the backend reads it via getattr(self._algo_module, "TOPO_NAME_TO_KIND", None).
   • kernel_args function is defined for the caller.
   • The actual algorithm function (e.g., allreduce_intercube_multidevice) is defined.
   • At the bottom of the module, kernel = allreduce_intercube_multidevice publishes the alias.

2. ccl.yaml install stage: kernbench.ccl.install.install_ipcq imports the same algorithm module while pushing the IPCQ neighbor table.

In short, the algorithm module's first act is "publish topology-kind constants, the TOPO_NAME_TO_KIND dict, the kernel_args function, and the kernel alias into the module namespace" — all as import-time side effects, no separate initialization call.

Context

AhbmCCLBackend (ADR-0047), at process-group creation, dynamically imports a module path obtained from ccl.yaml's defaults.algorithm (or a user-specified algorithm). The backend expects four things from the module:

• kernel: the collective's entry function.
• kernel_args(world_size, n_elem, cube_w=, cube_h=) -> tuple: a tuple packing the kernel's positional arguments.
• TOPO_NAME_TO_KIND (optional): a dict mapping topology.yaml's sips.topology string (e.g., "ring_1d", "torus_2d", "mesh_2d_no_wrap") to the integer kind constants.
• (Indirectly) IPCQ neighbor-table install: configure_sfr_intercube_multisip reads the module's TOPO_NAME_TO_KIND plus cube dimensions to decide the SFR.

The current corpus has one algorithm module: lrab_hierarchical_allreduce.py (248 lines). The name expands to "left-right alternating broadcast hierarchical allreduce". When future modules like ring_allreduce, tree_allreduce, or broadcast are added, they must follow this contract for the backend's dispatch path to keep working.

Without an ADR-level contract:

• A new algorithm author has to infer the signature from ADR-0047 D3's one-liner.
• The kernel-function argument order (especially t_ptr, n_elem, cube_w, cube_h, n_sips, sip_rank, sip_topo_kind, sip_topo_w, sip_topo_h, tl) is unclear without grep.
• It is conventional, but not documented, what kernel_args takes as inputs and what tuple it must return.

Decision

D1. The algorithm module exposes four public symbols

# src/kernbench/ccl/algorithms/<name>.py

from __future__ import annotations

# (required) topology-kind constants — referenced internally
SIP_TOPO_RING = 0
SIP_TOPO_TORUS = 1
SIP_TOPO_MESH = 2

# (optional) topology name → kind mapping. Used by the backend to
# translate ccl.yaml/topology's string SIP topology into an integer.
TOPO_NAME_TO_KIND = {
    "ring_1d": SIP_TOPO_RING,
    "torus_2d": SIP_TOPO_TORUS,
    "mesh_2d_no_wrap": SIP_TOPO_MESH,
}

# (required) kernel argument builder
def kernel_args(world_size: int, n_elem: int, *, cube_w: int = 4, cube_h: int = 4) -> tuple:
    return (n_elem, cube_w, cube_h, world_size)

# (required) kernel function (TLContext is injected via the `tl=...`
# keyword argument).
def my_allreduce_kernel(t_ptr, n_elem, cube_w, cube_h, n_sips,
                         sip_rank, sip_topo_kind, sip_topo_w, sip_topo_h, *, tl):
    ...

# (required) kernel alias — the backend accesses `module.kernel`
kernel = my_allreduce_kernel

• The kernel alias is the entry point the backend invokes. Whatever the function name is (e.g., allreduce_intercube_multidevice), it must be exposed via module.kernel = fn.
• Without kernel_args, the backend has no way to build the algorithm's argument list. See D2 for the signature.
• If TOPO_NAME_TO_KIND is absent, the backend falls back to sip_topo_kind = 0. An algorithm supporting only a single topology may omit it.

D2. kernel_args signature — (world_size, n_elem, *, cube_w, cube_h)

def kernel_args(world_size: int, n_elem: int, *,
                cube_w: int = 4, cube_h: int = 4) -> tuple:
    return (n_elem, cube_w, cube_h, world_size)

• Positional arguments: world_size (= number of ranks), n_elem (= element count of a single shard, f16-based).
• Keyword arguments: cube_w, cube_h (= cube-mesh dimensions). Default 4×4 — aligned with topology.yaml's sip.cube_mesh default.
• Return: a tuple in the order the kernel's positional arguments expect.

When the backend calls all_reduce:

kernel_args_tuple = self._algo_module.kernel_args(
    self._world_size, n_elem, cube_w=eff_cube_w, cube_h=eff_cube_h,
)
extra_args = (sip_rank, sip_topo_kind, sip_topo_w, sip_topo_h)
pending = self.ctx.launch(
    self._merged["algorithm"], kernel_fn, tensor,
    *kernel_args_tuple, *extra_args, _defer_wait=True,
)

So the kernel's full positional argument list becomes: (tensor_ptr, *kernel_args_tuple, sip_rank, sip_topo_kind, sip_topo_w, sip_topo_h), with tl=... injected as a keyword. The tuple length and order returned by kernel_args must match the kernel signature 1:1.

D3. Kernel signature — standardized 9 + tl arguments

Recommended signature:

def my_kernel(
    t_ptr: int,             # VA base of the row-wise-sharded tensor on this SIP
    n_elem: int,            # element count per cube tile (or per shard)
    cube_w: int,            # cube mesh width  (from kernel_args)
    cube_h: int,            # cube mesh height (from kernel_args)
    n_sips: int,            # equal to world_size (rank = SIP, ADR-0024)
    sip_rank: int,          # this SIP's rank
    sip_topo_kind: int,     # result of TOPO_NAME_TO_KIND lookup
    sip_topo_w: int,        # SIP mesh width (0 for ring_1d)
    sip_topo_h: int,        # SIP mesh height (0 for ring_1d)
    *, tl,                  # TLContext (auto-injected)
) -> None:

Even if kernel_args chose a different positional argument order, the kernel's last four positional arguments are always (sip_rank, sip_topo_kind, sip_topo_w, sip_topo_h) — the backend appends them as extra_args (ADR-0047 D5). A custom algorithm must accept these four, but a single-SIP algorithm may simply ignore them.

tl is injected via keyword — RuntimeContext.launch adds tl=tl_ctx just before invoking the kernel. The signature therefore exposes tl as keyword-only (*, tl) or as the trailing keyword parameter.

D4. Kernel body — freedom and constraints

Available inside the kernel: every tl.* primitive from ADR-0046 D3.

Common patterns:

• cube_id = tl.program_id(axis=1) — this PE's cube index.
• pe_addr = t_ptr + cube_id * nbytes — per-cube VA of the tile.
• acc = tl.load(pe_addr, shape=(n_elem,), dtype="f16") — load local data.
• tl.send(dir=...) / tl.recv(dir=..., shape=, dtype=) — IPCQ collective.
• acc = acc + recv — TensorHandle arithmetic operators (ADR-0046 D4).
• tl.store(pe_addr, acc) — store the result.

The kernel body is plain Python — branching and loops are fine. But:

• No SimPy yield or async (ADR-0046 D1).
• No direct access to TensorHandle .data — the Phase 1 timing model doesn't see data dependencies (ADR-0020's 2-pass separation).
• Kernel execution must be deterministic — the same input must produce the same op sequence. No random or external IO.

D5. SIP topology semantics — meaning of sip_topo_kind

The backend looks up topology.yaml's system.sips.topology string in the algorithm module's TOPO_NAME_TO_KIND and passes the integer as sip_topo_kind. The algorithm then branches:

if sip_topo_kind == SIP_TOPO_RING:
    acc = _inter_sip_ring(...)
elif sip_topo_kind == SIP_TOPO_TORUS:
    acc = _inter_sip_torus_2d(...)
elif sip_topo_kind == SIP_TOPO_MESH:
    acc = _inter_sip_mesh_2d(...)

Each topology branch communicates with peers via IPCQ direction names ("global_E", "W", "S", "N" …). Direction semantics are defined in ADR-0023/0025; configure_sfr_intercube_multisip installs the IPCQ neighbor table accordingly.

If a topology kind not supported by the algorithm appears, prefer an explicit raise ValueError(f"unsupported topology kind {sip_topo_kind}") over a silent no-op — fail fast on misconfiguration.

D6. The ccl.yaml algorithm entry

The algorithm module is paired with a ccl.yaml entry (ADR-0023 D10 + ADR-0047 D3):

defaults:
  algorithm: lrab_hierarchical_allreduce
  n_elem: 8

algorithms:
  lrab_hierarchical_allreduce:
    module: kernbench.ccl.algorithms.lrab_hierarchical_allreduce
    # optional: world_size override
    # optional: per-algorithm parameters consumed by configure_sfr_intercube_multisip

• module: the full Python module path; importlib.import_module consumes this string as-is.
• world_size (optional): when set, overrides the topology fallback (ADR-0047 D2).
• Algorithm-specific parameters are consumed by configure_sfr_intercube_multisip.

Workflow to add a new algorithm:

1. Write src/kernbench/ccl/algorithms/<name>.py following D1.
2. Add the entry under algorithms in ccl.yaml.
3. (If needed) extend kernbench.ccl.sfr_config with the SFR-install branch.
4. Add tests (e.g., tests/sccl/test_<name>.py, extending the ADR-0043 eval harness).

D7. Legacy "rank = flat PE index" mode

The world_size override in ccl.yaml, surfaced by ADR-0047 D2, is used by legacy "rank = flat PE index" tests. The algorithm module can assume n_sips=world_size ranks even in this mode — the backend maintains the rank↔(SIP, cube, PE) mapping, so no modal branching is needed inside the algorithm body.

In single-cube workloads (where cube_w=cube_h=1), the algorithm must skip mesh-based phases — see the single_cube = (cube_w == 1 and cube_h == 1) pattern in lrab_hierarchical_allreduce.py.

Alternatives Considered

A1. Organize the algorithm module as a class (class Allreduce: kernel(...))

Rejected. The Python module namespace already identifies an algorithm (see ADR-0047 D3's importlib.import_module). A class wrapper adds indirection without simplifying dispatch. Module-level free functions plus a kernel alias are clean and obvious.

A2. Type kernel_args with an explicit dataclass

Rejected (currently). Each algorithm normally has a different argument count; forcing one dataclass would hurt cross-algorithm interchange. The tuple return is simple and unpacks cleanly with the backend's *kernel_args_tuple. If an algorithm wants stronger internal typing, it may define its own NamedTuple.

A3. Move SFR installation inside the algorithm module

Rejected. SFR installation (configure_sfr_intercube_multisip) is a cross-module decision combining topology + algorithm; kernbench.ccl.sfr_config is a more natural home than the algorithm module itself. D6's "extend sfr_config if needed" workflow keeps responsibility boundaries clear.

A4. Auto-register algorithm names via a decorator (analogous to ADR-0045's @bench)

Rejected. Unlike benches, algorithms are already tied to ccl.yaml entries; an additional registry would be redundant. The string mapping in module is sufficient.

Consequences

• ADR-0047 D3's one-line contract expands to a D1–D7 author-facing guide; new algorithm signatures no longer need to be grep-derived.
• D3's standardized 9 + tl signature couples naturally with the backend's extra_args append (ADR-0047 D5). It is explicit that even single-SIP-only algorithms must accept the four sip_* trailing arguments.
• D5's fail-loud recommendation means a ccl.yaml topology that the algorithm doesn't support will surface as an explicit ValueError rather than a silent wrong result.
• D6's step-by-step addition workflow makes clear how far a new algorithm has to reach into sfr_config / tests / ccl.yaml.