[triton] 벤치마크에서 symmetric memory 해제

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

들어가며

Triton의 분산 matmul 벤치마크에서 symmetric memory pool이 실행 후 해제되지 않아, 반복 실행 시 메모리가 누수되는 문제가 있었습니다. 이 PR은 symm_mem_pool.release() 호출을 적절한 위치에 추가합니다.

핵심 코드 분석

1. SymmetricMemoryPool에 release 메서드 추가

class SymmetricMemoryPool:
    def release(self):
        if self._is_initialized:
            self.hdl.barrier(channel=0)
            self.hdl = None
            self.buf = None
            self.bufs = None
            self._is_initialized = False
            self.size = 0
            self.regions = {}

2. 벤치마크와 테스트에 cleanup 추가

# bench_mlp.py
proton.finalize()
triton_dist.cleanup_matmul()  # 추가

# distributed.py (테스트)
dist.barrier()
symm_mem_pool.release()  # 추가
dist.destroy_process_group()

왜 이게 좋은가

• 리소스 정리: 분산 실행 후 symmetric memory를 명시적으로 해제하여 GPU 메모리 누수를 방지합니다.
• 안전한 해제: barrier() 호출 후 해제하여 모든 rank가 사용을 완료한 후 정리합니다.

정리

분산 벤치마크/테스트의 symmetric memory pool 누수를 release() 메서드 추가와 적절한 cleanup 호출로 해결한 간결한 PR입니다.

참고 자료

• triton-lang/triton#8900

---

이 글은 AI의 도움을 받아 작성되었으며, 원본 PR의 코드 변경 사항을 기반으로 분석한 내용입니다.