Открыт FlashQLA — ускоренные fused linear attention‑ядра для GDN, оптимизированные под CP и обратный проход

По мере роста архитектур до масштабов 397A17B / 122A10B / 35B / 27B и увеличения контекстных окон свыше 256K накладные расходы блока GDN в сквозном обучении и инференсе перестали быть пренебрежимо малыми. В ответ на это команда представила FlashQLA — библиотеку линейных attention‑ядер, реализованную на TileLang и ориентированную на ускорение как прямого, так и обратного проходов GDN Chunked Prefill. FlashQLA достигает ускорения 2–3× по прямому проходу и 2× по обратному по сравнению с FLA Triton в ряде рабочих сценариев на архитектуре NVIDIA Hopper. Авторы подчёркивают, что такие выигрыши особенно важны в сценариях предобучения больших моделей и при развертывании агентных систем на устройствах края, где стоимость вычислений и эффективность использования GPU критичны.

Ключевой подход — применение разумного слияния операторов и целенаправленных оптимизаций производительности на уровне выполнения. Это включает преобразования, которые уменьшают количество отдельных CUDA‑ядерных вызовов и повышают плотность работы для тензорных и CUDA‑ядр, не жертвуя числовой точностью. Первая важная особенность — «gate‑driven automatic intra‑card context parallelism». Используя свойство экспоненциального затухания GDN‑врат, FlashQLA автоматически включает внутрикартовый CP в условиях tensor‑parallelism (TP), при длинных последовательностях и малом числе голов, что улучшает загрузку потоковых мультипроцессоров (SM) GPU.

Вторая особенность — аппаратно‑дружественная алгебраическая переработка прямого и обратного потоков GDN Chunked Prefill. Авторы перерассматривают выражения выполнения таким образом, чтобы снизить накладные расходы на Tensor Core, CUDA Core и SFU при сохранении численной точности, что даёт ощутимый выигрыш в производительности на целевой аппаратуре. Третья особенность — Fuse‑ядра, специализированные под warp и реализованные на TileLang. Такие слияния позволяют избавиться от множественных мелких кернелей и повысить эффективность вычислений на уровне warp. Код и результаты тестов выложены в репозитории: github.com/QwenLM/FlashQLA.

Авторы также разбирают существующую схему вычислений FLA GDN Chunked Prefill: для очереди чанка i, если не учитывать предобработку ворот и CP, каждый шаг этого потока в FLA соответствует отдельному кернелу. По их определению, такой поток содержит две основные проблемы эффективности, на которые направлены описанные оптимизации FlashQLA.