[Triton] Proton CuptiProfiler 다양한 버그 수정 및 개선

들어가며

Triton의 Proton 프로파일러에서 CUPTI(CUDA Profiling Tools Interface) 기반 프로파일링에 여러 버그가 있었다. Stream-captured 커널의 잘못된 correlation 처리, 비활성화된 프로파일러에서의 불필요한 작업, graph 리소스 추적의 불완전성 등이 이 PR에서 수정되었다.

핵심 코드 분석

Correlation 클래스에 clear() 추가

void clear() {
    corrIdToExternId.clear();
    externIdToState.clear();
    maxCompletedCorrelationId.store(0);
    maxSubmittedCorrelationId.store(0);
}

CUPTI가 unsubscribe 후 correlation ID를 리셋하므로, 그 전에 correlation 맵을 정리해야 한다.

Stream capture 커널 스킵

void handleApiEnterLaunchCallbacks(...) {
    auto &dataToEntry = threadState.dataToEntry;
    if (threadState.isStreamCapturing) // Stream captured 커널은 correlate하지 않음
        return;
    if (dataToEntry.empty()) // 프로파일러 비활성화 상태
        return;
}

기존에는 stream capture 확인이 함수 초입에서 이루어져 scope stack에 push조차 되지 않아 exit에서 불일치가 발생했다. 이제 scope push는 항상 수행하되, correlation 제출만 조건부로 처리한다.

Graph 리소스 콜백 개선

// Before: 커널 노드만 필터링하는 위치가 외부
if (graphData->nodeType == CU_GRAPH_NODE_TYPE_KERNEL)
    handleGraphResourceCallbacks(profiler, cbId, graphData);

// After: 내부에서 필터링 + GRAPH_DESTROY 콜백 추가
constexpr std::array<CUpti_CallbackId, 5> kGraphResourceCallbacks = {
    // ...
    CUPTI_CBID_RESOURCE_GRAPH_DESTROY_STARTING,  // 이전: GRAPHEXEC_CREATED
    CUPTI_CBID_RESOURCE_GRAPHEXEC_DESTROY_STARTING,
};

graph가 파괴될 때 관련 리소스를 정리하고, 커널이 아닌 노드(memcpy 등)는 조기에 무시한다.

오타 수정

// Before
"and t may cause memory leak"
// After
"and it may cause memory leak"

왜 이게 좋은가

• 안정성 향상: 프로파일러 활성화/비활성화 전환 시 메모리 누수와 deadlock이 방지된다.
• 정확한 프로파일링: stream capture된 커널이 잘못 프로파일링되는 문제가 해결되었다.
• 리소스 관리: graph 파괴 시 관련 메타데이터가 올바르게 정리된다.

정리

+47/-20의 변경이지만, 프로파일러의 edge case 처리를 크게 개선한 PR이다. CUPTI 콜백의 복잡한 생명주기를 올바르게 관리하는 것이 핵심이다.

참고 자료

• PR #9702

---

이 글은 AI(Claude)의 도움을 받아 작성되었습니다. 코드 분석 내용은 실제 PR diff를 기반으로 합니다.