HeyGen tối ưu hóa mô hình Avatar IV trên Google Cloud Trillium TPU

13 THÁNG 8, 2026

HeyGen là một nền tảng tạo video AI với các mô hình avatar. Avatar IV, hệ thống diffusion đứng sau các video talking-head của chúng tôi, vận hành trên hơn 18 tỷ tham số và khả dụng qua Web và API. Bằng cách hợp tác với đội ngũ tối ưu hóa hiệu năng hạ tầng AI chuyên dụng của Google Cloud, chúng tôi đã đưa Avatar IV lên một máy chủ Trillium (v6e) tám chip và giúp nó nhanh hơn 1,86 lần so với phiên bản hoạt động đầu tiên.

Một luồng dữ liệu không thể chờ đợi. Avatar IV kết xuất video talking-head theo từng khối (chunk), và nếu một khối bị trễ, video sẽ bị khựng. Mọi tối ưu hóa ở đây đều được thực hiện để đáp ứng thời hạn đó. Hợp tác với đội ngũ Google Cloud, chúng tôi đã chuyển pipeline sang máy chủ Trillium (v6e) tám chip và giúp nó nhanh gần gấp đôi so với phiên bản đầu tiên, với mọi thay đổi đều vượt qua các cổng kiểm soát chất lượng đầu ra tương đương. Có ba "bức tường" định hình công việc này: các tập lệnh all-to-all lộ diện trong mesh, các khối một phần trong lưới sparse attention, và sự phụ thuộc tuần tự trong vòng lặp softmax bên trong. Bài viết này đề cập đến cả ba yếu tố trên, cùng với các hợp đồng trình biên dịch và các cổng kiểm soát chất lượng đằng sau chúng.

Khối lượng công việc và quá trình chuyển đổi (Port)

Avatar IV biến một bức ảnh đơn lẻ và một bản âm thanh thành một người đang nói và cử động. Đằng sau sản phẩm này, ba mô hình thay phiên nhau xử lý trên mỗi khối video: một diffusion transformer kết xuất chuyển động dựa trên âm thanh, một transformer thứ hai thực hiện siêu phân giải (super-resolve), và một bộ giải mã VAE chuyển đổi các giá trị tiềm ẩn (latents) thành pixel. Đầu ra là video 720p hoặc 1080p ở tốc độ 25 khung hình/giây, được truyền phát ngay khi các khối hoàn tất, vì vậy quá trình phát lại bắt đầu trong khi các khối sau vẫn đang được kết xuất.

Pipeline của Avatar IV trên máy chủ Trillium. Hai diffusion transformer và một bộ giải mã VAE thay phiên nhau xử lý trên mỗi khối. Các trọng số được phân mảnh theo FSDP và các chuỗi được phân mảnh theo Ulysses trên cùng một mesh tám chip.

Avatar IV được viết cho GPU, giống như hệ sinh thái xung quanh nó. Quá trình chuyển đổi (port) được thực hiện thông qua torchax, một frontend PyTorch trên JAX: mã mô hình sản xuất chạy mà không cần sửa đổi, được điều phối trên các mảng JAX và được biên dịch bởi trình biên dịch XLA. Cùng một mã mô hình hiện nhắm mục tiêu vào cả hai stack, và kỹ thuật dành riêng cho TPU tập trung vào những nơi phần cứng có sự khác biệt. Mỗi biến thể attention của pipeline được điều phối đến một kernel Pallas được xây dựng cho hình dạng của nó. Chúng tôi cũng đã kiểm tra xem việc viết lại hoàn toàn bằng JAX thuần túy sẽ mang lại lợi ích gì so với cách tiếp cận này. Câu trả lời là gần như không có gì: XLA biên dịch toàn bộ pipeline từ đầu đến cuối trong cả hai trường hợp, vì vậy chi phí của frontend chỉ phải trả một lần tại thời điểm trace.

