XLA TPU 컴파일러의 조용한 수치 오류를 퍼징과 계층적 격리로 잡는다

무엇인가

이 논문은 구글의 XLA:TPU 컴파일러를 데이터센터 규모에서 검증한 방법론을 다룬다. XLA:TPU는 HLO→LLO→VLIW 번들로 내려가며 타일 단위 융합과 VMEM/VREG 스테이징을 적극 쓴다. 이런 공격적 최적화는 성능을 올리지만, 인덱스 오프셋 오류, 루프 누산기 재방문, 패킹 오류 같은 결함이 크래시 없이 텐서 출력만 바꿀 수 있다. 저자들은 컴파일 타임 크래시와 런타임 크래시는 상대적으로 쉽고, 조용한 수치 오류를 가장 어려운 대상으로 본다. 규모 근거로 약 20만 개 운영 그래프 표본에서 40% 이상이 융합 커널을 하나로 접어도 1만 노드를 넘고, 융합 내부 명령 수는 matmul 융합이 최대 약 700, element-wise 융합이 최대 5,900에 이른다고 제시한다.

어떻게 동작하나

먼저 사전적 검증은 하이브리드 컴파일러 퍼징이다. 1단계는 Csmith와 QuickCheck에서 영감을 받은 템플릿 기반 합성 프로그램 생성으로, HLO로 내리기 전 타입이 있는 함수형 IR인 Jaxpr에서 시작한다. 출력 타입 제약으로 텐서 랭크, 128의 배수 차원, 수치 타입을 두고 연산의 역 타입 규칙을 재귀적으로 적용해 유효한 프로그램을 만든다. conv, reshape, abs처럼 비단사 규칙도 처리하고, sequential composition, parallel concatenation, lifting 같은 타입 안전 조합자로 하위 트리를 붙인다. 템플릿으로 slice, scatter, gather 같은 연산 부분집합과 깊이를 제한하고, 하드웨어 실행을 참조 평가기와 비교한다. 2단계는 운영 오토튜닝을 퍼저로 재활용한다. 후보 window/tile 설정의 출력 텐서를 기본 설정과 엄격한 수치 허용오차로 비교하고, lowering에 민감한 플래그를 뒤집어 실제 플릿 융합을 압박한다. 3단계는 AI 에이전트 군집이다. 에이전트 하네스가 리드 에이전트와 하위 에이전트 트리, 격리된 소스 체크아웃, 저널링된 워크플로, 조건 감시를 제공하고, TPU 도구 계층으로 ISA 참조, HLO/LLO 덤프, 레이아웃과 DMA 디코딩, 단일 모듈 재컴파일, 프로파일 분석을 공유한다. 퍼징 하위 에이전트는 XLA:TPU 패스 구현을 읽고 맞춤 퍼저 플러그인을 만들며, MXU 누산기 에뮬레이션이 있는 시뮬레이터에서 주로 돌리고 물리 TPU는 최종 확인에 쓴다. 생성기 자체의 의미 위반, 예컨대 합에 비중립 초기값을 넣는 문제는 HLO 의미론에 대한 릿머스 테스트로 걸러 오탐을 줄이고, LLO 패스에서 찾은 버그는 상위 lowering에서 도달 가능함을 보이는 종단 간 HLO 재현기를 요구한다. 패치 하위 에이전트는 재현기 최소화, ISA 감사, 몇 줄로 제한된 원자적 수정, 대표 모델 티어에 assertion을 넣어 운영 경로 타격 여부 증명까지 맡고, 독립 적대적 리뷰 군집이 IEEE-754 경계와 부호와 모듈로 폴딩 같은 코너를 감사한다.

무엇과 다른가

사후적 쪽은 계층적 격리다. 큰 그래프에서 최종 출력만 비교하면 어느 연산이 문제인지 알 수 없고, 융합 커널 내부 값은 VREG와 VMEM에만 있어 보이지 않는다. 매크로 그래프 격리는 그래프를 독립 실행 가능한 하위 그래프로 분해한다. 경계 텐서를 HBM에 두도록 copy를 넣고 VMEM 오프셋을 대상 장치에 맞추며, 중복 제거한 하위 그래프를 병렬 워커에 샤딩한다. 각 워커는 물리 TPU에서 실행해 비융합 버전과 소프트웨어 참조 평가기 양쪽과 비교하고, 빠른 비융합 하드웨어 비교부터 시작해 통과하면 바로 단락한다. 오탐 억제를 위해 역방향 제약 전파를 쓴다. 구간 산술 기반 선형 시간 역방향 데이터플로 패스로 sqrt와 rsqrt의 음수 피연산자, 0 나눗셈, 큰 리덕션의 오버플로를 피하고, 리덕션 크기에 맞춰 범위를 동적으로 조정한다. 그래도 Inf나 NaN이 나오면 select와 pad로 주입된 음의 무한대 어텐션 마스크 같은 의도된 비유한 값인지 정적 의도 분석으로 추적한다. 참조 평가기도 부분 쓰기 버퍼의 수정 영역만 비교하고, 중첩 TPU 타일링 레이아웃과 EUP 초월함수 근사, 서브노멀 0 플러시를 모델링해 하드웨어와 소프트웨어 간 무해한 차이를 줄인다. 마이크로 융합 검사는 배리어를 넣어 중간 텐서를 HBM에 쓰는 순간 융합이 쪼개져 버그가 사라지는 하이젠버그를 피한다. 대신 텐서당 min, max, mean, 표준편차 16바이트 통계를 온디바이스로 계산해 메모리 예산 안에서 추적하고, 문제 융합에서는 TensorCore를 융합 진입점에서 멈춰 커널을 그대로 스텝 실행하며 내부 연산의 VREG 값을 하드웨어에서 읽는다. 참조 모듈은 동적 콜백으로 기대 중간 텐서 맵을 채우고, 참조와 대상 명령을 짝지어 최초 분기 지점을 찾는다.

