TLX로 GPU 커널 최적화하기: Triton 확장으로 워프 특화·파이프라이닝·공유메모리 직접 제어

핵심 메시지

• Triton은 연산을 타일 단위로 기술하게 해 저수준 복잡도를 감췄지만, 최신 GPU에서는 “누가 언제 어떤 일을 하는가”라는 스케줄 결정이 성능을 좌우하면서 컴파일러 휴리스틱만으로는 한계에 부딪힌다.
• TLX는 Triton을 대체하지 않고 확장한다. 스레드 단위 레이아웃은 여전히 감추되, 파이프라이닝·워프 특화·공유메모리 배치 같은 스케줄 결정을 사용자에게 명시적으로 넘긴다.
• 레이어 정규화 커널에서는 공유메모리에 입력을 한 번만 올려 두고 재사용하고, 클러스터 내 여러 CTA에 걸친 분산 공유메모리로 감축 연산을 나눠 레지스터 압박을 덜어낸다.
• GEMM 커널에서는 워프 특화 기반 파이프라이닝만으로 성능이 거의 두 배가 되고, 여기에 퍼시스턴트 루프·2-CTA 모드·데이터 파티셔닝을 얹어 텐서 코어의 빈 구간을 지워 나간다.
• 측정 결과는 BF16 GEMM이 cuBLAS에 근접하거나 일부 형태에서 소폭 앞서고, 플래시 어텐션은 계산 집약적 형태에서 최신 구현 대비 5~10% 차이가 남는 수준이다.

쉽게 이해하기

Triton은 강력한 전제 위에 서 있다. 연산을 타일 단위로 기술하기만 하면 메모리 병합이나 벡터화 같은 저수준 문제는 컴파일러가 알아서 처리해 준다는 것이다. 이 모델은 오랫동안 잘 작동했지만, 호퍼와 블랙웰 같은 최신 GPU에서는 사정이 달라졌다. 이제 성능은 무엇을 계산하느냐보다 누가 그 계산을 맡고, 언제 실행하며, 데이터가 파이프라인을 어떻게 흘러가느냐에 좌우된다. 이런 결정은 Triton에서 암묵적으로 남아 컴파일러 휴리스틱이 추론해야 하는데, GEMM처럼 단순한 커널에서는 잘 통해도 어텐션 커널처럼 복잡해지면 쉽게 어긋난다.

메타가 공개한 TLX(Triton low-level extension)는 여기서 출발한다. CUDA로 되돌아가지 않고 Triton을 확장해 최고 성능에 도달하자는 것이다. 사용자는 여전히 타일 단위로 사고한다. 스레드 레이아웃과 매핑은 그대로 감춰져 있고 동기화도 스레드가 아니라 타일 사이에서만 신경 쓰면 된다. 달라지는 것은 그 타일들이 GPU에서 어떤 순서로 실행되고 파이프라인을 통해 어떻게 제어되는가를 직접 기술한다는 점이다. 공유메모리도 직접 관리하지만 버퍼마다 정확한 오프셋을 지정할 필요는 없다. 어떤 버퍼가 겹쳐도 되고 어떤 버퍼가 안 되는지만 기술하면 실제 배치는 컴파일러가 맡는다.

첫 사례 연구는 레이어 정규화다. 파이토치 인덕터가 만들어 주는 Triton 커널은 평균, 분산, 정규화 세 단계에서 입력을 세 번 다시 읽는다. 한 행이 레지스터에 다 들어가지 않기 때문인데, 결국 L2 캐시에 기대는 셈이라 지연 시간이 남는다. TLX에서는 입력을 공유메모리에 한 번 올려 두고 이후 단계에서 거기서 레지스터로 되읽는다. 여기에 클러스터로 묶인 여러 CTA가 각자 부분 감축을 계산한 뒤 서로의 공유메모리 슬롯에 값을 비동기로 써 넣고 배리어로 도착을 확인하는 분산 공유메모리 방식을 더하면, 한 CTA가 감당해야 할 공유메모리와 레지스터가 함께 줄어든다. 마지막으로 전역 메모리에서 레지스터를 거치지 않고 곧장 공유메모리로 복사하는 비동기 적재까지 쓰면, 이 커널은 기존의 잘 알려진 고성능 구현보다 나은 성능을 냈다.

