Lograr una eficiencia extrema mediante la generación de kernels de GPU especializados

Tradicionalmente, los sistemas de inferencia en producción dependen de kernels genéricos para gestionar diversos modelos y cargas de trabajo. Esto no es óptimo porque las formas de las operaciones de la GPU se determinan mediante una combinación de parámetros estáticos del modelo y factores dinámicos en el momento de la solicitud; por ejemplo, mientras que un modelo define una de las dimensiones para una multiplicación de matrices, la otra dimensión fluctúa según el número específico de tokens de cada solicitud. Existe un interés creciente en la generación de kernels de GPU basada en agentes, y los esfuerzos recientes se muestran prometedores. Exploramos una pregunta fundamental: si la generación de kernels se puede automatizar, ¿por qué los modelos de tamaños enormemente diferentes (de 1000 millones a 1 billón de parámetros) deberían depender del mismo kernel? Al especializar los kernels para las formas específicas que se encuentran en el tiempo de ejecución, podemos lograr una eficiencia extrema.

En este blog, compartimos nuestros éxitos y aprendizajes al usar agentes para generar kernels de GPU. Creamos Proteus, un sistema diseñado para lograr una especialización extrema, el cual requiere un entorno de pruebas (harness) adaptado para una optimización, validación y gestión de contexto rigurosas.

Los entornos de pruebas convencionales suelen fallar aquí porque los agentes tienden a manipular las recompensas (reward-hack): siguen la letra de la ley en lugar del espíritu. Si le das un benchmark a un agente, puede que optimice el benchmark y no la operación prevista.

Para abordar esto, Proteus propone kernels, los verifica frente a una implementación de referencia controlada, mide el tiempo de los que tienen éxito y mejora de forma iterativa los mejores resultados. Aunque el proceso es sencillo, su éxito depende por completo de resolver dos desafíos fundamentales. La Figura 1 muestra la arquitectura simplificada de nuestro diseño. Usando nuestro entorno de pruebas Proteus, generamos kernels de Qwen 3.5 122B que fueron entre 1.8 y 5.2 veces más rápidos que los mejores disponibles en vLLM.

Validación

Originalmente considerábamos que la búsqueda de kernels era la parte difícil: ¿cómo explorar un gran espacio de programas sin quedarse estancado en una meseta sin mejorar? En la práctica, la primera pregunta era más básica. ¿Estamos midiendo lo que creemos que estamos midiendo?

Un modelo optimiza la puntuación que se le da. No necesita un exploit exótico: simplemente puede ser que la evaluación esté dando algo por sentado. Un ejemplo fueron los kernels para embeddings de posición rotatoria (RoPE), un paso común en las capas de atención. Un candidato podría reutilizar código compilado sobrante de un intento anterior y parecer más económico que una reconstrucción justa desde cero. Otro podría registrar un lote de lanzamientos de GPU en un gráfico (por ejemplo, un gráfico de CUDA) y reproducirlos como una sola unidad, mientras que la línea de base con la que comparamos seguía lanzando cada pieza por separado, por lo que ambas partes no estaban haciendo el mismo trabajo. Otro era fuerte en los tamaños de entrada que habíamos puesto en el conjunto de prueba visible y débil en los tamaños que no se le habían mostrado.

Así que dedicamos el trabajo de diseño inicial al verificador, no al prompt. Medimos el tiempo de ambas partes de la misma manera, incluso con más de un temporizador (por ejemplo, temporizador de eventos de CUDA, tiempo de reloj real y temporizador de CUPTI) cuando necesitamos una verificación cruzada. Limpiamos el estado compilado sobrante que no debería persistir y mantenemos la coherencia en el orden de configuración y desmontaje para que una parte no pueda omitir el trabajo por el que la otra aún paga. Volvemos a medir el tiempo de los ganadores antes de usarlos como punto de partida para la siguiente ronda. Mantenemos algunas pruebas que el candidato no puede ver, para que no pueda ajustarse solo al examen. Para evitar "trampas" en la evaluación con un rendimiento inflado artificialmente, implementamos comprobaciones de consistencia automatizadas para marcar aceleraciones teóricamente imposibles (por ejemplo, >100x) que superen el ancho de banda físico de la GPU y los límites de computación. Esto protege contra los mismos problemas de manipulación de recompensas (reward-hacking) vistos en casos anteriores de la industria, donde los agentes se optimizaban para las métricas del entorno de pruebas en lugar de obtener ganancias de rendimiento reales. Sin estas restricciones, generar más kernels principalmente producía más ruido.

