TPU attention 최적화 실전 가이드: collective overlap과 mask 제거를 trace로 검증하는 순서

요약: 대규모 생성형 AI 모델을 TPU에서 빠르게 돌리려면 커널을 “더 빠르게” 만드는 것만으로 부족합니다. trace에서 all-to-all collective가 critical path에 남아 있는지, sparse attention의 partial block이 mask predicate를 만들고 있는지, compiler layout copy가 숨어 있는지 순서대로 확인해야 합니다. HeyGen Avatar IV 사례는 이 절차를 그대로 보여줍니다.

이 글은 특정 회사의 성공담을 다시 요약하려는 글이 아닙니다. 실무 엔지니어가 TPU attention 병목을 볼 때 어떤 순서로 trace를 읽고, 어떤 변경을 실험하고, 어떤 기준으로 “성공”이라고 판단해야 하는지를 정리합니다. 배경 사례는 Google Cloud와 HeyGen이 공개한 Avatar IV 최적화입니다. Avatar IV는 18B+ diffusion stack으로 talking-head video를 chunk 단위로 생성하고, Trillium(v6e) 8칩 호스트에서 첫 TPU 동작 버전 대비 1.86배 빨라졌습니다.

많은 팀이 TPU 최적화를 시작할 때 kernel microbenchmark부터 봅니다. 틀린 출발은 아니지만, 실시간 파이프라인에서는 위험합니다. 개별 kernel이 10% 빨라져도 collective가 exposed 상태면 사용자 latency는 거의 안 줄 수 있습니다. 반대로 통신 시간 자체는 그대로인데 compute stream에서 숨기면 end-to-end latency가 크게 줄 수 있습니다.

1단계: chunk deadline과 critical path를 먼저 정한다

최적화의 단위부터 정해야 합니다. 텍스트 생성은 token latency, 이미지 생성은 step time, 비디오 생성은 chunk deadline이 기준이 됩니다. HeyGen Avatar IV는 영상 chunk가 늦으면 stream이 멈추기 때문에 time per generated video chunk를 기준으로 봤습니다.

실무에서도 같은 원칙이 필요합니다. 모델 전체 latency 평균, p50, GPU/TPU utilization만 보면 착시가 생깁니다. 사용자가 체감하는 deadline을 정하고, trace에서 그 deadline에 걸리는 path를 찾습니다.

확인할 항목은 다음입니다.

어떤 작업 단위가 사용자 지연 시간을 결정하는가
각 stage의 wall time과 critical path 기여도는 얼마인가
통신 시간이 compute stream에 노출되어 있는가
kernel microbenchmark 개선이 end-to-end에서 유지되는가
품질 gate를 통과하지 못하는 변경이 섞여 있지 않은가

이 기준이 없으면 최적화는 숫자 놀이가 됩니다. kernel 하나가 빨라졌는데 전체 파이프라인은 그대로인 상황이 흔합니다.

2단계: all-to-all collective가 숨을 곳을 찾는다

Ulysses sequence parallelism을 쓰면 self-attention 안에 all-to-all이 들어갑니다. sequence shard와 head shard를 교환하고 다시 되돌리는 통신입니다. trace에서 collective가 mesh bandwidth를 85~90%까지 쓰고 있다면, 대역폭을 더 늘리는 전략은 거의 막힙니다. 이때 질문은 “통신을 줄일 수 있나?”가 아니라 “통신을 어디 뒤에 숨길 수 있나?”입니다.

실험 순서는 다음이 좋습니다. 먼저 head를 독립 그룹으로 나눌 수 있는지 봅니다. 각 그룹이 all-to-all, attention, all-to-all을 별도로 실행하면 한 그룹의 통신이 다른 그룹의 attention compute와 겹칠 수 있습니다. 다음으로 compiler가 synchronous collective 대신 async start/done pair로 스케줄링하는지 trace에서 확인합니다. 마지막으로 group count를 조절합니다. 너무 적으면 exposed end가 길고, 너무 많으면 launch overhead와 interleaving 비용이 커집니다.

성공 기준은 wire time 감소가 아닙니다. collective footprint가 compute stream에서 줄었는지입니다. HeyGen 사례에서는 약 5배 축소됐지만, wire 자체가 사라진 것은 아닙니다. 실무 trace에서도 같은 식으로 봐야 합니다.

3단계: sparse attention의 partial block을 없앤다

Sparse attention이 느릴 때 많은 팀은 mask predicate 최적화부터 합니다. 하지만 먼저 봐야 할 것은 block geometry입니다. mask live region이 데이터의 자연 경계와 맞는데, tile size가 그 경계를 자르면서 partial block을 만드는 경우가 많습니다.

