Servir Dezenas de Modelos Fine-Tuned com vLLM no Amazon SageMaker AI e Amazon Bedrock

O Desafio da Subutilização de GPU com Múltiplos Modelos

Organizações que executam vários modelos de inteligência artificial customizados enfrentam um problema prático: quando modelos individuais não recebem tráfego suficiente para saturar um endpoint dedicado, há desperdício significativo de capacidade de GPU. Modelos recentes da família de Mistura de Especialistas (Mixture of Experts — MoE) amplificam esse desafio.

A AWS, em parceria com a comunidade vLLM, desenvolveu uma abordagem eficiente para servir múltiplos modelos adaptados usando Multi-Low-Rank Adaptation (Multi-LoRA), resolvendo o problema de subutilização. A solução permite que, por exemplo, cinco clientes utilizando apenas 10% de uma GPU dedicada sejam servidos a partir de uma única GPU compartilhada.

Como Funciona a Multi-LoRA para Modelos de Mistura de Especialistas

Entendendo os Modelos MoE

Modelos de Mistura de Especialistas contêm múltiplas redes neurais especializadas chamadas especialistas. Um roteador direciona cada token de entrada para os especialistas mais relevantes, cujas saídas são então agregadas. Essa arquitetura esparsa permite processar modelos maiores com menos recursos computacionais, pois apenas uma fração dos parâmetros totais é ativada por token.

Cada especialista é uma pequena rede de propagação direta que processa o estado oculto de um token em dois estágios. Primeiro, a projeção gate_up expande o estado oculto compacto (por exemplo, 4.096 dimensões) em um espaço intermediário maior (por exemplo, 11.008 dimensões). Essa expansão é necessária porque as características no espaço compacto estão fortemente interligadas — o espaço maior oferece lugar para separar, transformar e seletivamente ativar quais características importam. Em seguida, a projeção down comprime o resultado de volta à dimensão original, mantendo compatibilidade com o resto do modelo.

Ilustração de como os modelos MoE-LoRA funcionam com dimensão de estado oculto 4.096, dimensão de representação intermediária 11.008 e rank LoRA r = 32. Cada especialista possui duas projeções de peso: gate_up e down. Quando um adaptador LoRA é aplicado, adiciona duas operações de baixo rank: redução e expansão, para cada projeção. — Fonte: Aws

O Papel da Multi-LoRA

Multi-LoRA é uma abordagem popular para adaptar modelos sem retreinar pesos inteiros. Em vez disso, mantém os pesos originais congelados e injeta pequenos adaptadores treináveis nas camadas do modelo. No tempo de inferência, múltiplos modelos customizados compartilham a mesma GPU, com apenas os adaptadores sendo alternados por requisição.

Para uma projeção com pesos base W de dimensão h_in × h_out, LoRA treina duas matrizes pequenas: A com dimensão h_in × r e B com dimensão r × h_out, onde r é o rank LoRA (tipicamente 16-64). A saída adaptada torna-se y = xW + xAB. Cada adaptador LoRA adiciona duas operações a uma projeção: a operação de redução computa z=xA, reduzindo a entrada de h_in dimensões para r dimensões, e a operação de expansão projeta o resultado de r dimensões de volta para h_out dimensões multiplicando z por B.

Em uma configuração de Multi-LoRA servindo múltiplos adaptadores simultaneamente para diferentes usuários ou tarefas, o sistema deve gerenciar eficientemente essas quatro operações por especialista, por adaptador, por requisição — um gargalo chave de performance para modelos MoE.

Implementação Técnica no vLLM

A AWS utilizava a comunidade vLLM para criar um novo kernel chamado fused_moe_lora que integra operações LoRA ao kernel fused_moe existente. Este novo kernel realiza GEMMs (Multiplicações Gerais de Matrizes) de redução e expansão LoRA para as projeções gate_up e down.