Tính song song bị ép buộc bởi các phép toán. Hai transformer có tổng cộng hơn 36 GB trọng số bf16 so với 32 GB HBM trên mỗi chip Trillium, vì vậy các trọng số được phân mảnh FSDP trên tám chip. SparseCore của Trillium, một bộ đồng xử lý chạy song song với lõi chính, thực hiện việc thu thập trọng số (weight gathers) của từng lớp một cách bất đồng bộ. Trong các trace sản xuất, các thao tác thu thập đó ẩn sau quá trình tính toán: việc chạy chúng trên bộ đồng xử lý giúp giải phóng đơn vị ma trận khỏi việc di chuyển trọng số, vì vậy việc phân mảnh do các phép toán bộ nhớ ép buộc không gây tốn kém trên đường dẫn tới hạn (critical path). Tính song song chuỗi Ulysses kết hợp với việc phân mảnh trọng số trên cùng một mesh, chia nhỏ chính chuỗi video đó.

Sáu cột mốc quan trọng

Chiến lược phân mảnh đã được hoàn thiện trong phiên bản hoạt động đầu tiên. Mọi thứ sau đó là công việc về kernel và trình biên dịch, và chúng tôi theo dõi thời gian trên mỗi khối khi mỗi thay đổi được áp dụng. Biểu đồ dưới đây hiển thị hồ sơ đó dưới dạng sáu cột mốc:

Hình 1. Thời gian tương đối trên mỗi khối video được tạo, được chuẩn hóa theo phiên bản TPU hoạt động đầu tiên của chúng tôi (= 1,00×). Mỗi cột mốc bao gồm một vài thay đổi được triển khai cùng nhau.

Di chuyển từ trái sang phải trên biểu đồ, thời gian trên mỗi khối giảm xuống chỉ còn hơn một nửa giá trị ban đầu—tốc độ tăng 1,86 lần—với cùng một mô hình và các cổng chất lượng xuyên suốt. Cột mốc đầu tiên là mức giảm lớn nhất, đại diện cho quy trình đã biết được thực hiện đầy đủ: các kernel attention tùy chỉnh thay thế cho các kernel mặc định, bố cục song song chuỗi được khóa, cờ XLA được tinh chỉnh cho khối lượng công việc này thay vì để mặc định, và kích thước tile của kernel được khớp với hình dạng của nó. Ba bức tường theo sau, và các hợp đồng trình biên dịch sau đó, là những gì mà quy trình đó không thể đạt tới.

Kết quả là một pipeline truyền phát với hiệu năng tương đương với những gì chúng tôi thấy từ thiết lập sản xuất 8×H100, trong khi tiết kiệm chi phí hơn tới 25% trên mỗi phút video được tạo.

Ẩn đi các tập lệnh tập thể (Collective)

Tính song song chuỗi Ulysses đặt một cặp all-to-all bên trong mỗi self-attention: một để tráo đổi các mảnh chuỗi lấy các mảnh head, và một để tráo đổi ngược lại. Trên các trace của chúng tôi, các tập lệnh tập thể đó hoàn toàn lộ diện, đã di chuyển dữ liệu ở mức 85–90% băng thông bisection của mesh. Đường truyền không còn dư địa, điều này chỉ ra chính chương trình: một chuỗi nguyên khối gồm all-to-all, sau đó là attention, rồi all-to-all không cho XLA bất kỳ khoảng trống nào để lập lịch chuyển dữ liệu phía sau, vì vậy nó giữ cho nó ở trạng thái đồng bộ một cách chính xác.