El énfasis en el verificador también cambia el cuello de botella de la generación de kernels mediante agentes. En el trabajo de simple búsqueda de programas (es decir, optimización iterativa donde el sistema busca entre programas generando variantes repetidamente), los buenos candidatos son raros, por lo que escribirlos domina el costo. Podemos producir muchos borradores en paralelo, pero no podemos omitir la validación. Tenemos que diseñar la validación con cuidado, y la verificación debe ejecutarse en GPUs reales, de forma aislada y más de una vez. El sistema avanza tan rápido como puede confiar en un kernel, no tan rápido como puede escribir uno.

Gestión del contexto

Otro desafío es determinar qué se le permite ver al modelo de generación de kernels. Es un compromiso (trade-off). Si le das al modelo un prompt más grande, tendrá más información: el mejor kernel actual, fallos recientes, sugerencias del analizador de rendimiento (profiler), notas de ejecuciones anteriores. Eso puede ayudar. También cuesta más, porque pagamos por cada token que lee el modelo. Y a medida que el prompt crece, es más fácil que el siguiente intento se desvíe. Las señales útiles se mezclan con consejos obsoletos, sugerencias contradictorias y detalles que se aplican a un tamaño de entrada diferente o a una operación diferente. El modelo no siempre sabe en qué frases confiar, por lo que sigue las más llamativas o todas un poco.

Si le das demasiado poco, ocurre lo contrario. Cada intento comienza desde cero. Vuelven a aparecer los mismos callejones sin salida. Nada se transfiere de la última ejecución ni de una operación relacionada, y el bucle no avanza.

Queríamos una capa de conocimiento que ayudara con eso: recordar lo que funcionó, reutilizarlo más tarde y hacerlo sin intervención humana. Esa capa tiene un segundo compromiso, entre qué tan detallada es una lección almacenada y qué tan ampliamente se aplica.

Una nota muy específica (“en este kernel, con este tamaño de entrada, desenrollar este bucle”) puede ser exactamente lo que necesita el siguiente intento. También es fácil de usar incorrectamente en la siguiente operación, la siguiente GPU o un tamaño de entrada diferente. Una nota muy general (“hacer un mejor uso de la memoria integrada en el chip”) se aplica a casi todas partes y no le dice casi nada al modelo sobre qué hacer. Vimos ambos modos de fallo. Cuando las lecciones eran demasiado generales, reformulaban un fallo sin una acción. Cuando almacenábamos más detalles, a menudo estaban demasiado vinculados a una sola ejecución para ayudar a la siguiente. En una ejecución larga, la mayor parte de lo que el modelo leyó y escribió se gastó en recuperar y enrutar esa memoria en lugar de escribir kernels. La capa de memoria estaba haciendo mucho trabajo. No estaba mejorando al siguiente candidato. La Figura 2 muestra el desglose del costo de tokens de dicho sistema. El costo está dominado por las capas de conocimiento.

La versión de conocimiento que vale la pena conservar es más pequeña y logra un equilibrio entre generalidad y especificidad. Cuando el modelo está a punto de escribir un kernel, su prompt debe incluir solo contexto de alta confianza: conclusiones prácticas que asocien situaciones específicas con acciones (extraídas de un mapeo previo de modificación a impacto) y notas concisas de fallos de ejecuciones principales estrechamente relacionadas. Recuperadas a través de un filtrado jerárquico de etiquetas combinado con una búsqueda híbrida (palabras clave + semántica), las lecciones deben ser lo suficientemente específicas como para actuar en consecuencia y lo suficientemente acotadas como para aclarar dónde no se aplican. Las operaciones más profundas, como reorganizar y destilar aún más el almacén de lecciones, pertenecen a trabajos en segundo plano, no a un recorrido síncrono de múltiples saltos (multi-hop) sobre ejecuciones pasadas en cada intento. Si una conclusión no puede nombrar la situación y la acción, no vale la pena incluirla en el prompt. La Figura 3 muestra el desglose del costo de tokens después de corregir la capa de conocimiento; la mayoría de los tokens se gastan en la generación de candidatos después de esta corrección.

Caso de estudio: decodificación empaquetada de Gated DeltaNet

Un ejemplo concreto es el kernel de decodificación empaquetada (packed decode) en la ruta de Gated DeltaNet en Qwen 3.5 122B. La operación actualiza un estado recurrente y escribe la salida de decodificación a partir de entradas QKV empaquetadas, parámetros de puerta e índices de estado. Usamos esta tarea para ejercitar el bucle completo de Proteus en GPUs NVIDIA B200 con un backend de Triton: validar el contrato de la tarea, medir la implementación de referencia, pedir kernels candidatos a los agentes, ejecutar comprobaciones estáticas y compilaciones, verificar la corrección frente a la referencia controlada, realizar pruebas de rendimiento (benchmark) solo de los candidatos verificados y luego volver a medir los mejores candidatos.

