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