SAIL을 위한 Triton

Triton의 언어 기능, 컴파일러 설계 철학, 일반적인 설치 및 디버깅 팁에 대해서는 upstream Triton README와 공식 Triton 문서를 참조하십시오. 이 문서는 upstream Triton 대비 PPU 백엔드(backend)의 확장 및 최적화에 초점을 맞춥니다.

1. 개요 (Overview)

이 저장소는 T-Head Semiconductor가 독립적으로 개발한 PPU를 지원하기 위한 Triton의 **PPU 포크 (fork)**입니다.

• 대상 하드웨어 (Target hardware): 두 세대의 Tensor Core를 포함하는 T-Head PPU — PPU0010 (MMAv1) 및 PPU0015 (MMAv2).
• 런타임 (Runtime): T-Head SAIL SDK (PPU SDK) 환경이 필요합니다.
• 프로그래밍 인터페이스 (Programming interface): upstream Triton과 정확히 동일한 Python 프로그래밍 인터페이스를 사용하여 커널(kernel)을 작성합니다. PPU 백엔드(backend)가 PPU 하드웨어를 대상으로 하는 컴파일 및 실행을 투명하게 처리합니다.

2. 핵심 기능 및 최적화 (Core Features and Optimizations)

PPU 백엔드(backend)는 Triton의 표준 다단계 컴파일 흐름인 ttir → ttgir → llir → hgbin을 따르며, 각 단계에서 PPU 전용 패스(pass)를 삽입합니다. 핵심 구현은 third_party/ppu/backend/compiler.py에 위치합니다.

PPU에서의 주요 최적화:

• AIU 비동기 데이터 이동 (AIU asynchronous data movement): 글로벌 메모리(global memory)에서 공유 메모리(shared memory, TSM)로의 데이터 로드는 AIU (AI accelerator in compute Unit)에 의해 수행되는 비동기 복사로 자동 변환됩니다. 소프트웨어 파이프라이닝(software pipelining)과 결합하여, 루프 내부에서 멀티 버퍼링(multi-buffering)을 구현하며, 타일 프리페치(tile prefetch)와 MMA 연산을 중첩시켜 글로벌 메모리 액세스 지연 시간(latency)을 숨기고 연산 효율성을 향상시킵니다.

• Swizzled shared memory layout (스위즐드 공유 메모리 레이아웃): Triton은 워프(warp)의 수, 타일 형태(tile shape), 그리고 요소 비트 너비(element bit width)로부터 타일링 스킴(tiling scheme)과 스위즐 인코딩(swizzle encoding)을 자동으로 도출합니다. 이는 한편으로 공유 메모리 뱅크 충돌(bank conflicts)을 제거하고 다중 워프 동시 액세스 시 유효 대역폭(effective bandwidth)을 보장하며, 다른 한편으로는 TSM 내의 데이터 레이아웃을 MMA 피연산자(operand) 레이아웃과 일치시켜 추가적인 레이아웃 변환 오버헤드를 줄여줍니다.

• Tensor Core 가속 및 저정밀도 지원: Triton은 tl.dot을 PPU의 MMA 하드웨어 명령어로 컴파일하며, Tensor Core 가속을 최대한 활용하기 위해 타일 파티셔닝 입도(tile partitioning granularity)를 자동으로 도출합니다. 또한 FP8 (E5M2 / E4M3), FP16, BF16을 포함한 저정밀도(low-precision) / 혼합 정밀도(mixed-precision) 지원을 제공합니다.

3. AIU 사용 가이드

PPU Triton은 데이터 전송을 가속화하기 위해 AIU(AI accelerator in compute Unit)를 통해 글로벌 메모리에서 공유 메모리로 데이터를 이동시키는 새로운 aiu_load API로 Triton 언어를 확장합니다. 현재 tl.aiu_load, make_block_ptr, make_tensor_descriptor 인터페이스를 제공하며, 사용자는 이를 기존 Triton 커널에 쉽게 통합할 수 있습니다.

3.1 tl.aiu_load

def aiu_load(pointer, offsets=(0, 0), block_shape=(0, 0), shape=(0, 0), dtype=void, order=(1, 0), _semantic=None):
    """
    `pointer`로 정의된 위치의 메모리에서 로드된 값들을 가진 데이터 텐서를 반환합니다:
...

M x K 글로벌 데이터 텐서로부터 BLOCK_M x BLOCK_K 텐서를 공유 메모리로 로드해야 할 때 tl.aiu_load를 사용하십시오.

aiu_load 제약 사항:

• 텐서는 32바이트 정렬(32-byte aligned)되어야 합니다.
• BLOCK_M은 16의 배수여야 합니다.
• BLOCK_K 차원을 따른 데이터는 32바이트의 배수여야 합니다.
• 블록 텐서는 M 및 K 차원 모두에서 연속적(contiguous)이어야 합니다.
• 지원되는 데이터 타입:
  • PPU0010: b16
  • PPU0015: b16, b8

다음은 tl.aiu_load를 사용하는 matmul 커널의 예시입니다:

@triton.jit
def matmul_kernel_aiu(a_ptr, b_ptr, c_ptr,
                      stride_am, stride_ak,
...

3.2 make_block_ptr

make_block_ptr를 사용하면 aiu_load가 더 간단해집니다. 단, make_block_ptr / aiu_load / advance를 함께 사용해야 합니다.

def make_block_ptr(base: tensor, shape, strides, offsets, block_shape, order, _semantic=None):
    """
    부모 텐서 내 블록에 대한 포인터를 반환합니다.
