[triton] 캐시 테스트를 Device Agnostic하게 개선

PR 링크: triton-lang/triton#9528 상태: Merged | 변경: +12 / -7

들어가며

Triton은 CUDA뿐 아니라 AMD HIP 등 다양한 GPU 백엔드를 지원합니다. 그러나 일부 캐시 관련 테스트에서 device_caches[0]처럼 device index를 하드코딩하여 CUDA 이외의 환경에서 테스트가 실패하는 문제가 있었습니다. 이 PR은 나머지 캐시 테스트들을 device agnostic하게 만들어 모든 백엔드에서 올바르게 동작하도록 수정합니다.

핵심 코드 분석

Before:

assert len(kernel.device_caches[0][0]) == 4

After:

device = getattr(torch, device).current_device()
assert len(kernel.device_caches[device][0]) == 4

test_module_load_unload 함수에서도 device fixture를 활용하도록 변경했습니다:

Before:

def test_module_load_unload(fresh_knobs):
    # ...
    out = torch.randn(1, dtype=torch.float32, device='cuda')

After:

def test_module_load_unload(device, fresh_knobs):
    # ...
    out = torch.randn(1, dtype=torch.float32, device=device)

왜 이게 좋은가

Device index 하드코딩은 단일 GPU 환경의 CUDA에서만 우연히 동작하는 취약한 코드입니다. getattr(torch, device).current_device()를 사용하면 실행 환경에 맞는 정확한 device ID를 가져올 수 있어 multi-GPU 환경이나 HIP 백엔드에서도 올바르게 동작합니다. 이런 작은 수정이 CI 파이프라인의 안정성을 크게 높입니다.

정리

캐시 테스트 5곳에서 하드코딩된 device_caches[0]을 device_caches[current_device()]로, 그리고 'cuda' 문자열을 device fixture로 교체하여 device agnostic하게 개선했습니다.

참고 자료

• triton-lang/triton#9528

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