본문으로 건너뛰기

[Triton] TMA multicast 지원 추가

들어가며

NVIDIA Hopper GPU의 TMA(Tensor Memory Accelerator)는 글로벌 메모리에서 공유 메모리로 비동기 복사를 수행하는 하드웨어 유닛이다. Multicast 기능을 사용하면 하나의 TMA 연산으로 여러 CTA(Cooperative Thread Array)의 공유 메모리에 동시에 데이터를 복사할 수 있어 대역폭을 절약한다. 이 PR은 TMA multicast를 Triton에서 사용할 수 있도록 기반 인프라를 추가한다.

핵심 코드 분석

Before (AsyncTMACopyGlobalToLocalOp)

Arg<TTG_MemDescType, "", [MemWrite<SharedMemory>]>:$barrier,
Arg<TTG_MemDescType, "", [MemWrite<SharedMemory>]>:$result,
I1:$pred,
DefaultValuedAttr<TT_CacheModifierAttr, "...">:$cache,

multicast 옵션이 없어 단일 CTA에만 데이터를 복사할 수 있었다.

After

Arg<TTG_MemDescType, "", [MemWrite<SharedMemory>]>:$barrier,
Arg<TTG_MemDescType, "", [MemWrite<SharedMemory>]>:$result,
I1:$pred,
UnitAttr:$multicast,
DefaultValuedAttr<TT_CacheModifierAttr, "...">:$cache,

UnitAttr:$multicast 플래그가 추가되어 multicast 모드를 활성화할 수 있다. 또한 barrier 메모리 할당이 CTA 수에 맞게 변경되었다:

// Before
gpu::MemDescType::get({1}, rewriter.getI64Type(), ...);

// After
gpu::MemDescType::get({numCTAs}, rewriter.getI64Type(), ...);

Gluon API 수준에서도 multicast 파라미터가 전달된다:

# Python binding
def create_async_tma_copy_global_to_local(self, descPtr, coord,
    barrier, result, pred, multicast):  # multicast 추가

왜 이게 좋은가

  • 하드웨어 활용: CGA(Cooperative Grid Array) 내 여러 CTA가 같은 데이터를 필요로 할 때 메모리 대역폭을 크게 절약한다.
  • barrier 크기 정합: barrier를 numCTAs 크기로 할당하여 multi-CTA 시나리오에서의 정확한 동기화를 보장한다.
  • 증분 접근: 이 PR은 독립 실행 가능한 기반만 추가하고, tcgen05.mma와의 루프 통합은 후속 PR로 분리했다.

정리

+263/-61 변경으로 TMA multicast의 IR 표현, barrier 관리, Python binding, 테스트를 모두 포함한다. Hopper 이상 GPU에서의 multi-CTA 커널 성능 최적화의 기반이 되는 중요한 PR이다.

참고 자료

이 글은 AI(Claude)의 도움을 받아 작성되었습니다. 코드 분석 내용은 실제 PR diff를 기반으로 합니다.

댓글