[triton] CUDA 가변 인자 Pre-compiled Launcher로 커널 런치 오버헤드 제거

PR 링크: triton-lang/triton#6788 상태: Merged | 변경: +1408 / -1016

들어가며

Triton 커널을 실행할 때, Python에서 GPU 드라이버 API를 호출하는 "런처" 코드가 필요합니다. 기존에는 커널의 인자 시그니처에 맞는 C 코드를 문자열 치환으로 동적 생성한 후 컴파일하는 방식이었습니다. 이 PR은 가변 인자(variadic argument)를 사용하는 사전 컴파일된 런처로 전환하여, 동적 코드 생성의 복잡성과 오버헤드를 제거합니다.

핵심 코드 분석

1. 인자 메타데이터 타입 도입

typedef enum { ARG_CONSTEXPR = 0, ARG_KERNEL = 1, ARG_TUPLE = 2 } ArgType;

typedef struct {
  PyObject_HEAD;
  PyObject *nested_tuple;
  ArgType type;
} PyKernelArgObject;

각 커널 인자가 constexpr(컴파일 타임 상수), 실제 커널 인자, 또는 중첩 튜플인지 메타데이터로 구분합니다. 이를 통해 런타임에 인자를 파싱하고 적절히 처리할 수 있습니다.

2. 범용 런치 함수

static void _launch(int gridX, int gridY, int gridZ, int num_warps,
                    int num_ctas, int launch_cooperative_grid,
                    int shared_memory, int warp_size, hipStream_t stream,
                    hipFunction_t function, void **params) {
  if (num_ctas > 1) {
    // 클러스터 런치 (hipDrvLaunchKernelEx)
    HIP_LAUNCH_CONFIG config = { ... };
    HIP_CHECK(hipSymbolTable.hipDrvLaunchKernelEx(&config, function, params, 0));
  } else if (launch_cooperative_grid) {
    HIP_CHECK(hipSymbolTable.hipModuleLaunchCooperativeKernel(...));
  } else {
    HIP_CHECK(hipSymbolTable.hipModuleLaunchKernel(...));
  }
}

커널 시그니처에 관계없이 동일한 런치 함수를 사용하므로, 인자 수가 바뀌어도 C 코드 재컴파일이 불필요합니다.

3. 포인터 추출 로직

bool extractPointer(void *ptr, PyObject *obj) {
  if (obj == Py_None) {
    *dev_ptr = (hipDeviceptr_t)0;
    return true;
  }
  if (PyLong_Check(obj)) {
    *dev_ptr = (hipDeviceptr_t)PyLong_AsUnsignedLongLong(obj);
    return true;
  }
  PyObject *ret = PyObject_CallMethodNoArgs(obj, data_ptr_str);
  // ...
}

PyTorch 텐서, 정수, None 등 다양한 타입의 인자를 런타임에 처리합니다.

왜 이게 좋은가

기존 방식은 커널마다 고유한 C 런처 코드를 문자열 치환으로 생성했기 때문에, 캐시 미스 시 컴파일 비용이 발생했고 코드 유지보수도 어려웠습니다. 새 방식은 하나의 사전 컴파일된 런처로 모든 커널을 처리하므로, 빌드 복잡성이 크게 줄고 런치 오버헤드도 제거됩니다. AMD(HIP)와 NVIDIA(CUDA) 양쪽 모두에 적용되었으며, 런치 성능 회귀가 없음이 확인되었습니다.

정리

• Python 문자열 치환 기반 동적 런처를 C 가변 인자 사전 컴파일 런처로 대체
• PyKernelArgObject 메타데이터로 인자 타입 구분 및 런타임 파싱
• 클러스터, cooperative, 일반 런치 모드 모두 지원
• 코드 유지보수성 대폭 향상, 런치 성능 회귀 없음

참고 자료

• triton-lang/triton#6788

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