FlashAttention: внимание без матрицы n×n

Зачем это нужно

В уроке 4.4.3 вы посчитали: у Llama 3 8B на контексте 8192 матрица весов внимания занимает 4 ГиБ на слой, а на 131 072 токенах — терабайт. Там же было сказано, что все современные системы используют FlashAttentionFlashAttentionEN FlashAttentionТочный алгоритм внимания (Tri Dao и соавт., 2022), который не хранит матрицу n×n в видеопамяти: Q, K, V обрабатываются блоками в SRAM с online softmax, а в обратном проходе блоки пересчитываются. Память — O(n) вместо O(n²). Версии 2 и 3 улучшают параллелизм и используют возможности Hopper. и что F.scaled_dot_product_attention включает его сам. Теперь, когда вы знаете устройство GPU, можно понять, как он работает.

Идея умещается в одну фразу: посчитать внимание блоками, которые помещаются в быструю SRAMshared memory (SRAM)EN shared memoryБыстрая память внутри SM, общая для потоков одного блока; физически та же SRAM, что и кэш L1. На Ada — до 100 КБ на SM. В статьях о FlashAttention её называют просто SRAM., и ни разу не записать матрицу $T \times T$ в видеопамять. Трудность — в softmax: чтобы нормировать строку, нужно знать её максимум и сумму экспонент по всей строке, а блок видит только кусок. Решение называется online softmax. Разберём его, потом соберём из него весь алгоритм и напишем его на NumPy.

Где теряется время у обычного внимания

Обычное внимание для одной головы — четыре шага, и каждый ходит в видеопамять:

1. $S = Q K^\top / \sqrt{d}$ — матмул, запись матрицы $T \times T$.
2. Маска и softmax: чтение $S$, запись $P$ (тоже $T \times T$).
3. $O = P V$ — чтение $P$.
4. Для обратного прохода $P$ нужно хранить до самого конца шага.

Посчитаем для $T = 4096$, $d = 64$ в BF16. Одна матрица $T \times T$ — $4096^2 \cdot 2$ байт = 32 МиБ. Четыре прохода по таким матрицам — 134 МБ, а сами $Q, K, V, O$ — всего 2 МБ. Операций $4 T^2 d \approx 4.3$ GFLOP (два матмула по $2T^2 d$). Интенсивность:

$$I = \frac{4.3 \cdot 10^9}{136 \cdot 10^6} \approx 32\ \text{FLOP/байт}$$

— в четыре раза ниже точки перегиба RTX 4080 (136, урок 5.7.1). Внимание упирается в память: по roofline 0.19 мс на память против 0.044 мс на вычисления. Если бы $S$ и $P$ не покидали SRAM, по видеопамяти шли бы только $Q, K, V, O$, и интенсивность выросла бы до $4T^2 d / (4 T d \cdot 2) = T/2 = 2048$ — далеко в зоне вычислений.

Вторая беда — память для обратного прохода. Модель уровня GPT-2 small (12 голов) с батчем 8 на контексте 4096 хранит $P$ размером $8 \cdot 12 \cdot 4096^2 \cdot 2$ байт = 3.2 ГБ на каждый слой, 38.7 ГБ на 12 слоёв. В 32 ГБ вашей карты это не помещается, хотя сама модель весит меньше гигабайта.

Online softmax: нормировка за один проход

Обычный устойчивый softmax (урок 1.3.2) проходит по строке три раза: найти максимум $m$, посчитать сумму $\ell = \sum_j e^{x_j - m}$, поделить. online softmaxonline softmaxEN online softmaxSoftmax за один проход по данным с бегущими максимумом m и суммой ℓ: при новом максимуме старая сумма умножается на e^(m_старый − m_новый). Milakov и Gimelshein (NVIDIA, 2018); основа FlashAttention. (Milakov и Gimelshein, NVIDIA, 2018) находит $m$ и $\ell$ за один проход. Пусть мы прочитали часть строки и знаем для неё $m$ и $\ell$. Пришло новое число $x$:

