Iris — Multi-GPU Programming in Triton

Triton 은 단일 GPU 의 tile DSL 로 디자인됐다. 그러나 실제 LLM 학습 / inference 는 multi-GPU 가 default 다. tensor parallel, sequence parallel, expert parallel 같은 분산 패턴이 모두 그 위에 깔린다. 단일 GPU 추상으로는 이 자리를 표현할 수 없다.

강의의 frame 은 명확하다.

1. NCCL / RCCL 의 collective 는 black box 다 — allreduce, allgather 한 줄로 동작하지만 그 안에서 어떤 chunk 가 언제 어디로 가는지를 통제할 수 없다. fine-grained overlap 이 어렵다.
2. compute 와 통신을 같은 커널 안에서 짤 수 있어야 한다 — tensor parallel 의 GEMM + reduce-scatter, 또는 sequence parallel 의 ring-attention 같은 패턴은 본질적으로 compute-communication overlap 이다. NCCL 위에서는 별도 stream 으로 어색하게 표현된다.
3. Triton 의 tile 추상을 multi-GPU 까지 확장한다. tile 안에 "remote" 위치까지 표현. 한 program 이 다른 GPU 의 데이터를 자기 것처럼 indexing.

왜 NCCL 의 한 줄이 부족한가 — 강의에서 짚는 세 가지 자리.

NCCL 의 모델은 "여러 GPU 가 동시에 같은 collective 를 부르고 결과가 나올 때까지 기다린다". 이 모델은 scale 에서는 효율적이지만 compute 안에 communication 을 끼워 넣고 싶은 자리에는 어울리지 않는다.

Iris 의 핵심 자료형 두 개 — distributed tensor 와 remote tile load. Triton 에 익숙한 개발자가 자연스럽게 따라갈 수 있는 형태로 디자인.

# Iris 의 distributed GEMM (개념)
@triton.jit
def tp_gemm_kernel(
    A_ptr, B_ptr, C_ptr,
    M, N, K, world_size: tl.constexpr,
    BLOCK: tl.constexpr,
):
    pid = tl.program_id(0)
    rank = iris.rank()

    # 자기 shard 의 일부
    a_block = tl.load(A_ptr + offsets_a)

    # 다른 rank 의 B shard 를 가져온다
    for r in range(world_size):
        if r != rank:
            b_remote = iris.load_remote(
                r, B_ptr + offsets_b, mask=mask)
            acc += tl.dot(a_block, b_remote)

    # local 결과 reduce
    iris.atomic_add_remote(
        0, C_ptr + offsets_c, acc)
    iris.barrier(group=ALL)
    tl.store(C_ptr + offsets_c, final)

Iris 의 통신은 one-sided + tile-level barrier. NCCL 의 group call 이 아닌 RDMA 모델. 그래서 progress (실제 데이터가 도착했는지) 의 통제가 fine-grained 하다.

강의에서 가장 진하게 다룬 예제. Megatron-style TP 의 핵심 — "내 weight shard 와 모든 rank 의 input 의 곱을 누적". NCCL 로는 GEMM + allreduce 두 단계, Iris 로는 한 커널.

# NCCL 형 (PyTorch 표준)
y_local = torch.matmul(x, w_shard)
torch.distributed.all_reduce(y_local)
# Y = y_local now

# --- vs ---

# Iris 형 (단일 Triton kernel)
@triton.jit
def ag_gemm(x, w, y, ...):
    pid_m, pid_n = tl.program_id(0), tl.program_id(1)
    rank = iris.rank()

    acc = tl.zeros([BM, BN], dtype=tl.float32)
    for k in range(0, K, BK):
        x_tile = tl.load(x + offs_x)
        w_tile = tl.load(w + offs_w)
        acc += tl.dot(x_tile, w_tile)

    # 모든 rank 의 y[m,n] 에 atomic add
    for r in range(world_size):
        iris.atomic_add_remote(r, y + offs_y, acc)

    iris.barrier(group=ALL)

강의의 측정 (확인 필요): 큰 TP GEMM (예: M=N=K=8192, 8-way TP) 에서 Iris 가 NCCL 기반 baseline 보다 20-40% 더 빠른 wallclock 을 보였다. 작은 GEMM 에서는 NCCL 이 비슷하거나 더 빠름 — overhead 의 amortization 차이.

강의에서 명시적으로 강조 — Iris 는 NCCL/RCCL 을 대체하지 않는다. 두 도구의 자리가 다르다.

실전 워크플로 — 강의에서 권장하는 형태.

