, recorrido de 343 líneas de recuperación CUDA Top-K. Este kernel, CPU Oracle y conjunto de pruebas comparativas demuestran que el viaje de ida y vuelta estándar de Agentic RAG (rebotar consultas a través del bus PCIe) es el asesino silencioso de su canalización. Al mantener la búsqueda de similitudes residente en la memoria del dispositivo, esta arquitectura logra una aceleración de 8,6 veces sobre las líneas base de CPU optimizadas, incluso en una GTX 1080 de 7 años.
Esta es la Parte 3 de la serie “Inferencia agente de grado de producción”. Cada parte elimina un tipo de trabajo redundante de un proceso de LLM agente. La parte 1 eliminó el precarga redundante. La parte 2 acabó con las esperas redundantes: cómo varios microagentes comparten una GPU mediante la división del tiempo. La parte 3 (esta publicación) mantiene la recuperación de RAG en la GPU con un kernel CUDA Top-K personalizado. La parte 4 mantiene el estado del agente durante las transferencias para que el siguiente agente nunca tenga el problema de arranque en frío.
Conclusiones clave
El problema: en RAG agente, cada llamada a herramienta que necesita contexto activa una búsqueda de similitud. Una canalización predeterminada envía la consulta incrustada desde la GPU a Python, permite que la CPU puntúe N filas del corpus y elija la mejor K, luego envía la respuesta de regreso. Ese viaje de ida y vuelta es el impuesto silencioso. El cálculo está bien; el viaje es la cuenta. Todos sabemos que viajar nunca es barato, sin importar a dónde quieras ir (¡nunca mejor dicho!)
La solución fácil: cargar el corpus en VRAM una vez, luego mantener la puntuación de similitud, la selección Top-K y el paso de fusión en el dispositivo. Solo la pequeña incrustación por consulta (D flotantes) y los resultados K viajan a través de PCIe.
Los recibos: en la misma GTX 1080 de 7 años utilizada en las Partes 1 y 2, la ruta residente en la GPU ejecuta el salto de recuperación hasta 8,57 veces más rápido que una línea base de fuerza bruta de la CPU. En K=8 gana en las 15 configuraciones de barrido (N ∈ {10k, 50k, 100k, 500k, 1M}, D ∈ {384, 768, 1024}) con aceleraciones de 2,43× a 8,57×. Con K=32 gana en 13 de 15 configuraciones, alcanzando un máximo de 7,76×. En K=100, donde el selector V1 se mantiene intencionalmente simple, la CPU gana en 14 de 15 configuraciones. Esa última frase es la parte honesta (Bueno, incluso si hubiera mentido, fácilmente podrías haberlo descubierto).
El truco: las victorias no son victorias del “núcleo mágico”. Son las victorias de “dejamos de enviar el corpus a la RAM del host sin ningún motivo”. También es exactamente el tipo de decisión de “medir muchos candidatos, informar solo el mejor K al consumidor” que una estación base 5G y su teléfono han estado tomando cada pocos milisegundos desde que la retroalimentación de CSI se hizo realidad.
TL;DR: RAG agente predeterminado trata la GPU como una caja de servicio y la recuperación como una preocupación de Python. Cada llamada a la herramienta envía la consulta que incorpora D→H, permite que la CPU calcule N productos escalares, ordene los candidatos, elija el K superior y envíe índices y puntuaciones H→D. Para un agente que llama a una tienda de vectores diez veces por paso de razonamiento, ese viaje de ida y vuelta es el costo dominante: no el modelo, no la incorporación, es el viaje. CUDA-TopK-Retrieval mantiene el corpus residente en el dispositivo, ejecuta puntuación + Top-K parcial por bloque + una combinación multidireccional completamente en la GPU y expone una pequeña API de orquestador de C++ (upload_corpus_rowmajor una vez, search_resident por consulta). Los bytes que tocan el host por consulta colapsan a una longitud D incrustada hacia arriba y 2K resultados hacia abajo. En una GTX 1080, en un barrido de 45 configuraciones, la ruta residente de la GPU supera la línea base de ida y vuelta de la CPU en las 15 configuraciones K=8 (2,43× a 8,57×, con un máximo de N=1M, D=1024) y en 13 de 15 configuraciones K=32 (ambas pérdidas son mínimas N=10k para D=384 y D=768, donde el viaje de ida y vuelta en sí ya es barato; las victorias de big N K=32 suben a 7,76×). En K=100, el kernel V1 se mantiene deliberadamente simple (clasificación de burbujas de un solo carril por bloque con una combinación en serie) y la CPU gana en 14 de 15 configuraciones; ese techo es el remate honesto del artículo y una configuración limpia para la Parte 4.
Repositorio de GitHub: https://github.com/AnubhabBanerjee/cuda-topk-retrieval
(Confesión rápida antes de comenzar: llegué a esto con experiencia en ingeniería RAN 5G/6G. La selección de haces en una estación base parece sorprendentemente cercana a RAG Top-K: el UE califica un libro de códigos de haces candidatos según la potencia recibida e informa el mejor puñado por aire. Hay una sección completa sobre eso a continuación, sección 8, pero también es la razón por la que este núcleo existe en la forma que tiene).
Modelo mental de arquitectura: mantén esto abierto mientras lees.
agent.embed(consulta) → cudaMemcpy H→D (D flotantes) → row_dot_scores_kernel → parcial_topk_block_kernel (bloques P) → merge_partial_topk_kernel → cudaMemcpy D→H (índices K + puntuaciones K)
Todo lo que aparece a continuación es solo un comentario sobre una parte de esa línea.
1. Una confesión: cada paso de RAG en su agente es un pequeño viaje por carretera PCIe
En la Parte 2 de esta serie, aislamos con éxito el bucle de inferencia de nuestro agente LLM, manteniendo la generación de tokens funcionando de manera rápida y activa en el dispositivo. Diseñamos un sistema que evita el estancamiento. Pero en el momento en que le damos a ese agente una herramienta para buscar en una base de conocimiento externa (el núcleo de cualquier canal de generación aumentada de recuperación (RAG) de múltiples saltos), destruimos silenciosamente todo ese rendimiento ganado con tanto esfuerzo y nos estrellamos contra la pared. Si alguna vez ha conectado una canalización "agentica" a un almacén de vectores a través de un recuperador de Python, esto es lo que realmente sucede en cada llamada a la herramienta (con un poco de dramatización intencional):
Usted: "Agente, búsqueme los cinco puntos más relevantes para '¿cómo reclamo la deducción según la sección 80C?'"
Agente: "Claro. Incrustando la consulta en la GPU. ✅"
Agente: "Ahora rebotamos la consulta incrustada en el host".
(cudaMemcpy D→H, ~1024 flotantes) Recuperador de Python: "Entendido. Bucle NumPy. Producto escalar N veces. argpartition. Top-5".
(La CPU puntúa medio millón de filas de corpus, una fila a la vez, mientras una GPU de 9 TFLOP observa) Recuperador de Python: "Listo. Aquí están los índices y las puntuaciones".
Agente: "Genial. Regresándolos a la GPU ahora".
(cudaMemcpy H→D, 10 números) Agente: "Listo. ¿Cuál fue la pregunta otra vez?"
El agente tiene una GPU en perfecto estado. El corpus se encuentra en 4 GB de VRAM. La consulta incorporada ya estaba en la GPU; simplemente la generamos allí. Y luego, en cada salto de recuperación, enviamos la consulta de regreso al host, hacemos una similitud de fuerza bruta en NumPy / FAISS-on-CPU / un bucle hecho a mano y enviamos la respuesta de regreso.
Medidor de utilidad de su GPU: pasa la mayor parte del paso de recuperación inactivo. Su bus PCIe: recibe un entrenamiento al que no se registró. Latencia de llamada a la herramienta de su agente: dominada por algo que no es ni el modelo ni la incrustación. Ese es el chiste.
Ese es también el sucio secreto de cada demostración de RAG agente que supera la etapa de "diez trozos en la memoria" del juguete. El salto de recuperación rebota en la GPU y regresa cada vez, y cuanto mayor es el corpus, peor es el impuesto. En un millón de filas de incrustaciones de 1024 luces tenues, el viaje de ida y vuelta por sí solo (ni siquiera la puntuación, sí, sólo el viaje de ida y vuelta) consume la mayor parte del presupuesto del paso de recuperación en sí.
CUDA-TopK-Retrieval es lo que sucede cuando decide que el viaje de ida y vuelta es opcional y prefiere escribir 343 líneas de CUDA que dejar que el agente pase por la RAM del host cada vez que quiere un vecino.
Ahora imagine la verdadera carga de trabajo detrás de esto. No se trata de “cinco partes para una pregunta”. Se trata de múltiples microagentes especializados: cada uno ejecuta sus propios saltos RAG, cada uno necesita Top-K contra el mismo corpus y cada uno paga actualmente su propia factura de PCIe en cada llamada de herramienta. La parte 1 de esta serie acabó con el viaje de ida y vuelta previo al llenado. La parte 2 hizo que la GPU se pudiera compartir entre esos múltiples agentes. La parte 3 dice: ahora que comparten la tarjeta de manera justa, dejen de hacer que cada uno de ellos regrese al anfitrión para buscar a un vecino.
2. ¿Por qué existe la recuperación de Top-K? (un curso intensivo de un minuto)
Omita esto si ya tiene experiencia en esto. Para todos los demás que son nuevos en este campo, aquí hay una breve explicación amateur.
Un agente moderno no incluye toda la base de conocimientos en el mensaje. Se recupera. Para cada paso de razonamiento que necesita un contexto fundamentado, incrusta la consulta en un vector de dimensión fija (D flotantes, generalmente 384, 768 o 1024), califica ese vector en cada fila de un corpus de fragmentos preincrustados (N filas, también D flotantes cada una) y devuelve las K filas del corpus con la mayor similitud. Eso es todo. Esa es la búsqueda de vectores Top-K. Recuperación-Generación Aumentada es solo una forma educada de decir "Top-K más una plantilla de aviso".
Dos tipos de similitud aparecen en todas partes. El producto escalar es el más barato: una única suma múltiple fusionada por dimensión, trabajo N×D en total. El coseno es un producto escalar dividido por el producto de las normas L2, que se convierte en un producto escalar libre si prenormaliza el corpus una vez durante la ingesta. La mayoría de las tiendas de vectores de producción hacen el truco de prenormalización y lo llaman "coseno" mientras ejecutan matemáticas de productos escalares sin procesar en el momento de la consulta. El kernel CUDA-TopK-Retrieval admite ambos: simplemente multiplica por un puntero de norma por fila precalculado cuando el modo coseno está activado.
Las herramientas convencionales (FAISS, hnswlib, el lado Python de cuVS, su base de datos vectorial SaaS favorita) hacen este trabajo de puntuación + Top-K. La mayoría lo hace bien. El problema es dónde lo hacen. Casi todos los marcos de agentes del planeta llaman al recuperador desde Python, y en el momento en que Python está en el camino activo, el paso de recuperación ya no es una operación de GPU, es una operación PCIe con una GPU en un extremo.
La solución no es "un algoritmo mejor". Es “un viaje por carretera mucho más corto”.
3. La bombilla de “mantén el corpus en la GPU” (y por qué es más difícil de lo que parece)
El tono es simple
Cargue el corpus a la VRAM una vez durante la ingesta. Para cada consulta entrante, cudaMemcpy es un pequeño flotante D-dimensional integrado en el dispositivo. Inicie un núcleo de puntuación donde un subproceso CUDA por fila del corpus calcula el producto escalar. Inicie un kernel Top-K parcial donde cada bloque escanea un rango de filas separadas para emitir sus propios candidatos principales locales. Finalmente, inicie un kernel de fusión para recorrer los encabezados por bloque y emitir el Top-K global en el mejor orden.
CudaMemcpy envía exactamente 2K números al host: K índices y K puntuaciones.
Este es el paradigma de “tratar la recuperación de memoria como una primitiva de hardware, no como una llamada API de software”. La única razón por la que esto requiere más de un script PyTorch de 30 líneas para lograrlo es que tres tediosos casos extremos romperán inmediatamente el enfoque ingenuo.
Problema A: Top-K en una GPU es estructuralmente incómodo
Calificar los vectores es la parte fácil. Es simplemente multiplicación de matrices, y su GPU nació literalmente para hacer eso: es el lenguaje de amor del hardware. Sin embargo, en la selección es donde muere el romance. Pedirle a una GPU que realice una clasificación O(N log N) completa solo para obtener los K resultados principales es computacionalmente ofensivo; es como ordenar alfabéticamente toda su papelera de reciclaje solo para encontrar un solo recibo. Podrías probar una partición arg O(N), pero eso requiere un paseo por el árbol, lo que destroza la memoria de la GPU y se fusiona en un millón de lecturas no alineadas. La selección de torneos es rápida, suponiendo que quieras pasar el fin de semana depurando casos extremos. Y en el momento en que cedas y buscas una primitiva de empuje o clasificación de cachorros, felicitaciones: acabas de infectar tu canal C++ liviano e independiente con una dependencia de compilación masiva.
La arquitectura elige la respuesta aburrida a propósito. Se basa en una pequeña clasificación de burbujas O(K2) por bloque en un rango de filas separadas, impulsada por un único hilo por bloque y rematada con una combinación multidireccional en serie. Sobre el papel, esto suena terrible. En la práctica, funciona maravillosamente, por la razón exacta que se indica honestamente en los comentarios del núcleo:
// Escaneo por bloque de un solo subproceso que materializa una lista Top-K local para su partición de fila. // Esta no es la selección global más rápida, pero es fácil de razonar y coincide exactamente con la regla de ordenación de la CPU. __device__ void bubble_downward(float* const s, int* const ids, const int n) { // La clasificación Tiny O(K^2) es aceptable porque K tiene un límite (kMaxSupportedK) y se ejecuta en un solo carril por bloque. 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 = identificadores[j]; identificadores[j] = identificadores[j + 1]; identificadores[j + 1] = ti; } } } }
Este es el contrato de diseño: V1 prefiere la auditabilidad a la inteligencia. Todo el núcleo es lo suficientemente pequeño como para que un revisor pueda leerlo de un extremo a otro en una pausa para el café, alinearlo con el oráculo de la CPU y convencerse de que la salida de la GPU es correcta en bits. El día que alguien quiera un kernel 2×, este será reemplazado por un selector de torneos especializado en warp; la sección 9 del artículo lo promete explícitamente. Para K ≤ 32, la clasificación de burbujas de un solo carril está realmente bien. Para K=100, cae por un precipicio. Ese acantilado está documentado y evaluado. (Sección 9 nuevamente.)
Problema B: GPU y CPU deben estar de acuerdo, bit por bit, en el desempate
Bueno, como en cualquier partido de liga del grupo de la Copa Mundial de Fútbol, los empates ocurren con más frecuencia. Dos filas del corpus tienen la misma puntuación con una precisión de fp32. ¿Cuál gana?
Si el oráculo de la CPU y el núcleo de la GPU no están de acuerdo en el desempate, nunca se puede confiar en un punto de referencia, porque cada alarma de "desigualdad" ahora es ambigua: ¿la GPU obtuvo una puntuación incorrecta en una fila o simplemente rompió los empates de manera diferente? Pasarás una semana en un canal de Slack de las 3 a.m.
La solución es definir el comparador en una oración e implementarlo en dos lugares (una vez en el host y otra en el dispositivo) y hacer que esas dos implementaciones sean literalmente la misma expresión.
Del lado del anfitrión se ve así:
// "Mejor" relación lexicográfica para pares (puntuación, índice) bajo semántica de igualdad flotante. // Usamos un orden débil estricto para std::partial_sort: la puntuación más alta gana; En caso de empate exacto, gana el índice más pequeño. 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) { // Clave principal: puntuación de similitud (cuanto más alta, mejor para la recuperación). if (puntuación_lhs! = puntuación_rhs) { return puntuación_lhs > puntuación_rhs; } // Superficie de enlace determinista: prefiera la identificación de fila de corpus más pequeña para reflejar las claves primarias de base de datos estables. devolver idx_lhs < idx_rhs; }
En el lado del dispositivo se ve así:
// Réplica del lado del dispositivo del comparador de host para evitar problemas de vinculación entre TU para las rutas de código de __dispositivo__. __device__ bool device_is_better(const float score_lhs, const int idx_lhs, const float score_rhs, const int idx_rhs) { // Misma semántica de ordenamiento que topk::is_better_score_pair para superficies de unión bit a bit idénticas. if (puntuación_lhs! = puntuación_rhs) { return puntuación_lhs > puntuación_rhs; } devolver idx_lhs < idx_rhs; }
Esa es toda la política de desempate en cinco líneas, dos veces. La puntuación más alta gana; en caso de empate exacto, gana el índice de fila de corpus más pequeño. El oráculo de la CPU lo usa en std::partial_sort, la GPU lo usa en la clasificación de burbujas y en la combinación multidireccional, y el arnés de referencia no comenzará a cronometrar hasta que la salida de la GPU coincida exactamente con la salida de la CPU: los mismos índices en el mismo orden, puntuaciones dentro de una pequeña tolerancia fp32.
Ese único comparador es la razón por la que el artículo puede citar una aceleración. Sin él, "la GPU es 8 veces más rápida" es simplemente "la GPU es 8 veces más rápida al equivocarse de manera diferente".
Problema C: la VRAM es preciosa y el peor lugar para hacer malloc es la ruta activa
Asignar memoria de GPU por consulta es como firmar un contrato de arrendamiento de automóvil nuevo cada vez que necesita conducir hasta el supermercado. Es la forma más sencilla de convertir una búsqueda de 1 milisegundo en un atasco de 50 milisegundos.
En cambio, GpuTopkEngine::initialize compra el auto por adelantado. Ejecuta todas las llamadas de cudaMalloc durante el inicio del motor, dimensionando los buffers para la peor configuración posible. Una vez que el motor atiende consultas activamente, la ruta activa queda completamente libre de administración de memoria. Se trata simplemente de lanzamientos rápidos del kernel y pequeñas copias de datos. Sin fragmentación, sin negociación con el asignador, y cudaMalloc tiene prohibido permanentemente aparecer en sus seguimientos de rendimiento.
4. El proceso de cuatro etapas (la parte realmente interesante)
Paso 0: Inicio del motor: ocho cudaMallocs de tamaño (max_n, max_d, max_k) (GpuTopkEngine::initialize) Paso 1: Cargar el corpus una vez en VRAM (upload_corpus_rowmajor) Paso 2: Por consulta: H→D la incrustación (search_resident, primera línea) Paso 3: Calificar N filas en el dispositivo (row_dot_scores_kernel) Paso 4: Top-K parcial por bloque (partial_topk_block_kernel) Paso 5: Fusión multidireccional en Top-K global (merge_partial_topk_kernel) Paso 6: D→H los índices K + puntuaciones K (search_resident, últimas líneas)
Repasemos cada uno con el código real. Los fragmentos son incluso más cortos que los pequeños archivos fuente.
Paso 1: cargue el corpus una vez
Este es el paso aburrido que hace posible el resto del artículo. El corpus aumenta exactamente una vez por ingesta y permanece allí durante toda la vida útil del motor:
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(N) * static_cast(D); const cudaError_t st = cudaMemcpy(d_corpus_, host_corpus_rowmajor, corpus_bytes, cudaMemcpyHostToDevice); if (st! = cudaSuccess) { return st; } residente_n_ = N; residente_d_ = D; devolver cudaSuccess; }
Y esa es toda la API de ingesta. En 1024 dimensiones, un millón de vectores son exactamente 4 GB de datos float32, que se deslizan perfectamente en los 8 GB de VRAM de una GTX 1080 antigua. ¿Qué sucede cuando su corpus alcanza los 10 millones de vectores? Eso se convierte en un problema de sistemas distribuidos, no en un problema del kernel. Si sus datos exceden la VRAM, necesita una estrategia de fragmentación, que cubrimos en la Sección 9. Pero por ahora, estamos aquí para resolver el cuello de botella informático, no para inventar una nueva base de datos.
Paso 2: marque N filas en el dispositivo
Un hilo CUDA por fila de corpus. 256 hilos por bloque. Cada hilo acumula el producto escalar en D dimensiones y escribe un flotante en el búfer denso de puntuaciones[N]:
// Producto escalar de fila principal con normalización de coseno opcional; Las lecturas fusionadas a lo largo de D se sacrifican por claridad en v1. // Nota de microarquitectura: un hilo por fila es simple; un seguimiento puede colocar D en mosaico a través de deformaciones para aumentar la intensidad aritmética. __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 puntuaciones) { // Asigna cada subproceso CUDA a exactamente una fila del corpus para que la lógica de reducción sea fácil de auditar con respecto a la referencia de la CPU. const int fila = static_cast(blockIdx.x) * static_cast(blockDim.x) + static_cast(threadIdx.x); if (fila >= N) { retorno; } cuenta flotante = 0.0F; const int base = fila * D; for (int col = 0; col < D; ++col) { acc += corpus[static_cast(base + col)] * consulta[static_cast(col)]; } if (cosine_enabled! = 0) { const float denom = query_l2 * row_l2_fetch(row_l2, fila); puntuaciones[static_cast(row)] = denominación > 0.0F? (acc / denom): -std::numeric_limits::infinity(); } else { puntuaciones[static_cast(row)] = acc; } }
Un hilo por fila es el mapeo más simple posible. El comentario del código es honesto al respecto: un seguimiento puede colocar D en mosaicos a través de deformaciones para aumentar la intensidad aritmética. Para V1, esto le da al auditor una correspondencia uno a uno con el bucle de la CPU y le permite dormir tranquilamente por la noche.
Paso 3: cada bloque construye su propio Top-K local
Ahora viene la parte más incómoda. Elegir la K superior de N es conceptualmente "ordenar y dividir", pero una clasificación completa desperdicia la mayor parte del trabajo. Dividimos el rango de filas en bloques P (con un límite de 128), cada bloque recorre su porción inconexa con el tipo de burbuja pequeña de la Sección 3 y escribe su propia lista K superior local:
const int P = std::min(kMaxPartialBlocks, std::max(1, (static_cast(N) + 4095) / 4096)); parcial_topk_block_kernel<<
>>(d_scores_, static_cast(N), static_cast(K), P, d_partial_scores_, d_partial_indices_);
Un hilo por bloque. Sí, eso es un desperdicio en papel. También es por eso que un humano puede auditar este kernel en veinte minutos: la política if (threadIdx.x != 0 || blockIdx.x >= P) return; colapsa todo el razonamiento intrabloque hasta "el carril 0 de este bloque posee filas [inicio, fin)". Las s de cada bloque[]e identificaciones[]Las matrices viven en registros/memoria local, dimensionadas según el límite de tiempo de compilación kMaxSupportedK = 256.
Paso 4: fusionar los parciales en el Top-K global
Finalmente, un hilo en un bloque mueve los cursores P sobre las listas por bloque. Cada lista ya es la mejor primero. Elige la mejor cabeza; emitir; avance ese cursor; repetir K veces:
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 (cabezas[p] >= K) { continuar; } const float s = puntuaciones_parciales[static_cast(p * K + cabezas[p])]; const int idx = índices_parciales[static_cast(p * K + cabezas[p])]; if (mejor_p < 0 || dispositivo_es_mejor(s, idx, mejor_s, mejor_i)) { mejor_p = p; mejor_s = s; mejor_i = idx; } } out_scores[static_cast(out)] = best_s; out_indices[static_cast(out)] = best_i; cabezas[mejor_p] += 1; }
La fusión es despiadadamente eficiente: como máximo P * K lee y exactamente K escribe, ejecutado por un solo hilo. Para evitar el caos del punto flotante, el comparador device_is_better impone un determinismo estricto: si dos cabezas empatan en puntuación, gana el índice de la fila inferior del corpus, reflejando perfectamente el oráculo de la CPU. Finalmente, dos llamadas microscópicas de cudaMemcpy envían los índices ganadores K y las puntuaciones al anfitrión. El agente los ingiere y el bucle RAG se activa nuevamente.
Esa es la ruta activa completa: una transferencia de incrustación H -> D, tres lanzamientos del kernel y dos pequeñas copias de resultados D -> H. Sin bucles de host de Python, sin sobrecarga de marco y absolutamente sin vacaciones de PCIe.
5. Los recibos (es decir, los números)
Evalúémoslo ahora con respecto a la línea de base y veamos si valió la pena.
Una nota rápida sobre la metodología, antes de que llegue la policía de evaluación comparativa: cada comparación a continuación se ejecuta en la misma GPU que la Parte 1 y la Parte 2 (NVIDIA GeForce GTX 1080, Pascal sm_61, 8 GB), controlador 535.309.01, CUDA 12.2, CPU host Intel Core i7-8700K, indicadores del compilador -O3 -march=native –expt-relaxed-constexpr. Tres pruebas, un calentamiento, semilla 1, RNG fijo (std::mt19937 con std::normal_distribution), incrustaciones gaussianas, normalizado L2 en modo de producto escalar. El barrido completo es N ∈ {10k, 50k, 100k, 500k, 1M} × D ∈ {384, 768, 1024} × K ∈ {8, 32, 100} → 45 configuraciones, todas medidas a través de cudaEventElapsedTime con cudaDeviceSynchronize entre paréntesis de cada intervalo. Código en src/host/bench_main.cpp; números sin procesar en ejemplos/example-run-results/benchmark_run_results.csv.
Los dos caminos cronometrados:
Residente de GPU (tratamiento). Corpus ya está en el dispositivo. Cada iteración cronometrada: consulta cudaMemcpy H→D (D flotantes) + puntuación del kernel + Top-K parcial por bloque + fusión + cudaMemcpy K puntuaciones D→H + índices cudaMemcpy K D→H. De punta a punta. Viaje de ida y vuelta de CPU (línea de base). Modela el flujo de agente predeterminado: consulta cudaMemcpy D→H + puntuación de fuerza bruta de CPU + std::partial_sort con el mismo comparador + índices cudaMemcpy H→D + puntuaciones cudaMemcpy H→D. De punta a punta.
Ambas rutas se ejecutan dentro del mismo proceso, utilizan los mismos bytes de consulta y el mismo comparador. La única diferencia es dónde se realiza el trabajo. Si alguna vez ha sostenido la postura “PCIe está bien, comparamos los núcleos de forma aislada”, esto es lo que le cuesta cuando deja de fingir que el viaje de ida y vuelta es gratuito.
Titular (GTX 1080, tres pruebas, media en ms, proporciones calculadas a partir de medias por prueba):
Sí, estás leyendo bien los números.
Las primeras cinco filas son el punto del artículo: en K=8, la ruta residente en la GPU gana en cada configuración del barrido (las 15), por proporciones que van desde un cortés 2,43× en N=50k, D=384 hasta un ruidoso 8,57× en N=1M, D=1024. En K=32 gana en 13 de 15: las dos pérdidas son ambas en el N más pequeño (10k), para D=384 y D=768, donde el viaje de ida y vuelta en sí solo cuesta ~3–7 ms y los tres lanzamientos del núcleo de la GPU apenas tienen espacio para amortizar. Cuando se alcanzan tamaños de corpus agentes realistas (N ≥ 50k), K=32 también gana cómodamente, alcanzando un máximo de 7,76×. Las grandes aceleraciones no son aceleraciones del “núcleo mágico”; son aceleraciones del tipo “dejamos de enviar el corpus de regreso a la RAM del host sin ningún motivo”. La GPU siempre iba a ganar esta carrera; la única razón por la que perdió fue que seguíamos haciéndolo viajar innecesariamente.
Las dos últimas filas son donde este artículo se gana el derecho a llamarse honesto. En K = 100, la clasificación de burbujas de un solo carril por bloque se convierte en O (K²) = O (10,000) comparaciones secuenciales por bloque, y la combinación en serie recorre las posiciones de cabeza P × K. El std::partial_sort de la CPU está basado en montón, vectorizado por el compilador y efectivamente O(N log K), mucho más amigable con K=100. Entonces, la GPU pierde en 14 de 15 K=100 configuraciones, a veces 2×, a veces 46×. (Hay exactamente una configuración K=100 en la que la GPU aún gana: N=1M, D=1024, 1,44×, porque para entonces ya hay suficiente trabajo de puntuación para dominar el límite de selección. Una fila de quince no es una salvación; es una curiosidad.) Eso no es un error; es la declaración de diseño V1 (“auditabilidad sobre inteligencia”) que encuentra su primera consecuencia concreta. La solución está en la Sección 9, y es un selector de torneos especializado en warp, no una refactorización frenética.
Una advertencia más honesta sobre los números anteriores: en esta instantánea comprometida, los relojes de la GPU no estaban bloqueados. Eso significa que los milisegundos absolutos se mueven ligeramente con las térmicas y DVFS; las proporciones se mantienen estables. El repositorio envía scripts/lock_gpu_clocks.sh para cualquiera que quiera reproducir la tabla con relojes bloqueados en una GTX 1080. No hace falta mencionar que el hallazgo estructural no cambia.
6. “Está bien, pero ¿en qué se diferencia esto de FAISS/cuVS/hnswlib?”
Una pregunta muy razonable y que vale la pena responder directamente, porque el mundo de la búsqueda vectorial tiene muchas primitivas superpuestas y un lector de HPC preguntará esto en el primer comentario.
FAISS (índice de CPU). El valor predeterminado en la mayoría de los marcos de agentes. Excelente biblioteca. Vive en la CPU. Cada consulta que realiza un agente paga el viaje de ida y vuelta de PCIe que este artículo debe eliminar. Si ya está en IndexFlatIP y la recuperación depende de la CPU, usted es el público objetivo. FAISS (índice GPU). Resuelve el problema de residencia de la GPU, con un conjunto de kernel mucho más maduro que este repositorio. El objetivo de CUDA-TopK-Retrieval no es "superé en ingeniería a FAISS-GPU"; nunca lo fue y tampoco intenta serlo. El objetivo es mostrar, en 343 líneas, cómo se ve la primitiva de recuperación realmente genial y por qué los canales de agentes se sienten lentos cuando no está ahí. Si necesita un índice de producción hoy, utilice FAISS-GPU. Si desea comprender la pequeña ruta activa que marca la diferencia (una pequeña copia H→D, tres lanzamientos del kernel, dos pequeñas copias D→H), lea este kernel. NVIDIA cuVS/RAFT. La pila de búsqueda vectorial en GPU seria y de grado de producción. Más grande, más rápido, más algoritmos, más dependencias. Misma advertencia que FAISS-GPU: este kernel es la versión pedagógica/binaria única, no un competidor. hnswlib y amigos (vecino más cercano aproximado). Una forma de compensación completamente diferente: intercambian exactitud por tiempo de consulta sublineal en corpus enormes. CUDA-TopK-Retrieval es puntuación + selección exacta de fuerza bruta; La historia de la aceleración tiene que ver puramente con la residencia, no con faltar al trabajo.
El objetivo de este repositorio no es "construir su base de datos vectorial de producción sobre él". El punto es: el salto de recuperación del agente quiere permanecer en la GPU, y una vez que aceptas eso, incluso un pequeño núcleo escrito a mano supera una fuerza bruta alojada en la CPU en los valores K más realistas en una tarjeta de 7 años.
7. Entonces… ¿cómo lo pruebo realmente?
Clone el repositorio, luego compílelo con CMake (el indicador -DGGML_CUDA=ON refleja la ergonomía de compilación de llama.cpp de partes anteriores de la serie):
cmake-S. -B build -DGGML_CUDA=ON -DCMAKE_BUILD_TYPE=Liberar cmake –build build -j cd build && ctest –output-on-failure
Luego ejecute la demostración y el punto de referencia exactamente como lo hace el README:
./build/topk_demo # pequeña historia de humo (se requiere GPU) ./build/topk_bench –n 20000 –d 384 –k 32 –trials 3 –warmup 1 –seed 1 –metric 0
topk_demo es la pequeña historia del humo: 4096 filas de corpus, 128 atenuaciones, K = 8, imprime las identificaciones de los vecinos. topk_bench es el arnés que emite la línea TOPK_BENCH_JSON que consume el script de campaña de Python. Para el barrido completo de 45 configuraciones en hardware canónico:
python3 scripts/benchmark_campaign.py.example # barrido completo (se requiere GPU; escribe en ejemplos/benchmark-campaign-runs/run-*)
Antes de ejecutar la versión de publicación, bloquee los relojes de la GPU (el repositorio proporciona el script):
scripts sudo bash/lock_gpu_clocks.sh
Requisitos: Linux, kit de herramientas CUDA, una GPU NVIDIA (Pascal o más reciente) y paciencia para leer un archivo CMake una vez. Los artefactos aparecen en ejemplos/example-run-results/ para la ruta rápida o ejemplos/benchmark-campaign-runs/run–/ para el barrido completo, y el README es explícito en que está prohibido confirmar bases de datos .nsys-rep (solo líneas de tiempo PNG).
8. Giro de la trama: esto es solo la selección de un haz 5G disfrazado de CUDA
Probablemente debería confesar llegados a este punto: todavía no soy una “persona GPU” por formación. Llegué a través de las telecomunicaciones (5G NR con un pie arrastrándose firmemente hacia la investigación de 6G) y sigo notando que cada problema de infraestructura en la IA agente es un problema que ya se resolvió en la capa de radio, tal vez hace unos veinte años.
Para lectores sin experiencia en 3GPP: en una estación base 5G moderna, la antena no irradia por igual en todas las direcciones. Forma un libro de códigos de haces direccionales (lóbulos estrechos de energía de radio) y, en cualquier momento dado, su teléfono recibe servicio del único haz (o del pequeño puñado de haces) cuya potencia recibida es más alta en su dispositivo. Elegir el haz correcto y rápido es uno de los problemas de recuperación inalámbricos más estudiados. El UE mide L1-RSRP (una puntuación de potencia recibida por haz) a través de los haces candidatos que el gNB (estación base 5G) le ha indicado que mida, luego informa los mejores a través del canal de retroalimentación CSI. El gNB usa esos informes para decidir qué haces programar (bueno, esto fue lo más simple que pude hacer, ¡la verdadera historia es realmente fea!).
Esa es la búsqueda vectorial Top-K, disfrazada de radio. Los haces candidatos son el corpus. La medición instantánea del canal es la consulta. La puntuación es poder recibido. K es el número de mejores haces que devuelve el informe. El UE realiza la puntuación en la capa DSP de banda base: no envía las muestras de I/Q a una granja de CPU central y pregunta a un script de Python qué haz es mejor, porque hacerlo en un bucle por milisegundo derretiría la interfaz aérea.
Mire esto uno al lado del otro y dígame con seriedad que estos son problemas diferentes:
Un breve comentario sobre dos públicos muy diferentes
A mis primeros amigos de HPC y CUDA que leen esto: los escucho. Ninguno de los primitivos matemáticos aquí son novedosos. Todos sabemos que cuBLAS ejecuta matmuls más rápido, cuVS maneja Top-K a escala de centro de datos y una selección de torneos altamente ajustada aplastará este tipo de burbujas por bloque. Pero el objetivo aquí no es reinventar las bibliotecas empresariales de NVIDIA. El valor es el paquete de dependencia cero. Esta es una prueba arquitectónica de 343 líneas, completa con un estricto oráculo de CPU y un barrido de referencia de 45 configuraciones, diseñada para ejecutarse completamente en una GPU de consumo antigua de 8 GB. Es el tipo de artefacto de ingeniería de extremo a extremo que se construye para demostrar que realmente comprende los cuellos de botella de la memoria del hardware, en lugar de simplemente saber cómo llamar a una API de marco.
A mis amigos de las telecomunicaciones: si la “búsqueda de vectores Top-K” sonaba como un idioma extranjero hasta hace diez minutos, no se han quedado atrás: han llegado temprano. Durante veinte años nuestro mundo estuvo compuesto por FPGA, ASIC, PRB y diagramas de constelaciones. Optimizamos el espectro, no el silicio. Luego, los elementos de estudio de AI-RAN, NWDAF, NVIDIA Aerial y 3GPP Rel-20 sucedieron demasiado rápido en unos pocos meses, y la próxima década de carreras en telecomunicaciones ahora exige ser bilingüe entre el mundo del espectro y el mundo de la GPU. La intuición se traduce limpiamente. Ha estado haciendo Top-K del lado del receptor bajo duras limitaciones de tiempo real desde el primer libro de códigos MIMO. El mismo animal, sólo que en un zoológico nuevo.
9. Advertencias honestas (porque los comentarios van llegando)
Si vino aquí para encontrar el problema de este proyecto, felicidades, es el primer lector cuidadoso de este artículo. Directamente desde la sección LIMITACIONES del archivo README y los comentarios del código en línea:
K=100 es donde V1 pierde. La ruta parcial Top-K utiliza una selección de un solo carril por bloque para la auditabilidad; todavía no es un núcleo de selección de torneos especializado en warp. En K=100 domina el tipo de burbuja O(K²), y en el CSV la CPU avanza en 14 de 15 K=100 filas (a veces por 2×, a veces por 46×). La única victoria de la GPU con K=100 (N=1M, D=1024, 1,44×) es que el trabajo de puntuación finalmente es lo suficientemente grande como para superar el límite de selección, no que el selector mejore. Esto está documentado; la solución es un seguimiento conocido. Los relojes de la GPU no están bloqueados en el recibo confirmado. El entorno.json comprometido informa gpu_clocks_locked: false. Cambio absoluto de milisegundos con térmicas en una tarjeta de consumo; Los ratios del cuadro principal son duraderos. El repositorio envía scripts/lock_gpu_clocks.sh (modo de persistencia + bloqueo a relojes de aplicación de 1607 MHz para una GTX 1080 estándar) para cualquiera que necesite números de grado de publicación. Tolerancia numérica, no igualdad flotante exacta. Las comparaciones de puntuaciones de GPU y CPU utilizan una pequeña tolerancia de fp32 por puntuación; Los vínculos todavía se resuelven de manera determinista por índice. Esta es una necesidad del mundo real (las reducciones de fp32 se asocian de manera diferente en una GPU que en una CPU) y el arnés no comenzará a cronometrar hasta que los índices coincidan exactamente. Incrustaciones sintéticas. El punto de referencia utiliza vectores aleatorios gaussianos (std::normal_distribution, semilla 1) para aislar la señal de residencia versus ida y vuelta de los efectos del contenido y mantener las pruebas reproducibles bit por bit. Las incrustaciones reales producirán tiempos absolutos por prueba más ruidosos; la relación estructural entre PCIe-vacaciones y matemáticas en el dispositivo no se mueve. Una clase de arquitectura CUDA. Todos los números provienen de una GTX 1080 clase Pascal. En Ada/Hopper los milisegundos absolutos se reducirán para ambas rutas; El hallazgo estructural (el costo de ida y vuelta de PCIe domina una recuperación del lado de la CPU) se vuelve más importante en GPU más rápidas, no menos, porque el tiempo del kernel se reduce más rápido que el viaje de ida y vuelta. Corte RAG, no una base de datos vectorial completa. Esta es una similitud + porción Top-K. Sin compresión (PQ, OPQ), sin filtrado, sin fragmentación de múltiples GPU, sin simultaneidad entre consultas dentro de una instancia del motor. Es la primitiva de salto de recuperación a la que llama un agente, no un reemplazo de FAISS-GPU o cuVS.
Todo lo que figura en esta lista está en la hoja de ruta. Nada de eso cambia el resultado del titular. El objetivo de ponerlo por escrito es que no debería tener que buscarlo, y en el momento en que una publicación de blog de referencia oculta sus advertencias, es el momento en que sus números dejan de ser confiables.
10. Envoltura (y preparación para la parte final)
Si se gana la vida construyendo tuberías de agentes: vaya y mire a su perro perdiguero. Abra cualquier generador de perfiles en el que confíe. Llamada de una herramienta de un extremo a otro. Si la utilización de su GPU cae a cero mientras un proceso host de Python realiza una búsqueda de similitud, ya ha ganado la batalla del diagnóstico. La solución está en GitHub.
Si escribe CUDA para ganarse la vida: Sí, la clasificación de burbujas O(K2) es intencional. Un selector de torneos especializado en warp está en la hoja de ruta.
Si construyes infraestructura de telecomunicaciones para ganarte la vida: Sí, me atrapaste. Esta es exactamente la misma primitiva de recuperación de banda base que ha estado escribiendo en código DSP durante veinte años. La industria de la IA acaba de cambiar el vocabulario; las matemáticas no han cambiado.
Próximamente: Cómo evitar que sus agentes se arrojen traumas unos a otros
CUDA-TopK-Retrieval demuestra que puedes dejar de hacer rebotar cada salto de recuperación en la GPU. Pero si vuelve a leer la advertencia n.º 1 más las K=100 filas de la Sección 5, ya habrá detectado el siguiente límite: el trabajo por consulta es independiente entre consultas.
Cada salto de recuperación comienza en frío. El corpus está en el dispositivo, claro. Pero el estado del agente (las incorporaciones de sus decisiones anteriores, las claves y valores a los que naturalmente atendería) se abandona y se reconstruye en cada transferencia. La GPU se mantiene caliente; La memoria del agente permanece fría.
Eso está bien para un paso RAG de una sola vez. Se desmorona en el momento en que ejecutas la carga de trabajo para la que se creó esta serie: razonamiento de múltiples saltos entre un enjambre de agentes especializados. A esa escala, deja de preocuparse por "¿mantuvimos la recuperación en la GPU" y comienza a preocuparse por las preguntas que un núcleo de un solo disparo no puede responder:
Cuando el agente A pasa al agente B, ¿puede B reanudar con el contexto acumulado de A en lugar de comenzar en frío? ¿Qué tan pequeño puede ser el estado persistente por salto y seguir siendo útil? ¿Cuál es el costo de latencia de restaurar ese estado en la GPU del siguiente agente? ¿Cómo hacemos para que los traspasos no pierdan información?
Para obtener la respuesta a estas preguntas, tendrá que hacer exactamente lo que hace su CPU durante un cudaMemcpy: sentarse pacientemente y esperar la siguiente parte.
Nos vemos en la Parte 4, la final.
Descargo de responsabilidad: las ilustraciones de este artículo se generaron utilizando IA (Claude Opus 4.8). Son ilustrativos, no fotográficos, y cualquier etiqueta visible dentro de las imágenes está estilizada en lugar de autorizada; consulte el cuerpo del artículo y el código mismo para obtener nombres precisos de funciones, valores métricos y detalles de arquitectura.