[Triton] preload에 optional device 인자 추가

PR 링크: triton-lang/triton#8951 상태: Merged | 변경: +3 / -2

들어가며

Triton의 preload 메서드는 serialized specialization data를 사용하여 커널을 미리 컴파일하는 기능이다. 기존에는 항상 현재 활성화된 디바이스를 사용했는데, 멀티 GPU 환경에서는 특정 디바이스를 지정해야 하는 경우가 있다. 이 PR은 preload에 optional device 인자를 추가한다.

핵심 코드 분석

Before

def preload(self, specialization_data):
    import json
    import triton.language as tl
    device = driver.active.get_current_device()

항상 get_current_device()를 호출하여 현재 디바이스를 사용했다.

After

def preload(self, specialization_data, device=None):
    import json
    import triton.language as tl
    if device is None:
        device = driver.active.get_current_device()

device 인자가 None이면 기존과 동일하게 동작하고, 지정하면 해당 디바이스를 사용한다.

왜 이게 좋은가

1. 하위 호환성: 기존 코드는 수정 없이 그대로 동작
2. 멀티 GPU 지원: 여러 GPU에서 커널을 사전 로드할 때 디바이스를 명시적으로 지정 가능
3. API 일관성: 다른 Triton API들과 같은 패턴(optional device)을 따름

정리

3줄의 변경이지만, 멀티 GPU 환경에서의 사용성을 크게 향상시킨다. Optional parameter with None default는 하위 호환성을 유지하면서 기능을 확장하는 전형적인 Python 패턴이다.

참고 자료

• PR #8951

---

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