GPU-Resident Top-K for Agentic RAG: I Built a CUDA Kernel So My Retrieval Step of Resources Would Stop Boost Over GPU

Как замена платы за обмен данными в Python на пользовательскую архитектуру памяти графического процессора позволяет добиться детерминированных задержек в микросекундах для многошагового RAG.

Делиться

Наглядный, основанный на эмпирических данных, обзор алгоритма поиска Top-K в CUDA, состоящий из 343 строк кода. Этот набор тестов ядра, процессора и бенчмарков доказывает, что стандартный алгоритм Agentic RAG, использующий переадресацию запросов по шине PCIe, является скрытым убийцей вашего конвейера обработки данных. Благодаря хранению результатов поиска по сходству в оперативной памяти устройства, эта архитектура обеспечивает 8,6-кратное ускорение по сравнению с оптимизированными базовыми алгоритмами для ЦП даже на 7-летней видеокарте GTX 1080.

Это третья часть серии статей «Производственный уровень агентного вывода» . Каждая часть устраняет один из видов избыточной работы в конвейере агентного LLM. В первой части была устранена избыточная предварительная обработка. Во второй части было устранено избыточное ожидание — как несколько микроагентов совместно используют один графический процессор с помощью разделения по времени. В третьей части (эта статья) получение RAG-данных остается на графическом процессоре с помощью пользовательского ядра CUDA Top-K. В четвертой части состояние агента сохраняется при передаче управления, чтобы у следующего агента никогда не возникала проблема холодного запуска.

Основные выводы

Проблема: в агентном RAG каждый вызов инструмента, требующий контекста, запускает поиск сходства. Конвейер по умолчанию отправляет векторное представление запроса с графического процессора в Python, позволяет центральному процессору оценить N строк корпуса и выбрать лучшие K, а затем отправляет ответ обратно. Этот цикл обмена данными — это скрытая плата. Вычисления — это хорошо; расходы на поездки — это счет. Мы все знаем, что поездки никогда не бывают дешевыми, независимо от того, куда вы хотите поехать (игра слов!).

Простое решение: загрузить корпус в видеопамять один раз, а затем сохранить оценку сходства, выборку Top-K и этап слияния на устройстве. Через PCIe передаются только крошечные векторные представления для каждого запроса (D — числа с плавающей запятой) и K результатов.

Результаты: на той же 7-летней GTX 1080, что и в частях 1 и 2, путь, выполняемый на GPU, проходит этап извлечения данных до 8,57 раз быстрее, чем базовый вариант с использованием перебора на CPU. При K=8 он выигрывает во всех 15 конфигурациях сканирования (N ∈ {10k, 50k, 100k, 500k, 1M}, D ∈ {384, 768, 1024}) с ускорением от 2,43× до 8,57×. При K=32 он выигрывает в 13 из 15 конфигураций, достигая пика в 7,76× . При K=100 — где селектор V1 намеренно остается простым — CPU выигрывает в 14 из 15 конфигураций. Последнее предложение — это честная часть (Ну, даже если бы я солгал, вы бы легко это заметили).

Самое интересное: это не победы, достигнутые благодаря «волшебному ядру». Это победы, достигнутые благодаря тому, что «мы перестали отправлять корпус обратно в оперативную память хоста без всякой причины». Это также в точности тот тип решений, которые «измеряем множество кандидатов, сообщаем потребителю только лучшее значение K», которые базовая станция 5G и ваш телефон принимают каждые несколько миллисекунд с тех пор, как появилась обратная связь CSI.

Вкратце: В стандартном агентном RAG графический процессор рассматривается как сервер, а поиск — как задача Python. Каждый вызов инструмента передает векторное представление запроса D→H , позволяет ЦП вычислить N скалярных произведений, отсортировать кандидатов, выбрать K лучших и передать индексы и оценки H→D . Для агента, который обращается к векторному хранилищу десять раз за шаг рассуждения, это обращение является основной стоимостью — не модель, не векторное представление, а само перемещение. CUDA-TopK-Retrieval хранит корпус непосредственно на устройстве, выполняет оценку + частичный Top-K для каждого блока + многостороннее слияние полностью на графическом процессоре и предоставляет небольшой API оркестратора на C++ ( upload_corpus_rowmajor один раз, search_resident для каждого запроса). Количество байтов, затрагиваемых хостом для каждого запроса, сводится к одному векторному представлению длиной D вверх и 2K результатам вниз. На GTX 1080, при проверке 45 конфигураций, путь, используемый графическим процессором, превосходит базовый показатель времени выполнения на ЦП на всех 15 конфигурациях K=8 (в 2,43 раза против 8,57 раза, с пиком при N=1M, D=1024) и на 13 из 15 конфигураций K=32 (оба проигрыша наблюдаются при наименьшем N=10k для D=384 и D=768, где само время выполнения уже является дешевым; выигрыш при большом N K=32 достигает 7,76 раза). При K=100 ядро V1 намеренно остается простым — пузырьковая сортировка с одной полосой на блок с последовательным слиянием — и ЦП выигрывает на 14 из 15 конфигураций; этот потолок является честным финалом статьи и хорошей подготовкой к части 4.

(Небольшое признание, прежде чем мы начнем: я пришел к этому с опытом работы в области проектирования сетей RAN для 5G/6G. Выбор луча на базовой станции поразительно похож на алгоритм RAG Top-K — UE оценивает список потенциальных лучей по принимаемой мощности и передает по радио лучшие из них. Об этом есть целый раздел ниже — раздел 8 — но именно поэтому это ядро существует в таком виде.)

