Как на самом деле работает FlashAttention: мозаика, повторное использование SRAM и почему это в 10 раз быстрее

Как на самом деле работает FlashAttention: мозаика, повторное использование SRAM и почему это в 10 раз быстрее Когда Трансформеры обрабатывают длинные последовательности, внимание является бесспорным узким местом в вычислениях. Стандартное внимание масштабируется квадратично ($O(N^2)$) как по FLOP, так и по памяти. Но по мере того, как длина последовательности превышает 8 КБ, 32 КБ или 128 КБ токенов, графические процессоры обычно не исчерпывают вычислительную мощность в первую очередь. Им не хватает пропускной способности памяти. Несмотря на то, что современные графические процессоры, такие как NVIDIA A100, могут выполнять сотни терафлопс, они проводят большую часть своего времени в ожидании перемещения данных между медленной памятью с высокой пропускной способностью (HBM) и быстрой встроенной статической оперативной памятью (SRAM). Стандартное внимание материализует массивную матрицу оценок внимания $N \times N$ в HBM. Для последовательности из 16 000 токенов эта единственная промежуточная матрица занимает более 1 гигабайта памяти только для многократной записи и обратного чтения. FlashAttention меняет эту ситуацию, делая алгоритм внимания учитывающим ввод-вывод. Он не отвлекает внимания и не снижает математической точности. Он просто реорганизует выполнение с использованием тайлинга и онлайн-софтмакса, чтобы промежуточные матрицы никогда не касались медленного HBM. Стена памяти графического процессора: HBM против SRAM Чтобы понять, почему FlashAttention такой быстрый, вам нужно