[triton] ConSan: 상태 변경 시 커널 재컴파일을 보장하여 JIT 캐시 무효화

PR 링크: triton-lang/triton#8342 상태: Merged | 변경: +75 / -38

들어가며

Triton의 Concurrency Sanitizer(ConSan)는 GPU 커널에서 동시성 버그를 탐지하는 도구입니다. 하지만 ConSan을 활성화/비활성화해도 JIT 캐시에 이전 컴파일 결과가 남아있어, 도구의 상태 변경이 반영되지 않는 문제가 있었습니다. 이번 PR은 ConSan 상태를 컴파일 옵션에 포함시켜 캐시 미스를 강제합니다.

핵심 코드 분석

Before: 환경 변수 기반 설정

# 테스트에서 환경 변수로 ConSan 활성화
monkeypatch.setenv("TRITON_ENABLE_EXPERIMENTAL_CONSAN", "1")

After: instrumentation_mode 컴파일 옵션 도입

# 환경 변수 또는 knob으로 설정
os.environ["TRITON_INSTRUMENTATION_MODE"] = "consan"
knobs.refresh_knobs()
# 또는
knobs.compilation.instrumentation_mode = "consan"

테스트도 격리된 프로세스에서 실행하도록 변경:

def test_cache_miss_knob(device, monkeypatch):
    # ConSan 없이 먼저 실행
    run_in_process(run_failing_kernel, (device, False, "knob"))
    # ConSan 활성화 후 별도 프로세스에서 실행 - 캐시 미스 발생
    monkeypatch.setenv("CUDA_LAUNCH_BLOCKING", "1")
    result = run_in_process(run_failing_kernel, (device, True, "knob"))
    assert "device-side assert" in str(result.exc)

instrumentation_mode가 "consan" 또는 "proton"으로 설정되면 컴파일 옵션의 kwargs에 포함되어, 캐시 키가 변경되고 재컴파일이 트리거됩니다.

왜 이게 좋은가

1. 정확한 캐시 무효화: ConSan 상태가 컴파일 키에 포함되어 stale 캐시 문제가 해결됩니다.
2. Proton과 통합 준비: instrumentation_mode 기반으로 ConSan과 Proton이 동일한 메커니즘을 공유할 수 있습니다.
3. 테스트 격리: run_in_process를 사용하여 각 테스트가 독립적인 환경에서 실행됩니다.

정리

JIT 캐시 무효화는 도구(instrumentation) 개발에서 간과하기 쉬운 부분입니다. 컴파일 옵션에 도구 상태를 포함시키는 것은 간단하지만 필수적인 변경으로, 도구의 신뢰성을 보장합니다.

참고 자료

• triton-lang/triton#8342

---

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