Архитектурная ментальная модель — держите это открытым во время чтения.

agent.embed(query) → cudaMemcpy H→D (D floats) → row_dot_scores_kernel → partial_topk_block_kernel (P blocks) → merge_partial_topk_kernel → cudaMemcpy D→H (K indices + K scores)

Всё, что ниже приведено, — это лишь комментарии к одной части этой строки.

1. Признание: каждый шаг RAG в вашем агенте — это крошечное путешествие по PCIe-сети.

Во второй части этой серии мы успешно изолировали цикл вывода нашего агента LLM, обеспечив быструю и эффективную генерацию токенов на устройстве. Мы разработали систему, которая предотвращает задержки. Но как только мы даем этому агенту инструмент для поиска во внешней базе знаний — ядре любого многошагового конвейера генерации с дополненным поиском (RAG) — мы незаметно уничтожаем всю с трудом достигнутую производительность и сталкиваемся с препятствием. Если вы когда-либо подключали «агентный» конвейер к хранилищу векторов через Python-ретривер, вот что на самом деле происходит при каждом вызове инструмента (с небольшим намеренным драматизмом):

Вы: «Агент, найдите мне пять фрагментов, наиболее относящихся к вопросу „Как мне получить налоговый вычет в соответствии с разделом 80C?“»

Агент: «Конечно. Встраиваем запрос на графический процессор. ✅»

Агент: «Запрос на встраивание возвращается на хост».

(cudaMemcpy D→H, ~1024 float) Python retriever: “Got it. NumPy loop. Dot product N times. argpartition. Top-5.”

(Центральный процессор обрабатывает полмиллиона строк корпуса по одной строке за раз, в то время как графический процессор мощностью 9 ТФЛОП наблюдает за процессом.) Python retriever: «Готово. Вот индексы и оценки».

Агент: «Отлично. Сейчас перенаправляю их обратно на графический процессор».

(cudaMemcpy H→D, 10 чисел) Агент: «Готов. Какой был вопрос?»

У агента есть вполне исправный графический процессор. Корпус находится в 4 ГБ видеопамяти. Встраивание запроса уже было на графическом процессоре — мы просто сгенерировали его там. Затем на каждом этапе получения данных мы отправляем запрос обратно на хост, выполняем перебор сходства с помощью NumPy / FAISS-on-CPU / самодельного цикла и отправляем ответ обратно.

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

Это также грязный секрет каждой демонстрации агентного RAG, которая выходит за рамки упрощенного этапа «десять блоков в памяти». Каждый раз процесс извлечения данных проходит через графический процессор и обратно, и чем больше корпус, тем выше затраты. При миллионе строк 1024-мерных эмбеддингов один только обмен данными — даже не оценка, да, только обмен данными — съедает большую часть бюджета самого этапа извлечения.

CUDA-TopK-Retrieval — это то, что происходит, когда вы решаете, что обмен данными туда и обратно необязателен, и вы предпочитаете написать 343 строки кода CUDA, чем позволить агенту каждый раз, когда ему нужен сосед, «путешествовать» по оперативной памяти хоста.

Теперь представьте себе реальную рабочую нагрузку. Это не «пять фрагментов для одного вопроса». Это множество специализированных микроагентов — каждый из которых выполняет свои собственные RAG-переходы, каждому требуется Top-K для одного и того же корпуса, и каждый в настоящее время оплачивает свой собственный счет за PCIe при каждом вызове инструмента. В первой части этой серии был отменен цикл предварительного заполнения. Во второй части была сделана возможность совместного использования графического процессора этими несколькими агентами. В третьей части говорится: теперь, когда они справедливо используют карту, прекратите заставлять каждого из них возвращаться к хосту, чтобы найти соседнего агента.

2. Зачем нужен поиск Top-K? (краткий курс за одну минуту)

Если вы уже знакомы с этой темой, пропустите этот раздел. Для всех остальных, кто только начинает знакомиться с ней, вот краткое, любительское объяснение.

Современный агент не помещает всю базу знаний в подсказку. Он её извлекает. Для каждого шага рассуждения, требующего контекстного обоснования, он встраивает запрос в вектор фиксированной размерности (D чисел с плавающей запятой — обычно 384, 768 или 1024), оценивает этот вектор по каждой строке корпуса предварительно встроенных фрагментов (N строк, также D чисел с плавающей запятой в каждой) и возвращает K строк корпуса с наибольшим сходством. Вот и всё. Это поиск по K лучшим векторам. Генерация с расширенным поиском — это просто вежливый способ сказать «K лучших векторов плюс шаблон подсказки».

Везде встречаются два типа сходства. Скалярное произведение — более дешевый вариант: одно объединенное умножение-сложение на каждое измерение, всего N×D работы. Косинусное произведение — это скалярное произведение, деленное на произведение норм L2, которое становится бесплатным скалярным произведением, если предварительно нормализовать корпус при загрузке. Большинство производственных хранилищ векторной информации используют трюк с предварительной нормализацией и называют это «косинусом», выполняя при этом вычисления скалярного произведения во время запроса. Ядро CUDA-TopK-Retrieval поддерживает оба варианта — оно просто умножает на предварительно вычисленный указатель нормы для каждой строки, когда включен режим косинуса.

Основные инструменты (FAISS, hnswlib, Python-часть cuVS, ваши любимые SaaS-базы данных векторов) выполняют эту работу по оценке и поиску лучших K результатов. Большинство из них делают это хорошо. Проблема в том, где они это делают. Практически каждая агентская платформа в мире обращается к механизму извлечения данных из Python, и как только Python оказывается в центре внимания, этап извлечения данных перестает быть операцией GPU — это операция PCIe с GPU на одном конце.