$$m' = \max(m, x), \qquad \ell' = \ell \cdot e^{m - m'} + e^{x - m'}.$$

Читаем: «новый максимум — больший из старого и нового числа; старую сумму пересчитываем к новому максимуму, умножив на $e^{m - m'}$, и добавляем вклад нового числа». Почему это верно: $\ell = \sum_j e^{x_j - m}$, а нам нужно $\sum_j e^{x_j - m'}$. По правилу степеней $e^{x_j - m'} = e^{x_j - m} \cdot e^{m - m'}$, и общий множитель выносится за сумму. Если максимум не изменился, множитель равен 1.

Пример. Строка $x = (1, 3, 2)$:

Шаг	$x$	$m$	множитель $e^{m_{\text{старый}} - m}$	$\ell$
1	1	1	—	$e^0 = 1$
2	3	3	$e^{-2} = 0.1353$	$1 \cdot 0.1353 + e^0 = 1.1353$
3	2	3	$1$	$1.1353 + e^{-1} = 1.5032$

Проверка «в лоб»: $e^{-2} + e^0 + e^{-1} = 0.1353 + 1 + 0.3679 = 1.5032$. Вероятности $e^{x - 3} / 1.5032 = (0.090, 0.665, 0.245)$ — те же, что у логитов $(1, 2, 3)$ в уроке 1.3.2, только в другом порядке.

Главный шаг для внимания. Нам нужны не сами вероятности, а взвешенная сумма $o = \sum_j p_j v_j$. Её тоже можно копить на лету: держим ненормированный накопитель $a = \sum_j e^{x_j - m} v_j$ и пересчитываем его тем же множителем:

$$a' = a \cdot e^{m - m'} + e^{x - m'} \, v, \qquad o = a / \ell \text{ в самом конце}.$$

Для $x = (1, 3, 2)$ и значений $v = (10, 20, 30)$: после трёх шагов $a = 32.39$, $\ell = 1.5032$, $o = 21.547$. Точный ответ $0.090 \cdot 10 + 0.665 \cdot 20 + 0.245 \cdot 30 = 21.547$. Ни одной вероятности мы не хранили — только три числа $m$, $\ell$, $a$.

Прямой проход FlashAttention

Теперь вместо чисел — блоки. Разрежем $Q$ на блоки по $B_r$ строк, а $K$ и $V$ — на блоки по $B_c$ строк (это тайлингтайлинг (разбиение на блоки)EN tilingПриём быстрых ядер: данные режутся на тайлы, которые помещаются в быструю память, и каждый тайл используется многократно. Тайловый матмул с тайлом t читает из медленной памяти в t раз меньше, чем наивный. из прошлого урока). Для блока запросов $Q_i$ заводим в SRAM векторы $m$ и $\ell$ (по числу на строку) и накопитель $A$ размером $B_r \times d$. Затем идём по блокам ключей:

$$\begin{aligned} S &= Q_i K_j^\top / \sqrt{d} && \text{оценки блока } B_r \times B_c \\ m' &= \max\big(m,\ \mathrm{rowmax}(S)\big) && \text{бегущий максимум каждой строки} \\ \tilde P &= \exp(S - m') && \text{ненормированные вероятности блока} \\ \ell &\leftarrow \ell \cdot e^{m - m'} + \mathrm{rowsum}(\tilde P) \\ A &\leftarrow A \cdot e^{m - m'} + \tilde P \, V_j, \qquad m \leftarrow m' \end{aligned}$$

После последнего блока $O_i = A / \ell$ — один раз записываем в видеопамять. Это ровно формулы online softmax, только $x$ стал строкой блока $S$, а $v$ — блоком $V_j$.

flowchart LR
subgraph HBM["Видеопамять"]
  Q["Q: T×d"]
  KV["K, V: T×d"]
  O["O: T×d<br/>L: T чисел"]
end
subgraph SRAM["SRAM одного SM"]
  QI["блок Q_i"]
  KJ["блок K_j, V_j"]
  ST["S, P̃: B_r×B_c"]
  ACC["m, ℓ, A"]
