[triton] Triton FPSAN의 MMA 에뮬레이션 오버헤드 최적화 분석

PR 링크: triton-lang/triton#11159 상태: Merged | 변경: +40 / -54

들어가며

Triton의 FpSanitizer(FPSAN)는 부동소수점 연산의 정확성을 검증하기 위한 도구입니다. 하지만 기존 구현에서는 tcgen05.mma 연산 시 A와 D(결과) 텐서가 앨리어싱(aliasing)될 가능성을 가정하여, 불필요한 스크래치 메모리 라운드트립(roundtrip)을 수행하고 있었습니다. 이는 코드 생성(codegen) 오버헤드를 증가시키고 런타임 성능을 저하시키는 주요 원인이었습니다. 본 PR은 이 가정이 잘못되었음을 인지하고, 불필요한 복사를 제거하여 성능을 최적화했습니다.

코드 분석

1. FpSanitizer.cpp: 불필요한 스크래치 할당 제거

기존에는 createTmemOperandScratch 함수를 통해 모든 TMEM(Tensor Memory) 피연산자에 대해 별도의 스크래치 메모리를 할당하고 복사했습니다. PR에서는 이를 제거하고 createMmaOperandSource에서 직접 스크래치 정보를 가져오도록 변경했습니다.

Before:

auto info = createTmemOperandScratch(rewriter, loc, scratch, memdesc, memTy, scope);
if (!info) return std::nullopt;
return MmaOperandSource{info->ptr, Value(), tileTy, rowStride, stride};

After:

auto info = scratch.getOrCreate(memdesc, rewriter, scope);
if (!info || info->scaleSourceType) return std::nullopt;
return MmaOperandSource{info->ptr, Value(), tileTy, rowStride, info->tensorType.getShape().front()};

이 변경으로 인해 ttg.global_scratch_alloc 호출이 줄어들며, 불필요한 tt.store 및 tt.load 연산이 제거되었습니다.

2. TensorMemoryUtils.cpp: 레지스터 예산 최적화

getContextualMaxNReg 함수에 requestedRegisters를 확인하는 로직이 추가되어, 더 정확한 레지스터 예산을 계산할 수 있게 되었습니다.

if (auto requestedRegisters = partitions.getParentOp().getRequestedRegisters())
  return (*requestedRegisters)[idx];

왜 이게 좋은가

이번 최적화는 크게 두 가지 측면에서 이점을 제공합니다.

1. 성능 향상: 불필요한 메모리 복사를 제거함으로써 Attention FPSAN 컴파일 및 런타임 속도가 15-25% 향상되었습니다. GPU 연산에서 메모리 대역폭과 복사 오버헤드는 병목의 주범이므로, 이를 제거한 것은 매우 효과적인 최적화입니다.
2. 코드 생성 최적화: 불필요한 스크래치 할당이 사라지면서 생성되는 IR의 양이 줄어들고, 컴파일러가 더 효율적인 레지스터 할당을 수행할 수 있게 되었습니다.

교훈: 컴파일러 최적화에서 '보수적인 가정(Conservative assumption)'은 안전성을 보장하지만, 때로는 성능을 심각하게 저해합니다. 앨리어싱이 발생하지 않는다는 도메인 지식을 활용하여 가정을 완화하는 것이 성능 개선의 핵심임을 보여줍니다.

참고 자료

• Triton GitHub Repository
• MLIR PatternRewriter Documentation

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