본문으로 건너뛰기

[triton] Triton AMD 백엔드 최적화: Subtiling을 통한 GEMM 성능 향상

PR 링크: triton-lang/triton#9052 상태: Merged | 변경: +404 / -9

들어가며

Triton의 AMD 백엔드(gfx1250)에서 대규모 행렬 곱셈(GEMM) 연산을 수행할 때, 기존의 Warp-Specialized 커널은 공유 메모리(Shared Memory) 사용량과 레지스터 스필(Register Spills) 문제로 인해 블록 크기 확장에 한계가 있었습니다. 이번 PR은 'Subtiling' 기법을 도입하여 연산 단위를 더 작은 조각으로 나누어 처리함으로써, 메모리 효율성을 극대화하고 더 큰 블록 크기를 지원하도록 개선했습니다.

코드 분석

1. PersistentPartitionSubtiledArgs 도입

기존의 고정된 파라미터 전달 방식에서 벗어나, Subtiling을 위해 필요한 다양한 메타데이터를 관리하는 PersistentPartitionSubtiledArgs 클래스를 추가했습니다.

@aggregate
class PersistentPartitionSubtiledArgs:
    a_desc: ttgl.amd.gfx1250.tdm.tensor_descriptor
    # ... (생략) ...
    QUADRANT_M: ttgl.constexpr
    QUADRANT_N: ttgl.constexpr

이 구조체는 연산의 각 단계(Producer, Compute, Epilogue)에서 필요한 텐서 디스크립터와 버퍼 정보를 일관되게 전달하여 코드의 가독성과 유지보수성을 높였습니다.

2. Accumulator Subtiling 구현

가장 핵심적인 변경 사항은 split_accumulator_quadrant 함수입니다. 레지스터 압박을 줄이기 위해 누적기(Accumulator)를 4개의 서브타일로 분할합니다.

@gluon.jit
def split_accumulator_quadrant(acc):
    # Reshape [BLOCK_M, BLOCK_N] -> [2, SUBTILE_M, 2, SUBTILE_N]
    acc_4d = acc.reshape([2, SUBTILE_M, 2, SUBTILE_N])
    # Permute to [SUBTILE_M, SUBTILE_N, 2, 2] so split dimensions are at the end
    acc_4d = acc_4d.permute(1, 3, 0, 2)
    # Split along last dimension -> four tensors of [SUBTILE_M, SUBTILE_N]
    acc_00, acc_10 = acc_n0.split()
    acc_01, acc_11 = acc_n1.split()
    return acc_00, acc_01, acc_10, acc_11

이 과정은 큰 행렬을 한 번에 레지스터에 올리는 대신, 작게 쪼개어 처리함으로써 하드웨어의 레지스터 할당 효율을 극대화합니다.

왜 이게 좋은가

  1. 레지스터 스필 제거: 기존에는 큰 블록을 처리할 때 레지스터가 부족하여 메모리로 데이터가 넘치는(Spill) 현상이 발생했습니다. Subtiling을 통해 활성 레지스터 수를 제어함으로써 이를 해결했습니다.
  2. 공유 메모리 최적화: 연산 단위가 작아짐에 따라 공유 메모리 점유율이 최적화되어, 더 많은 워프(Warp)가 동시에 실행될 수 있는 환경을 조성합니다.
  3. 확장성: (128, 128, 256) 및 (256, 256, 128)과 같은 더 큰 블록 크기를 지원하게 되어, 대규모 연산에서 처리량(Throughput)이 크게 향상되었습니다.

일반적으로 GPU 커널 최적화에서 '공유 메모리 사용량 감소'와 '레지스터 압박 완화'는 성능의 병목을 해결하는 가장 강력한 수단입니다. 이번 변경은 Triton의 추상화 계층을 활용하여 복잡한 하드웨어 제어를 효율적으로 수행한 좋은 사례입니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글