end
Q -->|"один раз"| QI
KV -->|"по очереди j = 1, 2, …"| KJ
QI --> ST
KJ --> ST
ST --> ACC
ACC -->|"в конце: A / ℓ"| O

Прямой проход FlashAttention для одного блока запросов. Матрица оценок существует только кусками в SRAM

Сколько это занимает в SRAM. При $B_r = B_c = 64$ и $d = 64$ в BF16: блоки $Q_i$, $K_j$, $V_j$ — по 8 КиБ, блок $S$ во FP32 — 16 КиБ, накопитель $A$ во FP32 — 16 КиБ. Итого около 56 КиБ — помещается в 100 КБ shared memory SM у Ada. (В реальных ядрах $S$ и $A$ живут прямо в регистрах.)

Параллелизм. Блоки запросов независимы друг от друга, поэтому каждый обрабатывает свой блок потоковблок потоковEN thread block (CTA)До 1024 потоков, которые выполняются на одном SM и делят его shared memory. Ядро запускается сеткой (grid) из многих блоков. на своём SM. При батче 8, 12 головах и $T = 4096$ с $B_r = 128$ это $8 \cdot 12 \cdot 32 = 3072$ блока — работы хватит всем 76 SM.

Причинная маска. Если весь блок $K_j$ лежит правее последней строки блока $Q_i$, все его оценки замаскированы, и блок можно не считать вовсе. Маска нужна только на диагональных блоках. При $T = 4096$ и блоках по 128 из $32 \cdot 32 = 1024$ пар блоков считаются $32 \cdot 33 / 2 = 528$ — почти вдвое меньше работы.

Почему ответ точный

FlashAttention — не приближение вроде разреженного или линейного внимания (модуль 5.5). Это та же формула $\mathrm{softmax}(QK^\top/\sqrt{d})\,V$, переписанная в другом порядке. Online softmax — тождество: на каждом шаге $\ell$ и $A$ в точности равны суммам по уже просмотренным столбцам, приведённым к текущему максимуму. Отличия от наивной версии возникают только из-за другого порядка сложения чисел с плавающей точкой. В FP64 это $10^{-16}$, в FP32 — около $10^{-7}$, в BF16 — около $10^{-3}$. В практике ниже вы проверите это тестом.

Обратный проход: пересчитать дешевле, чем хранить

Для градиентов нужны вероятности $P$. Обычная реализация хранит их с прямого прохода — те самые 38.7 ГБ из примера выше. FlashAttention сохраняет только выход $O$ и одно число на строку:

$$L_i = m_i + \log \ell_i = \log \sum_j e^{S_{ij}}$$

— logsumexp строки (урок 1.3.2). В обратном проходе блоки $S = Q_i K_j^\top / \sqrt{d}$ считаются заново, и вероятности восстанавливаются без всякой нормировки: $P_{ij} = e^{S_{ij} - L_i}$. Дальше градиенты по формулам (их вывод — то же цепное правило, что в уроке 1.3.2 для $p - y$):

$$\begin{aligned} dV &= P^\top dO, \qquad dP = dO\, V^\top, \\ dS_{ij} &= P_{ij}\,(dP_{ij} - D_i), \quad D_i = \sum_j P_{ij}\, dP_{ij} = dO_i \cdot O_i, \\ dQ &= dS\, K / \sqrt{d}, \qquad dK = dS^\top Q / \sqrt{d}. \end{aligned}$$

Читаем вторую строку: градиент по оценке — это вероятность, умноженная на разницу «градиент по этой вероятности минус его среднее по строке». Среднее $D_i$ считается без матрицы: $\sum_j P_{ij} (dO_i \cdot V_j) = dO_i \cdot \sum_j P_{ij} V_j = dO_i \cdot O_i$.

