[triton] Triton AMD GPU 최적화: Warp-based Split-K 도입을 통한 MQA 성능 향상

PR 링크: triton-lang/triton#11178 상태: Merged | 변경: +388 / -306

들어가며

최신 LLM 추론 환경에서 Multi-Query Attention(MQA)은 메모리 대역폭을 절약하며 성능을 높이는 핵심 기술입니다. 하지만 시퀀스 길이가 짧은 경우, 기존의 CTA(Cooperative Thread Array)-based Split-K 방식은 오버헤드가 발생할 수 있습니다. 본 PR은 AMD gfx1250 아키텍처에서 MQA 디코딩 시, 각 Warp가 K 파티션을 소유하고 공유 메모리를 통해 최종 리덕션을 수행하는 'Warp-based Split-K'를 도입하여 성능을 개선합니다.

코드 분석

1. AttentionConfigBase의 유연성 확보

기존에는 SPLIT_K 모드가 고정적이었으나, 이제 SPLIT_K_MODE를 도입하여 'cta', 'warp', 'none' 중 선택할 수 있게 변경되었습니다.

Before:

self.SPLIT_K = ttgl.constexpr(SPLIT_K)
# ... (생략)
if SPLIT_K == 1:
    # ...
else:
    assert NUM_CTAS == SPLIT_K
    # ... (CTA 기반 고정)

After:

assert SPLIT_K_MODE in ['cta', 'warp', 'none']
self.SPLIT_K_MODE = ttgl.constexpr(SPLIT_K_MODE)
# ...
if SPLIT_K_MODE == 'warp':
    assert NUM_CTAS == 1
    assert SPLIT_K == NUM_WARPS
    # Warp 기반 분할 로직 수행
else:
    assert SPLIT_K_MODE == 'cta'
    # 기존 CTA 기반 로직 유지

2. 레이아웃 재구성

get_wmma_layout 호출 시 preshuffled 인자를 추가하고, SPLIT_K_MODE에 따라 wmma_layout을 동적으로 생성하여 하드웨어 가속기(WMMA)의 활용도를 높였습니다.

왜 이게 좋은가

1. 오버헤드 감소: 기존 CTA-based 방식은 여러 CTA 간의 동기화가 필요하여 시퀀스 길이가 짧을 때 오버헤드가 컸습니다. Warp-based 방식을 사용하면 단일 CTA 내에서 Warp 간 리덕션이 이루어지므로, 커널 실행 오버헤드가 대폭 감소합니다.
2. 하드웨어 활용 최적화: AMD gfx1250 아키텍처의 특성에 맞춰 Warp 단위로 작업을 분할함으로써, 레지스터와 공유 메모리 활용 효율을 극대화했습니다.
3. 교훈: GPU 커널 최적화 시, 작업의 크기(시퀀스 길이 등)에 따라 병렬화 단위를 CTA에서 Warp로 세분화하는 전략은 성능 향상의 핵심적인 레버가 됩니다.

결론

이번 변경은 특히 8k 이하의 시퀀스 길이를 가진 MQA 디코딩 작업에서 유의미한 성능 이득을 제공합니다. 하드웨어의 특성을 고려한 세밀한 스케줄링이 어떻게 실질적인 성능 개선으로 이어지는지 보여주는 좋은 사례입니다.

참고 자료

• https://triton-lang.org/main/index.html
• https://rocm.docs.amd.com/

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