Решение заключается не в «улучшении алгоритма», а в «значительно сокращении времени в пути».

3. Идея «просто храните корпус на графическом процессоре» (и почему это сложнее, чем кажется)

Суть предложения проста.

• Загрузите корпус в видеопамять сразу после его обработки.
• Для каждого входящего запроса cudaMemcpy для копирования крошечного D -мерного векторного представления числа с плавающей запятой на устройство.
• Запустите ядро оценки, в котором один поток CUDA на каждую строку корпуса вычисляет скалярное произведение.
• Запустите частичное ядро алгоритма Top-K, в котором каждый блок сканирует непересекающийся диапазон строк, чтобы выдать свои собственные локальные лучшие кандидаты.
• Наконец, запустите ядро слияния, чтобы пройтись по заголовкам блоков и выдать глобальный топ-K в порядке наилучшего первого порядка.

С cudaMemcpy вы отправляете на хост ровно 2K чисел: K индексов и K оценок.

Это парадигма «рассматривайте извлечение данных из памяти как аппаратный примитив, а не как вызов программного API». Единственная причина, по которой для этого требуется более 30 строк скрипта PyTorch, заключается в том, что три утомительных крайних случая немедленно нарушат наивный подход.

Проблема А: алгоритм Top-K на графическом процессоре имеет неудобную структуру.

Оценка векторов — это простая часть. Это всего лишь умножение матриц, и ваш графический процессор буквально создан для этого — это язык любви оборудования. Однако на этапе выбора романтика умирает. Заставлять графический процессор выполнять полную сортировку O(N log N) , чтобы получить только K лучших результатов, — это вычислительно обременительно; это как перебирать по алфавиту всю корзину для переработки, чтобы найти один-единственный чек. Можно попробовать O(N) argpartition , но это потребует обхода дерева, что приведет к тому, что память графического процессора превратится в миллион невыровненных операций чтения. Турнирный выбор — это быстро, если вы готовы потратить выходные на отладку граничных случаев. И в тот момент, когда вы сдадитесь и потянетесь к примитиву сортировки thrust или cub , поздравляю: вы только что заразили свой легковесный, автономный конвейер C++ огромной зависимостью от сборки.

Архитектура намеренно выбирает скучный вариант решения. Она основана на крошечной сортировке пузырьком O(K2) для каждого блока в диапазоне непересекающихся строк, управляемой одним потоком на блок, и завершается последовательным многосторонним слиянием. На бумаге это звучит ужасно. На практике же это работает прекрасно, именно по той причине, которая честно указана в комментариях к ядру:

 // Single-threaded per-block scan that materializes a local Top-K list for its row partition. // This is not the fastest global selection, but it is easy to reason about and matches the CPU ordering rule exactly. __device__ void bubble_downward(float* const s, int* const ids, const int n) { // Tiny O(K^2) sort acceptable because K is capped (kMaxSupportedK) and this runs on a single lane per block. for (int i = 0; i < n - 1; ++i) { for (int j = 0; j < n - 1 - i; ++j) { if (device_is_better(s[j + 1], ids[j + 1], s[j], ids[j])) { const float ts = s[j]; s[j] = s[j + 1]; s[j + 1] = ts; const int ti = ids[j]; ids[j] = ids[j + 1]; ids[j + 1] = ti; } } } }

Это проектный контракт: в версии V1 приоритет отдаётся возможности аудита, а не продуманности. Всё ядро достаточно компактно, чтобы рецензент мог прочитать его от начала до конца за перерывом на кофе, сравнить с результатами работы ЦП и убедиться в побитовой корректности выходных данных ГП. В тот день, когда кому-то понадобится ядро 2×, оно будет заменено на специализированный селектор турнира для варпов — это прямо обещано в разделе 9 статьи. Для K ≤ 32 однополосная пузырьковая сортировка действительно хороша. Для K=100 она резко падает. Этот обрыв задокументирован и протестирован. (Снова раздел 9.)

Задача B: GPU и CPU должны побитово согласовать решение по критерию разрешения ничьей.

Как и в любом матче группового этапа чемпионата мира по футболу, ничьи случаются довольно часто. Две строки корпуса имеют одинаковый счет с точностью до fp32. Кто из них победит?

Если ядро ЦП и ядро ГП расходятся во мнениях относительно определения победителя в случае равенства результатов, вы никогда не сможете доверять бенчмарку, потому что каждое сообщение о «несоответствии» становится неоднозначным: ГП неправильно оценил результат в строке или просто по-другому определил победителя в случае равенства результатов? Вы проведете неделю в канале Slack в 3 часа ночи.

Решение состоит в том, чтобы определить компаратор в одном предложении и реализовать его в двух местах — один раз на хосте, другой раз на устройстве — и сделать эти две реализации буквально одним и тем же выражением.

На стороне хоста это выглядит так:

 // Lexicographic "better" relation for (score, index) pairs under float equality semantics. // We use strict weak ordering for std::partial_sort: higher score wins; on exact tie, smaller index wins. bool is_better_score_pair(const float32_t score_lhs, const index_t idx_lhs, const float32_t score_rhs, const index_t idx_rhs) { // Primary key: similarity score (higher is better for retrieval). if (score_lhs != score_rhs) { return score_lhs > score_rhs; } // Deterministic tie surface: prefer the smaller corpus row id to mirror stable DB primary keys. return idx_lhs < idx_rhs; }

