Оптимизация CUDA-ядер разреженного внимания в llama.cpp b11140¶
4.0/10
В релизе b11140 фреймворка llama.cpp оптимизированы CUDA-ядра сканирования масок разреженного внимания (sparse-fa) для длинных контекстов и повторно активирована поддержка sparse-fa для dsv4 prefill. За счет устранения динамических условий в циклах время сканирования для одиночных декодов (batch 1) при длине контекста 49k столбцов сократилось с 46 до 17 мкс. В пакетных операциях разреженного внимания задержка ядра на контексте 49k снизилась с 586 до 244 мкс. Это нововведение снижает накладные расходы планирования и выборки индексов внимания в сценариях с экстремально длинным контекстом.
Метод оптимизации¶
Ядро `flash_attn_mask_to_sparse_indices` было шаблонизировано по параметру `ncols1`, что связало число итераций цикла запросов во время компиляции и превратило его в прямолинейный код, позволяя линиям варпа объединять операции загрузки памяти. Дополнительно проверка выхода за границы столбцов и проверка частичных финальных групп запросов были вынесены на сторону хоста через параметр `n_queries`, устранив динамические проверки границ и ранние выходы из цикла для конфигурации `ncols1 == 8`.
Прирост эффективности¶
На контексте 49k токенов время выполнения сканирования маски снизилось с 46 до 17 мкс (ускорение в 2,7x) для одиночного декодирования и с 586 до 244 мкс (ускорение в 2,4x) для батчевых операций, устраняя задержки уровня микросекунд при вычислении разреженного внимания.