[triton] NVIDIA canSkipBarSync 복원으로 MoE 커널 18GBps 성능 향상

PR 링크: triton-lang/triton#9246 상태: Merged | 변경: +55 / -52

들어가며

GPU 커널에서 shared memory 접근 간의 동기화를 위한 barrier는 필수적이지만, 불필요한 barrier는 파이프라인을 지연시킵니다. Triton의 Membar 분석 패스에는 특정 조건에서 barrier를 안전하게 생략할 수 있는 canSkipBarSync 필터가 있었는데, Blackwell GPU 지원 과정에서 커널 구조에 대한 가정이 맞지 않아 비활성화되었습니다. 이 PR은 더 보수적인 조건으로 이 최적화를 복원합니다.

핵심 코드 분석

Before (비활성화된 최적화):

// 컴파일 시에만 적용, 실제 빌드에서는 비활성화
bool NVIDIA::canSkipBarSync(Operation *before, Operation *after,
                            Allocation *allocation) {
  if (isa<ttng::InitBarrierOp>(before) && isa<ttng::InitBarrierOp>(after))
    return true;
  if (isa<ttng::InvalBarrierOp>(before) && isa<ttng::InvalBarrierOp>(after))
    return true;
  // ... (실제 빌드에서 호출되지 않음)
}

After (보수적으로 재설계):

bool NVIDIA::canSkipBarSync(Operation *before, Operation *after,
                            Allocation *allocation) {
  // elect predicate 사용으로 단일 스레드 실행 보장
  // InitBarrier/InvalBarrier는 uniform 연산이므로 barrier 불필요
}

// InitBarrierOp에 elect predicate 적용
Value pred = getElectWarp0OrThread0(*targetInfo, b);
::mlir::triton::PTXBuilder ptxBuilder;
const std::string ptx = "@$0 mbarrier.init.shared::cta.b64 [$1], " +
                        std::to_string(op.getCount()) + ";";

핵심 변경은 InitBarrierOp과 InvalBarrierOp가 이제 elect predicate를 사용하여 warp 내 단일 스레드만 실행하도록 한 것입니다. SM90+ GPU에서는 elect 명령어로 균일(uniform) 실행을 ptxas에 알려줘 더 나은 코드를 생성합니다.

왜 이게 좋은가

이 PR의 가치는 실측 데이터 기반 최적화 복원에 있습니다. 단순히 이전 코드를 되살리는 것이 아니라, 문제의 원인(커널 구조 가정)을 분석하고 보수적인 조건으로 재설계했습니다. persistent bf16 x mxfp4 MoE 워크로드에서 18 GBps 개선이라는 구체적 수치를 제시하며, 실제 프로덕션 워크로드에서의 효과를 검증했습니다.

정리

• canSkipBarSync 필터를 보수적 조건으로 재설계 후 실제 컴파일에 재활성화
• elect predicate로 barrier init/inval의 uniform 실행 보장
• SM90+ GPU에서 더 나은 ptxas 코드 생성
• persistent bf16 x mxfp4 MoE에서 18 GBps 성능 향상

참고 자료

• triton-lang/triton#9246

이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.