• data parallel gradient sync → NCCL/RCCL allreduce 한 줄. 단순하고 충분.
• tensor parallel 의 forward/backward GEMM → Iris kernel 한 개. compute-communication overlap 으로 throughput 회복.
• sequence parallel ring attention → Iris 가 가장 빛나는 자리. ring 단계마다 partial K/V tile 을 fine-grained 하게 가져온다.
• expert parallel all-to-all → 강의에서 다뤄졌는지 확인 필요. NCCL all_to_all 보다 token imbalance 에 robust.

• AMD ROCm 공식 라이브러리 — Iris 가 ROCm 의 multi-GPU Triton 표준 path 로 자리잡고 있다 (확인 필요). MI300X / MI325X / MI355X 에 default 로 들어감.
• vLLM AMD 백엔드 — TP serving 의 핫패스. Iris 기반 GEMM-allreduce fused kernel.
• SGLang / LMDeploy — 일부 fused multi-GPU op 가 Iris 위에서 짜여 있음 (확인 필요).
• Megatron-LM AMD fork — NCCL 기반 standard 외에 Iris 기반 TP 구현 옵션.

채택 곡선의 패턴 — AMD 진영의 "NVIDIA 따라잡기" 의 한 점. NVIDIA 의 NVSHMEM + Triton 등가물이 AMD 위에서 표준화된 길로 자리잡으려는 시도. 성공은 ROCm 의 전체 채택 곡선에 의존.

• node 간 RDMA 의 overhead — 같은 node 의 NVLink/Infinity Fabric 위에서는 빠르지만 multi-node 의 InfiniBand RDMA 는 latency 큼. tile 단위 통신이 오히려 NCCL 보다 느려질 수 있다.
• scale 의 한계 — NCCL 은 1000+ GPU 까지 검증, Iris 는 강의 시점에 single-node 또는 small multi-node 위주. 큰 cluster 에서의 동작은 확인 필요.
• SymmetricHeap 의 메모리 비용 — 모든 rank 에 동일 VA 가 mapped 되어야 하므로 메모리 fragmentation 비용이 있다.
• 디버깅의 어려움 — 한 rank 의 wrong store 가 다른 rank 의 silent corruption 으로 나타난다. 단일 GPU Triton 보다 디버깅 난이도 큼.
• NVIDIA 위 portability — Iris 는 AMD 출신. NVIDIA 위 NVSHMEM 백엔드가 가능하지만 강의 시점의 성숙도 확인 필요.
• AMD 외 채택의 inertia — 산업 default 가 NCCL. Iris 의 가치가 깊지만 학습 곡선 + AMD-bound 가 채택을 늦춘다.

• multi-node 검증 — 16 / 32 node 위에서 Iris 의 성능 곡선. RDMA overhead 의 amortization 가능성.
• NVIDIA NVSHMEM 백엔드 — 같은 추상이 NVIDIA 인프라 위에서도 동작하도록.
• autotuning 통합 — tile size, communication granularity 의 자동 결정.
• collective 의 자동 합성 — 사용자가 high-level 표현 (allreduce 한 줄) 을 적으면 Iris 가 fine-grained 으로 자동 lowering.
• compiler-level fusion — torch.compile 또는 FlexAttention 같은 high-level 표현에서 Iris 백엔드로 자동 lowering.
• FlashAttention with Iris — ring attention 의 fully-fused 버전.

손에 새기기 — 실습 시퀀스

1. local Iris setup — 2-GPU 기준으로 hello-world. SymmetricHeap allocate, simple put/get.
2. TP GEMM 직접 짜기 — 4096³ matmul 의 4-way TP 를 Iris 로. NCCL allreduce baseline 과 wallclock 비교.
3. ring attention — sequence parallel attention 의 ring 단계를 Iris 로. K/V tile 의 fine-grained 통신.
4. tile size sweep — Iris 의 tile size 와 communication granularity 의 trade-off.
5. NCCL 과 cross-check — 같은 input 의 결과가 두 backend 에서 비트 동일한지 확인. silent drift 점검.
6. profile 해석 — rocprof 로 timeline 을 dump. compute 와 communication 의 overlap 을 직접 본다.

• repo 의 정확한 위치 — github.com/ROCm/iris 추정. 강의 영상에 명시된 url 확인.
• NVIDIA backend 의 성숙도 — NVSHMEM 위 동작이 production-ready 인지.
• multi-node 측정 — 16+ node 위 측정 결과. RDMA latency 의 dominance 시점.
• EP all-to-all 패턴 — Mixture-of-Experts 에서 token routing. Iris 가 자연스럽게 표현 가능한지.
• 벤치마크의 baseline — § 05 의 20-40% 가 RCCL allreduce baseline 인지, 더 정교한 NCCL+overlap baseline 인지.
• 호환 PyTorch 버전 — Triton dependency 의 정확한 버전.
• FlashAttention-3 with Iris — sequence parallel ring attention 의 fully-fused 버전이 강의에서 다뤄졌는지.