Dois desafios técnicos demandaram soluções especializadas: vLLM não possuía um kernel para realizar LoRA em camadas MoE, pois os kernels Multi-LoRA densos existentes não gerenciam roteamento de especialistas. Além disso, MoE LoRA combina duas fontes de esparsidade — roteamento de especialistas e seleção de adaptadores — exigindo design de kernel especializado.

Com essa implementação inicial, foi possível executar Multi-LoRA com o modelo GPT-OSS 20B em uma GPU H200, alcançando 26 OTPS (Tokens de Saída Por Segundo) e 1.053 ms TTFT (Tempo Até o Primeiro Token) no dataset Sonnet com comprimento de entrada 1.600, comprimento de saída 600 e concorrência de 16.

Otimizações de Performance do vLLM

Otimizações em Nível de Execução

Análises de performance inicial revelaram que TTFT estava 10x pior que o modelo base. Usando ferramentas de profiling, identificou-se que o compilador Triton tratava variáveis dependentes do comprimento de entrada como constantes de tempo de compilação, causando recompilação do kernel fused_moe_lora a partir do zero para cada novo comprimento de contexto em vez de reutilizar uma versão em cache. Isso foi resolvido adicionando uma dica de compilador do_not_specialize, instruindo Triton a compilar o kernel uma vez e reutilizá-lo em todos os comprimentos de contexto.

O profiling também revelou que o kernel fused_moe_lora era lançado com a mesma sobrecarga elevada independentemente de a requisição usar apenas o modelo base, adaptadores somente de atenção ou adaptadores LoRA completos. Adicionou-se lógica de saída antecipada para pular o kernel fused_moe_lora em camadas sem adaptadores LoRA, evitando execução desnecessária.

Os kernels de redução e expansão executavam sequencialmente, criando gaps entre execuções. Implementou-se Programmatic Dependent Launch (PDL) para permitir que o kernel dependente comece a lançar antes do kernel primário terminar. Isso permite que o kernel de expansão pré-carregue pesos em memória compartilhada e cache L2 enquanto o kernel de redução executa.

Adicionou-se também suporte para decodificação especulativa com CudaGraph para LoRA, corrigindo um problema onde o vLLM capturava CudaGraphs diferentes para modelo base e adaptador. CudaGraphs são importantes para eficiência pois capturam sequências de operações GPU, reduzindo sobrecarga de kernel.

Com essas otimizações de execução, OTPS melhorou para 50/100 sem/com decodificação especulativa e TTFT melhorou para 150 ms para GPT-OSS 20B.

Otimizações em Nível de Kernel

Split-K é uma estratégia de decomposição de trabalho que melhora balanceamento de carga para matrizes estreitas. Redução LoRA computa xA onde x tem dimensão 1×h_in e A tem dimensão h_in×r. Cada um dos r elementos de saída requer somar h_in multiplicações. Kernels GEMM padrão atribuem diferentes grupos de threads a diferentes elementos de saída, mas cada grupo calcula sua soma sequencialmente. Com r na dezena e h_in na casa dos milhares, há poucos elementos de saída para paralelizar enquanto cada um requer uma longa soma sequencial. Split-K resolve isso dividindo a soma sobre a dimensão interna K de um GEMM entre múltiplos grupos de threads, que calculam somas parciais em paralelo e combinam resultados.

Cooperative Thread Array (CTA) swizzling reordena o agendamento para que grupos de threads trabalhando em colunas próximas executem simultaneamente, aumentando reutilização de cache L2. Esta técnica foi aplicada à operação de redução LoRA.

Removeu-se mascaramento desnecessário e operações de produto escalar dos kernels LoRA de redução e expansão. Kernels Triton carregam dados em blocos de tamanho fixo, mas dimensões de matrizes podem não se dividir uniformemente nesses tamanhos de bloco. Mascaramento previne acessos ilegais à memória, mas essas verificações condicionais executam em cada operação de carregamento, adicionando sobrecarga. Introduziu-se um parâmetro EVEN_K que verifica se K se divide uniformemente por BLOCK_SIZE_K — quando verdadeiro, carregamentos são válidos e mascaramento pode ser completamente ignorado.

