kernbench2/docs/adr/ADR-0029-hierarchical-allreduce.md

ADR-0029: Hierarchical All-Reduce — 3-level intra/inter-SIP 알고리즘

Status

Proposed

Context

목표

"Rank = SIP" 모델 (ADR-0024) 위에서 각 SIP 내부의 모든 PE를 참여시키는 3-level 계층 all-reduce 알고리즘을 정의한다. 각 레벨이 서로 다른 물리 연결(intra-cube ring, inter-cube NoC, inter-SIP UCIe)을 활용해 대역폭을 극대화한다.

왜 hierarchical인가

단순 ring/mesh/tree all-reduce는 SIP당 1 PE만 참여 (ADR-0024의 leader_only mapper). 이는 inter-SIP 단계는 잘 모델링하지만:

• Intra-SIP PE가 노는 시간이 발생. Leader PE가 inter-SIP 통신 중이면 나머지 7 PE / 16 cube는 유휴.
• Intra-cube/inter-cube 연결 대역폭 미활용. Cube NoC는 매우 빠르지만 단일 leader 사용 시 이 자원이 노출되지 않음.
• 실제 NCCL 등은 hierarchical: NVLink(intra-node) + InfiniBand(inter-node) 의 bandwidth 차이를 활용. KernBench 토폴로지도 동일 구조 (intra-cube / inter-cube / inter-SIP의 bandwidth·latency 차이).

현재 상태

• src/kernbench/ccl/algorithms/hierarchical_allreduce.py 이미 존재 (git log 10b33b4 — "Tensor indexing + hierarchical 3-level all-reduce kernel"). PE-level로 world_size = total PE를 가정하는 옛 모델 기반 구현.
• ADR-0024에 의해 launcher는 rank = SIP로 바뀜.
• Hierarchical 커널은 재해석 필요: 이제 각 worker(1 per SIP)가 자기 SIP의 모든 PE를 참여시키고, kernel은 intra-cube → inter-cube → inter-SIP 순으로 3-level reduce + broadcast.

풀어야 할 문제

1. ADR-0024 framework 위에 hierarchical 알고리즘 맞추기
   • Mapper: all_pes (ADR-0024 D5 제공)
   • Validator: multi_pe_sip_local (ADR-0024 D8 제공)
   • Kernel: 기존 hierarchical_allreduce.py 수정 — rank 계산 방식을 SIP 내 local (cube, pe)로 바꿈
2. PE-level neighbor graph 생성
   • Intra-cube: (sip, cube, pe) ↔ (sip, cube, pe±1 mod N_PE) (ring 내부)
   • Inter-cube: (sip, cube, 0) ↔ (sip, cube±1 mod N_CUBE, 0) (cube leader만)
   • Inter-SIP: (sip, 0, 0) ↔ (sip±1 mod N_SIP, 0, 0) (SIP leader만)
3. Tensor layout: 각 PE가 1 tile을 소유하고 시작 (multi_pe_sip_local validator가 이 layout 강제). DPPolicy(cube="column_wise", pe="column_wise")로 달성 가능.
4. PE-level topology 표현 부족 (ADR-0024 D6의 "책임 분산" 이슈 구체화)
   • Ring/mesh/tree 같은 단순 패턴은 rank-level topology_fn + mapper 조합으로 충분.
   • Hierarchical은 레벨마다 다른 peer 매핑이라 _build_pe_installs에서 multi-level 해석을 해야 함.
   • 장기적으로는 topology 모듈이 PE-level을 직접 표현하는 편이 명시적.

Non-problem (이 ADR 밖)

• Launcher / barrier / rank-to-SIP / mapper-validator registry → ADR-0024
• IPCQ direction addressing → ADR-0025
• DPPolicy 필드 정리 → ADR-0026
• Megatron TP → ADR-0027

---

Decision

D1. 알고리즘 구조 — 3-level reduce + 역순 broadcast

Level 1 (intra-cube, E/W ring):
  각 cube의 N_PE개 PE가 bidirectional ring reduce → cube 내 PE 0에 부분합 집중
