Глава 11a. FlashAttention: когда прорыв — это не новая формула

Иногда архитектура уже правильная,
а узкое место — в том, как она трогает память.

Глава 11 вывела attention и честно выставила счёт O(N²).
Но на практике трансформеры часто упираются не в «арифметику FLOPs», а в пропускную способность памяти.

FlashAttention — идеальный кейс для книги:
математика attention почти та же, меняется реализация под иерархию памяти GPU.

11a.1. Две скорости, которые путают интуицию

На GPU есть быстрая память рядом с вычислениями (SRAM / регистры уровня ядра) и большая медленная (HBM).

Грубая картинка:

Если алгоритм делает мало FLOPs на каждый байт, привезённый из HBM, вы «считаете» медленно даже при скромной арифметике.
Это memory-bound режим.

Наивный attention часто именно такой.

11a.2. Почему наивный attention бьёт по HBM

Стандартный путь (учебниковый):

1) посчитать S = QKᵀ          # матрица N×N
2) записать S в память
3) softmax(S) → A             # ещё N×N
4) записать A
5) P = A V

Даже если потом всё перемножается быстро, материализация S и A гоняет огромные промежуточные тензоры через медленную память.

N = 4096, fp16
N×N ≈ 32M элементов ≈ 64 MB  (только одна матрица)

И это ещё до батча, голов, слоёв, других активаций.
На обучении с backward цена памяти становится сюжетом №1.

Инженерный диагноз:

Узкое место не в том, что softmax «сложный».
Узкое место в том, что мы сохраняем то, что можно было бы
посчитать кусками и сразу потребить.

11a.3. Идея FlashAttention: tiling + online softmax

Смысл (без претензии на построчный разбор CUDA):

1. Разбить Q/K/V на блоки, которые помещаются в быструю память.
2. Считать attention по тайлам, не собирая полную N×N в HBM.
3. Использовать устойчивый online softmax: накопить знаменатели/максимумы так, чтобы результат совпал с обычным softmax, даже если шли кусками.
4. На обучении хранить/пересчитывать то, что нужно для backward, экономя активации.

Математически вы хотите тот же attention.
Инженерно вы меняете маршрут данных.

Это рифмуется со всей книгой:

Иногда «новая архитектура» не нужна. Нужна новая карта движения байтов.

11a.4. Что это меняет для архитектора модели

Практические следствия:

• более длинный контекст становится реалистичнее на том же железе;
• можно поднять batch / длину / число голов, упираясь позже;
• бенчмарки «FLOPs модели» ещё хуже предсказывают wall-clock — потому что выигрыш пришёл из IO;
• дальнейшие идеи (лучше tiling, разные attention-варианты) часто живут на стыке алгоритма и железа.

flowchart TD
    A[Self-Attention] --> B[Математика: softmax QK V]
    A --> C[Реализация: где живут промежуточные матрицы]
    C --> D[Наивно: материализовать N×N в HBM]
    C --> E[Flash: тайлы в SRAM + online softmax]
    D --> F[Memory-bound, короче контекст]
    E --> G[Тот же результат, выше tokens/s]

Важная оговорка честности:

FlashAttention не отменяет асимптотику O(N²) вычислений.
Он делает константы и память практически терпимыми гораздо дольше.

Если вам нужен принципиально другой asymptotics — это уже линейные/оконные/SSM-ветки из главы 16.
Если нужно выжать текущий transformer — сначала спросите про IO.

11a.5. Урок шире, чем attention

После этого кейса полезно завести привычку читать статьи двумя колонками:

Research Engineer, который умеет только менять YAML слоёв, и инженер, который умеет менять место хранения промежуточных тензоров, видят разные потолки одной и той же модели.

Если бы вы были автором статьи…

1. Профилируйте naive attention vs fused/Flash на фиксированных N ∈ {512, 2k, 8k}.
   Где разрыв в ms, где в peak memory?

2. Сформулируйте критерий: в каком режиме ваш новый блок memory-bound, а в каком compute-bound?
   Какой эксперимент это доказывает?

3. Придумайте “Flash” для своей боли вне NLP (например, большая cost-матрица в matching для DETR-подобных голов).
   Что нельзя материализовать целиком?

В послесловии этот урок вернётся как co-design модели и рантайма: будущее часто прячется не в новой букве блока, а в новой реализации старого.