Por fim, fundiu-se a adição dos pesos LoRA aos pesos do modelo base no kernel de expansão LoRA, reduzindo sobrecarga de lançamento de kernel. Essas otimizações de kernel alcançaram 144 OTPS e 135 ms TTFT para GPT-OSS 20B.

Ajuste de Configurações para Amazon SageMaker AI e Amazon Bedrock

Kernels Triton requerem ajuste de parâmetros como tamanhos de bloco (BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K) que controlam como a computação de matriz se divide entre grupos de threads. Parâmetros avançados incluem GROUP_SIZE_M, que controla ordem de grupos de threads para localidade de cache, e SPLIT_K, que paraleliza somas sobre a dimensão de matriz interna.

Descobriu-se que os kernels MoE LoRA usando configurações padrão otimizadas para MoE fundido padrão tinham performance ruim para serviço Multi-LoRA. Os padrões não contabilizavam a dimensão de grade adicional correspondente ao índice LoRA e a esparsidade composta de múltiplos adaptadores.

Adicionou-se suporte para usuários carregarem configurações ajustadas customizadas fornecendo um caminho de pasta. Clientes da AWS SageMaker AI e Bedrock agora têm acesso a essas configurações ajustadas, carregadas automaticamente e alcançando 171 OTPS e 124 ms TTFT para GPT-OSS 20B — melhorias subsequentes em relação às 144 OTPS e 135 ms TTFT da versão anterior.

Resultados e Impacto

Através da colaboração com a comunidade vLLM, a AWS implementou e disponibilizou em código aberto Multi-LoRA serving para modelos MoE incluindo GPT-OSS, Qwen3 MoE, DeepSeek, e Llama MoE. As otimizações resultaram em melhorias impressionantes: 454% de aumento em OTPS e 87% de redução em TTFT para GPT-OSS 20B comparando vLLM 0.15.0 versus vLLM 0.11.1rc3.

Tokens de Saída Por Segundo (OTPS) e Tempo Até o Primeiro Token (TTFT) para inferência Multi-LoRA do GPT-OSS 20B mostrando progressão da implementação inicial em vLLM 0.11.1rc3, vLLM 0.15.0, e vLLM 0.15.0 com ajuste de kernel customizado da AWS. Experimentos utilizaram 1.600 tokens de entrada e 600 tokens de saída com rank LoRA 32 e 8 adaptadores carregados em paralelo. — Fonte: Aws

Algumas otimizações, particularmente ajuste de kernel e CTA swizzling, também melhoraram performance para modelos densos — OTPS do Qwen3 32B melhorou 99%.

Para aproveitar este trabalho em deployments locais, utilize vLLM 0.15.0 ou versões posteriores. Otimizações específicas da AWS, disponíveis no Amazon SageMaker AI e Amazon Bedrock, entregam melhorias adicionais de latência em modelos, por exemplo 19% mais rápido em OTPS e 8% melhor em TTFT versus vLLM 0.15.0 para GPT-OSS 20B.

Como Começar

Para começar com hosting de modelos customizados na AWS, consulte a documentação de Amazon SageMaker AI hosting e Amazon Bedrock.

O repositório vLLM contém as implementações e pode ser explorado para entender melhor as otimizações técnicas por trás dessas melhorias de performance.

Fonte

Efficiently serve dozens of fine-tuned models with vLLM on Amazon SageMaker AI and Amazon Bedrock (https://aws.amazon.com/blogs/machine-learning/efficiently-serve-dozens-of-fine-tuned-models-with-vllm-on-amazon-sagemaker-ai-and-amazon-bedrock/)

Comments

Leave a Reply

Your email address will not be published. Required fields are marked *