FlashMLA: оптимизация ядер декодирования MLA для графических процессоров Hopper (День 1 недели открытых исходных кодов DeepSeek)

Общее введение

FlashMLA - это программа, разработанная DeepSeek Компания AI разработала эффективное ядро декодирования MLA (Multi-head Latent Attention), оптимизированное для графических процессоров NVIDIA Hopper Architecture и предназначенное для повышения производительности обработки последовательностей переменной длины. Проект был открыт на GitHub и доступен для разработчиков бесплатно. Он поддерживает вычисления с точностью BF16 и страничное KV-кэширование (размер блока 64) и демонстрирует высокую производительность на H800 SXM5, обеспечивая пропускную способность до 3000 ГБ/с в конфигурациях с интенсивным использованием памяти и до 580 TFLOPS в конфигурациях с интенсивным использованием вычислений. FlashMLA был вдохновлен FlashAttention 2&3 и проектом Cutlass. Компания DeepSeek AI продемонстрировала свои инновации в области технологий искусственного интеллекта с помощью этого проекта с открытым исходным кодом, который привлек большое внимание.

Список функций

• Эффективное декодирование MLA: Оптимизирован для графических процессоров Hopper, что значительно ускоряет обработку последовательностей переменной длины.
• Поддерживает точность BF16: Использование операций с плавающей точкой половинной точности для повышения эффективности вычислений при сохранении точности.
• Кэширование KV: Для эффективного управления памятью и повышения производительности вычислений используется механизм подкачки с размером блока 64.
• Высокая производительность: Обеспечивает пропускную способность памяти до 3000 ГБ/с и вычислительную мощность 580 TFLOPS на графическом процессоре H800.
• открытый исходный код: Предоставляется полный исходный код для поддержки модификации и интеграции разработчиками.

Использование помощи

Процесс установки

FlashMLA - это проект с открытым исходным кодом, основанный на GitHub, вам необходимо убедиться, что среда соответствует требованиям, и завершить установку перед использованием. Вот подробные шаги:

1. Экологическая готовность

• операционная система: Поддержка систем Linux (рекомендуется Ubuntu 20.04 или выше).
• требования к оборудованию: Требуется графический процессор с архитектурой NVIDIA Hopper (например, H800 SXM5).
• зависимость от программного обеспечения::
  • CUDA 12.6 или выше (инструкции по установке см. на сайте NVIDIA).
  • PyTorch 2.0 или выше (рекомендуется через pip install torch (Установка).
  • Python 3.8 или выше.
• Инструменты для осмотра: Убедитесь, что установлен Git для загрузки кода с GitHub.

2. Загрузка исходного кода

1. Откройте терминал и введите следующую команду, чтобы клонировать репозиторий FlashMLA:

   git clone https://github.com/deepseek-ai/FlashMLA.git

2. Перейдите в каталог проектов:

   cd FlashMLA
   

3. Установка зависимостей

Проект зависит от PyTorch и CUDA, которые можно установить с помощью следующих команд:

pip install -r requirements.txt

Если нет requirements.txt файл, то убедиться в том, что PyTorch установлен, не составит труда:

pip install torch torchvision

Убедитесь, что CUDA доступна:

python -c "import torch; print(torch.cuda.is_available())"

экспорт True Указывает на успешную настройку среды.

4. компиляция и тестирование

FlashMLA предоставляет предварительно скомпилированные CUDA-плагины, но убедитесь, что они соответствуют вашей локальной версии CUDA:

1. Перейдите в каталог с исходным кодом и запустите скрипт компиляции (если таковой имеется):

   python setup.py install
   

2. Проверьте успешность установки и запустите код примера:

   python example.py
   

Если ошибок нет, установка завершена.

Как использовать

Основной особенностью FlashMLA является обеспечение эффективной поддержки декодирования MLA для задач вывода моделей ИИ. Вот шаги:

Функция 1: Загрузка и запуск FlashMLA

1. модуль импорта::
   Внедрение основных функций FlashMLA в сценарии Python:

   from flash_mla import get_mla_metadata, flash_mla_with_kvcache
   

2. Подготовка к вводу данных::
   • cache_seqlens: Определяет длину последовательности кэша KV.
   • q_i: Тензор запросов.
   • kvcache_i: кэшированные данные KV.
   • block_table: Таблица блоков для кэша подкачки.
3. Получение метаданных::

   tile_scheduler_metadata, num_splits = get_mla_metadata(cache_seqlens, s_q * h_q // h_kv, h_kv)
   

4. работающее декодирование::

   o_i, lse_i = flash_mla_with_kvcache(q_i, kvcache_i, block_table, cache_seqlens, dv, tile_scheduler_metadata, num_splits, causal=True)
   

   экспорт o_i результат декодирования.lse_i значение логарифмической суммы.

Функция 2: Оптимизация обработки последовательностей переменной длины

• взять: При работе с входными последовательностями динамической длины FlashMLA уменьшает занимаемую память за счет подкачки KV-кэша.
• буровая установка::
  1. Настройте параметры подкачки: размер блока фиксирован и составляет 64, его можно изменить с помощью параметра cache_seqlens Контролирует длину последовательности.
  2. Укажите во время выполнения causal=True(c) Убедитесь, что механизм причинного внимания действует.
• эффектПропускная способность памяти на H800 составляет 3000 ГБ/с для решения масштабных задач вывода.

Функция 3: Тестирование производительности

• Методы испытаний::
  1. Отредактируйте пример сценария (например. example.py), увеличивая размер входных данных.
  2. Используйте следующий код для записи производительности:

     import time
     start = time.time()
     # 运行解码代码
     o_i, lse_i = flash_mla_with_kvcache(...)
     print(f"耗时: {time.time() - start} 秒")
     

• Ожидаемые результатыПочти 3000 ГБ/с для задач с интенсивным использованием памяти и 580 TFLOPS для задач с интенсивным использованием вычислений.

предостережение

• совместимость оборудования: Поддерживаются только графические процессоры Hopper, рекомендуется H800 или эквивалент.
• Советы по отладке: Если вы столкнулись с ошибками CUDA, проверьте соответствие версий или обратитесь за поддержкой к сообществу в GitHub Issues.
• производственная среда: Интегрируются непосредственно в существующие процессы вывода модели, обеспечивая соответствие форматов входных данных требованиям FlashMLA.

Выполнив описанные выше действия, пользователи смогут быстро приступить к работе с FlashMLA и насладиться улучшением производительности за счет эффективного декодирования. Полный код и документацию можно найти в репозитории GitHub, а параметры рекомендуется настраивать в соответствии с реальными требованиями проекта.