Level 2 (inter-cube within SIP, N/S ring, PE 0만 참여):
  N_CUBE개 cube-leader가 bidirectional ring reduce → SIP 내 (cube 0, PE 0)에
  SIP 전체 부분합 집중
Level 3 (inter-SIP, N_SIP peers, (cube 0, PE 0)만 참여):
  Ring 또는 pair exchange로 전역 합산 완료
Broadcast:
  역순 — Level 3 결과를 (cube 0, PE 0)에서 SIP 내 모든 cube-leader로, 다시
  각 cube 내 모든 PE로 전파

세부는 기존 hierarchical_allreduce.py의 커널 구현과 일치. ADR-0024 이후 변경점은 rank 계산 방식과 n_elem 해석뿐:

• 기존 (rank=PE 모델): rank = cube_id * pes_per_cube + local_pe, pe_addr = t_ptr + rank * nbytes
• 신규 (rank=SIP 모델): 커널은 SIP-local 좌표 (cube_id, local_pe)로만 동작. 텐서의 per-PE slice는 backend가 per-PE TensorArg로 전달 (ADR-0024 D3). 커널 내부 rank 계산 자체가 불필요해짐 — tl.program_id(0/1)로 충분.

D2. Framework integration — ADR-0024 infrastructure 재활용

ccl.yaml:

algorithms:
  hierarchical_allreduce:
    module: kernbench.ccl.algorithms.hierarchical_allreduce
    topology: hierarchical_3level        # NEW — D3 참고
    mapper: all_pes                      # ADR-0024 D5 built-in
    validator: multi_pe_sip_local        # ADR-0024 D8 built-in
    buffer_kind: tcm
    n_elem: 128

Framework 관점에서 hierarchical은 특별한 알고리즘이 아니라, 특정 topology / mapper / validator 조합. 본 ADR은 그 조합과 topology 패턴을 정의.

D3. hierarchical_3level topology (신규)

kernbench/ccl/topologies.py에 신규 추가:

def hierarchical_3level(rank: int, world_size: int, spec: dict) -> dict:
    """3-level hierarchical neighbor pattern.

    Returns a nested structure describing intra-cube + inter-cube + inter-SIP
    neighbors. Unlike ring_1d / mesh_2d which are rank → {dir: peer_rank},
    hierarchical is PE-level and requires spec for cube_mesh / pe_layout.
    """

반환 스키마 (초안):

{
    "intra_cube": {
        # 각 cube 내 ring neighbors: (cube, pe) → {"E": (cube, pe_e), "W": (cube, pe_w)}
        ...
    },
    "inter_cube": {
        # cube-leader 간 ring: (cube, 0) → {"N": (cube_n, 0), "S": (cube_s, 0)}
        ...
    },
    "inter_sip": {
        # SIP-leader 간: rank → {"parent": peer_rank} (또는 ring 방식)
        ...
    },
}

이 구조는 _build_pe_installs가 해석하여 각 PE의 neighbor table 엔트리 (4-direction)에 대응시킨다.

Rank-level topologies.py 현 API와의 관계: 기존 단순 패턴은 (rank → {dir: peer_rank}) 단일 레벨. Hierarchical은 multi-level이므로 기존 API와 schema가 다름. _resolve_topology는 알고리즘이 어떤 schema를 쓰는지 선언하고, builder가 그에 맞춰 해석하도록 확장 필요 (open question).

D4. PE-level neighbor graph — _build_pe_installs 확장

기존 (ring/mesh/tree): topology_fn이 반환한 (rank → {dir: peer_rank})를 각 참여 PE에 그대로 매핑 (leader_only일 경우 peer PE도 leader).

신규 (hierarchical): hierarchical_3level의 3단 구조를 per-PE neighbor table로 펼침:

