[Triton] TRITON_INTERPRET 모드에서 언어 패치 자동 정리

PR 링크: triton-lang/triton#8735 상태: Merged | 변경: +72 / -50

들어가며

Triton의 인터프리터 모드(TRITON_INTERPRET=1)는 GPU 없이 커널을 CPU에서 실행한다. 이를 위해 triton.language의 여러 함수를 인터프리터 버전으로 **패치(monkey-patch)**하는데, 기존에는 패치 후 원래 상태로 복원하지 않아 이후의 정상 실행에 영향을 줄 수 있었다.

핵심 코드 분석

Before

def _patch_attr(obj, name, member, builder):
    semantic = TritonSemantic(builder)
    new_member = lambda *args, **kwargs: member(*args, _semantic=semantic, **kwargs)
    setattr(obj, name, new_member)

setattr로 덮어쓰기만 하고 원본은 저장하지 않았다.

After

class _LangPatchScope:
    def __init__(self):
        self._changes: list[tuple[object, str, object]] = []

    def set_attr(self, obj, name, value):
        original = getattr(obj, name, None)
        self._changes.append((obj, name, original))
        setattr(obj, name, value)

    def restore(self):
        while self._changes:
            obj, name, original = self._changes.pop()
            setattr(obj, name, original)

_LangPatchScope가 모든 변경을 기록하고, restore()로 원래 값을 복원한다.

호출부

def _patch_lang(fn):
    scope = _LangPatchScope()
    # ... scope.set_attr(lang, "reduce", _new_reduce) 등 ...
    return scope

# 사용
patch_scope = _patch_lang(self.fn)
try:
    # 커널 실행
finally:
    patch_scope.restore()

try-finally로 예외 발생 시에도 반드시 복원된다.

왜 이게 좋은가

1. 부작용 제거: 인터프리터 실행 후에도 triton.language가 원래 상태로 복원됨
2. 테스트 격리: 같은 프로세스에서 인터프리터와 GPU 실행을 번갈아 할 수 있음
3. 예외 안전성: try-finally 패턴으로 어떤 상황에서도 정리가 보장됨

정리

Monkey-patching은 강력하지만 위험한 기법이다. 패치한 것은 반드시 되돌려야 하며, 이를 scope 기반으로 관리하는 것은 Python에서 흔히 사용되는 안전한 패턴이다. Context manager와 유사한 이 접근법은 리소스 정리가 필요한 모든 상황에 적용할 수 있다.

참고 자료

• PR #8735

---

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