Hacker News Digest

Новая версия Flash Attention 4 оптимизирована под архитектуру Blackwell от Nvidia и обещает ~20% прирост скорости по сравнению с предыдущим рекордсменом — закрытыми ядрами внимания в библиотеке cudnn. Хотя официального отчёта нет, исходный код уже доступен, что позволило разобрать его устройство. Главное изменение — не математические трюки (вроде быстрых приближённых экспонент или эффективного онлайн-softmax), а сложная асинхронная конвейеризация операций, напоминающая принципы параллельного программирования из высокопроизводительных систем вроде баз данных или веб-серверов.

Архитектура FA4 построена вокруг обработки «тайлов» — блоков данных, которые потоково считываются из глобальной памяти GPU. Один экземпляр ядра обрабатывает два тайла запросов, последовательно сканируя все ключи и значения, чтобы вычислить взвешенные выходные данные. Это напоминает векторized-сканирование в СУБД. Масштабирование достигается за счёт массового параллельного запуска таких программ по модели «одна программа — много данных». Подход требует глубокой асинхронности и эффективного использования warp-ов, но остаётся интуитивно понятным для инженеров, работавших с конкурентными системами.

• Обсуждение термина "reverse engineering" применительно к анализу исходного кода и его пониманию.
• Критика стиля и структуры блог-поста за избыточные отсылки к исследованиям и недостаток конкретики.
• Замечания о сложности написания эффективных GPU-кернелов для современного железа и упоминание тренда на "мегакернелы".
• Запрос рекомендаций по обучающим материалам для начинающих в GPU-программировании.
• Положительные отзывы о содержании поста и его развлекательном, доступном стиле.

Flash Attention на 5090 в CUDA C++

Цель — научиться писать attention-ядро на CUDA C++, чтобы использовать MXFP8/NVFP4 MMA для sm120, чего нет в Triton.
Код: learn-cuda/07_attention.

Бенчмарк (bs=1, heads=8, q=4096, kv=8192, BF16, 5090@400 W, CUDA 12.9, SOL 209.5 TFLOPS):

ядро	TFLOPS	%SOL
F.sdpa (Flash)	186.73	89.13
F.sdpa (CuDNN)	203.61	97.19
flash-attn	190.58	90.97
v1 (basic)	142.87	68.20
v2 (swizzle)	181.11	86.45
v3 (2-stage)	189.84	90.62
v4 (ldmatrix.x4)	194.33	92.76
v5 (pipe)	197.74	94.39

---

Алгоритм Flash Attention 2

Псевдокод:

scale = DIM**-0.5
for b, tile_Q:
    tile_O = 0
    tile_Q = load(Q[b, tile_Q])
    for tile_KV:
        tile_K = load(K[b, tile_KV])
        tile_S = tile_Q @ tile_K.T * scale
        online_softmax(tile_S)  # in-place
        tile_V = load(V[b, tile_KV])
        tile_O += tile_S @ tile_V
    store(O[b, tile_Q])

head_dim=128 помещается в регистры.

---

v1 — базовая версия

1. G2S: cp.async.ca.shared.global 128-битными транзакциями.
2. S2R: ldmatrix для Q, K, V → 8×8 фрагменты.
3. Softmax online:
   • m = max(m_prev, m_curr)
   • d = d_prev * exp(m_prev - m) + Σ exp(S - m)
   • O = O_prev * (d_prev/d) * exp(m_prev - m) + (exp(S - m)/d) @ V

---

v2 — swizzled shared memory

• 128-битные банки → конфликты при 8×8 tile.
• Swizzle K и V по 32-битным строкам; Q оставляем линейно.
• +40 % пропускной способности.

---

v3 — 2-stage pipeline

• Двойной буфер: пока вычисляем S/P@V, асинхронно грузим следующий KV.
• cp.async.commit_group() + cp.async.wait_group(1).
• +5 % к SOL.

---

v4 — ldmatrix.x4

• Одна инструкция ldmatrix.x4 загружает 4×8×8 фрагмента K/V за раз.
• Снижает инструкций на 25 %.
• +2 % к SOL.

---

v5 — улучшенный pipeline

• 3-4 стадии:
  1. prefetch KV
  2. compute S
  3. compute P@V
  4. write-back O
• __pipeline_wait_prior(N) + __pipeline_commit().
• +2 % к SOL.

---

Что дальше

• Использовать TMA (cp.async.bulk) и NVFP4/MXFP8 MMA.
• Поддержка head_dim > 128 (FlashMLA).

• Пользователи удивлены, что RTX 5090 даёт всего 209 TFLOPS BF16 — менее 10 % от серверного Blackwell B200 (2250 TFLOPS), но при цене ~$30-40 k за B200 соотношение цена/производительность почти сравнялось.
• Обсуждают, что NVIDIA с 4090 и далее искусственно ограничивает тензорные ядра игровых карт для ML-операций FP8/FP16.
• У 5090 выше TDP, чем у 4090, и можно ограничить мощность лишь до 70 % (4090 — до 50 %), что мешает апгрейду для ML-станций.
• Появились вопросы о поддержке Flash Attention на 5090/5080 и о нативной компиляции под Blackwell в PyTorch 2.7.
• Участники спорят, стоит ли вкладываться в Triton, если нужны фирменные типы NVFP4/MXFP8, которых там пока нет.