Giải pháp, được thực hiện cùng đội ngũ Google Cloud, là phương pháp tiêu chuẩn: pipeline hóa tập lệnh tập thể. Các head attention được chia thành một vài nhóm độc lập, mỗi nhóm chạy mô hình all-to-all, attention, all-to-all riêng, và việc truyền dữ liệu của mỗi nhóm giờ đây có attention "anh em" để ẩn phía sau, vì vậy XLA chuyển sang các cặp start/done bất đồng bộ. Công việc là thực hiện điều đó thông qua pipeline sản xuất và các cổng chất lượng. Trường hợp khó là sparse attention của giai đoạn siêu phân giải: mặt nạ (mask) của nó được định nghĩa trên một thứ tự token cụ thể, và việc nhóm head phải giữ nguyên thứ tự đó, để mọi nhóm vẫn thấy chính xác mặt nạ mà nó sẽ thấy nếu không bị chia tách. Trong các trace, dấu ấn của tập lệnh tập thể trên luồng tính toán đã giảm khoảng 5 lần, với thời gian attention không thay đổi. Thời gian trên đường truyền không giảm. Nó đã rời khỏi đường dẫn tới hạn, đó là tất cả những gì mà thời hạn quan tâm. Số lượng nhóm có một điểm tối ưu: nếu chia quá nhiều, chi phí overhead trên mỗi lần khởi chạy sẽ tiêu tốn những gì mà việc chồng lấp mang lại. Không chỉ một lần, một thay đổi tập lệnh tập thể đo được nhanh hơn khi đứng độc lập lại không duy trì được ở độ sâu pipeline đầy đủ, nơi các lần truyền dữ liệu xen kẽ với nhau và với mọi thứ khác trên đường truyền. Không có gì ở đây được tính là thành công cho đến khi nó thắng từ đầu đến cuối.

Hình 2. Cung cấp cho XLA thứ gì đó để chồng lấp. Sơ đồ, thời lượng không theo tỷ lệ. Một all-to-all nguyên khối không để lại cho trình biên dịch thứ gì để ẩn nó phía sau, vì vậy nó vẫn đồng bộ (trên). Chia các head thành các nhóm độc lập và các lần truyền dữ liệu tạo thành một pipeline: đường truyền vẫn mang một lần truyền tại một thời điểm, nhưng mỗi chặng giữa ẩn sau attention của nhóm khác, chỉ để lại điểm vào đầu tiên và điểm trả về cuối cùng lộ diện (dưới). Nhiều nhóm hơn sẽ thu nhỏ các đầu lộ diện đó hơn nữa, đó là cách dấu ấn của tập lệnh tập thể trên luồng tính toán giảm khoảng 5 lần, chứ không phải về 0.

Xóa bỏ mặt nạ (Mask)

Kernel đơn lẻ lớn nhất trong pipeline là sparse attention của giai đoạn siêu phân giải. Nó chạy trên hàng chục nghìn token với mô hình thưa thớt theo cửa sổ: mỗi khung hình chú ý đến một cửa sổ các khung hình lân cận cộng với một khung hình tham chiếu toàn cục. Mặt nạ đó được căn chỉnh theo khung hình về bản chất, nghĩa là các vùng hoạt động của nó khớp chính xác với ranh giới khung hình. Một kernel block-sparse đa năng không có lý do gì để biết về các khung hình. Họ kernel splash-attention Pallas mà chúng tôi bắt đầu từ đó chia chuỗi thành các bội số của 128, khớp với các làn vector của phần cứng, và phạm vi token của một khung hình không phải là bội số của 128. Vì vậy, các khối hoạt động nằm vắt qua ranh giới khung hình, và khoảng một trong năm khối trở thành khối một phần. Việc chia tile không khớp buộc phải sử dụng các cơ chế đắt đỏ xung quanh nó: các vị từ mặt nạ (mask predicates) trong vòng lặp bên trong bất cứ nơi nào một khối nằm vắt qua ranh giới, và một chuỗi được đệm để lưới khối chia hết toàn bộ chiều dài. Trên hết, kernel xử lý hai phần của mặt nạ, cửa sổ khung hình và khung hình tham chiếu, như hai lượt attention riêng biệt, ghi các giá trị trung gian độ chính xác đầy đủ mà một lượt gộp (merge pass) đọc lại để kết hợp chúng.

