HeyGen Avatar IV TPU 최적화: 18B 비디오 모델을 1.86배 빠르게 만든 병목 제거 방식

요약: HeyGen과 Google Cloud가 Avatar IV 비디오 생성 파이프라인을 Trillium TPU v6e 8칩 호스트로 옮기고, 첫 동작 버전 대비 1.86배 빠르게 만들었다고 공개했습니다. 단순 하드웨어 교체가 아니라 all-to-all collective, sparse attention mask, online softmax 의존성, compiler layout 계약을 하나씩 줄인 사례라서 대규모 생성형 AI 운영팀이 볼 만합니다.

HeyGen Avatar IV는 사진 1장과 오디오를 입력받아 말하는 사람 영상을 만드는 diffusion 기반 talking-head 비디오 모델입니다. Google Developers Blog에 따르면 Avatar IV는 18B+ 파라미터 규모의 diffusion stack이며 Web과 API로 제공됩니다. 파이프라인은 오디오 조건을 받아 모션을 렌더링하는 diffusion transformer, 초해상도 transformer, latent를 픽셀로 바꾸는 VAE decoder로 구성됩니다. 출력은 720p 또는 1080p, 25fps이고, 영상 전체를 기다리지 않고 chunk 단위로 스트리밍합니다.

중요한 점은 실시간성입니다. 텍스트 생성 모델은 몇 초 늦어도 사용자가 기다릴 수 있지만, 스트리밍 비디오는 chunk 하나가 늦으면 바로 재생이 끊깁니다. HeyGen 팀은 이 deadline을 기준으로 최적화를 진행했습니다. 결과적으로 8칩 Trillium(v6e) 호스트에서 첫 TPU 동작 버전 대비 chunk 생성 시간이 절반에 가까워졌고, 8×H100 프로덕션 구성과 비슷한 성능을 내면서 생성 영상 분당 비용은 최대 25% 효율적이라고 설명했습니다.

GPU 코드를 TPU로 옮긴다는 것의 실제 의미

많은 팀이 “GPU 모델을 TPU로 포팅한다”를 프레임워크 변경 정도로 생각합니다. 이번 사례는 반대입니다. HeyGen은 생산 모델 코드를 크게 바꾸지 않기 위해 torchax를 사용했습니다. torchax는 PyTorch frontend를 JAX 배열과 XLA 컴파일 경로에 연결합니다. 즉 기존 PyTorch 생산 코드를 유지하면서 TPU로 dispatch하고, TPU 차이가 큰 지점에만 Pallas kernel을 넣는 방식입니다.

이 접근의 장점은 명확합니다. 모델 코드를 두 벌로 유지하지 않아도 됩니다. HeyGen은 full native JAX rewrite가 추가 이득을 주는지도 확인했지만, XLA가 전체 파이프라인을 end-to-end로 컴파일하기 때문에 frontend 비용은 trace time에 한 번 지불되고 런타임 이득은 거의 없었다고 밝혔습니다.

그렇다고 자동 변환만으로 끝난 것은 아닙니다. 두 transformer의 bf16 weights만 해도 36GB 이상이고, Trillium 칩 하나의 HBM은 32GB입니다. 그래서 weights는 8칩에 FSDP로 shard해야 했습니다. 동시에 sequence는 Ulysses sequence parallelism으로 나눴습니다. 이 구조에서는 attention 안에 all-to-all 통신이 들어오고, 통신이 compute를 막으면 latency가 바로 커집니다. 결국 포팅의 핵심은 “돌아가게 하기”가 아니라 “통신·메모리·커널·컴파일러가 같은 deadline을 보게 만들기”였습니다.

첫 번째 병목: all-to-all collective를 숨기기

Ulysses sequence parallelism은 self-attention마다 두 번의 all-to-all을 만듭니다. 하나는 sequence shard를 head shard로 바꾸는 통신이고, 다른 하나는 다시 되돌리는 통신입니다. HeyGen의 trace에서는 이 collective가 mesh bisection bandwidth의 85~90%까지 쓰고 있었고, compute stream 위에 그대로 노출되어 있었습니다.

여기서 대역폭을 더 늘리는 것은 답이 아닙니다. 이미 wire가 거의 찼기 때문입니다. 해결책은 통신을 없애는 것이 아니라 숨기는 것이었습니다. attention head를 몇 개의 독립 그룹으로 나누고, 각 그룹이 all-to-all, attention, all-to-all을 따로 수행하게 만들었습니다. 그러면 한 그룹의 통신이 다른 그룹의 attention 계산 뒤에 숨을 수 있습니다. XLA는 이를 asynchronous start/done pair로 스케줄링할 수 있게 됩니다.

Google Cloud와 HeyGen은 이 방법으로 collective의 compute stream footprint를 약 5배 줄였다고 설명했습니다. wire time 자체가 줄어든 것이 아닙니다. critical path에서 사라진 것입니다. 실무에서 이 차이가 중요합니다. 네트워크 최적화는 “전송 시간이 줄었나”보다 “사용자 지연 시간에 남아 있나”로 봐야 합니다.

두 번째 병목: sparse attention mask를 삭제하기