...

• tl.make_block_ptr: 부모 텐서 내의 블록에 대한 포인터를 생성합니다.
• tl.aiu_load: make_block_ptr의 결과를 입력으로 받습니다.
• tl.advance: 블록 포인터 (block pointer)를 전진시킵니다.

def advance(base, offsets, _semantic=None):
    """
    블록 포인터를 전진시킵니다.
...

예시:

@triton.jit
def matmul_kernel_aiu(a_ptr, b_ptr, c_ptr,
                      stride_am, stride_ak,
...

3.3 make_tensor_descriptor

PPU Triton은 tl.make_tensor_descriptor에 대한 AIU 지원을 추가했습니다. 조건이 충족되면 로드 연산 (load op)이 자동으로 AIU 로드로 승격 (promote)될 수 있습니다.

def make_tensor_descriptor(
    base: tensor,
    shape: List[tensor],
...

Triton 커널에서 tl.make_tensor_descriptor를 사용할 때, 조건이 충족되면 로드 (load) 연산이 자동으로 AIU 로드로 승격될 수 있습니다.

3.4 추가 예시

더 많은 AIU 사용 사례는 저장소(repository)의 테스트 및 성능 예시를 참조하십시오:

• python/test/unit/ppu/aiu/ — aiu_load, block pointer (블록 포인터), tensor descriptor (텐서 디스크립터), dot (행렬 곱), addmm 및 다양한 데이터 타입 (FP16, FP8)과 순서 조합을 다루는 AIU 기능 테스트 (functional tests).
• python/test/unit/ppu/perf/ — 행렬 곱셈 (03-matrix-multiplication-aiu.py, 03-matrix-multiplication-mxfp4-aiu.py) 및 융합 어텐션 (fused attention, 06-fused-attention-aiu.py)과 같은 AIU 기반 성능 예시.

4. 빌드 및 설치 (Build and Installation)

4.1 필수 요구 사항 (Prerequisites)

• PPU SDK: PPU SDK를 설치하고 PPU_SDK 환경 변수를 통해 해당 경로를 지정하십시오 (기본값은 /usr/local/PPU_SDK입니다). SDK는 헤더, 라이브러리, 그리고 ppu-llc 및 llvm-irformatter와 같은 툴체인 실행 파일을 제공합니다.
• 런타임 드라이버 (Runtime driver): PPU 런타임 드라이버 libhggc.so를 설치하고, ldconfig 또는 LD_LIBRARY_PATH를 통해 찾을 수 있는지 확인하십시오.
• 빌드 의존성 (Build dependencies): Python 3.10 이상, C++ 컴파일러, CMake, Ninja 등 (upstream Triton README 참조).

export PPU_SDK=/usr/local/PPU_SDK           # 실제 설치 경로에 맞게 조정하십시오
export LD_LIBRARY_PATH=$PPU_SDK/lib:$LD_LIBRARY_PATH

4.2 소스에서 설치 (Install from source)

# 저장소 루트에서
pip install -r python/requirements.txt      # 빌드 시점 의존성
pip install -e .                            # 편집 가능 모드 (editable mode)로 컴파일 및 설치

또는 가상 환경 (virtualenv) 사용 시:

python -m venv .venv --prompt triton
source .venv/bin/activate
pip install -r python/requirements.txt
...

빌드 옵션 (커스텀 LLVM, ccache, MAX_JOBS를 이용한 메모리 사용량 제한 등)은 upstream Triton과 동일합니다. upstream Triton README의 "Building with a custom LLVM" 및 "Tips for building" 섹션을 참조하십시오.

4.3 설치 확인 (Verify the installation)

소스 코드로부터 빌드한 후, python -c "import triton; print(triton.__version__)" 명령어를 통해 임포트(import) 및 버전을 확인하십시오. 버전 번호가 올바르게 출력된다면, Triton이 성공적으로 설치되었으며 C++/MLIR 확장 기능(extensions)이 올바르게 로드된 것입니다.

그 다음, 엔드 투 엔드(end-to-end) 검증을 위해 AIU 섹션의 예제들을 실행할 수 있습니다.

4.4 환경 변수 (Environment variables)

일반적인 환경 변수에 대해서는 upstream Triton README를 참조하십시오. 아래는 추가적인 PPU 전용 환경 변수입니다:

• PPU_LLC_OPTIONS: ppu-llc에 전달되는 추가 옵션입니다.
• DISABLE_PPU_LLC_OPT: ppu-llc 최적화(optimizations)를 비활성화합니다.
• TRITON_LIBDEVICE_PATH: PPU libdevice 경로를 지정합니다.
• TRITON_DUMP_COMPILE_LOG: 컴파일 로그(compilation log)를 내보냅니다.