Со стороны устройства это выглядит так:

 // Device-side replica of the host comparator to avoid cross-TU linkage issues for __device__ code paths. __device__ bool device_is_better(const float score_lhs, const int idx_lhs, const float score_rhs, const int idx_rhs) { // Same ordering semantics as topk::is_better_score_pair for bitwise-identical tie surfaces. if (score_lhs != score_rhs) { return score_lhs > score_rhs; } return idx_lhs < idx_rhs; }

Вот вся политика разрешения ничьих в пяти строках, дважды. Побеждает более высокий балл; при точном совпадении побеждает индекс строки меньшего размера корпуса. CPU-оракул использует его в std::partial_sort , GPU — в сортировке пузырьком и в многостороннем слиянии, а тестовая среда не начнет замер времени, пока выходные данные GPU точно не совпадут с выходными данными CPU — одинаковые индексы в одном порядке, оценки в пределах небольшого допуска fp32.

Именно этот единственный компаратор позволяет статье вообще упоминать об ускорении. Без него фраза «GPU в 8 раз быстрее» превратилась бы просто в «GPU в 8 раз быстрее ошибается по-разному».

Проблема C: Видеопамять — ценный ресурс, и худшее место для выделения памяти malloc — это наиболее часто используемый участок видеопамяти.

Выделение памяти графического процессора для каждого запроса — это как каждый раз оформлять новый договор лизинга на автомобиль, когда вам нужно поехать в продуктовый магазин. Это самый простой способ превратить поиск, занимающий 1 миллисекунду, в пробку, растянувшуюся на 50 миллисекунд.

Вместо этого GpuTopkEngine::initialize берет управление на себя заранее. Он выполняет все вызовы cudaMalloc во время запуска движка, определяя размер буферов для наихудшей возможной конфигурации. Как только движок начинает активно обрабатывать запросы, «горячий путь» полностью освобождается от управления памятью. Остаются только быстрые запуски ядра и крошечные копии данных. Никакой фрагментации, никаких переговоров с распределителем памяти, и cudaMalloc навсегда исключается из трассировки производительности.

4. Четырехступенчатый конвейер (самая интересная часть)

 Step 0: Engine init — eight cudaMallocs sized for (max_n, max_d, max_k) (GpuTopkEngine::initialize) Step 1: Upload corpus once into VRAM (upload_corpus_rowmajor) Step 2: Per query — H→D the embedding (search_resident, first line) Step 3: Score N rows on device (row_dot_scores_kernel) Step 4: Per-block partial Top-K (partial_topk_block_kernel) Step 5: Multi-way merge into global Top-K (merge_partial_topk_kernel) Step 6: D→H the K indices + K scores (search_resident, last lines)

Давайте разберем каждый фрагмент кода. Приведенные фрагменты даже короче, чем крошечные исходные файлы.

Шаг 1 — Загрузите корпус один раз.

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

 cudaError_t GpuTopkEngine::upload_corpus_rowmajor(const float32_t* const host_corpus_rowmajor, const index_t N, const index_t D) { if (N > max_n_ || D > max_d_) { return cudaErrorInvalidValue; } const std::size_t corpus_bytes = sizeof(float) * static_cast<:size_t>(N) * static_cast<:size_t>(D); const cudaError_t st = cudaMemcpy(d_corpus_, host_corpus_rowmajor, corpus_bytes, cudaMemcpyHostToDevice); if (st != cudaSuccess) { return st; } resident_n_ = N; resident_d_ = D; return cudaSuccess; }

И это весь API для приема данных. При 1024 измерениях один миллион векторов — это ровно 4 ГБ данных float32 , идеально помещающихся в 8 ГБ видеопамяти старой GTX 1080. Что произойдет, когда ваш корпус достигнет 10 миллионов векторов? Это станет проблемой распределенных систем, а не проблемой ядра. Если ваши данные превышают объем видеопамяти, вам потребуется стратегия сегментирования, которую мы рассмотрим в разделе 9. Но сейчас мы здесь для того, чтобы решить проблему узкого места в вычислительных ресурсах, а не для того, чтобы изобрести новую базу данных.

Шаг 2 — Отметьте N строк на устройстве.

Один поток CUDA на строку корпуса. 256 потоков на блок. Каждый поток накапливает скалярное произведение по D измерениям и записывает одно число с плавающей запятой в плотный буфер scores[N] :

 // Row-major dot-product with optional cosine normalization; coalesced reads along D are sacrificed for clarity in v1. // Microarchitectural note: one thread per row is simple; a follow-up can tile D across warps to raise arithmetic intensity. __global__ void row_dot_scores_kernel(const float* const corpus, const float* const query, const float* const row_l2, const float query_l2, const int N, const int D, const int cosine_enabled, float* const scores) { // Map each CUDA thread to exactly one corpus row to keep the reduction logic easy to audit against the CPU reference. const int row = static_cast(blockIdx.x) * static_cast(blockDim.x) + static_cast(threadIdx.x); if (row >= N) { return; } float acc = 0.0F; const int base = row * D; for (int col = 0; col < D; ++col) { acc += corpus[static_cast<:size_t>(base + col)] * query[static_cast<:size_t>(col)]; } if (cosine_enabled != 0) { const float denom = query_l2 * row_l2_fetch(row_l2, row); scores[static_cast<:size_t>(row)] = denom > 0.0F ? (acc / denom) : -std::numeric_limits::infinity(); } else { scores[static_cast<:size_t>(row)] = acc; } }

