🌟 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