본문으로 건너뛰기

[Triton] FenceAsync에 비동기 읽기 의존성 추가 — st.shared와 copy_local_to_global 간 정합성 보장

PR 링크: triton-lang/triton#9610 상태: Merged | 변경: +91 / -30

들어가며

Triton의 NVIDIA 백엔드에는 서로 다른 메모리 프록시 간 데이터 이동 시 fence를 삽입하는 ProxyFenceInsertion 패스가 있다. GPU에서 shared memory에 동기적으로 쓴 값(st.shared)을 비동기적으로 글로벌 메모리로 복사(copy_local_to_global)할 때, 두 연산이 서로 다른 하드웨어 프록시를 사용하므로 반드시 fence가 필요하다.

기존 코드는 이 의존성을 감지하지 못해 fence를 삽입하지 않았고, 이는 잠재적인 데이터 정합성 버그를 유발할 수 있었다. 이 PR은 비동기 읽기 연산의 소스를 정확히 추적하는 로직을 추가하여 문제를 해결한다.

핵심 코드 분석

Before: 비동기 읽기 연산 단순 분류

기존에는 어떤 연산이 비동기 프록시 읽기인지만 판단했을 뿐, 구체적으로 어떤 operand가 shared memory에서 읽히는지 추적하지 않았다.

bool isAsyncProxyRead(Operation *op) {
  return isa<triton::nvidia_gpu::WarpGroupDotOp,
             triton::nvidia_gpu::TCGen5MMAOp,
             triton::nvidia_gpu::TCGen5MMAScaledOp,
             triton::nvidia_gpu::AsyncTMACopyGlobalToLocalOp,
             triton::nvidia_gpu::AsyncTMAScatterOp,
             triton::nvidia_gpu::AsyncTMAReduceOp>(op);
}

After: 연산별 소스 operand 추적 함수 추가

새로운 isAsyncProxyReadSource 함수는 특정 value가 해당 연산에서 shared memory 소스로 사용되는지를 정확히 판별한다.

bool isAsyncProxyReadSource(Operation *op, Value value) {
  auto memDescType = dyn_cast<triton::gpu::MemDescType>(value.getType());
  if (!memDescType ||
      !isa<triton::gpu::SharedMemorySpaceAttr>(memDescType.getMemorySpace()))
    return false;
  if (auto asyncTMACopyLocalToGlobalOp =
          dyn_cast<triton::nvidia_gpu::AsyncTMACopyLocalToGlobalOp>(op)) {
    return value == asyncTMACopyLocalToGlobalOp.getSrc();
  }
  if (auto warpGroupDotOp =
          dyn_cast<triton::nvidia_gpu::WarpGroupDotOp>(op)) {
    return value == warpGroupDotOp.getA() || value == warpGroupDotOp.getB();
  }
  // ... 각 연산별 소스 매칭
}

또한 연산 목록 자체도 수정되었다. AsyncTMACopyGlobalToLocalOp(글로벌→로컬)을 제거하고 AsyncTMACopyLocalToGlobalOp(로컬→글로벌), TMEMCopyOp를 추가했다.

bool isAsyncProxyRead(Operation *op) {
  return isa<
      triton::nvidia_gpu::WarpGroupDotOp,
      triton::nvidia_gpu::TCGen5MMAOp,
      triton::nvidia_gpu::TCGen5MMAScaledOp,
      triton::nvidia_gpu::TMEMCopyOp,            // 새로 추가
      triton::nvidia_gpu::AsyncTMACopyLocalToGlobalOp, // 방향 수정
      triton::nvidia_gpu::AsyncTMAScatterOp,
      triton::nvidia_gpu::AsyncTMAReduceOp>(op);
}

왜 이게 좋은가

  1. 정합성 버그 수정: shared memory에 동기적으로 쓴 값을 비동기 글로벌 복사가 올바르게 읽으려면 fence가 필수다. 이 패치 없이는 stale 데이터를 읽을 수 있다.
  2. 세밀한 의존성 추적: 단순히 "이 연산은 비동기 읽기다"가 아닌, "이 연산의 이 operand가 shared memory 소스다"까지 추적하여 불필요한 fence 삽입을 방지한다.
  3. 확장성: 새로운 비동기 연산이 추가될 때 isAsyncProxyReadSource에 케이스만 추가하면 되는 구조다.

정리

이 PR은 Triton NVIDIA 백엔드의 fence 삽입 패스에서 비동기 읽기 연산에 대한 의존성 추적이 불완전했던 문제를 수정한다. st.shared → copy_local_to_global 경로에서 누락되던 fence를 올바르게 삽입하도록 연산별 소스 operand 매칭 로직을 추가했다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글