초해상도 단계의 sparse attention은 tens of thousands tokens를 처리합니다. 각 프레임은 주변 프레임 window와 하나의 global reference frame을 봅니다. 문제는 mask의 live region이 frame boundary에 맞춰져 있는데, 기존 block-sparse kernel의 tile은 128 배수였습니다. 프레임 token span이 128 배수가 아니면 live block 일부가 boundary를 걸칩니다. HeyGen은 약 5개 중 1개 block이 partial이었다고 설명합니다.

Partial block은 비쌉니다. inner loop에 mask predicate가 들어가고, sequence padding도 필요합니다. 또 frame window와 reference frame을 두 번의 attention pass로 처리하고 중간 결과를 merge해야 했습니다.

해결은 mask path를 빠르게 만드는 것이 아니었습니다. block-size constraint를 128 배수에서 16 배수로 낮춰 frame token span을 정확히 나누게 만들었습니다. 그러면 block은 mask 안에 완전히 들어가거나 완전히 밖에 있게 됩니다. partial block이 사라지고, mask predicate와 padding, second pass도 사라집니다.

이건 성능 최적화에서 자주 나오는 패턴입니다. 느린 분기를 미세하게 고치는 것보다, 데이터 레이아웃을 바꿔 분기 자체를 없애는 편이 큽니다.

세 번째 병목: online softmax의 직렬 의존성

Flash-style attention은 query row마다 running maximum을 들고 다니며 key block을 순회합니다. 새 block에서 더 큰 logit이 나오면 accumulator를 rescale합니다. 이 방식은 안정적이지만 hot loop 안에 직렬 의존성을 만듭니다. 앞 block 결과가 다음 block 계산에 영향을 주기 때문입니다.

HeyGen과 Google Cloud는 running max 대신 사전 계산한 upper bound를 넣었습니다. Cauchy-Schwarz inequality로 query norm과 최대 key norm을 곱하면 가능한 logit 상한을 구할 수 있습니다. key를 mean-centering하면 bound가 더 타이트해집니다. softmax는 일정한 값을 빼도 결과가 변하지 않으므로 수학적으로 가능한 접근입니다.

물론 모든 head에 적용할 수는 없습니다. bound가 너무 느슨하면 exponential이 underflow 쪽으로 밀릴 수 있습니다. 그래서 head별 eligibility를 확인하고, 생산 데이터에서 98~99% head가 qualify했다고 설명합니다. 나머지는 같은 kernel 안에서 online path로 fallback합니다.

이 접근은 “수학적 보장 + 품질 게이트 + fallback”이 함께 있어야 운영 가능합니다. 증명만 있거나 벤치마크만 있으면 위험합니다.

compiler contract와 품질 게이트가 성능을 완성한다

마지막 포인트는 compiler와의 계약입니다. HeyGen은 attention 입력을 만드는 normalization, rotary embedding, projection, head packing 같은 작은 연산 체인을 하나의 Pallas kernel로 fusion했습니다. 그리고 그 출력 layout을 뒤따르는 all-to-all collective가 원하는 물리 layout과 맞췄습니다. 결과적으로 projection과 collective 사이의 repack chain이 사라졌습니다.

또 XLA scheduler가 custom kernel 내부 비용을 모르면 주변 작업을 잘못 스케줄링할 수 있습니다. HeyGen은 in-house kernel에 정직한 cost estimate를 붙여 latency-hiding scheduler가 올바르게 판단하게 했습니다. compiled program도 model weight처럼 release artifact로 다뤘습니다. 캐시하고, 버전 관리하고, 환경 간 promote하며, zero recompile과 bit-exact output을 게이트로 삼았습니다.

품질 검증은 두 단계였습니다. reduction order가 유지되는 변경은 byte-identical hash를 요구했습니다. reduction order가 바뀌는 변경은 bf16 reassociation이 만드는 좁은 similarity band 안에 들어와야 했고, 그보다 낮으면 모델 소유자가 frame-by-frame blind review를 했습니다. 한 후보는 end-to-end로 빨랐지만 output 품질이 기준 아래라 버렸습니다. 이 대목이 중요합니다. 생성형 AI 성능 최적화는 “빨라졌다”가 아니라 “같은 품질로 빨라졌다”여야 합니다.

실행 체크리스트

• 실시간 생성 모델은 평균 latency보다 chunk deadline miss를 먼저 봅니다.
• GPU→TPU 포팅은 framework 변경보다 통신·layout·kernel 병목 추적이 핵심입니다.
• all-to-all이 bandwidth 한계에 가까우면 줄이는 대신 overlap 가능성을 봅니다.
• sparse attention은 mask 분기를 최적화하기 전에 block geometry로 mask를 없앨 수 있는지 확인합니다.
• online softmax 최적화는 수학적 bound, head별 fallback, 품질 게이트를 같이 설계합니다.
• compiler가 원하는 layout을 ABI처럼 취급해 repack copy를 줄입니다.
• custom kernel에는 scheduler가 이해할 수 있는 cost estimate를 붙입니다.
• compiled executable을 release artifact로 관리하고 recompile 여부를 배포 게이트에 넣습니다.
• 성능 개선은 항상 byte hash 또는 유사도 band로 품질을 같이 검증합니다.

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