Google опубликовала разбор оптимизации механизма внимания в видеодиффузионных моделях — затратного этапа при создании длинных роликов высокого разрешения. Ядро на JAX и Pallas использует разреженное внимание и выравнивает маску внимания по тайлам, которые вычисляет TPU.
В тестах на одном чипе TPU v6e финальная версия работала 32.76 миллисекунды против 78.70 миллисекунды у плотного внимания Splash — заявленное ускорение составило 2.40 раза. Измерения проводились на синтетических входных данных и не учитывали маршрутизацию, перестановку токенов и обмен данными между устройствами, поэтому они не подтверждают ускорение генерации видео в целом.
Почему это важно
Для команд, оптимизирующих генерацию видео на TPU, этот результат показывает: разреженная маска может работать медленнее плотного внимания, если ядро не умеет пропускать ненужные тайлы и эффективно обрабатывать нужные. В источнике не приводятся данные о влиянии оптимизации на общее время генерации видео.
