6 мин

74

Кратко: Три изменения в HNSW-поиске ускоряют KNN-поиск до 29% при больших k и дают более 20% прироста при параллельной нагрузке. Без изменений API, без перестроения индексов и без новых настроек — просто более быстрый поиск.

Как мы ускорили KNN-поиск в Manticore

KNN-поиск в Manticore построен поверх hnswlib , открытой реализации HNSW. Исторически большая часть нашей работы над KNN была связана с собственными функциями расстояния, например для бинарной квантизации, а не с основным поисковым циклом hnswlib. Мы также добавили предфильтрацию с ACORN-1 и раннее завершение , но сам цикл оставался прежним: hnswlib по-прежнему обходил соседей, вычислял расстояния и обновлял набор кандидатов.

На этот раз изменения затрагивают сам основной поисковый цикл hnswlib: то, как он обходит соседей, вызывает функции расстояния и работает с иерархией памяти CPU. В сочетании с новыми реализациями функций расстояния под AVX-512 в колонковой библиотеке они убирают три источника накладных расходов: неэффективный доступ к памяти, повторные загрузки данных и косвенные вызовы функций.

Специализация функции расстояния на этапе компиляции

Раньше функция расстояния хранилась в HNSW-индексе как указатель на функцию и выбиралась во время выполнения; её вызывали для каждого кандидата. Когда HNSW приходится проверять много кандидатов, это даёт множество косвенных вызовов на один запрос. Косвенные вызовы мешают компилятору встроить функцию расстояния в поисковый цикл и ухудшают предсказание ветвлений.

В новом коде выбор функции расстояния переносится на этап компиляции с помощью шаблонов C++. В начале поиска один оператор switch выбирает нужную специализацию шаблона по метрике расстояния и настройкам квантизации. После этого весь внутренний цикл — обход соседей, вычисление расстояния и обновление набора кандидатов — выполняется как единая функция с полностью встроенным расчётом расстояния. Благодаря этому компилятор может оптимизировать распределение регистров, планирование инструкций и разворачивание циклов с учётом вычисления расстояния.

Двухпроходная обработка соседей

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

Новая реализация делит эту работу на два прохода. На первом проходе код перебирает всех соседей текущего узла, пропускает уже посещённых и собирает непосещённых соседей в небольшой массив для пакетной обработки. Для каждого добавленного соседа сразу выдаётся prefetch-инструкция для его векторных данных. На втором проходе код проходит по этому массиву и вычисляет расстояния. К этому моменту предзагрузка из первого прохода уже успевает подтянуть данные в кэш.

Во втором проходе код идёт по компактному последовательному массиву идентификаторов кандидатов, а не по самой структуре графа. Сами векторы по-прежнему лежат в разных местах памяти, но к этому моменту нужные данные уже предварительно загружены в кэш.

Для нефильтрованных запросов (без условия WHERE в KNN-поиске) новый код также выбирает быстрый путь и полностью убирает проверку фильтра для каждого кандидата.

Пакетная обработка

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

При обработке двух кандидатов вектор запроса загружается один раз и повторно используется для обоих вычислений расстояния. Благодаря этому исчезают лишние загрузки.

Так сокращаются повторные загрузки вектора запроса, а цикл оценки может обрабатывать кандидатов парами; если кандидатов нечётное число, последний обрабатывается отдельно. Функции для обработки двух кандидатов за раз реализованы для скалярного произведения, L2 и бинарной квантизации.

Поддержка AVX-512

Новый код вычисления расстояний под AVX-512 обрабатывает 16 чисел с плавающей точкой за итерацию вместо 8 в AVX2. Для скалярного произведения и расстояния L2 основной цикл использует fused multiply-add (_mm512_fmadd_ps): умножение и накопление выполняются одной инструкцией. Для векторных данных с бинарной квантизацией расширение AVX-512 VPOPCNTDQ ускоряет операции подсчёта бит, используемые при вычислении расстояния.

Теперь Manticore поставляется в трёх сборках библиотеки: базовой, с поддержкой AVX2 и с поддержкой AVX-512. При запуске Manticore определяет возможности CPU и автоматически загружает подходящую библиотеку. Ничего дополнительно настраивать не нужно.

Результаты тестов

Бенчмарки ниже запускались на датасете dbpedia-openai-1M-1536-angular (1 млн векторов, 1536 измерений, косинусное расстояние) на AMD Ryzen 7 9700X (Zen 5, 8 физических ядер / 16 логических ядер). Во всех тестах использовалась 1-битная квантизация, oversampling и рескоринг были отключены. В многопоточных запусках каждый воркер выполнял свой набор запросов; мы измеряли QPS каждого воркера отдельно и усредняли значения. Каждый результат — среднее по 6 независимым запускам. Раннее завершение также отключили, чтобы измерить влияние этих оптимизаций именно на обход HNSW.

