Великий парадокс VRAM: почему мы платим миллионы за память, чтобы делать всё, лишь бы ей не пользоваться

Если вы посмотрите на эволюцию видеокарт для машинного обучения, вы увидите одну тенденцию: гонку за объемом видеопамяти. Размеры LLM пухнут, KV кэш сжирает терабайты, батчи становятся всё больше. Нам нужно больше VRAM. Еще больше VRAM.

Но если вы когда-нибудь писали собственные ядра на triton, вы знаете одну жестокую тайну, о которой не задумываются дата саентисты высокоуровневых фреймворков.

Самая дорогая часть вашей видеокарты это самое медленное, узкое и отвратительное место во всей системе. И вся современная ML оптимизация (Kernel Fusion, FlashAttention, PagedAttention) сводится к одному правилу: делай что угодно, хоть считай одно и то же дважды, только не трогай VRAM.

Давайте разберем эту аппаратную иронию.

Физика GPU: Работа против cклада

Нужно посмотреть на иерархию памяти внутри GPU. Грубо говоря, у нас есть три уровня:

1. Регистры (Registers). Это руки процессора. Доступ мгновенный (доли наносекунды). Вычисления происходят только здесь. Но их ничтожно мало.

2. SRAM (Shared Memory). Это ваш рабочий стол. Она находится прямо внутри вычислительного блока. Доступ невероятно быстрый (терабайты в секунду с минимальной задержкой). Размер пара сотен килобайт на блок.

3. VRAM (HBM). Та самая память на 24, 80 или 192 ГБ. Это огромный холодный ангар на другом конце завода.

Проблема современных GPU (так называемая стена памяти) заключается в том, что вычислительная мощность (FLOPs) растет экспоненциально, а скорость чтения из VRAM линейно. Современный тензорный чип может перемножать матрицы с безумной скоростью, но он 90% времени простаивает в ожидании, пока погрузчик из VRAM довезет ему очередную партию матриц.

Наивный PyTorch: Смерть от логистики

Что происходит, когда вы пишете стандартный код на Pytorch без оптимизации?

x = x + bias           # Операция 1y = F.gelu(x)          # Операция 2z = layer_norm(y)      # Операция 3

Фреймворк выполняет это последовательно.

1. Он берет X и Bias из VRAM, складывает их в регистрах, и записывает результат обратно в VRAM.

2. Затем он запускает ядро GELU. Он снова читает только что сохраненный результат из VRAM, применяет нелинейность в регистрах, и снова записывает в VRAM.

3. То же самое для LN.

Ваш дорогущий графический процессор в этот момент не занимается математикой. Он работает грузчиком. Вы упираетесь в пропускную способность памяти (Memory Bound). Ваша Арифметическая интенсивность (количество FLOPs на каждый прочитанный байт) катастрофически мала.

Искусство избегания: Зачем мы пишем triton ядра

И вот тут на сцену выходят разработчики оптимизированных ядер. Наша задача устроить слияние ядер.

Мы пишем код на triton, который делает следующее:

1. Загрузи блок данных из VRAM в быстрый SRAM один раз.

2. Перенеси данные в регистры.

3. Прибавь bias, примени GELU, посчитай статистики для LayerNorm, примени LayerNorm всё это не выходя из SRAM/регистров!

4. Запиши финальный результат обратно в VRAM один раз.

Мы только что ускорили выполнение куска сети в 3-4 раза. Математики меньше не стало (FLOPs те же). Мы просто перестали использовать VRAM как промежуточную мусорку.

Пик иронии: FlashAttention

Самый яркий пример этой философии FlashAttention. Эта технология навсегда изменила LLM и позволила нам использовать контексты на миллионы токенов.

В чем магия FlashAttention? Вы не поверите, но на обратном проходе (backward pass) FlashAttention делает БОЛЬШЕ математических вычислений (FLOPs), чем стандартный Attention. Он буквально заново пересчитывает часть прямого прохода.

С точки зрения классического программирования это безумие: зачем считать одно и то же дважды?
Но с точки зрения GPU это гениально. Пересчитать матрицу внимания в сверхбыстрых регистрах (SRAM) стоит в разы дешевле по времени, чем прочитать огромную квадратную матрицу N x N из медленной VRAM.

FlashAttention пожертвовал вычислительными тактами, чтобы вообще не материализовывать матрицу внимания в VRAM. VRAM здесь главный злодей, и алгоритм платит налог на дополнительные вычисления, лишь бы избежать контакта с ней.

Итог

В этом и заключается великая ирония искусственного интеллекта.

Мы требуем от корпораций производить карты с бесконечной VRAM. Мы готовы отдавать по 10 миллиона рублей за узел с H200, только чтобы у нас было куда сложить веса наших 70B моделей и KV-кэш. Мы боготворим объем.

Но настоящая инженерия, которая заставляет эти модели работать в реальном времени, происходит в крошечных кэшах по 256 килобайт. Как только веса загружены на карту, VRAM становится нашим злейшим врагом. Нейросеть живет не в терабайтах HBM памяти. Она дышит и мыслит в микроскопических транзисторах SRAM и регистров.

VRAM это просто огромный, медленный паркинг. А мы гонщики, которые изо всех сил стараются не выезжать за пределы гоночного трека процессора. И чем лучше мы умеем не пользоваться VRAM, тем круче мы как инженеры.