[triton] Triton, TMA 전송 시 L2 캐시 정책 지원 추가: 성능 최적화 분석

PR 링크: triton-lang/triton#12000 상태: Merged | 변경: +845 / -86

들어가며

최신 GPU 아키텍처에서는 메모리 계층 구조의 효율적인 활용이 성능의 핵심입니다. 특히, Tensor Memory Accelerator (TMA)는 GPU의 로컬 메모리와 전역 메모리 간의 데이터 전송을 가속하는 중요한 메커니즘입니다. 하지만 기존 Triton의 TMA 연산은 L2 캐시 정책을 명시적으로 제어할 수 없어, 잠재적인 성능 최적화 기회를 놓치고 있었습니다. 이번 PR은 Triton의 TMA 연산에 L2 캐시 정책 지원을 추가하여, GPU의 캐시 계층 구조를 더욱 정교하게 활용하고 전반적인 메모리 접근 성능을 향상시키는 것을 목표로 합니다.

이 PR은 async_load, async_load_im2col, async_store, async_gather, async_scatter와 같은 핵심 TMA API에 cache_policy 인자를 도입합니다. 이를 통해 개발자는 L2 캐시의 동작 방식을 명시적으로 제어할 수 있게 되며, 이는 특히 데이터 재사용성이 높은 워크로드에서 상당한 성능 향상을 가져올 수 있습니다.

코드 분석

이번 PR의 핵심 변경 사항은 Triton IR(Intermediate Representation)의 TMA 관련 연산 정의와 이를 지원하는 C++ 구현에 집중되어 있습니다. 주요 변경 사항을 파일별로 살펴보겠습니다.

1. include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUOps.td

이 파일은 Triton의 NVIDIA GPU 관련 연산들을 정의하는 TableGen 파일입니다. TMA 연산들의 속성(attribute) 정의가 수정되어 새로운 cachePolicy 인자를 지원하도록 변경되었습니다.

주요 변경점:

• async_tma_copy_global_to_local 연산:

  • 기존의 cacheModifier와 evictionPolicy 속성이 제거되었습니다.
  • OptionalAttr<AnyAttr>:$cachePolicy 속성이 새로 추가되었습니다. 이 속성은 tt.cache_policy 또는 NVIDIA 고유의 ttng.cache_policy 객체를 받을 수 있으며, L2 캐시 정책에 집중합니다.
  • 설명 주석이 업데이트되어 새로운 cachePolicy 속성의 동작 방식과 제약 사항(L1 정책, 캐시 수정자, 사전 로드 크기 등 지원 불가)을 명확히 했습니다.

  Before:

  -    DefaultValuedAttr<TT_CacheModifierAttr, 
  

참고 자료

• https://docs.nvidia.com/cuda/parallel-thread-execution/#tensor-tiled-mode

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