Giải pháp không phải là làm cho đường dẫn có mặt nạ nhanh hơn. Đó là xóa bỏ mặt nạ. Chúng tôi đã nới lỏng ràng buộc kích thước khối của kernel từ bội số của 128 xuống bội số của 16, kích thước tile bf16 mịn nhất mà phần cứng hỗ trợ dọc theo chiều chuỗi: đủ mịn để kích thước khối có thể chia hết phạm vi token trên mỗi khung hình, và vẫn đủ lớn để giữ cho đơn vị ma trận được cung cấp dữ liệu. Giờ đây, mọi khối đều nằm hoàn toàn bên trong mặt nạ hoặc hoàn toàn bên ngoài nó. Các khối hoạt động là đầy đủ theo cấu trúc, và các khối của khung hình tham chiếu chỉ đơn giản là nhiều khối đầy đủ hơn trong cùng một softmax trực tuyến, vì vậy một lượt duy nhất bao phủ toàn bộ mặt nạ. Với điều đó, các vị từ mặt nạ, lượt thứ hai và phần đệm đều biến mất.

Hình 3. Mặt nạ, bị xóa bỏ theo cấu trúc. Hình học minh họa. Vùng hoạt động của mặt nạ (mỗi khung hình chú ý đến một cửa sổ các khung hình cộng với một tham chiếu) khớp với ranh giới khung hình. Kích thước khối chung nằm vắt qua các ranh giới đó, và các khối một phần mà chúng tạo ra buộc phải có logic mặt nạ và phần đệm trên toàn bộ kernel (trái). Các khối chia hết phạm vi token của khung hình là đầy đủ hoặc bị bỏ qua, không bao giờ là một phần, vì vậy cả hai đều biến mất theo cấu trúc (phải).

Vòng thứ hai xây dựng lại thân kernel xung quanh cùng một lưới: vòng lặp bên trong dense-attention nhanh hơn chạy trên bố cục thưa thớt đã căn chỉnh theo khung hình, với các phép giảm softmax được tái cấu trúc để cắt giảm lưu lượng thanh ghi trong vòng lặp nóng. Vòng căn chỉnh đã nâng kernel từ khoảng một nửa giới hạn mà hình dạng attention này có thể đạt được trên phần cứng lên gần ba phần tư, và thân được xây dựng lại đã đóng lại ở mức khoảng 86%. Cùng nhau, hai vòng này đã cắt giảm hơn mười phần trăm thời gian của giai đoạn siêu phân giải.

Việc căn chỉnh đã loại bỏ công việc giữa các khối. Bên trong chúng, một sự phụ thuộc tuần tự vẫn còn: softmax trực tuyến.

Giải phóng Softmax

Attention kiểu Flash mang theo một giá trị tối đa chạy trên mỗi hàng truy vấn, thay đổi tỷ lệ bộ tích lũy của nó bất cứ khi nào một khối khóa mới làm tăng nó. Việc ghi chép đó là một sự phụ thuộc tuần tự trong vòng lặp bên trong nóng nhất của các kernel attention của pipeline.

Cùng với đội ngũ Google Cloud, chúng tôi đã thay thế giá trị tối đa chạy bằng một giới hạn trên được tính toán trước. Theo bất đẳng thức Cauchy–Schwarz, logit lớn nhất có thể của một truy vấn bị giới hạn bởi tích của chuẩn của nó và chuẩn khóa lớn nhất. Một mảng nhỏ các chuẩn được tính toán trước, được cung cấp cho kernel thông qua scalar prefetch, cho phép mỗi hàng suy ra giới hạn của nó khi nó bắt đầu, và giá trị tối đa trực tuyến không còn cần thiết nữa. Việc căn giữa giá trị trung bình các khóa trước tiên sẽ thắt chặt giới hạn, vì các khóa chia sẻ một thành phần chung lớn, và nó bất biến với softmax: đầu ra không thay đổi trong số học chính xác, và các cổng chất lượng được mô tả bên dưới sẽ bắt được những gì làm tròn gây ra trong thực tế. Với giá trị tối đa được cố định trước, việc thay đổi tỷ lệ và chuỗi tuần tự của nó hoàn toàn biến mất khỏi vòng lặp bên trong.

