# CCL Algorithm Author Guide 이 문서는 kernbench에서 CCL (Collective Communication Library) 알고리즘을 직접 작성하는 사람을 위한 step-by-step 가이드이다. 시스템 내부 설계와 컴포넌트 구조는 [ADR-0023](adr/ADR-0023-ipcq-pe-collective.md)에 있다. 본 가이드는 알고리즘 작성자가 **자신이 만져야 할 곳**과 **만지지 않아도 될 곳**을 명확히 분리하고, 가장 짧은 경로로 첫 알고리즘을 동작시키는 것을 목표로 한다. --- ## 0. 5분 요약 | 만지는 것 | 위치 | |----------|------| | 알고리즘 모듈 (kernel + 선택적 neighbors) | `src/kernbench/ccl/algorithms/.py` | | 알고리즘 등록 | `ccl.yaml` | | 호스트 bench (PE 수, 메모리 init, launch, 검증) | `benches/.py` | | (선택) 단위 테스트 | `tests/test_.py` | | 만지지 않는 것 | 위치 | |---------------|------| | TLContext API | `src/kernbench/triton_emu/tl_context.py` (ADR-0022 spec) | | 프레임워크 (topology generators, helpers, mock testing) | `src/kernbench/ccl/` | | PE_IPCQ / PE_DMA 컴포넌트 | `src/kernbench/components/builtin/` | | backend 구현 (install_ipcq) | `src/kernbench/runtime_api/distributed.py` 및 `kernbench/ccl/install.py` | 흐름: 1. 알고리즘 모듈에 `kernel` 작성 2. `ccl.yaml`에 entry 등록 3. 호스트 bench에서 `install_ipcq` + `launch` 4. (선택) mock runtime으로 단위 테스트 (수 ms) 5. `kernbench run --bench --verify-data`로 SimPy 검증 --- ## 1. Hello World — 가장 단순한 send/recv 각 PE가 자기 데이터를 E 방향 이웃에 한 번 보내고, W 방향에서 한 번 받는 가장 단순한 알고리즘이다. 실제 동작 코드는 [`src/kernbench/ccl/algorithms/hello_send.py`](../src/kernbench/ccl/algorithms/hello_send.py) 에 있다. ### Step 1: kernel 작성 새 파일 `src/kernbench/ccl/algorithms/hello_send.py`: ```python """Hello world: 자기 데이터를 다음 rank에 보내고 이전 rank에서 받기.""" def kernel(t_ptr, n_elem, tl): # 글로벌 rank는 program_id(0/1)에서 계산 (ADR-0022) local_pe = tl.program_id(axis=0) cube_id = tl.program_id(axis=1) pes_per_cube = tl.num_programs(axis=0) rank = cube_id * pes_per_cube + local_pe nbytes = n_elem * 2 # f16 pe_addr = t_ptr + rank * nbytes # 자기 슬라이스를 로드해서 E로 보낸다. src = tl.load(pe_addr, shape=(n_elem,), dtype="f16") tl.send(dir="E", src=src) # W 방향에서 받아서 그대로 자기 슬라이스에 store한다. recv = tl.recv(dir="W", shape=(n_elem,), dtype="f16") tl.store(pe_addr, recv) ``` 핵심 포인트: - **글로벌 rank는 `program_id(axis=0)` + `program_id(axis=1)`에서 계산.** TL에는 `tl.rank` / `tl.world_size` 같은 약속되지 않은 확장이 없다. 호스트가 `world_size` 같은 알고리즘 파라미터가 필요하면 `torch.launch`의 일반 인자로 전달한다. - **`tl.send`는 `TensorHandle`을 받는다.** 핸들의 `addr`/`space`/`shape`/`dtype`/`nbytes`를 PE_IPCQ가 읽어 PE_DMA에 IpcqDmaToken을 발행한다. - **`tl.recv`는 `shape`와 `dtype`이 필수.** 반환된 TensorHandle은 IPCQ ring slot을 가리키며, `tl.store(pe_addr, recv)`처럼 dst 핸들로 그대로 사용할 수 있다. Phase 2 dma_write replay가 (slot, hbm) 복사를 수행하므로 numpy `.data`를 직접 만질 필요가 없다. ### Step 2: ccl.yaml 등록 `ccl.yaml`의 `algorithms` 섹션에 entry를 추가한다. (defaults.algorithm은 호스트 bench가 `install_ipcq(algorithm=...)`로 명시 전달해도 되므로 꼭 바꿀 필요는 없다.) ```yaml algorithms: hello_send: module: kernbench.ccl.algorithms.hello_send topology: ring_1d buffer_kind: tcm ``` ### Step 3: 호스트 bench 작성 새 파일 `benches/ccl_hello.py`: ```python """Hello-world ring rotation bench (각 PE가 W 이웃의 데이터를 1번 받음).""" import numpy as np from kernbench.ccl.algorithms import hello_send from kernbench.policy.placement.dp import DPPolicy ALGORITHM = "hello_send" N_ELEM = 8 WORLD_SIZE = 8 def run(torch): plan = torch.install_ipcq(algorithm=ALGORITHM) a = torch.zeros( (1, WORLD_SIZE * N_ELEM), dtype="f16", dp=DPPolicy( sip="replicate", cube="replicate", pe="column_wise", num_sips=1, num_cubes=1, ), name="hello_in", ) store = torch.engine.memory_store base = a._handle.va_base or a._handle.shards[0].pa nbytes = N_ELEM * 2 for r in range(WORLD_SIZE): store.write("hbm", base + r * nbytes, np.full((N_ELEM,), float(r + 1), dtype=np.float16)) torch.launch(ALGORITHM, hello_send.kernel, a, N_ELEM) # rank r은 rank (r-1)%ws의 데이터를 가져야 한다. for r, (sip, cube, pe) in enumerate(plan["rank_to_pe"]): result = store.read("hbm", base + r * nbytes, shape=(N_ELEM,), dtype="f16") prev = float(((r - 1) % WORLD_SIZE) + 1) ok = np.allclose(result, prev) print(f" [{'OK ' if ok else 'FAIL'}] rank {r} got {float(result.mean()):.1f}, " f"expected {prev:.1f}") ``` ### Step 4: 단위 테스트 (선택, 강력 추천) `tests/test_hello_send.py`: ```python import numpy as np from kernbench.ccl.algorithms.hello_send import kernel from kernbench.ccl.testing import run_kernel_in_mock def test_hello_send_4_ranks(): n_elem = 8 inputs = [np.full((n_elem,), float(r + 1), dtype=np.float16) for r in range(4)] outputs = run_kernel_in_mock( kernel_fn=kernel, world_size=4, topology="ring_1d", inputs=inputs, kernel_args=(n_elem,), ) # rank r은 rank (r-1) % 4의 데이터를 받아야 함 for r in range(4): assert np.array_equal(outputs[r], inputs[(r - 1) % 4]) ``` `run_kernel_in_mock`는 SimPy 없이 순수 Python으로 모든 rank를 동시 실행하므로 **ms 단위로 끝난다**. 알고리즘 logic 정합성만 검증. ### Step 5: 시뮬 검증 ```bash kernbench run --topology topology.yaml --bench ccl_hello --verify-data ``` Phase 1에서 SimPy 시뮬레이션 + MemoryStore 데이터 이동, Phase 2에서 op_log 정합성 replay. 호스트 bench의 `print` 검증이 모든 rank에 대해 OK여야 한다. --- ## 2. Ring All-Reduce — 두 번째 알고리즘 조금 더 복잡한 예제. Ring all-reduce는 N-1 라운드 동안 각 PE가 자기 데이터를 E로 보내고 W에서 받아 누적한다. 최종적으로 모든 PE가 글로벌 sum을 갖는다. 실제 동작 코드는 [`src/kernbench/ccl/algorithms/ring_allreduce.py`](../src/kernbench/ccl/algorithms/ring_allreduce.py) 참조. 핵심 흐름: ```python """Ring all-reduce.""" def kernel(t_ptr, n_elem, world_size, tl): # rank local_pe = tl.program_id(axis=0) cube_id = tl.program_id(axis=1) pes_per_cube = tl.num_programs(axis=0) rank = cube_id * pes_per_cube + local_pe nbytes = n_elem * 2 pe_addr = t_ptr + rank * nbytes # HBM의 자기 슬라이스를 가리키는 TensorHandle. greenlet 모드에선 .data가 # 채워지지만 커널은 .data를 직접 만질 필요가 없다. acc = tl.load(pe_addr, shape=(n_elem,), dtype="f16") current = acc # 첫 라운드 send 출처 for _step in range(world_size - 1): tl.send(dir="E", src=current) recv = tl.recv(dir="W", shape=(n_elem,), dtype="f16") # TensorHandle 연산자 오버로드 → MathCmd → PE_MATH 디스패치. # Phase 1은 타이밍만, Phase 2 DataExecutor가 실제 numpy 누적을 수행한다. acc = acc + recv current = recv # 다음 라운드는 직전에 받은 슬롯을 다시 forward # 최종 누적값을 자기 슬라이스에 store. 출처는 acc(=PE-local scratch addr) # 이고 dst는 HBM. op_log dma_write가 (scratch, hbm) 복사 정보를 기록하므로 # Phase 2가 검증 시점에 HBM[pe_addr]에 정답을 채워준다. tl.store(pe_addr, acc) ``` 네 가지 포인트: 1. **누적은 TensorHandle 연산자**: `acc + recv`는 `MathCmd`를 emit하고 PE_MATH로 디스패치된다 — 실제 하드웨어 경로를 거치므로 latency 모델이 정확하다. ADR-0020 D3대로 Phase 1은 타이밍만 시뮬레이션하고, Phase 2 `DataExecutor`가 op_log를 재실행하면서 numpy 누적을 수행한다. 2. **`current = recv`로 forward**: 매 라운드의 send 출처를 직전에 받은 슬롯 핸들로 갱신해야 같은 데이터가 ring을 순회하면서 누적이 한 번씩 일어난다. `current = acc`로 두면 누적값이 다시 송출되어 결과가 부풀려진다. 3. **`tl.store(pe_addr, acc)` 한 번이면 끝**: 중간에 store→reload 패턴은 금지다. acc는 PE-local scratch에 살고, op_log가 (src=scratch, dst=hbm) 메타데이터를 기록한다. Phase 2가 math를 먼저 실행해 scratch를 채운 뒤 dma_write 스냅샷으로 HBM에 복사한다. 4. **`world_size`는 호스트가 명시 전달**: TL은 topology slot 수만 안다 (예: `num_programs(axis=0)`은 cube당 PE 수). 실제 참여하는 CCL group 크기는 bench가 알고 호스트→kernel 인자로 넘긴다. `ccl.yaml` 등록 + 호스트 bench는 [`benches/ccl_allreduce_tcm.py`](../benches/ccl_allreduce_tcm.py) 참조. mock 단위 테스트는 [`tests/test_ccl_mock_runtime.py`](../tests/test_ccl_mock_runtime.py) 를 그대로 따라하면 된다 (`kernel_args=(n_elem, world_size)` 인자 형태). --- ## 3. neighbors() override — Custom topology 대부분의 알고리즘은 builtin topology(`ring_1d`, `mesh_2d`, `tree_binary`, `ring_1d_unidir`, `none`)로 충분하다. builtin을 변형하거나 새로 만들고 싶으면 알고리즘 모듈에 `neighbors()`를 정의한다. ### 시그니처 ```python def neighbors(rank: int, world_size: int, neighbor_map: dict[str, int]) -> dict[str, int] | None: """builtin topology가 만든 neighbor_map을 override. Args: neighbor_map: ccl.yaml의 topology 필드가 만든 builtin 매핑. 예: ring_1d → {"E": (rank+1)%ws, "W": (rank-1)%ws} mutable dict — 직접 수정 가능. Returns: dict: neighbor_map을 override한 결과 (또는 수정한 그 dict) None: override 안 함, neighbor_map 그대로 사용 """ return None ``` ### Pattern A: builtin을 base로 일부만 수정 ```python def neighbors(rank, world_size, neighbor_map): # 짝수 rank만 W 방향 사용 (홀수 rank는 W 제거) if rank % 2 == 1: neighbor_map.pop("W", None) return neighbor_map ``` ### Pattern B: 완전히 새로 작성 (skip-connection ring) ```python def neighbors(rank, world_size, neighbor_map): # neighbor_map은 무시하고 새로 작성 return {"E": (rank + 2) % world_size} ``` ### Pattern C: builtin 사용, override 없음 `neighbors()` 함수를 정의하지 않거나 None을 반환: ```python def neighbors(rank, world_size, neighbor_map): return None # 명시적으로 builtin 사용 ``` --- ## 4. PE 커널 API 레퍼런스 (ADR-0023 D4) ### IPCQ API | API | 설명 | Blocking? | |-----|------|-----------| | `tl.send(dir, src=TensorHandle)` | direction으로 데이터 send | Yes (peer slot full 시 wait) | | `tl.send(dir, src_addr=..., nbytes=..., shape=..., dtype=..., space=...)` | 동일, keyword 형태 | Yes | | `tl.recv(dir, shape=..., dtype=...)` | 특정 방향에서 blocking recv | Yes | | `tl.recv(shape=..., dtype=...)` | 4방향 round-robin recv (방향 미지정) | Yes | | `tl.recv_async(dir, shape=..., dtype=...) → RecvFuture` | non-blocking recv | No | | `tl.wait(future)` | non-blocking future 완료 대기 → TensorHandle | Yes | ### 기존 TL API (ADR-0020/0022, 그대로 사용 가능) | API | 설명 | |-----|------| | `tl.load(addr, shape, dtype) → TensorHandle` | DMA read; greenlet 모드에서 `.data`에 ndarray | | `tl.store(addr, handle)` | DMA write — handle.data가 있으면 MemoryStore에 propagate | | `tl.composite(op, ...)` | GEMM/Math compute 비동기 submit | | `tl.program_id(axis=0)` | cube 내 local PE id | | `tl.program_id(axis=1)` | cube id (ADR-0022) | | `tl.num_programs(axis=0/1)` | topology 슬롯 수 (참여 ranks 수가 아님) | ### `recv` 두 가지 모드 기본은 `return_slot` (zero-copy): IPCQ slot 주소가 그대로 handle.addr에 들어온다. slot 데이터를 별도 위치로 복사하고 싶으면 `dst_addr` + `dst_space`를 명시: ```python recv = tl.recv( dir="W", shape=(8,), dtype="f16", dst_addr=my_scratch_addr, dst_space="hbm", ) # 이제 recv.addr == my_scratch_addr (copy_to_dst 모드) ``` --- ## 5. Helpers (`kernbench.ccl.helpers`) 알고리즘 코드를 짧게 유지하기 위한 헬퍼들: ```python from kernbench.ccl.helpers import chunked, ring_step, tree_step ``` ### `chunked(base_addr, n_chunks, n_elem, dtype="f16") → list[Chunk]` 총 `n_elem` 개의 element를 `n_chunks` 등분한 view 리스트를 반환. 각 `Chunk`는 `addr`, `n_elem`, `nbytes` 필드를 가진다. ```python chunks = chunked(t_ptr, n_chunks=4, n_elem=64, dtype="f16") # chunks[0..3] 각각 16 element view, addr이 연속 ``` ### `ring_step(rank, step, world_size) → (send_idx, recv_idx)` Ring algorithm의 step별 chunk 인덱스 (reduce-scatter / all-gather): ```python for step in range(world_size - 1): send_idx, recv_idx = ring_step(rank, step, world_size) tl.send(dir="E", src_addr=chunks[send_idx].addr, nbytes=chunks[send_idx].nbytes, shape=(chunks[send_idx].n_elem,), dtype="f16") recv = tl.recv(dir="W", shape=(chunks[recv_idx].n_elem,), dtype="f16") # accumulate ... ``` ### `tree_step(rank, world_size) → {"parent": int|None, "children": list[int]}` Binary tree의 parent/children rank: ```python info = tree_step(rank, world_size) if info["parent"] is None: print(f"rank {rank} is the root") for child in info["children"]: ... ``` --- ## 6. 단위 테스트 — Mock Runtime `kernbench.ccl.testing.run_kernel_in_mock`은 SimPy를 거치지 않고 알고리즘을 빠르게 검증할 수 있다. ### 기본 사용법 ```python from kernbench.ccl.testing import run_kernel_in_mock from kernbench.ccl.algorithms.my_algo import kernel import numpy as np def test_my_algo(): n_elem = 16 inputs = [np.arange(n_elem, dtype="f16") + r for r in range(4)] expected = sum(inputs) outputs = run_kernel_in_mock( kernel_fn=kernel, world_size=4, topology="ring_1d", inputs=inputs, kernel_args=(n_elem, 4), # kernel의 (t_ptr 이후) 추가 positional 인자 ) for r in range(4): assert np.allclose(outputs[r], expected, rtol=1e-3) ``` ### 동작 - 4개 rank의 kernel을 greenlet으로 동시 실행 - `tl.send/recv`를 in-memory FIFO로 즉시 처리 (DMA, latency 무시) - 각 rank가 마지막에 store한 데이터를 ndarray로 반환 ### 한계 - latency / 성능 측정 불가 (시뮬레이션이 아님) - PE_DMA, fabric, BW 모델 안 함 - 정합성 검증만 가능 - 한 cube 안에서 동작하는 가정 — `program_id(axis=1)`은 항상 0 --- ## 7. 디버깅 ### CCL trace ```bash KERNBENCH_CCL_TRACE=1 kernbench run --topology topology.yaml \ --bench ccl_allreduce_tcm --verify-data ``` 각 rank의 send/recv 시점이 stdout에 출력된다: ``` [ccl t=346.4 send] sip0.cube0.pe1 dir=E nbytes=64 seq=0 [ccl t=360.4 recv] sip0.cube0.pe2 dir=W nbytes=64 ... ``` ### Pointer dump `kernbench.ccl.diagnostics.pointer_dump(engine)`는 모든 PE_IPCQ의 ring buffer 상태(`my_head`, `my_tail`, `peer_head_cache`, `peer_tail_cache`)를 multi-line 문자열로 반환한다. hang이 발생하면 어느 rank가 어떤 상태에서 막혔는지 한눈에 보인다. ### Deadlock detection 매칭되지 않는 send/recv 등으로 SimPy 스케줄이 비면 engine이 `IpcqDeadlock`을 던지며 pointer dump를 메시지에 포함시킨다 (ADR-0023 D14 F3). 별도 wait-for graph 시각화는 미래 작업. --- ## 8. 흔한 실수 ### 1. install 안 된 direction 사용 ccl.yaml의 `topology: ring_1d`는 E/W만 install한다. N/S 사용 시: ```python tl.send(dir="N", ...) # → IpcqInvalidDirection 예외 ``` 해결: `topology: mesh_2d`로 바꾸거나, `neighbors()` override로 N/S 추가. ### 2. send만 호출하고 recv 없음 ```python def kernel(..., tl): for _ in range(100): tl.send(dir="E", ...) # peer 측 recv 없음 → ring buffer 가득 차면 backpressure → deadlock ``` 해결: 모든 send에 짝이 되는 recv가 있어야 한다. 안 그러면 `IpcqDeadlock`이 발생한다. ### 3. dtype/shape 불일치 기본 모드에서는 dtype/shape mismatch를 검증하지 않는다. 작성자가 직접 보장하거나, PE_IPCQ 노드 attrs에 `strict_validation: true`를 설정해 D14 F2 strict 모드로 mismatch를 즉시 잡을 수 있다. ### 4. round-robin recv의 fairness 가정 `tl.recv()` (방향 미지정)는 round-robin으로 가져오지만, 도착한 첫 슬롯을 반환한다. **도착 순서를 알 수 없으므로** 알고리즘이 도착 방향에 의존하면 안 된다. 필요하면 `tl.recv(dir="N", ...)`처럼 명시. ### 5. CCL 그룹 크기 가정 `tl.num_programs(axis=0/1)`은 토폴로지 슬롯 개수이지 CCL group 크기가 아니다. 참여하는 rank 수(`world_size`)는 호스트 bench가 알고 있고, kernel 인자로 명시 전달해야 한다. ### 6. 호스트가 send-source 메모리를 도착 전에 덮어씀 PE_DMA가 송신 시점에 src 데이터를 토큰에 스냅샷해서 in-flight 데이터의 의미가 보존된다. 그래도 하나의 PE 안에서 같은 주소를 여러 step에 걸쳐 갱신할 때는 direct send 후 다른 step에서 같은 주소를 store해도 안전하다 (token snapshot 덕분). 하지만 `tl.send`가 PE_DMA 큐에 enqueue되기 전에 주소를 덮어쓰면 잘못된 데이터가 스냅샷된다 — `tl.send`를 먼저, 메모리 변경을 나중에 하는 게 권장. --- ## 9. 다음 단계 - `mesh_2d` / `tree_binary` 같은 다른 topology 활용 - recursive halving/doubling 등 더 빠른 알고리즘 - `buffer_kind` (tcm/hbm/sram) / `backpressure` (poll/sleep) 모드별 latency 비교 - `ccl_ring_allreduce_multicube.py`, `ccl_ring_allreduce_multisip.py`처럼 큰 scale의 ring 검증 새 알고리즘이나 패턴을 추가했다면 PR로 기여해주세요. --- ## 참고 - [ADR-0023](adr/ADR-0023-ipcq-pe-collective.md): IPCQ + PE-level collective 설계 - [ADR-0022](adr/ADR-0022-program-id-2d-grid.md): 2D grid program_id (axis=0/1) - [ADR-0020](adr/ADR-0020-data-execution-two-pass.md): 2-pass data execution - [ADR-0021](adr/ADR-0021-pe-pipeline-refactor.md): PE pipeline refactor 기존 알고리즘 예제: - [`src/kernbench/ccl/algorithms/hello_send.py`](../src/kernbench/ccl/algorithms/hello_send.py) — 가장 단순한 send/recv - [`src/kernbench/ccl/algorithms/ring_allreduce.py`](../src/kernbench/ccl/algorithms/ring_allreduce.py) — ring all-reduce - [`src/kernbench/ccl/algorithms/mesh_allreduce.py`](../src/kernbench/ccl/algorithms/mesh_allreduce.py) — 2D mesh all-reduce - [`src/kernbench/ccl/algorithms/tree_allreduce.py`](../src/kernbench/ccl/algorithms/tree_allreduce.py) — binary tree all-reduce