[triton] AMD GFX950에서 Padded Layout Async Copy의 OOM 버그 수정

PR 링크: triton-lang/triton#9442 상태: Merged | 변경: +37 / -0

들어가며

AMD GFX950(CDNA4)에서는 async copy를 위해 padded shared memory layout을 사용하여 LDS bank conflict를 줄입니다. 그러나 wrap 값이 0이 되는 작은 타일 크기(예: 16x16)에서 pipelining이 잘못된 레이아웃을 생성하여 OOM(Out Of Memory) 크래시가 발생했습니다.

핵심 코드 분석

Before:

unsigned wrap = std::min(contigDim, 128u) / padding;
unsigned requiredDim = warpSize / contigLanes * wrap;
if (nonContigDim < requiredDim) {
    return {};
}
// wrap == 0인 경우 division by zero 또는 잘못된 레이아웃 생성

After:

unsigned wrap = std::min(contigDim, 128u) / padding;
// wrap == 0 means padding > contigDim, which is not a valid configuration
if (wrap == 0) {
    return {};
}
unsigned requiredDim = warpSize / contigLanes * wrap;
if (nonContigDim < requiredDim) {
    return {};
}

이 수정과 함께 16x16 타일의 pipelining을 검증하는 테스트도 추가되었습니다:

// CHECK-NOT: ttg.padded_shared
tt.func @pipeline_padded_layout_gfx950(
    %arg0: !tt.ptr<f16>, %arg1: !tt.ptr<f16>) {
    // 16x16 MFMA with padded layout async copy
    // CHECK: ttg.async_wait

왜 이게 좋은가

padding > contigDim인 경우는 padding을 적용할 수 없는 무효한 설정입니다. 기존에는 이 경우가 처리되지 않아 wrap = 0이 되고, 이후 연산에서 비정상적인 레이아웃이 생성되어 과도한 shared memory를 요청하게 되었습니다. 단 3줄의 early return으로 이 문제를 해결하며, padded layout을 사용하지 않는 fallback 경로가 자동으로 적용됩니다.

정리

Padded layout 생성에서 wrap == 0(padding > contigDim)인 경우를 조기 반환하여 작은 타일 크기에서의 OOM 크래시를 방지하고, 해당 경우의 테스트를 추가했습니다.

참고 자료

• triton-lang/triton#9442

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