[Triton] 커널 끝에 cross-CTA barrier 추가 — 클러스터 메모리 정합성 보장

PR 링크: triton-lang/triton#9413 상태: Merged | 변경: +94 / -7

들어가며

NVIDIA Hopper GPU의 클러스터(cluster)는 여러 CTA(Cooperative Thread Array)를 하나의 그룹으로 묶어 상호 간 distributed shared memory 접근을 가능하게 한다. 그러나 한 CTA에서 발생한 메모리 연산이 다른 CTA에서 보이려면 적절한 barrier가 필요하다.

이 PR은 커널 종료 시점에 미처리(outstanding) 읽기 또는 쓰기가 남아있고, TMEM deallocation이 없는 경우 클러스터 수준의 barrier를 삽입한다.

핵심 코드 분석

Before: 커널 종료 시 cross-CTA barrier 없음

기존에는 커널 본문 내의 동기화만 처리했고, 커널 종료 시점의 미처리 연산에 대한 cross-CTA 동기화가 없었다.

After: 커널 끝에 조건부 barrier 삽입

// ClusterMembarAnalysis에 추가된 로직
void insertEndOfKernelBarrier(triton::FuncOp funcOp) {
  // 커널의 마지막 블록에서 미처리 읽기/쓰기 확인
  auto &lastBlock = funcOp.getBody().back();
  auto returnOp = cast<triton::ReturnOp>(lastBlock.getTerminator());

  if (hasOutstandingReadOrWrite() && !hasTMEMDeallocation()) {
    // 클러스터 수준 barrier 삽입
    OpBuilder builder(returnOp);
    builder.create<triton::nvidia_gpu::ClusterArriveOp>(
        returnOp.getLoc(), /*relaxed=*/false);
    builder.create<triton::nvidia_gpu::ClusterWaitOp>(
        returnOp.getLoc());
  }
}

ClusterArriveOp + ClusterWaitOp 조합은 클러스터 내 모든 CTA가 해당 지점에 도달할 때까지 대기하며, 이전의 모든 메모리 연산이 다른 CTA에서 관측 가능함을 보장한다.

TMEM deallocation이 이미 있는 경우에는 barrier가 불필요한데, deallocation 자체가 암묵적으로 동기화를 포함하기 때문이다.

왜 이게 좋은가

1. 정합성 보장: 클러스터 모드에서 CTA 간 shared memory 접근이 커널 종료 후에도 올바르게 관측되도록 보장한다.
2. 조건부 삽입: 미처리 연산이 없거나 TMEM deallocation으로 이미 동기화된 경우에는 barrier를 삽입하지 않아 불필요한 성능 저하를 방지한다.
3. 방어적 프로그래밍: 컴파일러가 자동으로 필요한 barrier를 삽입하므로, 사용자가 수동으로 동기화를 관리할 필요가 없다.

정리

이 PR은 Triton의 클러스터 membar 분석에 커널 종료 시점의 cross-CTA barrier 삽입 로직을 추가한다. 미처리 메모리 연산이 있을 때만 조건부로 삽입되어 정합성과 성능을 동시에 확보한다.

참고 자료

• NVIDIA Hopper Architecture — Thread Block Clusters
• Triton Cluster Membar Analysis

---

이 글은 AI(Claude)의 도움을 받아 작성되었습니다. 핵심 코드와 explaination은 실제 PR diff를 기반으로 합니다.