예를 들어 frame-based video attention에서 mask는 frame boundary에 맞습니다. 그런데 tile이 128 배수이고 frame token span이 128 배수가 아니면 block이 frame 경계를 걸칩니다. partial block이 생기면 inner loop에 mask predicate가 들어갑니다. padding도 필요해지고, window attention과 reference attention을 별도 pass로 돌린 뒤 merge해야 할 수 있습니다.

해결 순서는 이렇습니다.

live block 중 partial block 비율을 측정합니다.
mask가 어떤 semantic boundary에 맞춰져 있는지 확인합니다.
hardware vector lane을 해치지 않는 선에서 block size constraint를 낮춥니다.
block이 mask 안/밖으로 완전히 분리되는 tile size를 찾습니다.
mask predicate, padding, second pass가 실제로 사라졌는지 trace와 kernel code로 확인합니다.

HeyGen은 block-size constraint를 128 배수에서 16 배수로 낮춰 partial block을 없겼습니다. 중요한 점은 16이 마법 숫자가 아니라 “hardware가 지원하는 finest bf16 tiling과 frame token span을 동시에 만족한 값”이라는 것입니다. 다른 모델에서는 답이 다릅니다.

4단계: online softmax 직렬 의존성을 분리한다

Flash-style attention의 online softmax는 안정적인 대신 running max 의존성을 만듭니다. 각 key block을 보며 최대값이 갱신될 수 있고, 갱신되면 accumulator를 rescale해야 합니다. 이 구조는 hot loop 안에서 block 간 직렬 체인을 만듭니다.

이를 없애려면 상한을 사전에 줄 수 있어야 합니다. Cauchy-Schwarz inequality를 쓰면 query norm과 최대 key norm으로 logit upper bound를 만들 수 있습니다. key를 mean-centering하면 bound가 더 타이트해질 수 있고, softmax 결과는 정확한 산술에서 변하지 않습니다.

하지만 이 최적화는 안전장치 없이 적용하면 안 됩니다. bound가 실제 최대값보다 너무 높으면 exponential이 underflow 쪽으로 밀립니다. 따라서 head별로 eligibility를 검사하고, bound가 느슨한 head는 online path로 fallback해야 합니다. 또한 reduction order가 바뀌면 byte-identical output이 깨질 수 있으므로 similarity band를 정해야 합니다.

5단계: layout copy와 compiler mispricing을 잡는다

마지막으로 compiler와의 계약을 봅니다. XLA 같은 compiler는 물리 layout을 정하고, layout mismatch가 있으면 copy를 넣습니다. attention 입력을 만드는 normalization, rotary embedding, projection, head packing 사이에 작은 op가 많으면 예상치 못한 repack chain이 생길 수 있습니다.

해결은 producer kernel이 consumer가 원하는 layout으로 바로 쓰게 만드는 것입니다. HeyGen은 작은 op chain을 Pallas kernel로 fusion하고, 그 출력 buffer를 all-to-all collective의 입력 buffer layout에 맞췄습니다. 한 사례에서는 five-stage repack chain이 사라졌습니다.

또 custom kernel은 compiler가 내부 FLOPs와 bytes를 모를 수 있습니다. 그러면 latency-hiding scheduler가 주변 작업 비용을 잘못 계산합니다. custom kernel에 cost estimate를 붙이면 kernel code를 바꾸지 않고도 스케줄링이 좋아질 수 있습니다. 이건 많은 팀이 놓치는 영역입니다. “커널은 빨라졌는데 end-to-end가 이상하다”면 compiler가 커널을 어떻게 가격 매기는지 봐야 합니다.

실행 체크리스트

사용자 deadline 기준으로 trace를 끊습니다. token, image step, video chunk 중 무엇인지 정합니다.
kernel microbenchmark보다 end-to-end critical path를 먼저 봅니다.
all-to-all이 exposed 상태인지, async overlap이 되는지 trace에서 확인합니다.
head grouping은 group count를 바꿔 full pipeline에서 비교합니다.
sparse attention은 partial block 비율부터 측정합니다.
mask predicate를 고치기 전에 block geometry로 mask를 삭제할 수 있는지 봅니다.
online softmax 최적화는 upper bound, fallback, 품질 gate를 같이 설계합니다.
layout mismatch로 생기는 repack chain을 trace에서 찾습니다.
custom kernel에는 compiler scheduler가 볼 수 있는 cost estimate를 제공합니다.
compiled executable을 캐시·버전 관리하고 zero recompile을 배포 조건에 넣습니다.

참고: Google Developers Blog, “HeyGen x Google Cloud: Bringing Avatar IV to TPUs”, 2026-08-13.