두 번째 사례는 GEMM이다. 아무 최적화도 없는 상태에서는 데이터 적재와 행렬 곱이 완전히 순차적으로 실행돼 겹치는 구간이 전혀 없다. 워프 하나를 적재 전담으로, 다른 워프를 행렬 곱 전담으로, 또 다른 워프를 에필로그 전담으로 나누고 배리어로 신호를 주고받게 하면 이 구간들이 겹치기 시작한다. 발표자는 이 파이프라이닝만으로 성능이 거의 두 배가 됐다며 가장 중요한 기법으로 꼽았다. 그다음은 하나의 CTA가 여러 타일을 연달아 처리하는 퍼시스턴트 루프다. 타일 사이 경계에서 생기던 텐서 코어의 빈 구간이 메워지고, 에필로그도 다음 타일의 연산과 겹치게 된다.

블랙웰에서만 가능한 기법도 소개됐다. 2-CTA 모드는 두 CTA가 한 쌍으로 묶여 B 피연산자를 절반씩만 보관하게 해 공유메모리를 절약하고, 그렇게 아낀 공간으로 타일 크기를 키운다. 다만 타일이 커지면 이번엔 텐서 메모리가 부족해져 파이프라인 깊이가 얕아지므로, 타일을 다시 잘게 나눠 두 서브타일을 핑퐁 방식으로 번갈아 처리하는 데이터 파티셔닝으로 균형을 맞춘다. 여기에 클러스터 실행 제어를 쓰면 각 SM이 아직 시작되지 않은 작업을 실시간으로 가져가는 동적 스케줄링이 가능해져, 다른 작업이 SM 일부를 점유하는 상황에서도 부하가 고르게 분산된다. 이 모든 항목은 튜닝 가능한 설정으로 노출돼 있고, 발표자들은 메타의 Triton 브랜치에서 언어 명세와 튜토리얼을 함께 공개하고 있다고 밝혔다.

주요 인사이트

• 추상화의 층위를 통째로 낮추는 것과 특정 결정만 사용자에게 돌려주는 것은 다르다. TLX는 스레드 레이아웃과 스위즐링은 계속 감춘 채 스케줄 결정만 노출해 Triton의 생산성을 지키려 한다.
• 같은 문제를 겨냥한 도구라도 성격이 갈린다. 공유메모리 스위즐링까지 손수 조율해 뱅크 충돌을 최소화하려면 더 낮은 층위의 도구가 낫고, 컴파일러의 추론에 맡기고 싶다면 TLX 쪽이 편하다는 것이 발표자의 설명이다.
• cuBLAS 같은 고도로 튜닝된 라이브러리는 후속 원소별 연산을 융합할 통로가 없다. TLX로 작성한 GEMM을 인덕터의 템플릿으로 넣으면 편향이나 양자화가 섞인 연산까지 하나의 커널로 융합할 수 있다.
• 레지스터가 모자라면 컴파일러는 전역 메모리로 값을 밀어낸다. 그 스필을 공유메모리로 직접 옮기는 것만으로도 커널 성능이 달라지며, TLX는 그 판단을 수동으로 내릴 수 있게 한다.
• 성능 차이의 원인을 컴파일러 탓으로 돌리기는 어렵다. 발표자는 비교 대상 구현이 느렸던 이유를 묻는 질문에 컴파일러보다 프로그래밍 방식의 차이일 가능성이 높다고 답했고, 그 모호함은 TLX에도 똑같이 적용된다고 인정했다.

자주 묻는 질문

관련 AI 소식