Hình 4. Loại bỏ chuỗi khỏi softmax. Softmax trực tuyến thay đổi tỷ lệ bộ tích lũy của nó bất cứ khi nào một khối khóa-giá trị làm tăng giá trị tối đa chạy, vì vậy mọi khối đều chờ đợi khối trước (trên). Một giới hạn trên mỗi hàng có thể chứng minh được, suy ra từ các chuẩn được tính toán trước thông qua bất đẳng thức Cauchy–Schwarz, cố định giá trị tối đa trước để các khối truyền qua vòng lặp bên trong (dưới). Các head mà giới hạn quá lỏng sẽ quay lại đường dẫn trực tuyến bên trong cùng một kernel.

Không phải head nào cũng phù hợp: một giới hạn xa hơn giá trị tối đa thực tế sẽ đẩy các hàm mũ về phía underflow. Tính đủ điều kiện được kiểm tra trên mỗi attention head, và các head có giới hạn quá lỏng sẽ quay lại đường dẫn trực tuyến tiêu chuẩn bên trong cùng một kernel. Trên dữ liệu sản xuất của chúng tôi, 98–99% các head đủ điều kiện. Trước khi bất cứ thứ gì được triển khai, một cuộc kiểm toán chung về toán học đã mở rộng các đảm bảo của giới hạn cho các mặt nạ thưa thớt theo cửa sổ của chúng tôi, nơi mỗi hàng truy vấn thấy một lát cắt khác nhau của các khóa. Khi sự phụ thuộc tuần tự đã biến mất, hình học khối tối ưu của kernel thay đổi, vì vậy chúng tôi đã tinh chỉnh lại hình dạng khối cùng với softmax mới thay vì kế thừa chúng từ cái cũ.

Trên biểu đồ, đây là một trong những mức giảm dốc nhất sau quy trình mở đầu.

Bố cục là ABI

Các bức tường chia sẻ một đòn bẩy yên tĩnh hơn: các hợp đồng rõ ràng với trình biên dịch.

Ví dụ rõ ràng nhất là bố cục. Đầu vào của attention được tạo ra bởi một chuỗi các thao tác nhỏ: chuẩn hóa, rotary embeddings, phép chiếu, đóng gói head. Trình biên dịch gán một bố cục vật lý cụ thể cho toán hạng của all-to-all theo sau, và bất kỳ sự không khớp nào cũng được vá bằng các bản sao. Chúng tôi đã hợp nhất chuỗi đó thành một kernel Pallas duy nhất ghi đầu ra của nó chính xác theo bố cục mà tập lệnh tập thể muốn. Bộ đệm đầu ra của kernel là bộ đệm đầu vào của tập lệnh tập thể. Trong một trường hợp, một chuỗi đóng gói lại năm giai đoạn giữa các phép chiếu và tập lệnh tập thể đã biến mất hoàn toàn.

Hình 5. Bố cục là ABI. Trình biên dịch gán cho all-to-all một bố cục vật lý cho toán hạng của nó, và bất kỳ sự không khớp nào cũng được vá bằng các bản sao (trên). Kernel hợp nhất ghi đầu ra của nó trực tiếp theo bố cục đó, vì vậy bộ đệm đầu ra của kernel là bộ đệm đầu vào của tập lệnh tập thể và chuỗi đóng gói lại biến mất (dưới).

Các cờ (flags) cũng là các hợp đồng. Các kernel self-attention nặng chỉ đạt được tốc độ đã tinh chỉnh dưới một bộ lập lịch lệnh thay thế giúp chồng lấp công việc vector của softmax với đơn vị ma trận tốt hơn. Vì vậy, mỗi kernel yêu cầu nó theo tên (XLA_TPU_FORCE_LP_LLO_SCHEDULER) thay vì tin tưởng vào các mặc định.