Один поток на строку — это простейшее возможное отображение. Комментарий к коду об этом честно говорит: последующий поток может замостить D по варпам, чтобы повысить интенсивность арифметических операций. Для V1 это обеспечивает аудитору однозначное соответствие с циклом ЦП и позволяет ему спокойно спать по ночам.

Шаг 3 — Каждый блок строит свой собственный локальный Top-K

Теперь начинается самая неудобная часть. Выбор K лучших элементов из N концептуально представляет собой «сортировку и нарезку», но полная сортировка отнимает большую часть работы. Мы разделяем диапазон строк на P блоков (ограниченных 128), каждый блок проходит по своему непересекающемуся срезу с помощью крошечной пузырьковой сортировки из раздела 3 и записываем свой собственный локальный список из K лучших элементов:

 const int P = std::min(kMaxPartialBlocks, std::max(1, (static_cast(N) + 4095) / 4096)); partial_topk_block_kernel<<
>>(d_scores_, static_cast(N), static_cast(K), P, d_partial_scores_, d_partial_indices_);

Один поток на блок. Да, на бумаге это выглядит расточительно. Именно поэтому человек может проверить это ядро за двадцать минут — политика if (threadIdx.x != 0 || blockIdx.x >= P) return; сводит всю внутриблочную логику к «полоса 0 этого блока владеет строками [start, end) ». Массивы s[] и ids[] каждого блока находятся в регистрах/локальной памяти, их размер определяется ограничением kMaxSupportedK = 256 , заданным на этапе компиляции.

Шаг 4 — Объедините частичные значения в глобальный Top-K.

Наконец, один поток на одном блоке проходит P курсоров по спискам для каждого блока. Каждый список уже имеет приоритет наилучшего варианта. Выбираем наилучший начальный элемент; выводим результат; продвигаем этот курсор на один уровень вперед; повторяем K раз:

 for (int out = 0; out < K; ++out) { int best_p = -1; float best_s = -std::numeric_limits::infinity(); int best_i = std::numeric_limits::max(); for (int p = 0; p < P; ++p) { if (heads[p] >= K) { continue; } const float s = partial_scores[static_cast<:size_t>(p * K + heads[p])]; const int idx = partial_indices[static_cast<:size_t>(p * K + heads[p])]; if (best_p < 0 || device_is_better(s, idx, best_s, best_i)) { best_p = p; best_s = s; best_i = idx; } } out_scores[static_cast<:size_t>(out)] = best_s; out_indices[static_cast<:size_t>(out)] = best_i; heads[best_p] += 1; }

Слияние выполняется с невероятной эффективностью: максимум P * K операций чтения и ровно K операций записи, выполняемых одним потоком. Для предотвращения хаоса, связанного с вычислениями с плавающей запятой, компаратор device_is_better обеспечивает строгий детерминизм — если два варианта имеют одинаковый результат, побеждает вариант с меньшим индексом строки корпуса, идеально повторяя алгоритм ЦП. Наконец, два микроскопических вызова cudaMemcpy передают K выигрышных индексов и результатов обратно на хост. Агент их обрабатывает, и цикл RAG запускается снова.

Это весь процесс "горячего" доступа: одна передача эмбеддинга H -> D , три запуска ядра и два крошечных копирования результата D -> H Никаких циклов Python на хосте, никаких накладных расходов фреймворка и абсолютно никаких отключений PCIe.

5. Чеки (т.е., цифры)

Теперь давайте сравним это с базовым показателем и посмотрим, стоило ли это затраченных усилий.

Краткое замечание по методологии, прежде чем появятся эксперты по бенчмаркингу: каждое сравнение ниже выполняется на том же графическом процессоре, что и в Части 1 и Части 2 (NVIDIA GeForce GTX 1080, Pascal sm_61, 8 ГБ), драйвер 535.309.01, CUDA 12.2, процессор Intel Core i7-8700K, флаги компилятора -O3 -march=native --expt-relaxed-constexpr . Три испытания, одна разминка, начальное значение генератора случайных чисел 1, фиксированный генератор случайных чисел ( std::mt19937 со std::normal_distribution ), гауссовы вложения, L2-нормализация в режиме скалярного произведения. Полный цикл измерений: N ∈ {10k, 50k, 100k, 500k, 1M} × D ∈ {384, 768, 1024} × K ∈ {8, 32, 100} → 45 конфигураций , все измерения выполнены с помощью cudaEventElapsedTime , при этом cudaDeviceSynchronize ограничивает каждый интервал. Код находится в src/host/bench_main.cpp ; исходные данные — в examples/example-run-results/benchmark_run_results.csv .

Проверены два маршрута по времени:

• Процессор находится на графическом процессоре (обработка). Корпус уже находится на устройстве. Каждая итерация с таймером: запрос cudaMemcpy H→D (D чисел с плавающей запятой) + ядро оценки + частичный Top-K для каждого блока + слияние + k оценок cudaMemcpy D→H + k индексов cudaMemcpy D→H. От начала до конца.
• Обращение ЦП туда и обратно (базовый вариант). Моделирует стандартный поток работы агента: запрос cudaMemcpy D→H + перебор результатов ЦП + std::partial_sort с тем же компаратором + индексы cudaMemcpy H→D + оценки cudaMemcpy H→D . От начала до конца.

Оба пути выполняются внутри одного и того же процесса, используют одни и те же байты запроса и один и тот же компаратор. Единственное различие заключается в том, где происходит работа. Если вы когда-либо придерживались позиции «PCIe — это нормально, мы тестируем ядра изолированно», вот во что вы в конечном итоге ввязываетесь, когда перестаёте делать вид, что обмен данными в обоих направлениях бесплатен.

