[triton] Blackwell GPU Cluster Launch Control 지원으로 Persistent Kernel 워크로드 밸런싱 구현

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

들어가며

NVIDIA Blackwell(SM100+) GPU에는 Cluster Launch Control(CLC)이라는 새로운 하드웨어 기능이 존재합니다. CLC를 사용하면 실행 중인 워커가 아직 시작되지 않은 클러스터의 작업을 취소하고 대신 처리할 수 있어, SM 가용성이 변하는 환경에서 persistent kernel의 로드 밸런싱을 크게 개선할 수 있습니다. 이 PR은 Triton의 Gluon API에 CLC 지원을 추가합니다.

핵심 코드 분석

1. MLIR Op 정의

CLC 기능을 위한 세 가지 핵심 MLIR 연산이 추가되었습니다.

def TTNG_CLCTryCancelOp : TTNG_Op<"clc_try_cancel", []> {
  let summary = "Issue CLC try_cancel to cancel a pending cluster";
  let arguments = (ins
    Arg<TTG_MemDescType, "", [MemWrite<SharedMemory>]>:$result,
    Arg<TTG_MemDescType, "", [MemWrite<SharedMemory>]>:$mbarrier,
    I1Attr:$multicast
  );
}

def TTNG_CLCIsCanceledOp : TTNG_Op<"clc_is_canceled", [Pure]> {
  let arguments = (ins I128:$clcResult);
  let results = (outs I1:$is_canceled);
}

def TTNG_CLCGetProgramIdOp : TTNG_Op<"clc_get_program_id", [Pure]> {
  let arguments = (ins I128:$clcResult, TT_ProgramDim:$dim);
  let results = (outs I32:$result);
}

try_cancel은 비동기적으로 CLC 요청을 발행하고, is_canceled로 취소 성공 여부를 확인하며, get_program_id로 취소된 클러스터의 CTA ID를 가져옵니다.

2. Python Gluon API 바인딩

.def("create_clc_try_cancel",
     [](GluonOpBuilder &self, Value result, Value mbarrier, bool multicast) {
       self.create<ttng::CLCTryCancelOp>(result, mbarrier, multicast);
     })
.def("create_clc_is_canceled",
     [](GluonOpBuilder &self, Value clcResult) -> Value {
       return self.create<ttng::CLCIsCanceledOp>(clcResult);
     })
.def("create_clc_get_program_id",
     [](GluonOpBuilder &self, Value clcResult, int dim) -> Value {
       return self.create<ttng::CLCGetProgramIdOp>(clcResult, dim);
     })

Gluon의 OpBuilder에 CLC 연산을 직접 생성할 수 있는 Python 인터페이스가 추가되어, 사용자가 Python 레벨에서 CLC 기반 persistent kernel을 작성할 수 있습니다.

왜 이게 좋은가

Persistent kernel은 GPU 전체 SM에 상주하며 작업을 반복 처리하는 패턴으로, 대규모 행렬 연산에서 launch 오버헤드를 제거하는 데 핵심입니다. 그러나 SM 간 작업 분배가 균일하지 않으면 일부 SM이 유휴 상태가 되는 문제가 있습니다. CLC는 이 문제를 하드웨어 레벨에서 해결합니다.

PR에 포함된 튜토리얼에서는 8192x8192x8192 FP16 행렬 곱에서 cuBLAS 성능의 92.5%를 달성했다고 합니다. 이는 Triton 사용자가 하드웨어 수준의 고급 기능을 활용하여 라이브러리 수준의 성능에 근접할 수 있음을 보여줍니다.

정리

• Blackwell SM100+ GPU의 CLC 기능을 MLIR Op + Python API로 완전 지원
• try_cancel / is_canceled / get_program_id 3개 연산 추가
• Persistent kernel에서 동적 작업 재분배로 로드 밸런싱 개선
• FP16 matmul에서 cuBLAS 대비 92.5% 성능 달성

참고 자료

• triton-lang/triton#9361
• NVIDIA PTX ISA - Cluster Launch Control

이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.