Пересчёт добавляет FLOPs, но убирает чтение огромной матрицы — а внимание, как мы выяснили, упирается в память. По рис. 2 статьи FlashAttention (внимание GPT-2 medium на A100): у обычной реализации 66.6 GFLOP и 40.3 ГБ чтений и записей видеопамяти, у FlashAttention — 75.2 GFLOP и 4.4 ГБ; время — 41.7 мс против 7.3 мс. Операций больше на 13%, трафик меньше в 9 раз, время меньше в 5.7 раза.

Проверка формул обратного прохода на PyTorch (CPU, 10 секунд)

import math
import torch

torch.manual_seed(0)
T, d = 64, 16
Q, K, V = (torch.randn(T, d, dtype=torch.float64, requires_grad=True) for _ in range(3))
mask = torch.tril(torch.ones(T, T, dtype=torch.bool))

# эталон: autograd по обычному причинному вниманию
S = (Q @ K.T / math.sqrt(d)).masked_fill(~mask, float("-inf"))
O = torch.softmax(S, dim=-1) @ V
dO = torch.randn_like(O)
O.backward(dO)

# «как во FlashAttention»: из прямого прохода сохранены только O и L
with torch.no_grad():
    S = (Q @ K.T / math.sqrt(d)).masked_fill(~mask, float("-inf"))
    L = torch.logsumexp(S, dim=-1)              # одно число на строку
    P = torch.exp(S - L[:, None])               # пересчёт вероятностей, без нормировки
    D = (dO * O).sum(dim=-1)                    # D_i = dO_i · O_i
    dV = P.T @ dO
    dS = P * (dO @ V.T - D[:, None])
    dQ, dK = dS @ K / math.sqrt(d), dS.T @ Q / math.sqrt(d)

for name, mine, ref in [("dQ", dQ, Q.grad), ("dK", dK, K.grad), ("dV", dV, V.grad)]:
    print(name, f"{(mine - ref).abs().max().item():.1e}")

dQ 1.1e-15
dK 8.9e-16
dV 6.7e-16

Здесь матрица $T \times T$ строится целиком ради простоты; в настоящем ядре те же формулы применяются к блокам.

FlashAttention-2, -3 и дальше

FlashAttention-2 (Tri Dao, 2023) не меняет математику, только раскладку работы по железу. Ненормированный накопитель делится на $\ell$ один раз в конце, а не на каждом блоке, — меньше операций вне тензорных ядер. Блоки запросов распределяются по SM (параллелизм по длине последовательности, а не только по батчу и головам) — длинные последовательности с маленьким батчем загружают всю карту. Внутри блока работа делится между варпами по строкам $Q$, и им не нужно обмениваться частичными суммами через shared memory. Итог — в 1.7–3 раза быстрее первой версии. Именно FlashAttention-2 стоит внутри F.scaled_dot_product_attention в PyTorch и работает на вашей RTX 4080.

FlashAttention-3 (2024) написан под Hopper (H100, H800): асинхронные копирования и специализация варпов (одни варпы грузят данные, другие считают), чередование матмулов и softmax, чтобы тензорные ядра не ждали экспонент, и FP8 с поблочным масштабированием. По статье — до 740 TFLOPS в FP16 (75% пика H100) и около 1.2 PFLOPS в FP8. FlashAttention-4 в том же репозитории нацелен на Hopper и Blackwell. На Ada эти версии не запускаются.

Где вы это уже видели. Блочное линейное внимание из урока 5.5.3 (Lightning Attention, ядра flash-linear-attention) — та же идея: блок в SRAM, внутри — матмулы, между блоками — бегущее состояние. NSA из урока 5.5.2 — FlashAttention, который считает только выбранные блоки. FlashMLA из урока 5.2.2 — FlashAttention для decode с MLA.

Практика

Реализуйте online-нормировку и блочное внимание в стиле FlashAttention-2 с причинной маской и пропуском блоков из будущего. Тесты сравнят результат с наивным вниманием до $10^{-6}$ на длинах, которые не делятся на размер блока, проверят logsumexp и посчитают экономию памяти.

Online softmax и блочное causal attentionnumpy · нужен интернет8 тестов

import math

import numpy as np