Заголовок (GTX 1080, три испытания, среднее значение в мс, соотношения рассчитаны на основе средних значений по каждому испытанию):

Да, вы правильно читаете цифры.

Первые пять строк отражают суть статьи: при K=8 путь, использующий графический процессор, выигрывает в каждой из 15 конфигураций в ходе исследования , с коэффициентами от вежливых 2,43× при N=50k, D=384 до громких 8,57× при N=1M, D=1024. При K=32 он выигрывает в 13 из 15 случаев — оба проигрыша приходятся на наименьшее N (10k), для D=384 и D=768, где само время кругового пути составляет всего ~3–7 мс, а три запуска ядра графического процессора едва успевают окупиться. К моменту достижения реалистичных размеров корпуса агентов (N ≥ 50k) K=32 также уверенно выигрывает, достигая пика в 7,76× . Значительное ускорение достигается не за счет «магического ядра», а за счет того, что мы «прекратили отправлять корпус обратно в оперативную память хоста без всякой причины». Графический процессор всегда был обречен на победу в этой гонке; единственная причина, по которой он когда-либо проиграл, заключалась в том, что мы постоянно заставляли его выполнять ненужные запросы.

Последние две строки — это то, где эта статья заслуживает называться честной. При K=100 однополосная пузырьковая сортировка на блок становится O(K²) = O(10 000) последовательных сравнений на блок, а последовательное слияние проходит P × K позиций головки. std::partial_sort процессора основана на куче, векторизована компилятором и фактически имеет сложность O(N log K) — гораздо более удобна для K=100. Таким образом, графический процессор проигрывает в 14 из 15 конфигураций K=100, иногда в 2 раза, иногда в 46 раз. (Есть ровно одна конфигурация K=100, где графический процессор все еще выигрывает — N=1M, D=1024, 1,44× — потому что к этому моменту работы по оценке достаточно, чтобы превзойти потолок выбора. Одна строка из пятнадцати — это не спасение; это курьез.) Это не ошибка; Это первое конкретное следствие принципа проектирования V1 («возможность проверки важнее оригинальности»). Исправление находится в разделе 9, и это специализированный селектор турниров для варпов, а не лихорадочная переработка кода.

Ещё одно важное замечание к приведённым выше цифрам: в этом зафиксированном снимке тактовые частоты графического процессора не были заблокированы . Это означает, что абсолютные значения в миллисекундах немного изменяются в зависимости от температуры и DVFS; соотношения остаются стабильными. В репозитории есть scripts/lock_gpu_clocks.sh для тех, кто хочет воспроизвести таблицу с заблокированными тактовыми частотами на GTX 1080. Само собой разумеется, структурные результаты не меняются.

6. «Хорошо, но чем это отличается от FAISS / cuVS / hnswlib?»

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

• FAISS (индекс ЦП). Используется по умолчанию в большинстве агентских фреймворков. Отличная библиотека. Работает на ЦП. Каждый запрос, выполняемый агентом, влечет за собой обмен данными по PCIe, для устранения которого и предназначена эта статья. Если вы уже используете IndexFlatIP и производительность ЦП ограничена при получении данных, вы — целевая аудитория.
• FAISS (индекс GPU). Решает проблему размещения на GPU, используя гораздо более зрелый набор ядер, чем в этом репозитории. Суть CUDA-TopK-Retrieval не в том, чтобы «превзойти FAISS-GPU по инженерным решениям» — так никогда не было и не пытается быть. Суть в том, чтобы показать в 343 строках, как выглядит действительно крутой примитив извлечения данных и почему агентные конвейеры работают медленно, когда его нет. Если вам сегодня нужен производственный индекс, используйте FAISS-GPU. Если вы хотите понять небольшой «горячий путь», который имеет значение — одно крошечное копирование H→D, три запуска ядра, два небольших копирования D→H — прочтите это ядро.
• NVIDIA cuVS / RAFT. Серьезный, предназначенный для промышленного использования стек векторного поиска на графическом процессоре. Больше, быстрее, больше алгоритмов, больше зависимостей. Как и в случае с FAISS-GPU: это ядро является учебной/однобинарной версией, а не конкурентом.
• hnswlib и аналогичные библиотеки (метод приблизительного ближайшего соседа). Компромисс здесь совершенно иной — точность приносится в жертву сублинейному времени выполнения запросов к огромным корпусам данных. CUDA-TopK-Retrieval — это точный метод грубой силы с оценкой и выбором; ускорение достигается исключительно за счет места хранения, а не за счет экономии на работе.

Цель этого репозитория не в том, чтобы «создать на его основе свою базу данных производственных векторов». Суть в том, что этап получения данных агентом должен оставаться на графическом процессоре, и как только вы это примете, даже крошечное ядро, написанное вручную, превосходит по эффективности перебор, размещенный на ЦП, при большинстве реалистичных значений K на 7-летней видеокарте.

7. Итак… как же мне это попробовать?

Клонируйте репозиторий, затем выполните сборку с помощью CMake (флаг -DGGML_CUDA=ON повторяет принципы сборки файла llama.cpp из предыдущих частей серии):

 cmake -S . -B build -DGGML_CUDA=ON -DCMAKE_BUILD_TYPE=Release cmake --build build -j cd build && ctest --output-on-failure

