ywkang
cfc2d74ec4
The unified ccl_allreduce bench previously carried two execution models
in one worker with ``if world_size == n_sips:`` branching:
- TP mode (rank = SIP, ADR-0024/0027): proper ProcessGroup semantics.
- Legacy rank = PE mode: single-driver worker allocating one big tensor
distributed across all PEs via _derive_dp, with kernel-level SPMD via
program_id.
The second model is unnecessary — intra-SIP PE-level collectives are
expressed inside the kernel (tl.send/tl.recv with program_id, IPCQ) and
do not need a host-side ProcessGroup. Removing it lets the bench be a
clean reference implementation of the TP launcher.
benches/ccl_allreduce.py:
- Config resolved once in run() via _resolve_cfg -> _BenchCfg dataclass.
- rank != n_sips now raises RuntimeError explicitly.
- _worker / _allocate_rank_tile / _init_with_rank_value / _report each
have one concern; duplicated init + verification paths collapsed.
- _derive_dp and the second verify+print block deleted.
- 166 lines -> 91 lines.
ccl.yaml:
- mesh_allreduce_4 (world_size: 4) and tree_allreduce_7 (world_size: 7)
algorithm entries removed (rank = PE only).
- Algorithm kernel files (kernbench.ccl.algorithms.mesh_allreduce,
tree_allreduce) kept as-is for direct-dispatch future use.
tests/test_ccl_allreduce_matrix.py:
- Matrix shrinks from 7 cases to 3: ring × {tcm, hbm, sram} at ws =
topology SIP count (= 2). mesh_2x2, tree_binary_7, ring_multi_cube,
and the three ring_*_8 cases removed.
tests/test_ccl_performance.py:
- _run_8rank renamed to _run_ring; world_size: 8 override dropped; now
exercises rank = SIP ring all-reduce.
tests/test_mp_spawn.py, tests/test_ccl_ddp_launcher.py:
- Monkeypatch target updated from bench.worker to bench._worker
(signature now takes BenchCfg instead of (rank, world_size)).
555 passed, 1 intentional skip. Tests that directly call
install_ipcq(world_size_override=N) for kernel-level sanity
(test_ccl_hello_world_guide, test_recv_copy_to_dst, test_tl_recv_async,
test_ccl_deadlock_detection) are unchanged — they never went through
the bench and still exercise the kernel-only path.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
109 lines
2.9 KiB
Python
109 lines
2.9 KiB
Python
"""End-to-end matrix tests for the unified ``ccl_allreduce`` bench.
|
||
|
||
Only covers the rank = SIP TP launcher path (ADR-0024 + ADR-0027). Each
|
||
case writes a tmp ``ccl.yaml`` that selects a specific (algorithm,
|
||
buffer_kind) pair; ``world_size`` is always derived from topology SIP
|
||
count (2 in the shipped topology).
|
||
|
||
The legacy rank = PE single-driver path was removed; intra-SIP PE-level
|
||
collectives are expressed inside the kernel via ``tl.program_id`` and do
|
||
not require a host-side ``ProcessGroup``.
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import os
|
||
import textwrap
|
||
|
||
import pytest
|
||
|
||
import kernbench.cli.main as cli_main
|
||
|
||
|
||
CCL_YAML_TEMPLATE = textwrap.dedent("""\
|
||
defaults:
|
||
algorithm: {algorithm}
|
||
buffer_kind: {buffer_kind}
|
||
backpressure: sleep
|
||
n_slots: 4
|
||
slot_size: 4096
|
||
vc_chunk_size: 256
|
||
ipcq_credit_size_bytes: 16
|
||
|
||
algorithms:
|
||
{algorithm}:
|
||
module: {module}
|
||
topology: {topology}
|
||
buffer_kind: {buffer_kind}
|
||
""")
|
||
|
||
|
||
def _write_ccl_yaml(
|
||
tmp_path,
|
||
*,
|
||
algorithm: str,
|
||
module: str,
|
||
topology: str,
|
||
buffer_kind: str,
|
||
) -> str:
|
||
body = CCL_YAML_TEMPLATE.format(
|
||
algorithm=algorithm,
|
||
module=module,
|
||
topology=topology,
|
||
buffer_kind=buffer_kind,
|
||
)
|
||
(tmp_path / "ccl.yaml").write_text(body)
|
||
return str(tmp_path)
|
||
|
||
|
||
CASES = [
|
||
# Ring all-reduce across SIPs (ws == topology SIP count = 2),
|
||
# one case per IPCQ buffer location.
|
||
pytest.param(
|
||
"ring_allreduce_tcm", "kernbench.ccl.algorithms.ring_allreduce",
|
||
"ring_1d", "tcm",
|
||
id="ring_tcm",
|
||
),
|
||
pytest.param(
|
||
"ring_allreduce_hbm", "kernbench.ccl.algorithms.ring_allreduce",
|
||
"ring_1d", "hbm",
|
||
id="ring_hbm",
|
||
),
|
||
pytest.param(
|
||
"ring_allreduce_sram", "kernbench.ccl.algorithms.ring_allreduce",
|
||
"ring_1d", "sram",
|
||
id="ring_sram",
|
||
),
|
||
]
|
||
|
||
|
||
@pytest.mark.parametrize("algorithm,module,topology,buffer_kind", CASES)
|
||
def test_ccl_allreduce_matrix(
|
||
tmp_path, capsys, monkeypatch,
|
||
algorithm, module, topology, buffer_kind,
|
||
):
|
||
"""Each (algorithm × buffer_kind) combo passes through the unified
|
||
rank = SIP bench and yields ``ws OK`` where ``ws == topology SIP count``."""
|
||
project_root = os.path.abspath(
|
||
os.path.join(os.path.dirname(__file__), "..")
|
||
)
|
||
yaml_dir = _write_ccl_yaml(
|
||
tmp_path,
|
||
algorithm=algorithm,
|
||
module=module,
|
||
topology=topology,
|
||
buffer_kind=buffer_kind,
|
||
)
|
||
monkeypatch.chdir(yaml_dir)
|
||
rc = cli_main.main([
|
||
"run",
|
||
"--topology", os.path.join(project_root, "topology.yaml"),
|
||
"--bench", "ccl_allreduce",
|
||
"--verify-data",
|
||
])
|
||
assert rc == 0
|
||
out = capsys.readouterr().out
|
||
assert "FAIL" not in out, f"unexpected FAIL in output:\n{out}"
|
||
assert f"{algorithm}" in out and "OK" in out, (
|
||
f"expected pass line for '{algorithm}' in output:\n{out}"
|
||
)
|