# ── Готовый эталон: обычное внимание с матрицей T×T ───────────────────────────
def naive_attention(Q, K, V, causal=False):
    """softmax(Q Kᵀ / √d) V «в лоб»: строит всю матрицу оценок S размером T×T.

    Q, K, V: (T, d). При causal=True позиция i видит только позиции j ≤ i.
    """
    T, d = Q.shape
    S = Q @ K.T / math.sqrt(d)                       # (T, T) — та самая матрица n×n
    if causal:
        S = np.where(np.tril(np.ones((T, T), dtype=bool)), S, -np.inf)
    S = S - S.max(axis=1, keepdims=True)             # устойчивый softmax (урок 1.3.2)
    P = np.exp(S)
    P = P / P.sum(axis=1, keepdims=True)
    return P @ V


# ── Ваша часть ───────────────────────────────────────────────────────────────
def online_normalizer(xs):
    """Online softmax: за ОДИН проход по xs найти m = max(xs) и l = Σ exp(x - m).

    xs может быть итератором, который можно пройти только один раз.
    Правило обновления для нового x:
        m_new = max(m, x);   l = l · exp(m - m_new) + exp(x - m_new);   m = m_new
    Возвращает (m, l).
    """
    m, l = -math.inf, 0.0
    # TODO: один цикл for x in xs
    ...
    return m, l


def flash_attention(Q, K, V, block_q=16, block_k=16, causal=False):
    """Блочное внимание в стиле FlashAttention-2 без матрицы T×T.

    Внешний цикл — блоки строк Q по block_q, внутренний — блоки K, V по block_k.
    Для блока строк хранятся только m (бегущий максимум), l (бегущая сумма) и acc (ненормированный выход).
    При causal=True блоки K целиком «из будущего» пропускаются, а на диагонали ставится маска.
    T не обязано делиться на размеры блоков.
    Возвращает (O, L): O — выход (T, d), L — logsumexp каждой строки оценок (T,): L = m + log(l).
    """
    T, d = Q.shape
    scale = 1.0 / math.sqrt(d)
    O = np.zeros((T, V.shape[1]))
    L = np.zeros(T)
    for qs in range(0, T, block_q):
        qe = min(qs + block_q, T)
        q = Q[qs:qe]
        m = np.full(qe - qs, -np.inf)
        l = np.zeros(qe - qs)
        acc = np.zeros((qe - qs, V.shape[1]))
        # TODO: цикл по блокам K/V: оценки s, маска, m_new, p, alpha, обновление l и acc
        ...
        # TODO: нормировка O[qs:qe] и L[qs:qe]
    return O, L


def count_blocks(T, block_q, block_k, causal=False):
    """Сколько пар блоков (Q-блок, K-блок) вычисляет flash_attention.
    Без маски — все ceil(T/bq)·ceil(T/bk); с причинной маской — только те, где ks ≤ qe − 1."""
    # TODO
    ...


def extra_memory(T, bytes_per_elem=2):
    """Память под внимание одной головы сверх Q, K, V и O, которую надо хранить до обратного прохода.

    Обычное внимание хранит матрицу вероятностей P: T·T чисел по bytes_per_elem байт.
    FlashAttention хранит только L (logsumexp строк) во FP32: T чисел по 4 байта.
    Возвращает (naive_bytes, flash_bytes).
    """
    # TODO
    ...


if __name__ == "__main__":
    print(online_normalizer(iter([1.0, 3.0, 2.0])))   # ждём (3.0, 1.5032...)
    rng = np.random.default_rng(0)
    Q, K, V = rng.standard_normal((3, 100, 16))
    O, L = flash_attention(Q, K, V, 32, 32)
    print("max |flash - naive| =", np.abs(O - naive_attention(Q, K, V)).max())

ЗапуститьПроверитьПодсказкаСбросить

После решения поэкспериментируйте: замерьте naive_attention и flash_attention на $T = 2000$ через time.perf_counter. На NumPy блочная версия медленнее — питоновские циклы и маленькие матмулы. Выигрыш FlashAttention — от иерархии памяти GPU, а не от алгоритма как такового. В следующем уроке вы перенесёте этот код на Triton, и тогда блоки окажутся в настоящей SRAM.