Затем запустите демонстрацию и тест производительности точно так, как указано в файле README:

 ./build/topk_demo # tiny smoke story (GPU required) ./build/topk_bench --n 20000 --d 384 --k 32 --trials 3 --warmup 1 --seed 1 --metric 0

topk_demo — это небольшой пример анализа данных с помощью Smoke: 4096 строк корпуса, 128 измерений, K=8, выводит идентификаторы соседей. topk_bench — это тестовая среда, которая генерирует строку TOPK_BENCH_JSON обрабатываемую скриптом кампании на Python. Полный анализ 45 конфигураций на стандартном оборудовании:

 python3 scripts/benchmark_campaign.py.example # full sweep (GPU required; writes under examples/benchmark-campaign-runs/run-*)

Перед запуском тестов, пригодных для публикации, заблокируйте тактовую частоту графического процессора (скрипт находится в репозитории):

 sudo bash scripts/lock_gpu_clocks.sh

Требования: Linux, инструментарий CUDA, графический процессор NVIDIA (Pascal или более новая версия) и терпение, чтобы один раз прочитать файл CMake. Артефакты размещаются в папке examples/example-run-results/ для быстрого доступа или examples/benchmark-campaign-runs/run--/ для полного доступа, а в файле README четко указано, что добавление баз данных .nsys-rep запрещено — только временные шкалы PNG.

8. Неожиданный поворот сюжета — это всего лишь выбор луча 5G в костюме CUDA.

Наверное, стоит признаться: я по образованию не специалист по графическим процессорам. Я пришел из телекоммуникационной отрасли — 5G NR с незримым намеком на исследования в области 6G — и постоянно замечаю, что каждая инфраструктурная проблема в агентном ИИ — это проблема, которая уже была решена на радиоуровне, может быть, лет двадцать назад.

Для читателей, не знакомых с 3GPP: в современной базовой станции 5G антенна излучает неравномерно во всех направлениях. Она формирует кодовую книгу направленных лучей — узких лепестков радиоэнергии — и в любой момент времени ваш телефон получает сигнал от одного луча (или небольшого количества лучей), мощность которого на вашем устройстве самая высокая. Быстрый выбор правильного луча — одна из наиболее изученных проблем поиска в беспроводной связи. UE измеряет L1-RSRP (показатель мощности принимаемого сигнала для каждого луча) по всем лучам-кандидатам, которые gNB (базовая станция 5G) указала ему измерить, а затем сообщает о лучших лучах через канал обратной связи CSI. gNB использует эти отчеты, чтобы решить, какие лучи следует запланировать (Ну, это было максимально упрощенное объяснение, на самом деле все гораздо сложнее!).

Это векторный поиск Top-K в радиоформате. Кандидатами являются лучи из корпуса. Мгновенное измерение канала — это запрос. Оценка — это принимаемая мощность. K — количество лучших лучей, которые передаются в отчете. UE выполняет оценку на уровне цифровой обработки сигналов в полосе частот — он не отправляет I/Q-сэмплы обратно в центральный процессорный центр и не запрашивает у скрипта на Python, какой луч является лучшим, потому что выполнение этого в цикле с интервалом в миллисекунду привело бы к перегрузке радиоинтерфейса.

Сравните эти два примера и скажите мне совершенно серьезно:

Небольшое замечание, адресованное двум совершенно разным аудиториям.

Моим друзьям, которые в первую очередь разбираются в высокопроизводительных вычислениях и CUDA и читают это : я вас понимаю. Ни один из математических примитивов здесь не является новым. Мы все знаем, cuBLAS быстрее обрабатывает умножения матриц, cuVS справляется с алгоритмом Top-K в масштабах центров обработки данных, а тщательно настроенный турнирный отбор превзойдёт эту сортировку пузырьком по блокам. Но цель здесь не в том, чтобы заново изобрести корпоративные библиотеки NVIDIA. Ценность заключается в отсутствии зависимостей в пакетах. Это архитектурное доказательство из 343 строк — в комплекте со строгим оракулом ЦП и бенчмарком из 45 конфигураций — разработанное для работы исключительно на устаревшей потребительской видеокарте с 8 ГБ памяти. Это своего рода комплексный инженерный артефакт, который вы создаёте, чтобы доказать, что вы действительно понимаете узкие места в аппаратном обеспечении, связанные с памятью, а не просто знаете, как вызывать API фреймворка.

Моим друзьям из телекоммуникационной отрасли: если еще десять минут назад «векторный поиск Top-K» казался вам чем-то совершенно непонятным, вы не отстаете — вы опережаете всех. Двадцать лет наш мир состоял из FPGA, ASIC, PRB и диаграмм созвездий. Мы оптимизировали спектр, а не кремний. Затем AI-RAN, NWDAF, NVIDIA Aerial и исследования 3GPP Rel-20 появились слишком быстро, всего за несколько месяцев, и следующее десятилетие карьеры в телекоммуникационной отрасли теперь требует двуязычия между миром спектра и миром GPU. Интуиция понятна. Вы занимались векторным поиском Top-K на стороне приемника в условиях жестких ограничений реального времени с момента появления первого свода правил MIMO. То же самое, просто в новом мире.

9. Честные оговорки (потому что комментарии не заставят себя долго ждать)

Если вы пришли сюда, чтобы узнать, что не так с этим проектом, — поздравляю, вы первый внимательный читатель этой статьи. Прямо из раздела «ОГРАНИЧЕНИЯ» файла README и комментариев к коду:

