[Triton] Proton에서 module_unload를 kernel_unload로 이름 변경 및 콜백 구현

2026년 2월 24일

들어가며

Triton의 Proton 프로파일러에서 커널 모듈이 해제될 때의 hook 이름이 module_unload에서 kernel_unload로 변경되었다. 기존 kernel_load_start_hook, kernel_load_end_hook과의 일관성을 위한 이름 변경이며, Proton에서 실제 unload 콜백을 구현하여 메타데이터를 정리한다.

핵심 코드 분석

Before

# knobs.py
module_unload_hook: HookChain[InitHandleHook] = HookChain()

# compiler.py
if knobs.runtime.module_unload_hook is not None:
    knobs.runtime.module_unload_hook(self.module, ...)

After

# knobs.py
kernel_unload_hook: HookChain[InitHandleHook] = HookChain()

# compiler.py
if knobs.runtime.kernel_unload_hook is not None:
    knobs.runtime.kernel_unload_hook(self.module, ...)

Proton에서 실제 cleanup 로직이 구현되었다:

# hook.py
def destroy_handle(module, function, name, metadata_group, hash):
    for hook in reversed(HookManager.active_hooks):
        hook.destroy_handle(module, function, name, metadata_group, hash)

// InstrumentationProfiler.cpp
void InstrumentationProfiler::destroyFunctionMetadata(uint64_t functionId) {
    functionScopeIdNames.erase(functionId);
    functionScopeIdContexts.erase(functionId);
    functionNames.erase(functionId);
    functionMetadata.erase(functionId);
}