🌟 DeepSelect: набор TopK-ядер для разреженного внимания и сэмплирования

DeepSeek выложила DeepSelect, высокопроизводительную реализацию ядра TopK, которое используется в механизме разреженного внимания DeepSeek Sparse Attention и на котором работают модели V3.2, V4 и V4.1.

Относительно torch.topk из PyTorch ядро даёт ускорение от 2 до 20 раз.

Быстрейший алгоритм зависит от типа данных, размера батча, размера словаря и величины k, поэтому библиотека заточена под два сценария:

🟢для индексатора в разреженном внимании, работает с bfloat16 при любых размерах батча и словаря;

🟢для сэмплирования, рассчитан на float32 и словарь около 128 тысяч элементов.

В обоих случаях k ограничен четырьмя тысячами, поддерживаются строки переменной длины через задание верхней границы для каждой строки, а проверка на NaN включена всегда и по умолчанию обрывает исполнение при обнаружении.

Для выжимания скорости DeepSeek советуют отключать сортировку индексов и не запрашивать значения, если нужны только позиции (последнее даёт около 10%).

⚠️ Есть вещи, которые придётся учесть:

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

🟠выходные тензоры библиотека выделяет сама и тоже выравнивает, из-за чего они могут оказаться несмежными;

🟠репозиторий свежий, инструмент только выложен и толком вне Deepseek ещё не обкатан.


📌Лицензирование: MIT License


🖥Github


#AI #ML #Kernels #GPU #DeepSelect #Deepseek