def _build_pe_installs_hierarchical(rank, world_size, sip, pes, topo, spec):
    """Hierarchical 전용 PE neighbor table 빌더."""
    result = []
    for (cube, pe) in pes:
        entries = []
        # Level 1: intra-cube ring (E/W)
        for d, peer in topo["intra_cube"][(cube, pe)].items():
            entries.append(NeighborTableEntry(direction=d, ...))
        # Level 2: inter-cube ring (N/S) — cube leader (pe == 0)만
        if pe == 0:
            for d, peer in topo["inter_cube"][(cube, 0)].items():
                entries.append(NeighborTableEntry(direction=d, ...))
        # Level 3: inter-SIP — SIP leader (cube == 0 and pe == 0)만
        if cube == 0 and pe == 0:
            for d, peer_rank in topo["inter_sip"][rank].items():
                # peer_rank → peer SIP의 (0, 0)
                entries.append(NeighborTableEntry(
                    direction=d, peer_sip=peer_rank, peer_cube=0, peer_pe=0, ...))
        result.append(PeInstallSpec(cube=cube, pe=pe, neighbors=tuple(entries)))
    return tuple(result)

build_install_plans에서 algorithm_config의 topology에 따라 적절한 builder 선택 (기존 simple builder vs hierarchical builder).

D5. Kernel 재해석 — SIP-local 좌표로

src/kernbench/ccl/algorithms/hierarchical_allreduce.py를 ADR-0024 D3에 맞춰 수정:

def kernel_args(*, n_elem: int, world_size: int, pes_per_cube: int,
                cubes_per_sip: int, num_sips: int, **kw) -> tuple:
    """world_size (= num_sips), pes_per_cube, cubes_per_sip를 스칼라로."""
    return (n_elem, pes_per_cube, cubes_per_sip, num_sips)

def kernel(t_ptr, n_elem, pes_per_cube, cubes_per_sip, num_sips, tl):
    """SIP-local 좌표 기반.

    이전 (rank=PE 모델):
        rank = cube_id * pes_per_cube + local_pe
        pe_addr = t_ptr + rank * nbytes
    현재 (rank=SIP 모델):
        per-PE tensor slice는 backend가 TensorArg로 전달 → t_ptr은 이미 local.
        intra-cube ring은 tl.program_id(0) 사용.
        inter-cube ring은 pe_id == 0 조건으로 제한.
        inter-SIP reduce는 cube_id == 0 and pe_id == 0 조건으로 제한.
    """
    local_pe = tl.program_id(axis=0)
    cube_id = tl.program_id(axis=1)

    # Level 1: intra-cube ring
    for _ in range(intra_rounds(pes_per_cube)):
        tl.send(dir="E", src=acc)
        recv = tl.recv(dir="W", shape=(n_elem,), dtype="f16")
        acc = acc + recv

    # Level 2: inter-cube (cube leader only)
    if local_pe == 0:
        for _ in range(inter_cube_rounds(cubes_per_sip)):
            tl.send(dir="N", src=acc)
            recv = tl.recv(dir="S", shape=(n_elem,), dtype="f16")
            acc = acc + recv

    # Level 3: inter-SIP (SIP leader only)
    if local_pe == 0 and cube_id == 0:
        for _ in range(inter_sip_rounds(num_sips)):
            tl.send(dir="parent", src=acc)
            recv = tl.recv(dir="parent", shape=(n_elem,), dtype="f16")
            acc = acc + recv

    # Broadcast (reverse chain)
    # ...
    tl.store(t_ptr, acc)

kernel_args는 ADR-0024 D4의 keyword-only signature 계약을 따른다.

D6. Validator — multi_pe_sip_local

ADR-0024 D8의 built-in 그대로 활용. ccl.yaml에서 validator: multi_pe_sip_local 지정 시 backend가 각 SIP에 cubes × pes_per_cube개 shard가 있는지 검증.

D7. Bench — 기본 all-reduce bench 확장

benches/ccl_allreduce.py의 worker는 ccl.yaml이 hierarchical_allreduce를 선택하면 자동으로:

