[triton] Triton 커널 최적화: 불필요한 텐서 메모리 할당 제거하기

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

들어가며

Triton은 고성능 GPU 커널을 작성하기 위한 강력한 도구입니다. 하지만 커널의 메타데이터를 계산하는 과정에서 예상치 못한 메모리 오버헤드가 발생할 수 있습니다. 이번에 살펴볼 PR은 reduce_forward 호출 시 발생하는 마스크(mask) 텐서 처리 문제를 다룹니다. 특히, 특정 차원의 stride가 0인 경우, 기존 코드는 이를 브로드캐스팅하여 거대한 중간 텐서를 생성함으로써 메모리 부족(OOM)이나 불필요한 연산 비용을 초래했습니다. 이 글에서는 해당 문제를 어떻게 해결했는지 코드 레벨에서 분석합니다.

코드 분석: python/triton_kernels/triton_kernels/reduce.py

문제의 원인: 암시적 브로드캐스팅

기존 코드에서는 m = (Mask != 0)을 통해 마스크를 즉시 평가했습니다. 만약 Mask 텐서가 특정 차원에서 stride가 0이라면, 이는 메모리상에서 실제로는 작은 크기임에도 불구하고 논리적으로는 매우 큰 텐서로 확장됩니다.

# Before
-        m = (Mask != 0)
-        total_loads = m.sum()
-        total_adds = (m.sum(dim=dim) - 1).clamp(min=0).sum()

해결책: Stride 기반의 축소

수정된 코드는 텐서를 직접 확장하는 대신, stride가 0인 차원을 식별하여 narrow 연산으로 텐서 크기를 1로 고정합니다. 이후 계산된 결과에 broadcast_repeats를 곱해주는 방식으로 논리적인 값을 보존합니다.

# After
+        # Avoid materializing broadcasted mask elements, then account for their repetitions below.
+        broadcast_repeats = tuple(size if stride == 0 else 1 for size, stride in zip(m.shape, m.stride()))
+        for d, size in enumerate(broadcast_repeats):
+            if size > 1:
+                m = m.narrow(d, 0, 1)
+        m = (m != 0)
+        total_loads = m.sum() * broadcast_repeats[0] * broadcast_repeats[1] * broadcast_repeats[2]
+        total_adds = ((m.sum(dim=dim) * broadcast_repeats[dim] - 1).clamp(min=0).sum()
+                      * broadcast_repeats[dim0] * broadcast_repeats[dim1])

왜 이게 좋은가

1. 메모리 효율성: 불필요하게 큰 텐서를 메모리에 올리지 않음으로써, 대규모 배치 처리 시 발생할 수 있는 메모리 폭증을 방지합니다.
2. 연산 최적화: m.sum() 연산이 실제 데이터 크기가 아닌 축소된 텐서에 대해 수행되므로, CPU/GPU 상에서의 연산 속도가 비약적으로 향상됩니다.
3. 일반적 교훈: 텐서 연산 시 stride가 0인 차원은 브로드캐스팅을 의미합니다. 이를 무비판적으로 평가(materialize)하면 메모리 오버헤드가 발생하므로, 메타데이터 계산 시에는 논리적 크기를 유지하면서 물리적 메모리 점유를 최소화하는 전략이 필수적입니다.

리뷰 과정에서 yongjik님은 total_adds 계산 시에도 동일하게 거대한 텐서가 생성될 수 있음을 지적했고, 이를 반영하여 전체 연산 로직을 수정함으로써 안정성을 확보했습니다.

결론

이번 최적화는 Triton 커널의 메타데이터 계산 로직을 더욱 견고하게 만들었습니다. 텐서의 물리적 레이아웃(stride)을 이해하고 이를 활용해 연산량을 줄이는 방식은 고성능 라이브러리 개발에서 매우 중요한 패턴입니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.Tensor.narrow.html
• https://pytorch.org/docs/stable/generated/torch.Tensor.stride.html

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