Процессор Zen 5 был выбран потому, что AVX-512 на нём работает с полной 512-битной шириной: без разбиения на две 256-битные части и без сильного снижения частоты, характерного для некоторых более старых процессоров Intel. Это помогает отделить изменения в алгоритме от особенностей конкретных CPU при выполнении AVX-512.

Только алгоритмические улучшения

Первый график показывает эффект алгоритмических изменений отдельно от набора SIMD-инструкций: двухпроходной обработки, оценки кандидатов парами и диспетчеризации на этапе компиляции. Для этого новая сборка AVX2 сравнивается с предыдущей сборкой AVX2. Обе используют один и тот же набор SIMD-инструкций, поэтому разница объясняется именно изменениями в поисковом цикле.

В однопоточном режиме прирост постепенно увеличивается с +3% при k=10 до +24% при k=1000: вычисление расстояний начинает доминировать в нагрузке поиска. Когда больше потоков конкурирует за пропускную способность памяти, выигрыш на поток снижается: до +9-10% при 4 или 8 потоках и до +2-5% при 16 потоках.

Тест с 16 потоками — это SMT: на каждом физическом ядре работают два потока. Вычисление расстояний упирается в память, поэтому когда два потока делят L1/L2-кэши одного ядра, выигрыш от предзагрузки и пакетной обработки частично теряется из-за конкуренции за общие ресурсы. Алгоритмические улучшения всё равно помогают, но запас по производительности становится меньше.

Что даёт большая ширина SIMD (AVX-512 vs AVX2)

Второй график показывает эффект AVX-512 отдельно: сборка AVX-512 сравнивается с новой сборкой AVX2, при этом в обеих используются одни и те же алгоритмические улучшения.

При k=10 AVX-512 немного медленнее AVX2 — примерно на 2% независимо от числа потоков. Это особенность именно AVX-512: одни только алгоритмические улучшения такого регресса не показывают, значит дело не в фиксированных накладных расходах, одинаковых для всех запросов. Начиная с k=30 AVX-512 выходит вперёд при любом количестве потоков.

Интересно, что выигрыш от AVX-512 растёт вместе с числом потоков. Хотя в этом бенчмарке oversampling отключён, KNN-запрос в Manticore по умолчанию использует LIMIT 20, а стандартное значение oversampling=3.0 умножает k перед повторной оценкой после квантизированного поиска, поэтому HNSW внутри работает с k=60. При k=60 AVX-512 даёт прирост относительно новой AVX2-сборки: +1,2% в одном потоке, +2,6% при 4 потоках, +3,4% при 8 потоках и +6,5% при 16 потоках.

Суммарный эффект (новый AVX-512 vs старый AVX2)

Третий график показывает суммарный эффект: AVX-512 со всем новым кодом сравнивается с предыдущей сборкой AVX2. Именно такой результат увидит пользователь, который обновится с предыдущей версии Manticore на новую, если его CPU поддерживает AVX-512.

Кривая для одного потока растёт с +0,5% при k=10 до +29% при k=1000. В многопоточном режиме все кривые доходят до +22-24% при k=1000. Улучшение видно при любом числе потоков: алгоритмические и SIMD-выигрыши складываются по-разному при разном уровне параллелизма, но итоговый результат остаётся стабильно высоким при средних и больших k.

Почему выигрыш растёт вместе с k

На всех трёх графиках форма одинаковая: небольшой выигрыш при низких k и большой при высоких. Причина в том, что при небольших k запросы тратят большую долю времени на обход структуры графа — посещение узлов, проверку битов посещения и работу с очередью кандидатов. Эта работа зависит от структуры графа, а не от k. С ростом k увеличивается число кандидатов, для которых нужно посчитать расстояние, и всё больше времени уходит именно на эти вычисления. Оптимизации нацелены на вычисление расстояний и связанные с ним циклы, поэтому их эффект растёт вместе с долей такой работы.

Что это значит для вас

Эти улучшения не требуют никаких действий. Они доступны в недавнем релизе Manticore Search 27.1.5 ; изменений в API нет, новых параметров конфигурации нет, и перестраивать таблицы не нужно.

Эти улучшения хорошо сочетаются с ранним завершением KNN : раннее завершение сокращает число вычислений расстояния на запрос, а новые оптимизации ускоряют каждое вычисление.

Наибольший эффект заметен в таких сценариях:

• Векторы высокой размерности (больше арифметики на одно вычисление расстояния, выше выигрыш от SIMD)

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

• Запросы с oversampling (oversampling увеличивает k и переводит запросы в диапазон, где выигрыш максимален)

Дополнительные материалы

• Документация по раннему завершению KNN — как Manticore определяет сходимость HNSW и останавливает поиск заранее

• Документация по фильтрации KNN — как работает предфильтрация с ACORN-1

• Справочник по векторному KNN-поиску — полный синтаксис, параметры, квантизация и рескоринг