# Worker 예
dp = DPPolicy(cube="column_wise", pe="column_wise")
tensor = torch.zeros((1, intra_sip_pes * n_elem), dp=dp, name="in")
# tensor는 각 SIP의 모든 PE에 1 tile씩 분산 (multi_pe_sip_local validator 통과)
dist.all_reduce(tensor, op="sum")

Worker 코드 자체는 알고리즘 종류를 모름 (ccl.yaml 선택에 의존). 단, DPPolicy가 hierarchical 요구와 일치해야 함 — cube/pe="column_wise" 같은 SIP-내 분산을 하는 DPPolicy여야 multi_pe_sip_local 검증 통과. 이 DPPolicy 선택은 bench 설정 또는 sample bench에서 결정.

---

Dependencies

• ADR-0024: Launcher, all_pes mapper, multi_pe_sip_local validator, registry + import path. 본 ADR 구현의 전제.
• ADR-0025: IPCQ direction addressing — cube/pe/SIP 간 다중 direction을 동시 사용하므로 정확한 direction 매칭 필수.
• ADR-0023: IPCQ protocol (neighbor table, send/recv, credit return).
• 기존 hierarchical_allreduce.py: 본 ADR은 그 커널의 재해석 + 주변 framework integration.

---

Non-goals

• ADR-0024 framework 변경: 재활용만.
• Alternative reduce topology (tree-in-tree 등): 3-level ring이 첫 구현.
• Dynamic level count: 현재 SIP/cube/PE 3단 고정. 2단 (SIP + PE, cube skip) 또는 4단 이상은 future.
• Bandwidth-optimal schedule tuning: reduce round 수 / chunk size 조정 같은 tuning은 별도.
• Pipelined hierarchical: 여러 chunk를 파이프라인으로 겹쳐서 돌리는 NCCL-style 최적화는 future.

---

Open questions

🟠 중간 영향 — 구현 시 결정 필요

• topologies.py 스키마 확장: 기존 ring_1d 등은 단일 레벨 (rank → {dir: peer}). hierarchical_3level은 multi-level. _resolve_topology가 둘을 모두 반환할 수 있도록 schema를 일반화할지, 아니면 hierarchical 전용 return type을 두고 builder가 분기할지.

  • Option A: 모든 topology를 neighbor-list 형태로 단일화 ([{direction, peer_sip, peer_cube, peer_pe}, ...])
  • Option B: topology 모듈이 kind 필드 제공, builder가 분기
  • 권장: Option A (single source of truth, ADR-0024 Open Q의 "PE-level topology 일원화" 방향과 일치)

• hierarchical_3level vs algorithm별 topology 모듈: 향후 mesh-based hierarchical 등 variant이 생기면? hierarchical_3level 같은 이름이 이미 topology-specific. 변형은 새 key 추가 (hierarchical_mesh_3level 등) 또는 알고리즘 모듈에서 topology 생성 override.

🟡 Nice-to-have

• Reduce round 수 최적화: Bidirectional ring은 ceil((N-1)/2) round. Non-power-of-2 group size에서 idle PE 발생 가능.
• Non-uniform topology 대응: cube_mesh가 w != h일 때 inter-cube ring balance.
• Single SIP 케이스: world_size = 1 (SIP 1개)일 때 Level 3 skip. Degenerate case 검증.

🟢 Framework evolution 시사점 (ADR-0024로부터 이관)

• PE-level topology 일원화 (중장기): 현 설계는

  • topology (rank graph 또는 level-separated)
  • mapper (per-SIP PE set)
  • _build_pe_installs (actual edges)

  의 3단 분산. Hierarchical이 이 분산을 가장 스트레스 받는 케이스. 중장기로는 topologies.py가 PE-level neighbor list를 직접 반환하고 mapper는 단순히 "어느 PE가 참여하느냐"만 결정, _build_pe_installs는 flat mapping으로 단순화되는 방향이 자연스러움. 본 ADR에서 Option A를 채택하면 이 방향으로 이미 정합.

---

Test strategy

T1. Topology generator

tests/test_hierarchical_topology.py (new):

