[triton] Triton: Blackwell 아키텍처를 위한 TMEM Load-Reduce 연산 퓨전 최적화

PR 링크: triton-lang/triton#10551 상태: Merged | 변경: +495 / -22

들어가며

최신 NVIDIA Blackwell(sm103+) 아키텍처는 하드웨어 수준에서 메모리 로드와 동시에 연산을 수행할 수 있는 강력한 기능을 제공합니다. 특히 tcgen05.ld.red PTX 명령어는 Tensor Memory(TMEM)에서 데이터를 읽어오는 동시에 행 축(row reduction)에 대한 연산을 수행할 수 있습니다. 이번 Triton PR은 기존에 분리되어 있던 ttng.tmem_load와 tt.reduce 연산을 하나의 퓨전된 명령어로 통합하여, 메모리 대역폭 효율을 극대화하고 연산 지연 시간을 줄이는 최적화를 도입했습니다.

코드 분석

1. TargetFeatures.h 및 TensorMemoryUtils.h

먼저, Blackwell 아키텍처의 특정 컴퓨트 능력(Compute Capability)을 확인하고, 퓨전이 가능한 조건인지 검증하는 로직이 추가되었습니다.

// include/triton/Dialect/TritonNvidiaGPU/IR/TargetFeatures.h
bool supportLdRed() const {
  return computeCapability >= 103 && computeCapability / 10 != 12;
}

또한, supportsTMemLoadReduce 함수를 통해 레지스터 레이아웃이 퓨전 연산에 적합한지(packed format 여부, N 차원이 스레드 간에 분할되지 않았는지 등)를 사전에 검증합니다.

2. FuseTMemLoadReduce.cpp (신규 패스)

이 PR의 핵심은 FuseTMemLoadReducePattern 클래스입니다. 이 패스는 tt.reduce 연산자를 찾아 그 이전의 ttng.tmem_load를 추적한 뒤, 두 연산을 하나로 합칩니다.

// lib/Dialect/TritonNvidiaGPU/Transforms/FuseTMemLoadReduce.cpp
class FuseTMemLoadReducePattern : public OpRewritePattern<triton::ReduceOp> {
  // ... 매칭 로직 ...
  // ttng.tmem_load를 찾아 redOp 속성을 부여하여 tcgen05.ld.red 생성 유도
  tmemLoad.setRedOp(redOpKind);
}

이 과정에서 stripConvertLayout 함수를 사용하여 중간에 끼어든 ttg.convert_layout 연산들을 건너뛰고 실제 로드 연산자를 찾아내는 것이 핵심입니다.

왜 이게 좋은가

이 최적화는 메모리 로드와 연산을 단일 명령어로 묶음으로써 다음과 같은 이점을 제공합니다.

1. 명령어 오버헤드 감소: 별도의 로드 및 리덕션 명령어를 실행하는 대신, 하드웨어 유닛이 로드와 동시에 연산을 완료하므로 명령어 발행(issue) 횟수가 줄어듭니다.
2. 데이터 이동 최소화: 레지스터로 데이터를 가져온 뒤 다시 연산하는 과정에서 발생하는 데이터 이동을 줄여, 특히 대규모 행렬 연산에서 성능 향상을 가져옵니다.

실제 벤치마크 결과, fused-attention 워크로드에서 N_CTX가 8192 이상일 때 FP16 기준 최대 10% 이상의 성능 향상을 보였습니다. 이는 메모리 집약적인 어텐션 연산에서 하드웨어 가속기를 최대한 활용한 결과입니다.

교훈

컴파일러 최적화에서 '연산 퓨전(Operator Fusion)'은 단순히 연산을 합치는 것을 넘어, 하드웨어가 제공하는 특수 명령어(Specialized Instruction)를 타겟팅할 때 가장 큰 효과를 발휘합니다. 특히 GPU 아키텍처가 발전함에 따라 하드웨어 수준에서 지원하는 복합 연산(Load-Reduce, Load-Store 등)을 컴파일러가 자동으로 매핑해주는 것이 성능 최적화의 핵심임을 다시 한번 확인시켜 줍니다.

참고 자료

• Triton GitHub PR #5164
• NVIDIA PTX ISA Documentation

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