1. При K=100 V1 проигрывает. Частичный путь Top-K использует выборку по одной полосе на блок для удобства проверки; это еще не специализированное ядро турнирного отбора для варпов. При K=100 доминирует пузырьковая сортировка O(K²), и на CSV-файле CPU опережает на 14 из 15 строк при K=100 (иногда в 2 раза, иногда в 46 раз). Единственная победа GPU при K=100 — N=1M, D=1024, 1,44× — заключается в том, что работа по оценке наконец-то стала достаточно масштабной, чтобы преодолеть потолок отбора, а не в улучшении самого селектора. Это задокументировано; исправление является известным последующим решением.
2. В зафиксированном чеке тактовые частоты графического процессора не зафиксированы. В зафиксированном файле environment.json указано gpu_clocks_locked: false . На потребительской видеокарте абсолютные значения в миллисекундах изменяются в зависимости от температуры; соотношения в заголовочной таблице являются постоянными. В репозитории есть scripts/lock_gpu_clocks.sh (режим сохранения + блокировка тактовых частот приложения на уровне 1607 МГц для стандартной GTX 1080) для тех, кому нужны данные, пригодные для публикации.
3. Числовая погрешность, а не точное равенство чисел с плавающей запятой. При сравнении результатов GPU и CPU используется небольшая погрешность fp32 для каждого показателя; при равных результатах разрешение по-прежнему осуществляется детерминированно по индексу. Это необходимо в реальных условиях — уменьшение fp32 происходит по-разному на GPU и CPU — и тестовые данные не начнут измерение времени, пока индексы не совпадут точно.
4. Синтетические эмбеддинги. В бенчмарке используются гауссовские случайные векторы ( std::normal_distribution , начальное значение 1) для выделения сигнала, зависящего от времени пребывания и времени кругового пути, из влияния содержимого и обеспечения воспроизводимости результатов испытаний побитово. Реальные эмбеддинги приведут к более шумным абсолютным значениям времени для каждого испытания; структурное соотношение между временем простоя PCIe и вычислениями на устройстве не изменится.
5. Один класс архитектуры CUDA. Все цифры получены с одной видеокарты GTX 1080 класса Pascal. На Ada/Hopper абсолютное время в миллисекундах будет сокращаться для обоих путей; структурный вывод (стоимость обмена данными по PCIe доминирует при извлечении данных на стороне ЦП) становится более важным на более быстрых графических процессорах, а не менее, поскольку время выполнения ядра сокращается быстрее, чем время обмена данными.
6. Это RAG-срез, а не полноценная векторная база данных. Это срез, объединяющий сходство и Top-K. Без сжатия (PQ, OPQ), без фильтрации, без сегментирования на нескольких GPU, без параллельного выполнения запросов внутри одного экземпляра движка. Это примитив перехода между этапами поиска, вызываемый агентом, а не замена FAISS-GPU или cuVS.

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

10. Заключение (и подготовка к заключительной части)

Если вы занимаетесь разработкой конвейеров обработки данных с использованием агентов: пожалуйста, проверьте свой инструмент поиска. Откройте любой профилировщик, которому вы доверяете. Замерьте время выполнения одного вызова инструмента от начала до конца. Если загрузка вашего графического процессора падает до нуля, пока хост-процесс Python выполняет поиск сходства, вы уже выиграли битву за диагностику. Решение уже есть на GitHub.

Если вы зарабатываете на жизнь написанием кода на CUDA: Да, сортировка пузырьком O(K2) — это намеренное решение. В планах разработки — специализированный селектор турниров для варпов.

Если вы зарабатываете на жизнь созданием телекоммуникационной инфраструктуры: Да, вы меня раскусили. Это тот самый примитив извлечения базовой полосы, который вы писали в коде DSP на протяжении двадцати лет. Индустрия ИИ просто изменила терминологию; математические методы остались прежними.

Далее: Как предотвратить ситуацию, когда ваши агенты вываливают друг на друга свои травмы

CUDA-TopK-Retrieval доказывает, что можно прекратить перенаправлять каждый шаг поиска на графический процессор. Но если вы перечитаете предостережение №1, а также строки с K=100 в разделе 5, вы уже заметили следующий предел: работа над каждым запросом независима друг от друга.

Каждый этап извлечения данных начинается с нуля. Корпус данных, конечно, находится на устройстве. Но состояние агента — эмбеддинги его предыдущих решений, ключи и значения, которые он бы естественным образом обработал при каждом переходе — теряется и восстанавливается заново. Графический процессор остается «горячим», а память агента — «холодной».

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

• Когда агент А передает управление агенту В, может ли В продолжить работу, используя накопленный контекст А, вместо того, чтобы начинать с нуля?
• Насколько малым может быть сохраняемое состояние на каждом этапе, чтобы оно всё ещё оставалось полезным?
• Какова задержка при восстановлении этого состояния на графическом процессоре следующего агента?
• Как нам обеспечить передачу информации без потери данных при переходе от одного источника к другому?

Чтобы получить ответы на эти вопросы, вам придётся делать то же самое, что делает ваш процессор во время выполнения cudaMemcpy : терпеливо ждать следующей части.

Увидимся в 4-й, заключительной части.

Примечание: Иллюстрации в этой статье созданы с помощью ИИ (Claude Opus 4.8). Они носят иллюстративный, а не фотографический характер, и любые подписи, видимые на изображениях, являются стилизованными, а не авторитетными — для получения точных названий функций, значений метрик и сведений об архитектуре обратитесь к тексту статьи и самому коду.

Анубхаб Банерджи Посмотреть все в Анубхаб Банерджи

Источник: towardsdatascience.com