[triton] ConSan 컴파일 타임 19분에서 34초로 단축 - 대규모 최적화

PR 링크: triton-lang/triton#9366 상태: Merged | 변경: +1079 / -753

들어가며

Triton의 Concurrency Sanitizer(ConSan)는 GPU 커널의 동시성 버그를 감지하는 도구입니다. 그러나 attention forward 커널 컴파일 시 cubin 생성에 19분이 소요되는 등 컴파일 타임이 비현실적이었습니다. 이 PR은 여러 최적화를 통해 컴파일 시간을 34초로 단축합니다.

핵심 코드 분석

1. Warp-local layout으로 IR 크기 32배 축소

기존에는 thread-local layout(모든 스레드가 텐서의 복사본 보유)을 사용했지만, warp-local layout(각 warp 내에서 분산)으로 전환하여 생성 IR 크기를 32배 줄였습니다.

2. createMultiColumnMask 재작성

Before: make_range + masking으로 컴퓨팅 (기가바이트 단위 IR 생성)

After:

// constant tensor로 직접 생성

이 단일 함수가 기가바이트 단위의 IR을 생성하고 있었습니다. 상수 텐서로 직접 대체하여 IR 크기를 극적으로 줄였습니다.

3. Global 상수 캐싱

// assert 문자열 등 동일한 상수를 여러 번 생성하지 않도록 캐싱

4. LLVM IR Canonicalize 패스

def CanonicalizeLLVMIR : Pass<"canonicalize-llvm-ir", "mlir::LLVM::LLVMFuncOp"> {
  let summary = "Canonicalize LLVM IR";
}

LLVM dialect에 누락된 select %false|%true, %a, %b 정규화 패턴을 추가하여 불필요한 select 연산을 제거합니다.

왜 이게 좋은가

컴파일 시간 결과:

• llir: 99.6초 -> 17.0초 (5.9배)
• llvmir: 27.7초 -> 3.7초 (7.5배)
• cubin: 19.15분 -> 34.36초 (33.4배)

이 개선으로 ConSan이 개발 워크플로우에서 실제로 사용 가능한 도구가 됩니다. 특히 warp-local layout 전환은 레지스터 사용량도 줄여 런타임 성능에도 긍정적 영향을 미칩니다.

정리

• Thread-local에서 warp-local layout으로 전환 (IR 32배 축소)
• createMultiColumnMask를 constant tensor로 재작성
• Global 상수 캐싱으로 중복 문자열 제거
• CanonicalizeLLVMIR 패스 추가
• 헬퍼 함수 중복제거로 코드 크기 감소
• 전체 컴파일 타임 19분 -> 34초

참고 자료

• triton-lang/triton#9366

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