[Triton] AMD gfx950/gfx1250에 AsyncCopy 기본 활성화 — 파이프라인 성능 향상

PR 링크: triton-lang/triton#9087 상태: Merged | 변경: +18 / -11

들어가며

GPU에서 행렬 곱셈 커널의 성능은 메모리 로딩과 연산의 파이프라인 병렬화에 크게 좌우된다. Triton의 AMD 백엔드는 ttg.async_copy_global_to_local을 사용하여 글로벌 메모리에서 LDS(Local Data Share)로 비동기 복사를 수행할 수 있다.

이 PR은 gfx950(CDNA4)과 gfx1250 아키텍처에서 이 비동기 복사를 기본적으로 활성화한다. 기존에는 환경변수 TRITON_HIP_USE_ASYNC_COPY를 명시적으로 설정해야 했다.

핵심 코드 분석

Before: 환경변수로만 제어

# knobs.py
class amd_knobs(base_knobs):
    use_async_copy: env_bool = env_bool("TRITON_HIP_USE_ASYNC_COPY")

# compiler.py — 환경변수가 설정되지 않으면 False
use_async_copy = knobs.amd.use_async_copy

After: 아키텍처별 기본값 + 환경변수 오버라이드

# knobs.py — env_bool에서 env_opt_bool로 변경 (None 허용)
class amd_knobs(base_knobs):
    use_async_copy: env_opt_bool = env_opt_bool("TRITON_HIP_USE_ASYNC_COPY")

# compiler.py — 아키텍처 기반 기본값 함수 추가
def is_async_copy_enabled(arch):
    return (arch in ["gfx950", "gfx1250"]) \
        if knobs.amd.use_async_copy is None \
        else knobs.amd.use_async_copy

# 사용처
use_async_copy = is_async_copy_enabled(options.arch)

핵심은 env_bool(True/False)에서 env_opt_bool(True/False/None)로 타입을 변경한 것이다. 환경변수가 설정되지 않으면 None이 되어 아키텍처별 기본값이 적용되고, 명시적으로 설정하면 사용자 의도가 우선된다.

테스트 업데이트: shared memory 초과 설정 스킵

비동기 복사는 레지스터 버퍼 하나를 추가 LDS 버퍼로 대체하므로 LDS 사용량이 증가한다. 이에 따라 일부 큰 블록 설정이 shared memory를 초과하게 되어 테스트에 skip 조건이 추가되었다.

if (BLOCK_M == 256 or BLOCK_N == 256) and BLOCK_K == 256:
    pytest.skip("Config requires too much shared memory")

왜 이게 좋은가

1. 자동 최적화: 지원되는 아키텍처에서 별도 설정 없이 비동기 복사가 활성화되어 파이프라인 효율이 향상된다.
2. 하위 호환성: 환경변수로 여전히 오버라이드 가능하므로 기존 워크플로우가 깨지지 않는다.
3. 3단 논리(None/True/False): "사용자가 의견이 없으면 최적 기본값, 의견이 있으면 그대로" 패턴으로 유연성과 편의성을 모두 확보한다.

정리

이 PR은 AMD gfx950/gfx1250에서 비동기 메모리 복사를 기본 활성화한다. env_bool → env_opt_bool 변환으로 아키텍처별 기본값과 사용자 오버라이드를 모두 지원하는 3단 설정 패턴을 도입했다.

참고 자료

• AMD CDNA4 Architecture
• Triton AMD Backend Compiler

---

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