• hierarchical_3level(rank, world_size, spec) → 각 level의 neighbor set이 예상 구조인지 (intra-cube는 ring, inter-cube는 cube-leader만 참여, inter-SIP은 SIP-leader만 참여)
• 2 SIP × 4 cubes × 4 PEs 같은 작은 토폴로지로 수작업 검증 가능
• Symmetry: rank r의 E neighbor가 peer에서 W로 역포인팅

T2. Install plan — hierarchical × all_pes

tests/test_ccl_install_plan.py (확장):

• build_install_plans(algorithm="hierarchical_allreduce", mapper="all_pes", validator="multi_pe_sip_local") 호출 시
  • 각 SIP의 모든 PE가 participating_pes에 포함
  • PE 0 (cube leader)만 inter-cube neighbor를 가짐
  • (cube 0, pe 0) (SIP leader)만 inter-SIP neighbor를 가짐
  • Non-leader PE는 intra-cube neighbor만

T3. Kernel unit — mock runtime

tests/test_hierarchical_mock_runtime.py (new):

• run_kernel_in_mock (kernbench.ccl.testing)을 확장해 multi-level 지원
• 2 SIP × 2 cubes × 4 PEs (총 16 PE) 토폴로지에서 초기 tile을 rank+1로 채우고 hierarchical all-reduce 실행
• 모든 PE의 최종 결과가 sum(1..16)인지

T4. E2E — 실제 SimPy backend

tests/test_ccl_allreduce_matrix.py (확장):

• hierarchical @ ws=SIP_count: multi_pe_sip_local layout + 3-level 알고리즘 전체 stack 통과 검증

T5. Validator enforcement

• multi_pe_sip_local validator가 wrong layout (예: leader_only 스타일 1 shard per rank) 입력에 raise

T6. 회귀

기존 ring/mesh/tree 알고리즘 모두 그대로 통과. 본 ADR은 그들을 건드리지 않음.

---

Consequences

Positive

• Intra-SIP PE 활용도 증가: Inter-SIP 통신 중에도 intra-cube / inter-cube reduce가 진행되어 전체 PE 가동률 향상.
• Multi-level bandwidth 활용: cube NoC, UCIe 모두 작동 → 더 정확한 HW 모델.
• ADR-0024 framework 검증: all_pes mapper + multi_pe_sip_local validator의 첫 non-trivial use case. Framework 설계 타당성 확인.
• 기존 커널 재활용: hierarchical_allreduce.py 큰 구조 유지, SIP-local 좌표만 재해석.

Negative

• topologies.py schema 확장 필요: Single-level vs multi-level 표현. 해결안(Option A)은 기존 ring/mesh/tree의 마이그레이션 비용 유발.
• Validator / mapper 조합 요구: 사용자가 DPPolicy를 multi_pe_sip_local에 맞춰 선택해야 함 (bench 설정 복잡도 증가).

Neutral

• 본 ADR 구현 전까지 hierarchical_allreduce.py는 deprecated 상태 유지 또는 ADR-0024 matrix test에서 제외. 현재 파일을 곧바로 삭제하지는 않음.

---

Affected files

File	Change
src/kernbench/ccl/topologies.py	D3: hierarchical_3level topology 함수 추가. (Option A 채택 시) 기존 topology 출력 format 통일
src/kernbench/ccl/install_plan.py	D4: hierarchical builder 분기 (또는 단일 builder가 level 개수로 dispatch)
src/kernbench/ccl/algorithms/hierarchical_allreduce.py	D5: SIP-local 좌표로 kernel 재작성, kernel_args keyword-only signature
ccl.yaml	D2: hierarchical_allreduce 엔트리 추가 (mapper: all_pes, validator: multi_pe_sip_local, topology: hierarchical_3level)
tests/test_hierarchical_topology.py (new)	T1
tests/test_ccl_install_plan.py	T2 확장
tests/test_hierarchical_mock_runtime.py (new)	T3
tests/test_ccl_allreduce_matrix.py	T4: hierarchical row 추가