[triton] AMD 백엔드에서 Floating-Point Sanitizer(FPSan) 지원 활성화

PR 링크: triton-lang/triton#9455 상태: Merged | 변경: +178 / -36

들어가며

FPSan(Floating-Point Sanitizer)은 부동소수점 연산의 오류(NaN, Inf 전파 등)를 감지하는 Triton의 계측 도구입니다. 기존에는 CUDA 백엔드에서만 지원되었는데, 이 PR은 AMD GPU(HIP 백엔드)에서도 FPSan을 활성화합니다. 핵심 과제는 AMD의 warp size(64)가 NVIDIA(32)와 다르다는 점에서 발생하는 레이아웃 불일치 처리입니다.

핵심 코드 분석

Before (CUDA 전용):

def _require_cuda_backend(device):
    if is_hip():
        pytest.skip("fpsan tests currently cover the CUDA backend only")

After (CUDA + HIP 지원):

THREADS_PER_WARP = triton.runtime.driver.active.get_current_target().warp_size

def _require_cuda_backend(device):
    if not (is_cuda() or is_hip()):
        pytest.skip("fpsan tests require CUDA or HIP")
    if is_hip() and not _hip_device_supports_fpsan():
        pytest.skip("fpsan is not supported on this HIP device")

Warp size를 매개변수화:

@gluon.jit
def _binop_kernel(x_ptr, y_ptr, out_ptr, n_elements, OP, BLOCK,
                  THREADS_PER_WARP: gl.constexpr):
    layout = gl.BlockedLayout(
        size_per_thread=[2],
        threads_per_warp=[THREADS_PER_WARP],  # 32 or 64
        warps_per_cta=[4], order=[0])

AMD 전용 MLIR 테스트도 추가:

module attributes {ttg.target = "hip:gfx942",
                   "ttg.threads-per-warp" = 64 : i32} {
  tt.func public @dot_emulation() -> tensor<16x16xf32, #blocked> {

왜 이게 좋은가

AMD에서 FPSan을 지원함으로써 크로스-벤더 부동소수점 정합성 검증이 가능해집니다. THREADS_PER_WARP를 런타임에서 가져와 constexpr로 전달하는 패턴은 vendor-agnostic 커널 작성의 좋은 예시입니다. ptxas 13.x 컴파일러 버그 workaround(-Ofc mid)를 fpsan에도 확장 적용한 것은 AMD/NVIDIA 양쪽에서 안정적인 계측 코드를 생성합니다.

정리

AMD GPU(CDNA3/CDNA4/GFX1250)에서 FPSan을 활성화하고, warp size 차이를 매개변수화하여 기존 테스트가 양 벤더에서 동작하도록 개선했습니다.

참고 자료

• triton-lang/triton#9455

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