🖥 На вашей машине (4×GPU)

1. Три реализации внимания в PyTorch. F.scaled_dot_product_attention выбирает ядро сам, но его можно заставить взять конкретное: math (обычное, с матрицей $T \times T$), flash (FlashAttention-2) или mem_eff (memory-efficient attention из библиотеки xFormers, похожая идея). Скрипт меряет прямой и обратный проход и пиковую память:

"""sdpa_bench.py — три реализации внимания в PyTorch: время и память прямого + обратного прохода."""
import time

import torch
import torch.nn.functional as F
from torch.nn.attention import SDPBackend, sdpa_kernel

B, H, d = 4, 16, 64
backends = {"math": SDPBackend.MATH, "flash": SDPBackend.FLASH_ATTENTION, "mem_eff": SDPBackend.EFFICIENT_ATTENTION}

for T in [1024, 2048, 4096, 8192, 16384]:
    q, k, v = (torch.randn(B, H, T, d, device="cuda", dtype=torch.bfloat16, requires_grad=True) for _ in range(3))
    row = []
    for name, backend in backends.items():
        try:
            with sdpa_kernel(backend):                       # разрешить только эту реализацию
                torch.cuda.empty_cache()
                torch.cuda.reset_peak_memory_stats()
                torch.cuda.synchronize()
                for i in range(4):                           # 1 прогрев + 3 замера
                    if i == 1:
                        torch.cuda.synchronize()             # дождаться конца прогрева, иначе он попадёт в замер
                        t0 = time.perf_counter()
                    out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
                    out.sum().backward()
                torch.cuda.synchronize()
                dt = (time.perf_counter() - t0) / 3
                mem = torch.cuda.max_memory_allocated() / 2**30
                row.append(f"{name} {dt * 1e3:7.1f} мс {mem:5.1f} ГиБ")
        except (torch.OutOfMemoryError, RuntimeError) as e:  # OOM или «нет такого ядра для этих входов»
            row.append(f"{name} {'OOM' if 'out of memory' in str(e).lower() else 'нет ядра':>19}")
        q.grad = k.grad = v.grad = None
    print(f"T = {T:5d} | " + " | ".join(row))

Чего ждать (оценка, проверьте у себя). У math время и память растут как $T^2$: при $T = 16\,384$ одна матрица оценок — это $4 \cdot 16 \cdot 16384^2 \cdot 2$ байт ≈ 34 ГБ, так что будет OOM, а при 8192 — на грани. У flash память растёт линейно, а время квадратично, но в несколько раз медленнее, чем у math. Проверьте и обратное: при $T = 1024$ разница небольшая — матрица оценок всего батча занимает около 128 МБ (у одной головы — 2 МБ, она целиком помещается в L2 на 64 МБ), и обычная реализация не так страдает от видеопамяти.

2. Найдите ядро в профиле. Оберните один вызов в torch.profiler из урока 5.7.1: у flash вы увидите одно ядро с flash_fwd в имени (и flash_bwd в обратном проходе), у math — цепочку из матмула, softmax, маски и второго матмула.

3. Отдельная библиотека. Пакет flash-attn (FlashAttention-2) тоже работает на Ada: pip install flash-attn --no-build-isolation (собирается из исходников долго — десятки минут на 48 ядрах; нужен CUDA Toolkit той же версии, что у PyTorch). Сравните flash_attn_func(q, k, v, causal=True) (формат (B, T, H, d)) со встроенным flash. Для обучения своей GPT это не нужно — SDPA уже выбирает FlashAttention-2.

Настройка карт — на странице «Ваша лаборатория (4×GPU)».

Китайский угол: FlashMLA и SageAttention