La Figura 4 se lee de izquierda a derecha. El nodo de la línea de base fija el benchmark en 0.025 ms. El candidato 0000 es la semilla segura: reprodujo la estructura de decodificación empaquetada y pasó la validación, pero fue más lento que la referencia, por lo que Proteus lo mantuvo como un elemento principal medido en lugar de tratarlo como una victoria. A partir de ahí, Proteus dejó de optimizar un kernel genérico para cada forma (shape) y dividió la búsqueda en rutas específicas para cada forma.

La ruta de reparación de Batch-1 produjo un kernel específico para la forma en el Candidato 012, alcanzando 1.5 veces de aceleración en la forma de decodificación de un solo lote (single-batch). Los resultados más sólidos se obtuvieron en la ruta de decodificación de servicio (serving-decode): el Candidato 030 encontró la latencia de kernel medida más baja con 0.018 ms, y el Candidato 036 produjo la mejor aceleración de forma con 1.6 veces. Ese fragmento de servicio ganador se especializó para el diseño (layout) Batch=4, Key=128, Value=128 y procesó la dimensión de valor en fragmentos de 64 de ancho, por lo que es un kernel seguro para esa forma específica en lugar de un reemplazo universal.

El desvío final en la línea de tiempo muestra por qué la traza es importante. Los intentos de generaciones posteriores en C++ (en lugar de Triton) se enfrentaron a fallos de compilación y generación, y la ejecución larga terminó tras agotar el presupuesto de intentos de esa rama. El artefacto útil, por lo tanto, no es solo el candidato más rápido. Es la ruta completa que se muestra en la figura: se rechazaron los fallos semánticos, se midieron los kernels correctos pero más lentos y las mejoras de rendimiento reales se mantuvieron asociadas a la forma que las hace seguras para componer en un kernel de producción.

¿En qué estamos trabajando ahora?

El bucle evolutivo que construimos es estricto a la hora de escribir. A menudo llama al modelo siguiendo un patrón fijo: tomar el mejor kernel actual, intentar una pequeña edición, comprobar y repetir. Eso le quita la autonomía que el agente necesita. No puede cambiar fácilmente de estructura, cambiar de idioma o abandonar un diseño sin salida.

El bucle sigue siendo necesario. No para crear el kernel, sino para darle al agente una pista confiable sobre el siguiente paso. Esa pista tiene que provenir de dos lugares.

Primero, la comunicación con la capa de conocimiento: algunas conclusiones que sean lo suficientemente específicas como para actuar en consecuencia, y delimitadas para que sepamos dónde no se aplican. Sin eso, cada intento empieza desde cero.

Segundo, los resultados de un verificador de confianza: la corrección y los tiempos que el agente no midió por sí mismo. Esos números son pistas para el siguiente intento. También son las únicas puntuaciones en las que deberíamos creer. Si el agente mide el tiempo de su propio trabajo, volveríamos a encontrarnos con cachés residuales, comparaciones que no coinciden y pruebas que puede ver.

Por lo tanto, la división que queremos es más estrecha que «agente frente a bucle». Hay que dar al agente autonomía sobre cómo se escribe un kernel. Mantener el bucle como el canal para la memoria y la evaluación. El agente propone. El bucle devuelve lo que se le permite ver y si la última propuesta realmente ganó.

Proteus construyó kernels especializados para partes de Qwen 3.5 122B en su ruta Gated DeltaNet (un bloque de estilo de atención lineal) que se ejecuta en GPU NVIDIA B200. Las aceleraciones en los kernels individuales estuvieron en el rango de 1,8x a 5,2x.

La lección es que la generación es el paso fácil. La validación y la gestión del contexto son la parte difícil. Ahí es donde se necesita un tiempo de diseño cuidadoso e innovación.

La generación agéntica de kernels de GPU ha liberado el increíble potencial de la especialización extrema, pero la construcción de bancos de pruebas confiables y listos para producción sigue siendo un desafío complejo. Estamos abordando los retos más difíciles en la intersección de la AI y los sistemas, y buscando ingenieros audaces que se unan a nosotros para dar forma al futuro de la inferencia eficiente. Si te apasiona superar los límites de lo posible, ¡estamos contratando!

(Esta entrada del blog ha sido traducida utilizando herramientas basadas en inteligencia artificial) Publicación original