[triton] 컴파일된 커널 모듈 명시적 unload 지원

PR 링크: triton-lang/triton#9444 상태: Merged | 변경: +70 / -1

들어가며

Triton은 커널을 컴파일하면 GPU 드라이버 모듈(CUmodule 또는 hipModule_t)을 로드하지만, 이를 명시적으로 해제하는 방법이 없었습니다. 장시간 실행되는 서비스에서는 모듈이 계속 축적되어 GPU 메모리가 누수됩니다. 이 PR은 Python __del__ 메서드와 CUDA/HIP unload_module 함수를 추가합니다.

핵심 코드 분석

1. CompiledKernel에 del 추가

Before:

class CompiledKernel:
    def __init__(self, src, metadata_group, hash):
        self.module = None
        self.function = None
        self._run = None
    # __del__ 없음

After:

def __del__(self):
    if self.module is not None:
        if knobs.runtime.module_unload_hook is not None:
            knobs.runtime.module_unload_hook(
                self.module, self.function, self.name,
                self.metadata_group, self.hash)
        driver.active.utils.unload_module(self.module)
        self.module = None

2. CUDA/HIP 드라이버에 unload 함수 추가

// NVIDIA
static PyObject *unloadModule(PyObject *self, PyObject *args) {
  CUmodule mod;
  if (!PyArg_ParseTuple(args, "K", &mod)) return NULL;
  Py_BEGIN_ALLOW_THREADS;
  CUDA_CHECK_AND_RETURN_NULL_ALLOW_THREADS(cuModuleUnload(mod));
  Py_END_ALLOW_THREADS;
  return Py_None;
}

// AMD
static PyObject *unloadModule(PyObject *self, PyObject *args) {
  hipModule_t mod;
  if (!PyArg_ParseTuple(args, "K", &mod)) return NULL;
  HIP_CHECK_AND_RETURN_NULL(hipSymbolTable.hipModuleUnload(mod))
  return Py_None;
}

3. module_unload_hook 추가

class runtime_knobs(base_knobs):
    # hook to unload module when kernel is freed
    module_unload_hook: HookChain[InitHandleHook] = HookChain()

왜 이게 좋은가

• 메모리 누수 방지: 더 이상 사용하지 않는 커널 모듈을 GPU 메모리에서 해제할 수 있습니다.
• Hook 지원: module_unload_hook으로 사용자가 unload 이벤트를 감시하거나 추가 정리 작업을 수행할 수 있습니다.
• 크로스 플랫폼: CUDA와 HIP 모두 지원합니다.

정리

CompiledKernel의 수명주기 관리를 완성한 PR입니다. __del__에서 GPU 드라이버 모듈을 해제하고, hook을 통해 확장 가능한 정리 메커니즘을 제공합니다.

참고 자료

• triton-lang/triton#9444

---

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