Las organizaciones y las personas que ejecutan múltiples modelos de IA personalizados, especialmente las familias de modelos recientes de Mezcla de Expertos (MoE), pueden enfrentar el desafío de pagar por la capacidad de GPU inactiva cuando los modelos individuales no reciben suficiente tráfico para saturar un punto final de computación dedicado. Para resolver este problema, nos asociamos con la comunidad vLLM y desarrollamos una solución eficiente para el servicio de adaptación de rango bajo múltiple (Multi-LoRA) de modelos MoE de código abierto populares como GPT-OSS o Qwen. Multi-LoRA es un enfoque popular para ajustar modelos. En lugar de volver a entrenar los pesos completos del modelo, multi-LoRA mantiene los pesos originales congelados e inyecta pequeños adaptadores entrenables en las capas del modelo. Con multi-LoRA, en el momento de la inferencia, varios modelos personalizados comparten la misma GPU, y solo los adaptadores se intercambian por solicitud. Por ejemplo, cinco clientes, cada uno de los cuales utiliza solo el 10 % de una GPU dedicada, pueden recibir servicio desde una sola GPU con LoRA múltiple, convirtiendo cinco GPU infrautilizadas en una GPU compartida de manera eficiente.
En esta publicación, explicamos cómo implementamos la inferencia multi-LoRA para modelos Mixture of Experts (MoE) en vLLM, describimos las optimizaciones a nivel de kernel que realizamos y le mostramos cómo puede beneficiarse de este trabajo. Usamos GPT-OSS 20B como nuestro ejemplo principal a lo largo de esta publicación.
Puede utilizar estas mejoras hoy en sus implementaciones locales de vLLM con la versión 0.15.0 o posterior. El servicio Multi-LoRA ahora funciona para familias de modelos MoE, incluidas GPT-OSS, Qwen3-MoE, DeepSeek y Llama MoE. Nuestras optimizaciones también ayudan a mejorar el alojamiento multi-LoRA para modelos densos, por ejemplo, Llama3.3 70B o Qwen3 32B. Las optimizaciones específicas de Amazon ofrecen mejoras de latencia adicionales con respecto a vLLM 0.15.0, por ejemplo, tokens de salida por segundo (OTPS) un 19 % más altos (es decir, qué tan rápido el modelo genera resultados) y un 8 % menos de tiempo hasta el primer token (TTFT) (es decir, cuánto tiempo hay que esperar antes de que el modelo comience a generar resultados) para GPT-OSS 20B. Para beneficiarse de estas optimizaciones, aloje sus modelos personalizados de LoRA en Amazon SageMaker AI o Amazon Bedrock.
Implementación de inferencia multi-LoRA para modelos MoE en vLLM
Antes de sumergirnos en nuestra implementación inicial de inferencia multi-LoRA para modelos MoE en vLLM, queremos brindar información general sobre los modelos MoE y el ajuste de LoRA que es importante para comprender la lógica detrás de nuestras optimizaciones. Los modelos MoE contienen múltiples redes neuronales especializadas llamadas expertos. Un enrutador dirige cada token de entrada a los expertos más relevantes, cuyas salidas luego se agregan. Esta arquitectura dispersa procesa modelos más grandes con menos recursos computacionales porque solo una fracción de los parámetros totales del modelo se activa por token; consulte la Figura 1 a continuación para obtener una visualización.
Cada experto es una pequeña red de retroalimentación que procesa el estado oculto de un token en dos etapas. Primero, la proyección gate_up expande el estado oculto compacto (por ejemplo, 4096 atenuaciones) a un espacio intermedio más grande (por ejemplo, 11008 atenuaciones). Esta expansión es necesaria porque las características en el espacio compacto están estrechamente entrelazadas: el espacio más grande le da a la red espacio para separarlas, transformarlas y controlar selectivamente cuáles son importantes. En segundo lugar, la proyección hacia abajo comprime el resultado a la dimensión original. Esto ayuda a mantener la salida compatible con el resto del modelo y actúa como un cuello de botella, lo que obliga a la red a conservar solo las funciones más útiles. En conjunto, este patrón de “expandir y luego comprimir” permite a cada experto aplicar transformaciones enriquecidas mientras mantiene un tamaño de salida consistente. vLLM utiliza un kernel fusionado_moe para ejecutar estas proyecciones como operaciones de multiplicación de matriz general de grupo (Grupo GEMM): un GEMM por experto asignado a un token determinado. El ajuste fino de Multi-LoRA mantiene congelados los pesos del modelo base W, por ejemplo, W_gate_up para la proyección gate_up, y entrena dos pequeñas matrices A y B que juntas forman un adaptador. Para una proyección con pesos base W de forma h_in × h_out, LoRA entrena A de forma h_in × r y B de forma r × h_out, donde r es el rango de LoRA (normalmente 16-64). La salida ajustada se convierte en y = xW + xAB. Cada adaptador LoRA agrega dos operaciones a una proyección. La operación de reducción calcula z=xA, reduciendo la entrada desde las dimensiones h_in hasta las dimensiones r. La operación de expansión toma ese resultado de dimensión r y lo proyecta nuevamente a las dimensiones h_out multiplicando z por B. Esto se ilustra a la derecha de la Figura 1.
Figura 1: Ilustración de cómo funcionan los modelos MoE-LoRA con un ejemplo de dimensión de estado oculto 4096, dimensión de representación intermedia 11008 y rango LoRA r = 32.
Cada experto tiene dos proyecciones de peso: gate_up y down. Cuando se aplica un adaptador LoRA, agrega dos operaciones de bajo rango, es decir, encoger y expandir, a cada proyección. Esto significa que cada experto requiere cuatro operaciones del kernel LoRA en total: reducir y expandir para gate_up, y reducir y expandir para bajar. En una configuración de servicio multi-LoRA, donde se sirven múltiples adaptadores LoRA simultáneamente para diferentes usuarios o tareas, el sistema debe gestionar de manera eficiente estas cuatro operaciones por experto, por adaptador y por solicitud. Esto lo convierte en un cuello de botella clave en el rendimiento de los modelos MoE. Las cuatro operaciones involucran matrices, donde una dimensión (el rango r de LoRA) es 100-300 veces más pequeña que la otra (por ejemplo, estado oculto y dimensión de representación intermedia). Los núcleos GEMM estándar están diseñados para matrices aproximadamente cuadradas y funcionan mal en matrices delgadas, razón por la cual las optimizaciones del núcleo que se describen más adelante en esta publicación son necesarias. Además de tener que optimizar para matrices delgadas, agregar soporte multi-LoRA para modelos MoE presentó dos desafíos técnicos. Primero, vLLM carecía de un kernel para realizar LoRA en capas MoE porque los kernels multi-LoRA densos existentes no manejan el enrutamiento experto. En segundo lugar, MoE LoRA combina dos fuentes de escasez: enrutamiento de expertos (tokens asignados a diferentes expertos) y selección de adaptadores (solicitudes que utilizan diferentes adaptadores LoRA). Esta escasez de compuestos requiere un diseño de núcleo especializado. Para abordar estos desafíos, creamos un kernel fusionado_moe_lora que integra las operaciones LoRA en el kernel fusionado_moe. Este nuevo kernel realiza GEMM de reducción y expansión de LoRA para las proyecciones gate_up y down. El kernel fused_moe_lora sigue la misma lógica que el kernel fused_moe y agrega una dimensión adicional a la cuadrícula para los adaptadores LoRA activados correspondientes.
Con esta implementación fusionada en vLLM, podríamos ejecutar un servicio multi-LoRA con GPT-OSS 20B en una GPU H200, alcanzando 26 OTPS y 1053 ms TTFT en el conjunto de datos Sonnet (un punto de referencia basado en poesía) con una longitud de entrada de 1600, una longitud de salida de 600 y una simultaneidad de 16. Para reproducir estos resultados, consulte nuestro PR en la versión 0.11.1.rc3 de Repositorio vLLM GitHub. En el resto de este blog, mostraremos cómo optimizamos el rendimiento a partir de estos números de habilitación básicos.
Mejora del rendimiento de inferencia multi-LoRA en vLLM
Después de finalizar nuestra implementación inicial, utilizamos NVIDIA Nsight Systems (Nsys) para identificar cuellos de botella y descubrimos que el kernel fusionado_moe_lora era el componente de mayor latencia. Luego utilizamos NVIDIA Nsight Compute (NCU) para perfilar el rendimiento de la computación y la memoria para las cuatro operaciones del kernel: gate_up_shrink, gate_up_expand, down_shrink y down_expand. Estos hallazgos nos llevaron a desarrollar optimizaciones de ejecución, optimizaciones a nivel de kernel y configuraciones ajustadas para estos cuatro kernels.
Optimizaciones de ejecución
Con nuestra implementación inicial, el TTFT multi-LoRA era 10 veces mayor (peor) que el TTFT del modelo base (es decir, la versión de lanzamiento público de GPT-OSS 20B). Nuestro perfil reveló que el compilador Triton trataba las variables dependientes de la longitud de entrada como constantes de tiempo de compilación, lo que provocaba que el kernel fusionado_moe_lora se volviera a compilar desde cero para cada nueva longitud de contexto en lugar de reutilizarse. Esto es visible en la Figura 2: las llamadas cuModuleLoadData antes de cada ejecución del kernel fusionado_moe_lora indican que la GPU está cargando un binario del kernel recién compilado en lugar de reutilizar uno almacenado en caché, y las grandes brechas entre los tiempos de inicio del kernel muestran que la GPU permanece inactiva durante la recompilación. Esta sobrecarga impulsó la regresión TTFT 10× sobre el modelo base. Resolvimos esto agregando una sugerencia del compilador do_not_specialize para estas variables, indicando a Triton que compilara el kernel una vez y lo reutilizara en todas las longitudes de contexto.
Figura 2: Resultados de la creación de perfiles para el kernel fused_moe_lora antes de nuestras optimizaciones de ejecución.
La creación de perfiles también reveló que nuestro kernel fusionado_moe_lora se lanzó con la misma sobrecarga elevada independientemente de si la solicitud utilizó solo el modelo base, adaptadores de solo atención (adaptadores LoRA con pesos solo en las capas de atención) o adaptadores LoRA completos (adaptadores con pesos en las capas de atención y MoE). Para ayudar a resolver esto, agregamos una lógica de salida temprana para omitir el kernel fusionado_moe_lora en capas sin adaptadores LoRA, lo que ayuda a evitar la ejecución innecesaria del kernel.
Los núcleos de reducción y expansión se ejecutan en serie, lo que creó burbujas entre las ejecuciones de dos núcleos en nuestra implementación inicial. Para superponer la ejecución del kernel, implementamos el lanzamiento dependiente programático (PDL). Con PDL, un kernel dependiente puede comenzar a iniciarse antes de que finalice el kernel primario, lo que permite que el kernel expandido precapture pesos en la memoria compartida y la caché L2 mientras se ejecuta el kernel reducido. Cuando se completa el kernel de reducción, el kernel de expansión ya ha cargado sus pesos y puede comenzar a calcular inmediatamente.
También agregamos soporte para decodificación especulativa con CudaGraph para LoRA, solucionando un problema en vLLM que capturaría diferentes CudaGraphs para el modelo base y el adaptador. Los CudaGraphs son importantes para la eficiencia ya que se utilizan para capturar secuencias de operaciones de GPU para ayudar a reducir la sobrecarga del kernel de GPU, por ejemplo, kernels como una sola unidad. Como resultado, CudaGraphs puede reducir los gastos generales de la CPU y las latencias de inicio del kernel. Con nuestras optimizaciones de ejecución, OTPS mejoró a 50/100 sin/con decodificación especulativa y TTFT mejoró a 150 ms para GPT-OSS 20B usando la configuración predeterminada. Durante el resto del blog, informaremos los números con decodificación especulativa.
Optimizaciones del kernel
Split-K es una estrategia de descomposición del trabajo que ayuda a mejorar el equilibrio de carga para matrices delgadas. La contracción LoRA calcula xA donde x tiene una dimensión de 1×h_in y A tiene una dimensión h_in×r. Cada uno de los r elementos de salida requiere sumar h_in multiplicaciones. Los núcleos GEMM estándar asignan diferentes grupos de subprocesos (lotes de subprocesos de GPU que comparten memoria rápida en el chip) a diferentes elementos de salida, pero cada grupo de subprocesos calcula su suma h_in de forma secuencial. Con r en decenas y h_in en miles, hay pocos elementos de salida para paralelizar, mientras que cada uno requiere una suma secuencial larga. Split-K soluciona esto dividiendo la suma sobre la dimensión interna K de un GEMM (en este ejemplo, K=h_in) en múltiples grupos de subprocesos, que calculan sumas parciales en paralelo y luego combinan sus resultados. Estos resultados parciales requieren una suma atómica para producir la suma final. Dado que realizamos una suma atómica pura sin lógica adicional, utilizamos la libertad del compilador Triton para optimizaciones configurando el parámetro sem="relaxed" para la operación de suma atómica.
El programador de GPU asigna varios grupos de subprocesos al mismo elemento de salida y ejecuta grupos de subprocesos para diferentes elementos de salida al mismo tiempo. Para lora_shrink, cada elemento de salida requiere leer una columna de A, que abarca las filas h_in. Con h_in en miles, cada columna toca líneas de caché distribuidas en una gran región de memoria. Las columnas cercanas comparten las mismas filas y se superponen en la caché, por lo que los grupos de subprocesos que trabajan en columnas vecinas pueden beneficiarse de la reutilización de los datos cargados de cada uno. La rotación cooperativa de subprocesos (CTA) reordena la programación para que los grupos de subprocesos que trabajan en columnas cercanas se ejecuten al mismo tiempo, lo que aumenta la reutilización de la caché L2. Aplicamos CTA swizzling a la operación lora_shrink.
También eliminamos operaciones innecesarias de enmascaramiento y productos escalares de los núcleos LoRA de reducción y expansión. Los núcleos Triton cargan datos en bloques de tamaño fijo, pero es posible que las dimensiones de la matriz no se dividan uniformemente en estos tamaños de bloques. Por ejemplo, si BLOCK_SIZE_K es 64 pero la dimensión de la matriz K es 100, el segundo bloque intentaría leer 28 ubicaciones de memoria no válidas. El enmascaramiento ayuda a prevenir estos accesos ilegales a la memoria al verificar si cada índice está dentro de los límites antes de cargarlo. Sin embargo, estas comprobaciones condicionales se ejecutan en cada operación de carga, lo que añade gastos generales incluso cuando los elementos son válidos. Introdujimos un parámetro EVEN_K que verifica si K se divide uniformemente por BLOCK_SIZE_K. Cuando es verdadero, las cargas son válidas y el enmascaramiento se puede omitir por completo, lo que ayuda a reducir tanto la sobrecarga de enmascaramiento como los cálculos innecesarios de productos escalares.
Por último, fusionamos la adición de los pesos de LoRA con los pesos del modelo base en el núcleo expandido de LoRA. Esta optimización ayuda a reducir la sobrecarga de lanzamiento del kernel. Estas optimizaciones del kernel nos ayudaron a alcanzar 144 OTPS y 135 ms TTFT para GPT-OSS 20B.
Ajuste de las configuraciones del kernel para Amazon SageMaker AI y Amazon Bedrock
Los núcleos Triton requieren ajustar parámetros como el tamaño de los bloques (BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K), que controlan cómo se divide el cálculo de la matriz entre los grupos de subprocesos. Los parámetros avanzados incluyen GROUP_SIZE_M, que controla el orden de los grupos de subprocesos para la localidad de caché, y SPLIT_K, que paraleliza las sumas en la dimensión de la matriz interna.
Descubrimos que los núcleos MoE LoRA que utilizan configuraciones predeterminadas optimizadas para MoE fusionado estándar tuvieron un rendimiento deficiente para el servicio multi-LoRA. Estos valores predeterminados no tuvieron en cuenta la dimensión de cuadrícula adicional correspondiente al índice LoRA y la escasez compuesta de múltiples adaptadores. Para abordar este cuello de botella, agregamos soporte para que los usuarios carguen configuraciones personalizadas proporcionando una ruta de carpeta. Para obtener más información, consulte la documentación de ajuste de vLLM LoRA. Ajustamos las cuatro operaciones de fusionado_moe_lora (reducción de puerta_arriba, expansión de puerta_arriba, reducción de reducción, expansión descendente) simultáneamente, ya que comparten el mismo parámetro BLOCK_SIZE_M. Los clientes de Amazon SageMaker AI y Bedrock ahora tienen acceso a estas configuraciones optimizadas, que se cargan automáticamente y alcanzan 171 OTPS y 124 ms TTFT para GPT-OSS 20B.
Resultados y conclusión
A través de nuestra colaboración con la comunidad vLLM, implementamos y abrimos servicios multi-LoRA para modelos MoE, incluidos GPT-OSS, Qwen3 MoE, DeepSeek y Llama MoE. Luego aplicamos optimizaciones, por ejemplo, obteniendo un 454 % de mejoras en OTPS y un 87 % menos de TTFT para GPT-OSS 20B en vLLM 0.15.0 frente a vLLM 0.11.1rc3. Algunas optimizaciones, particularmente el ajuste del kernel y el cambio de CTA, también mejoraron el rendimiento para modelos densos, por ejemplo, Qwen3 32B OTPS mejoró en un 99%. Para aprovechar este trabajo en sus implementaciones locales, utilice vLLM 0.15.0 o posterior. Las optimizaciones específicas de Amazon, disponibles en Amazon Bedrock y Amazon SageMaker AI, ayudan a ofrecer mejoras de latencia adicionales en todos los modelos, por ejemplo, OTPS un 19 % más rápido y TTFT un 8 % mejor frente a vLLM 0.15.0 para GPT-OSS 20B. Para comenzar con el alojamiento de modelos personalizados en Amazon, consulte la documentación sobre alojamiento de IA de Amazon SageMaker y Amazon Bedrock.
Figura 3: Tokens de salida por segundo (OTPS) y tiempo hasta el primer token (TTFT) para la inferencia multi-LoRA GPT-OSS 20B: 1/ Implementación inicial en vLLM 0.11.1rc3; 2/ con vLLM 0.15.0; 3/ con vLLM 0.15.0 y ajuste del kernel personalizado de AWS. Los experimentos utilizaron 1600 tokens de entrada y 600 tokens de salida con adaptadores LoRA de rango 32 y 8 cargados en paralelo.
Expresiones de gratitud
Nos gustaría agradecer a los contribuyentes y colaboradores de la comunidad vLLM: Jee Li, Chen Wu, Varun Sundar Rabindranath, Simon Mo y Robert Shaw, y a los miembros de nuestro equipo: Xin Yang, Sadaf Fardeen, Ashish Khetan y George Karypis.
Sobre los autores
Danielle Maddix Robinson
Danielle Maddix Robinson es científica aplicada sénior en la investigación de IA/ML de AWS. Obtuvo su doctorado en Ingeniería Computacional y Matemática del ICME de Stanford. Su experiencia abarca análisis numérico, álgebra lineal y optimización para PDE numéricas y aprendizaje profundo. Desde que se unió a AWS en 2018, ha desarrollado pronósticos de series temporales y modelos básicos tabulares, incluidos Chronos y Mitra. Lidera el esfuerzo de optimización de la inferencia Multi-LoRA y la investigación sobre el aprendizaje automático limitado por la física para la informática científica.
Canción Xiang
Xiang Song es científico aplicado principal en la investigación de IA/ML de AWS, donde trabaja en los desafíos de los sistemas ML en capacitación e inferencia distribuidas. Xiang también ha sido desarrollador líder de los esfuerzos de aprendizaje automático de gráficos de AWS, incluidos los marcos de aprendizaje profundo GraphStorm. Recibió su doctorado en arquitectura y sistemas informáticos en la Universidad de Fudan, Shanghai, en 2014.
Haipeng Li
Haipeng Li es científico aplicado en AWS, donde se especializa en optimización de inferencia de modelos de lenguaje grande (LLM). Su trabajo se centra en mejorar la eficiencia del modelo, desarrollar núcleos de inferencia de alto rendimiento y reducir la latencia y los costos computacionales para implementar LLM a escala. A Haipeng le apasiona cerrar la brecha entre la investigación de IA de vanguardia y las soluciones prácticas y listas para producción que mejoran el rendimiento de los servicios de aprendizaje automático basados en la nube.
Yu Gong
Yu Gong trabaja en el avance de los sistemas de inteligencia artificial para su implementación en el mundo real. Como científico aplicado en Amazon, diseña arquitecturas escalables y marcos de inferencia de rendimiento optimizado, centrándose en la optimización del kernel de GPU y el diseño de sistemas con reconocimiento de hardware. Su trabajo une los sistemas de producción e investigación de IA para brindar servicios inteligentes eficientes y confiables.
George Novak
George Novack es ingeniero de software en AWS, donde se concentra en escribir núcleos de GPU eficientes y crear sistemas de inferencia rápidos y confiables.
Mani Kumar Adari
Mani Kumar Adari es ingeniero de software principal en AWS. Su enfoque actual es construir una infraestructura escalable para la inferencia y el ajuste de modelos de lenguaje grandes en AWS. Desde que se unió a Amazon en 2015, Mani ha trabajado en sistemas de inteligencia artificial conversacional y de comprensión del lenguaje natural, incluidos Amazon Lex y Amazon Connect, antes de pasar a la infraestructura de inteligencia artificial generativa en AWS Bedrock.
Florian Saupe
Florian Saupe es director técnico principal de productos en la investigación de IA/ML de AWS y trabaja en optimización de inferencia de modelos, capacitación distribuida a gran escala y resistencia a fallas. Antes de unirse a AWS, Florian dirigió la gestión técnica de productos para la conducción automatizada en Bosch, fue consultor de estrategia en McKinsey & Company y trabajó como científico en robótica y sistemas de control, campo en el que tiene un doctorado.