[Triton] Proton GlobalScratchAllocOp 폐기 — TritonGPU 공용 op으로 통합

PR 링크: triton-lang/triton#8976 상태: Merged | 변경: +54 / -54

들어가며

Triton의 Proton 프로파일러는 프로파일링 데이터를 저장하기 위해 글로벌 스크래치 메모리를 할당한다. 이전에는 Proton이 자체 proton_gpu.global_scratch_alloc op을 사용했는데, 이것이 TritonGPU dialect에 이미 존재하는 ttg.global_scratch_alloc과 기능이 중복되었다.

이 PR은 Proton 전용 op을 폐기하고 TritonGPU의 공용 op을 사용하도록 마이그레이션하며, backend 속성을 추가하여 할당 정책을 구분한다.

핵심 코드 분석

Before: Proton 전용 op

// Proton dialect의 전용 op
%0 = proton_gpu.global_scratch_alloc {
    alignment = 128 : i32,
    nbytes = 384 : i32
} : !tt.ptr<i8>

After: TritonGPU 공용 op + backend 속성

// TritonGPU 공용 op에 backend="proton" 명시
%0 = ttg.global_scratch_alloc {
    alignment = 128 : i32,
    backend = "proton",
    nbytes = 384 : i32
} : !tt.ptr<i8>

TritonGPU의 op 정의에 backend 속성이 추가되었다:

// Before
let arguments = (ins
  I32Attr:$nbytes,
  I32Attr:$alignment
);

// After
let arguments = (ins
  I32Attr:$nbytes,
  I32Attr:$alignment,
  DefaultValuedAttr<StrAttr, "\"default\"">:$backend
);

할당 로직에서 backend별로 분기한다:

// GlobalScratchMemoryAllocation.cpp
if (auto alloc = dyn_cast<triton::gpu::GlobalScratchAllocOp>(op)) {
  if (alloc.getBackend() != "default")
    return;  // proton 등 다른 backend는 별도 처리
  nbytes = alloc.getNbytes();
  align = alloc.getAlignment();
}

왜 이게 좋은가

1. 코드 중복 제거: 동일 기능의 op이 두 dialect에 중복 정의되던 것을 하나로 통합했다.
2. 확장 가능한 설계: backend 문자열 속성으로 향후 다른 서브시스템(profiler, debugger 등)도 같은 op을 재사용할 수 있다.
3. 기존 패스 호환성: backend != "default"인 op은 기본 할당 패스가 건너뛰므로, 각 backend가 자체 할당 정책을 구현할 수 있다.

정리

이 PR은 Proton의 global_scratch_alloc op을 TritonGPU의 공용 op으로 통합하고, backend 속성으로 할당 정책을 구분하는 설계를 도입했다. 코드 중복 제거와 확장성 확보를 동시에 달성한다.

참고 자료

• TritonGPU GlobalScratchAllocOp
• Triton Proton GPU Dialect

---

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