FlashMLA (DeepSeek, github.com/deepseek-ai/FlashMLA, MIT) — ядро внимания для MLA (урок 5.2.2). В README прямо сказано, что оно вдохновлено FlashAttention 2 и 3 и библиотекой CUTLASS: те же блоки в SRAM и online softmax, но под сжатый KV-кэш MLA. По README — до 660 TFLOPS на H800 SXM5. Работает только на Hopper и Blackwell, на вашей карте — только чтение кода.

SageAttention (группа машинного обучения Университета Цинхуа, github.com/thu-ml/SageAttention, Apache-2.0) — FlashAttention с квантованием: в первой версии $Q$ и $K$ переводятся в INT8, во второй — в INT4, а $P$ и $V$ — в FP8, с поправками против выбросов. Это ускоряет инференс без дообучения модели. Статьи приняты на ICLR 2025, ICML 2025 и NeurIPS 2025. Ключевое для вас: библиотека поддерживает Ada, и на графиках README для RTX 4090 ядро SageAttention2++ быстрее FlashAttention-2 примерно в 3–4 раза (в зависимости от длины и варианта квантования). Это хороший опыт на вашей машине: подменить F.scaled_dot_product_attention на sageattn при генерации и сравнить скорость и качество.

По-китайски тему хорошо объясняет курс Ли Хунъи «Машинное обучение» 2026 Spring (лекции о FlashAttention и KV-кэше).

Термины. Online softmax по-китайски так и называют — «онлайн-softmax». Тайлинг — «разбиение на блоки», пересчёт — «повторное вычисление», SRAM — «память на чипе» (калька с on-chip memory), а видеопамять, как в прошлом уроке, — «память отображения».

Проверь себя

Проверь себя6 вопросов·

1 Online softmax прочитал строку $(2, 5)$. Чему равна бегущая сумма $\ell$ (при $m = 5$)? Округлите до двух знаков.

2 Чем результат FlashAttention отличается от обычного внимания?

Это приближение: часть оценок отбрасываетсяНичем, кроме ошибок округления из-за другого порядка сложенияSoftmax заменён на линейную функциюРезультат точный только без причинной маски

3 Что FlashAttention хранит после прямого прохода для обратного? Выберите все.Можно выбрать несколько ответов

Входы $Q$, $K$, $V$Выход $O$Logsumexp каждой строки $L_i = m_i + \log \ell_i$Матрицу вероятностей $P$ размером $T \times T$

4 Сколько ГБ занимает матрица вероятностей $P$ одного слоя при батче 8, 12 головах, $T = 4096$ в BF16? Округлите до десятых (1 ГБ = $10^9$ байт).

5 Причинное внимание, $T = 2048$, блоки $Q$ и $K$ по 128. Сколько пар блоков вычисляет FlashAttention с пропуском блоков из будущего?

6 Обратный проход FlashAttention делает больше FLOPs, чем обычный (пересчитывает $S$). Почему он всё равно быстрее?

Внимание упирается в память, а пересчёт избавляет от чтения и записи матриц $T \times T$ в видеопамятьПересчёт идёт на CPU параллельно с GPUТензорные ядра считают пересчёт бесплатноОбратный проход FlashAttention не считает градиент по $K$

Проверить ответы

Итоги

• Обычное внимание memory-bound: матрицы $T \times T$ несколько раз ходят через видеопамять, а для обратного прохода $P$ хранится целиком — гигабайты на слой.
• Online softmax копит максимум $m$, сумму $\ell$ и ненормированный выход за один проход, пересчитывая старое множителем $e^{m - m'}$.
• FlashAttention применяет это к блокам $Q$, $K$, $V$ в SRAM: матрица оценок существует только кусками, память — $O(T)$, причинная маска позволяет пропустить почти половину блоков. Результат точный.
• В обратном проходе блоки пересчитываются из $Q$, $K$ и сохранённого logsumexp $L$: больше FLOPs, но в разы меньше трафика. FlashAttention-2 лучше раскладывает работу по SM и варпам, FlashAttention-3 использует асинхронность и FP8 на Hopper.

Дальше: Пишем ядра на Triton — от сложения векторов до тайлового матмула и forward FlashAttention на настоящей SRAM вашей видеокарты.