어떻게 쓰나

운영 배포 결과로 저자들은 142개 이상의 컴파일러 버그 해결을 이끌었고 103개 추가 발견이 검토 중이라고 밝힌다. 구체적 사례로 복소 tan와 tanh 버그가 있다. |y|가 44.36을 넘을 때 sinh(y)가 무한대가 되어 무한대 곱하기 0이 NaN이 되고, tanh에서는 분모가 2^126을 넘을 때 TPU의 f32 역수 나눗셈이 서브노멀로 언더플로한 뒤 0+0i로 플러시되는 문제였다. 에이전트는 tan(z)를 -i tanh(-y+xi) 항등식으로 바꾸고 분모가 2^126 이상이면 실수부를 sign(a)로 포화해 고쳤다. int8 singleton scatter에서는 N이 128보다 클 때 int8 iota가 256 모듈로로 감겨 인덱스 5에 쓴 값이 261도 덮어썼고, 32비트 미만 인덱스를 int32로 넓혀 수정했다. axis-0 Top-K에서는 [N,B]를 [M,128,B]로 접으며 lane 오프셋을 idx + l·M으로 잘못 계산한 것을 idx·128+l로 바로잡고, is_stable=True에서는 strided split을 우회했다. cum-sum 상수 폴딩은 초기 f32와 최종 f32만 비교해 중간 int32 절단인 1.75를 건너뛴 문제였고, 모든 중간 convert 타입을 순차 평가하도록 패치했다.

전제와 한계

가장 설득력 있는 사례는 quantization-aware training 융합이다. f32 텐서를 f16, bf16, fp8로 내렸다가 f32로 되돌려 r=x-f32(f16(x)) 잔차를 잡는 코드가 있다. 최적화된 융합 HLO에는 두 캐스트가 남아 있지만 XLA 기본 정밀도 정책이 중간값을 32비트 VREG에 두고 두 캐스트를 건너뛰어, 뺄셈이 x-x=0.0으로 자기 상쇄됐다. TPU에서는 비영 원소가 0.0%였고 참조 플랫폼에서는 93.7~100%였다. 이 융합을 전통적으로 디버깅하려고 쪼개면 중간 텐서가 16비트 HBM으로 흘러 반올림이 복원되어 참조와 비트 단위로 같아지는 하이젠버그가 난다. 계층적 격리는 원형 블록에서 먼저 최대 상대 오차 순위로 QAT 융합을 100% 상대 오차, 131,072개 원소 전부 0.0으로 표면화했고, 마이크로 융합 검사에서 두 convert가 입력을 반올림하지 않은 채 VREG≡x로 남아 뺄셈이 0.0이 된 것을 확인했다. 해당 융합에만 명시적 16비트 레지스터 반올림을 강제하자 QAT 잔차가 0.00% 오차로 비트 단위 복원됐고 모델의 다른 32비트 누산에는 영향이 없었다.

저자들이 뽑은 교훈은 실무에 바로 적용할 만하다. 성능 오토튜너는 경계 조건을 이미 공격적으로 탐색하므로 자연스러운 퍼저가 되고, 합성 퍼징은 새 패스와 구조적 코너를 빠르게 찌르며 실제 플릿 그래프는 운영 토폴로지와 메모리 압력을 보완한다. 또한 조용한 수치 오류는 CI의 정적 단위 테스트만으로는 잡기 어렵고, 오탐 5%만으로도 개발자 신뢰가 무너지므로 역제약 전파와 참조 평가기 충실도에 퍼저만큼 투자해야 한다. 실패 설정을 독립 실행 가능한 단위 테스트로 자동 변환하는 일은 무시되는 경고와 머지되는 수정을 가르는 차이였고, 이 테스트가 AI 에이전트에게 버그를 고치게 하는 실행 가능한 명세가 됐다. AI 에이전트는 큰 크래시로 쏠리기 쉬워 조용한 수치 오류 쪽으로 의도적으로 조향해야 하고, 수백 줄 패치보다 작은 원자적 패치와 최소 재현 테스트가 리뷰어 부담을 크게 줄였다. JAX와 XLA로 TPU를 쓰며 커스텀 융합, QAT 캐스트, 좁은 정수 인덱스 scatter, Top-K, cumsum 최적화를 건드린다면 출력 텐서 차분 검사와 격리 재현, 융합 내부 중간값 확인을 기본 절차로 두는 것이 좋다.

한계와 전제도 분명하다. 저자들은 비결정성 버그, 즉 다중 버퍼 VMEM 파이프라인의 비동기 DMA 경쟁이나 초기화되지 않은 VMEM/VREG 패딩 읽기로 같은 입력에서도 실행마다 출력이 달라지는 문제를 주 초점이 아니라고 밝힌다. 방법론이 이런 결함을 가끔 드러낼 수는 있지만 주 대상은 아니다. 또 운영 오토튜닝과 TPU 시뮬레이터 및 물리 TPU 자원에 의존하고, AI 에이전트는 인간 저자와 컴파일러 메인테이너의 승인을 전제로 움직인다. 이 논문은 표준 벤치마크 데이터셋에서 정확도나 성능을 비교한 표를 제시하지 않는다. 대신 약 20만 개 플릿 그래프 표본의 분포, 운영 배포에서 해결한 142개 버그와 103개 검토 중 발견, 네 가지 구체적 버그 사례, QAT 사례로 방법의 효과를 보인다.