Tối Ưu Hóa Attention Thời Gian Không Gian Cho Video Diffusion Trên TPU

Để giải quyết bottlenek độ trễ bình phương của self-attention trong các mô hình video diffusion độ phân giải cao, các nhà phát triển đã triển khai Sparse VideoGen (SVG) để định tuyến động các đầu attention tới các mặt nạ thưa thớt không gian hoặc thời gian có cấu trúc cao. Dịch mảnh giảm thành tốc độ phần cứng thực tế trên TPU đòi hỏi tối ưu hóa kernel Splash Attention bằng cách bỏ qua các ô nhớ trống, giới hạn việc mặt nạ tọa độ chính xác chỉ ở các ô biên giới, và sắp xếp lại bố trí bộ nhớ token theo thứ tự thời gian-chính để truy cập liên tục. Bằng cách căn chỉnh các mặt nạ thưa thớt này với quy trình thực thi tile phần cứng thực sự, các tối ưu hóa kết hợp đã giảm đáng kể các phép toán ma trận lãng phí và đạt được tốc độ tăng lên tới 1.69x cho suy diễn video độ phân giải 1440p.