Быстрые ядра внимания и слияние операторов реорганизуют вычисления в трансформерах, чтобы ограничить перемещение данных и запуски ядер на акселераторе. Эти оптимизации в принципе оставляют неизменными веса модели и её математическую функцию, хотя порядок операций и точность могут приводить к небольшим численным различиям. Опубликованное Machine Learning Mastery руководство подробно описывает эти техники с примерами кода.
Слой трансформера опирается на стандартные операции: линейные проекции, внимание, нормализацию, функции активации и остаточные сложения. Прямая реализация запускает отдельное вычислительное ядро («kernel») для каждой операции и записывает каждый промежуточный тензор в память акселератора. Вычисления остаются корректными, но время выполнения может тратиться больше на перемещение значений и запуск ядер, чем на полезные вычисления.
Слияние операторов
Блок SwiGLU, используемый многими моделями с одним декодером, иллюстрирует эту проблему. Прямая реализация запускает матричные умножения для проекций, ядро для функции SiLU, другое для поэлементного умножения, а затем последнее матричное умножение. На каждой границе промежуточный тензор может быть записан в память с высокой пропускной способностью (HBM), а затем прочитан. Слияние заключается в объединении этих этапов в одно ядро. Оно экономит время, когда перемещения памяти и запуски ядер имеют большой вес, и гораздо меньше, когда большое матричное умножение уже доминирует в выполнении. Выбор ядра остаётся эмпирическим: необходимо профилировать всю рабочую нагрузку.
Другие слияния распространены в трансформерах, например, добавление остаточного соединения с последующей нормализацией RMSNorm или LayerNorm или совмещение деквантования с матричным умножением. Однако у слияния есть пределы: операции с несовместимыми организациями данных или очень объёмное промежуточное состояние плохо комбинируются. Слишком большое ядро может увеличить нагрузку на регистры, снизить коэффициент занятости и стать медленнее нескольких хорошо настроенных ядер.
Внимание типа FlashAttention и постраничное внимание
Классическое внимание вычисляет матрицу оценок размером T² на голову, где T — длина последовательности. На GPU её создание, запись в глобальную память и последующее чтение потребляют много пропускной способности, что может стать узким местом при длительном предварительном заполнении («prefill»). FlashAttention обрабатывает запросы, ключи и значения блоками. Для каждого он вычисляет часть оценок, обновляет статистики softmax и накапливает соответствующий выход в памяти на чипе, не сохраняя полную матрицу. PyTorch предоставляет эту оптимизацию через функцию scaled_dot_product_attention(). В руководстве уточняется, что жёсткое указание бэкенда может привести к сбою кода, если ни один бэкенд не поддерживает входные данные или оборудование.
Кэш KV, который сохраняет уже вычисленные ключи и значения, растёт на один токен за раз, причём конечная длина непредсказуема. Непрерывное выделение памяти на запрос оставляет зарезервированное и неиспользуемое пространство. Постраничный кэш разбивает хранилище на блоки фиксированного размера и использует таблицу блоков, связывающую логические позиции с физическими блоками. Ядру внимания тогда необходимо переводить каждую позицию через эту таблицу. Страничная организация не уменьшает количество логических элементов, но ограничивает фрагментацию и зарезервированную ёмкость.
В приведённом примере vLLM использует модель Qwen/Qwen2.5-1.5B-Instruct с блоками по 16 токенов (block_size=16) и бэкендом «TRITON_ATTN». Согласно руководству, vLLM проверяет совместимость бэкенда с моделью и оборудованием при инициализации и сообщает об ошибке в случае несовместимости.
Слияние и CUDA Graphs с torch.compile
Функция torch.compile() захватывает последовательность операций с помощью TorchDynamo, а затем компилирует её с помощью TorchInductor, своего бэкенда по умолчанию. TorchInductor может сгенерировать одно ядро Triton для совместимых операций, которые стандартный режим запускал бы отдельно. Разрывы графа, вызванные, например, потоком управления, зависящим от данных, или преобразованиями тензоров в скаляры Python, ограничивают области, которые компилятор может оптимизировать.
Режим «reduce-overhead» предписывает TorchInductor использовать CUDA Graphs, когда захваченный граф соответствует критериям. CUDA Graph записывает последовательность операций GPU и воспроизводит её с помощью единой отправки на процессор. Он не создаёт дополнительного слияния, а дополняет слияние TorchInductor. Он требует стабильных адресов памяти и совместимых форм и может потреблять больше памяти. Руководство рекомендует проверять результаты с допусками, соответствующими типу данных, профилировать выполнение и измерять сквозную производительность.
