[Triton] Proton에서 선택적 커널 메타데이터 기록 및 커스텀 메트릭 지원

PR 링크: triton-lang/triton#9158 상태: Merged | 변경: +156 / -17

들어가며

Triton의 Proton 프로파일러는 모든 커널의 launch metadata를 기록한다. 하지만 특정 커널(예: attention)의 metadata 수집이 GPU->CPU sync를 유발하여 CUDA graph를 깨뜨릴 수 있다. 이 PR은 LaunchHook에 커널 이름 기반 include/exclude 필터를 추가하고, flops/bytes 외의 임의 메트릭(예: n_tokens)도 기록할 수 있게 한다.

핵심 코드 분석

Before: 모든 커널에 대해 고정 메트릭만 수집

class LaunchHook(Hook):
    metrics = [f"flops{width}" for width in flops_width] + ["bytes"] + ["flops"]

    def enter(self, metadata: LazyDict) -> None:
        enter_state(COMPUTE_METADATA_SCOPE_NAME)
        lazy_metadata = metadata.get()
        exit_state()
        fn_metrics = {k: lazy_metadata[k] for k in LaunchHook.metrics
                      if k in lazy_metadata}
        op_name.set(lazy_metadata["name"])
        # ...

After: 필터 + 임의 메트릭 지원

class LaunchHook(Hook):
    _reserved_metadata_keys = {"name", "function", "stream"}

    def configure(self, *, include: Optional[str] = None,
                  exclude: Optional[str] = None) -> None:
        self._include_re = re.compile(include) if include else None
        self._exclude_re = re.compile(exclude) if exclude else None

    def _matches_kernel_name(self, kernel_name: str) -> bool:
        if self._include_re and self._include_re.match(kernel_name) is None:
            return False
        if self._exclude_re and self._exclude_re.match(kernel_name) is not None:
            return False
        return True

    @staticmethod
    def _extract_metrics(lazy_metadata: dict) -> dict:
        # reserved field 제외, 지원되는 타입만 추출
        return {
            k: v for k, v in lazy_metadata.items()
            if k not in LaunchHook._reserved_metadata_keys
            and LaunchHook._is_supported_metric_value(v)
        }

    def enter(self, metadata: LazyDict) -> None:
        kernel_name = metadata.data.get("name")
        if not self._matches_kernel_name(kernel_name):
            enabled.set(False)
            return
        # ...
        enabled.set(True)
        fn_metrics = LaunchHook._extract_metrics(lazy_metadata)

    def exit(self, metadata: LazyDict) -> None:
        if not enabled.get():
            return
        libproton.exit_op(id.get(), op_name.get())

커스텀 Hook 인스턴스 전달

# Before: 문자열만 지원
proton.start(..., hook="triton")

# After: Hook 인스턴스도 지원
hook = LaunchHook()
hook.configure(include=r"_p_matmul.*")  # matmul 커널만 프로파일링
proton.start(..., hook=hook)

왜 이게 좋은가

1. 선택적 프로파일링: attention 같은 GPU sync가 필요한 커널은 제외하고 matmul만 프로파일링 가능하다.
2. 임의 메트릭: n_tokens 등 도메인 특화 메트릭을 launch_metadata에 추가하면 자동으로 기록된다.
3. CUDA graph 호환성: GPU->CPU sync가 필요한 커널을 exclude하여 CUDA graph를 깨뜨리지 않는다.
4. Fast path: 커널 이름이 필터에 매칭되지 않으면 metadata.get()(lazy evaluation)을 호출하지 않아 오버헤드가 없다.

정리

이 PR은 Proton LaunchHook에 커널 이름 기반 include/exclude regex 필터와 임의 메트릭 자동 수집을 추가했다. 이를 통해 CUDA graph를 깨뜨리는 커널을 제외하고, flops/bytes 외의 커스텀 메트릭도 프로파일링에 포함할 수 있다.

참고 자료

• triton-lang/triton#9158

---

이 글은 AI를 활용하여 PR의 핵심 변경사항을 분석하고 정리한 것입니다. 실제 코드의 맥락은 원본 PR을 참고해 주세요.