InteligĂȘncia artificial

SGLang: Execução Eficiente de Programas de Modelos de Linguagem Estruturados

mm
SGLang: Efficient Execution of Structured Language Model Programs

Os grandes modelos de linguagem (LLMs) estão sendo cada vez mais utilizados para tarefas complexas que exigem múltiplas chamadas de geração, técnicas de prompt avançadas, controle de fluxo e entradas/saídas estruturadas. No entanto, os sistemas eficientes para programar e executar essas aplicações estão faltando. SGLang, um sistema recém-introduzido, visa abordar isso fornecendo execução eficiente de programas de modelos de linguagem complexos. SGLang compreende uma linguagem de frontend e um tempo de execução. O frontend simplifica a programação com primitivos para geração e controle de paralelismo, enquanto o tempo de execução acelera a execução por meio de otimizações novas, como RadixAttention para reutilização do cache KV e máquinas de estado finito comprimidas para decodificação de saída estruturada mais rápida. Experimentos demonstram que SGLang alcança até 6,4 vezes mais taxa de transferência em comparação com sistemas de inferência de ponta em vários grandes modelos de linguagem e multimodais, abordando tarefas como controle de agente, raciocínio lógico, benchmarks de aprendizado de poucos disparos, decodificação JSON, pipelines de geração aumentada por recuperação e chat de várias voltas.

Os recentes avanços nas capacidades dos LLMs expandiram sua utilidade, permitindo que eles lidem com uma gama mais ampla de tarefas gerais e funcionem como agentes autônomos. Nesses aplicativos, os LLMs se envolvem em planejamento de múltiplas rodadas, raciocínio e interação com ambientes externos. Isso é facilitado por meio do uso de ferramentas, múltiplas modalidades de entrada e várias técnicas de prompt, como aprendizado de poucos disparos, autoconsistência, esqueleto de pensamento e árvore de pensamento. Esses novos casos de uso exigem chamadas de geração de LLMs múltiplas e frequentemente dependentes, indicando uma tendência de usar estruturas de múltiplas chamadas para concluir tarefas complexas.

Essa mudança marca uma transição do uso simples de chat para um uso mais sofisticado e programático dos LLMs, onde os programas agendam e controlam os processos de geração dos LLMs. Esses programas são referidos como “Programas de Modelo de Linguagem” (LM Programs). Técnicas de prompt avançadas e fluxos de trabalho de agente caem dentro do escopo dos programas LM. Existem duas propriedades comuns dos programas LM: (1) Os programas LM normalmente envolvem chamadas de LLMs múltiplas intercaladas com controle de fluxo para concluir tarefas complexas e melhorar a qualidade geral. (2) Os programas LM recebem entradas estruturadas e produzem saídas estruturadas, permitindo a composição de programas LM e a integração em sistemas de software existentes.

Neste artigo, vamos mergulhar mais fundo no framework SGLang, explorando sua arquitetura, analisando seu desempenho e comparando-o com frameworks de ponta. Vamos começar.

Introdução ao SGLang

Apesar do uso generalizado de programas LM, os sistemas atuais para expressá-los e executá-los permanecem ineficientes. SGLang identifica dois desafios principais associados ao uso eficiente de programas LM:

  • Complexidade de programação: Desenvolver programas LM é tedioso e difícil devido à natureza não determinística dos LLMs. Isso envolve manipulação extensiva de strings, ajuste experimental de prompts, parsing de saída frágil, tratamento de múltiplas modalidades de entrada e implementação de mecanismos de paralelismo. Essa complexidade reduz significativamente a legibilidade de programas simples.
  • Ineficiência de execução: A execução de programas LM é ineficiente devido à computação redundante e ao uso de memória. Os sistemas de inferência de ponta, otimizados para reduzir a latência e melhorar a taxa de transferência, carecem de conhecimento direto da carga de trabalho, resultando em ineficiências significativas. Um exemplo notável é a reutilização do cache KV, que consiste em tensores intermediários reutilizáveis essenciais para a inferência gerativa. Os sistemas atuais carecem de mecanismos eficazes para facilitar a reutilização do cache KV em chamadas de LLMs múltiplas que compartilham um prefixo comum, levando a computações desnecessárias e memória desperdiçada. Além disso, a decodificação de saída estruturada, como o modo JSON, é subótima, pois os sistemas existentes apenas decodificam um token de cada vez.

Para abordar esses desafios, SGLang introduz uma Linguagem de Geração Estruturada para LLMs. A ideia central é explorar sistematicamente a estrutura de múltiplas chamadas nos programas LM para execução eficiente. Como mostrado na figura a seguir, SGLang tem duas partes: uma linguagem de frontend e um tempo de execução.

O frontend simplifica a programação de programas LM, e o tempo de execução acelera a execução deles. Essas partes podem trabalhar juntas para melhor desempenho ou funcionar independentemente.

SGLang é uma linguagem específica de domínio incorporada em Python, fornecendo primitivos para geração (por exemplo, extend, gen, select) e controle de paralelismo (por exemplo, fork, join). É compatível com o controle de fluxo e bibliotecas do Python, permitindo que os usuários desenvolvam fluxos de trabalho de prompt avançados facilmente com a sintaxe nativa do Python. SGLang inclui um interpretador e um compilador. O interpretador gerencia o estado do prompt como um fluxo e envia operações primitivas para o fluxo para execução assíncrona, garantindo controle adequado sobre a sincronização e o paralelismo intra-programa. Além disso, os programas SGLang podem ser rastreados e compilados para otimizações adicionais.

  • RadixAttention: Essa técnica permite a reutilização automática do cache KV em chamadas de geração múltiplas. Nos sistemas de inferência existentes, o cache KV de um pedido é descartado após o processamento, impedindo a reutilização em chamadas múltiplas e retardando a execução. SGLang mantém um cache LRU do cache KV dentro de uma árvore radix, gerenciando o cache KV como um cache tradicional e usando a árvore radix para correspondência, inserção e evição eficientes. Isso permite que o tempo de execução lidere com vários padrões de reutilização de forma eficiente.
  • Máquina de Estado Finito Comprimida: Essa técnica permite a decodificação de saída estruturada mais rápida. Os sistemas existentes seguem as restrições apenas para o próximo token, podendo decodificar apenas um token de cada vez. Em vez disso, SGLang analisa as restrições e constrói uma máquina de estado finito comprimida para representá-las, comprimindo um caminho de vários tokens em um caminho de um passo sempre que possível, permitindo a decodificação de vários tokens de uma vez para uma velocidade mais rápida.
  • API Speculative Execution: Para modelos de API, como o OpenAI, SGLang introduz a execução especulativa de API para otimizar programas de múltiplas chamadas.

Usando SGLang, várias aplicações de LLM foram implementadas, incluindo controle de agente, raciocínio lógico, benchmarks de aprendizado de poucos disparos, decodificação JSON, pipelines de geração aumentada por recuperação, chat de várias voltas e processamento de multimodalidade. O desempenho foi testado em modelos, incluindo Llama-7B/70B, Mistral-8x7B, LLaVA-v1.5-7B (imagem) e LLaVA-NeXT-34B (vídeo) em GPUs NVIDIA A10G e A100. Os resultados experimentais mostram que SGLang alcança até 6,4 vezes mais taxa de transferência em uma ampla gama de cargas de trabalho, modelos e configurações de hardware, em comparação com os sistemas de programação e inferência existentes, incluindo Guidance, vLLM e LMQL.

SGLang: Modelo de Programação e Metodologia

O modelo de programação SGLang é introduzido por meio de um exemplo em execução, descrevendo sua linguagem primitiva e modos de execução, e delineando oportunidades de otimização de tempo de execução. Esse modelo simplifica operações tediosas em fluxos de trabalho de múltiplas chamadas (por exemplo, manipulação de strings, chamadas de API, especificação de restrições, paralelismo) fornecendo primitivos flexíveis e compostos. SGLang é uma linguagem específica de domínio incorporada em Python.

A função multi_dimensional_judge recebe três argumentos: `s`, `path` e `essay`. s gerencia o estado do prompt, path é o caminho do arquivo de imagem e essay é o texto do ensaio. Novas strings e primitivos SGLang podem ser anexados ao estado s para execução usando o operador +=. Primeiro, a função adiciona a imagem e o ensaio ao prompt. Em seguida, verifica se o ensaio está relacionado à imagem usando select, armazenando o resultado em s[“related”]. Se relacionado, o prompt é bifurcado em três cópias para avaliação paralela de diferentes dimensões, usando gen para armazenar os resultados em f[“judgment”]. Em seguida, combina os julgamentos, gera um resumo e atribui uma nota de letra. Finalmente, retorna os resultados no formato JSON, seguindo um esquema definido por uma restrição de expressão regular regex. SGLang simplifica significativamente esse programa, pois um programa equivalente usando uma interface semelhante à do OpenAI exigiria 2,1 vezes mais linhas de código devido à manipulação manual de strings e controle de paralelismo.

SGLang fornece primitivos para controle de estado do prompt, geração e paralelismo, que podem ser usados com a sintaxe e bibliotecas do Python. Aqui estão os primitivos:

gen: Chama um modelo para gerar e armazena os resultados em uma variável com o nome especificado em seu primeiro argumento. Ele suporta um argumento `regex` para restringir a saída a seguir uma gramática definida por uma expressão regular (por exemplo, um esquema JSON).

  • select: Chama um modelo para escolher a opção de maior probabilidade de uma lista.
  • += ou extend: Anexa uma string ao prompt.
  • [nome_da_variável]: Recupera os resultados de uma geração.
  • fork: Cria bifurcações paralelas do estado do prompt.
  • join: Reune o estado do prompt.
  • imagem e vídeo: Recebem entradas de imagem e vídeo.

A forma mais simples de executar um programa SGLang é por meio de um interpretador, onde um prompt é tratado como um fluxo assíncrono. Primitivos como extend, gen e select são submetidos ao fluxo para execução assíncrona. Essas chamadas não bloqueadoras permitem que o código Python continue executando sem esperar que a geração termine, semelhante ao lançamento de kernels CUDA de forma assíncrona. Cada prompt é gerenciado por um executor de fluxo em uma thread em segundo plano, permitindo o paralelismo intra-programa. A recuperação dos resultados da geração bloqueará até que eles estejam prontos, garantindo a sincronização correta. Alternativamente, os programas SGLang podem ser compilados como gráficos computacionais e executados com um executor de gráfico, permitindo mais otimizações. Este artigo usa o modo interpretador por padrão e discute os resultados do modo compilador no Apêndice D. SGLang suporta modelos de peso aberto com seu próprio tempo de execução SGLang (SRT), bem como modelos de API, como o OpenAI e o Anthropic.

Os sistemas de programação para LLMs podem ser classificados como de alto nível (por exemplo, LangChain, DSPy) e de baixo nível (por exemplo, LMQL, Guidance, SGLang). Os sistemas de alto nível fornecem prompts pré-definidos ou gerados automaticamente, como o otimizador de prompts do DSPy. Os sistemas de baixo nível normalmente não alteram os prompts, mas permitem a manipulação direta dos prompts e primitivos. SGLang é um sistema de baixo nível semelhante ao LMQL e Guidance. A tabela a seguir compara suas características.

SGLang se concentra mais na eficiência de tempo de execução e vem com seu próprio tempo de execução co-projetado, permitindo otimizações novas. As linguagens de alto nível (por exemplo, DSPy) podem ser compiladas para linguagens de baixo nível (por exemplo, SGLang). A integração de SGLang como um backend no DSPy para melhor eficiência de tempo de execução é demonstrada posteriormente.

O exemplo acima ilustra as operações RadixAttention com uma política de evição LRU em nove pontos no tempo, mostrando a evolução dinâmica da árvore radix em resposta a diferentes solicitações. Essas solicitações incluem duas sessões de chat, um lote de consultas de aprendizado de poucos disparos e amostragem para prompting de autoconsistência. Cada aresta da árvore carrega um rótulo denotando uma substring ou uma sequência de tokens. Os nós são coloridos para refletir diferentes estados: verde para nós novos, azul para nós em cache acessados durante o ponto no tempo e vermelho para nós que foram evitados.

Etapa 1: A árvore radix está inicialmente vazia.

Etapa 2: O servidor processa uma mensagem de usuário de entrada “Olá” e responde com a saída do LLM “Oi”. O prompt do sistema “Você é um assistente útil”, a mensagem do usuário “Olá!” e a resposta do LLM “Oi!” são consolidados na árvore como uma aresta única vinculada a um novo nó.

Etapa 3: Uma nova solicitação de prompt chega, e o servidor encontra o prefixo do prompt (ou seja, a primeira rodada da conversa) na árvore radix e reutiliza seu cache KV. A nova rodada é anexada à árvore como um novo nó.

Etapa 4: Uma nova sessão de chat começa. O nó da Etapa 3 é dividido em dois nós para permitir que as duas sessões de chat compartilhem o prompt do sistema.

Etapa 5: A segunda sessão de chat continua. No entanto, devido aos limites de memória, um nó da Etapa 4 deve ser evitado. A nova rodada é anexada após o nó restante da Etapa 4.

Etapa 6: O servidor recebe uma consulta de aprendizado de poucos disparos, processa-a e a insere na árvore. O nó raiz é dividido porque a nova consulta não compartilha nenhum prefixo com os nós existentes.

Etapa 7: O servidor recebe um lote de consultas de aprendizado de poucos disparos adicionais. Essas consultas compartilham o mesmo conjunto de exemplos de poucos disparos, então um nó da Etapa 6 é dividido para permitir o compartilhamento.

Etapa 8: O servidor recebe uma nova mensagem da primeira sessão de chat. Ele evita todos os nós da segunda sessão de chat, pois eles são os menos recentemente usados.

Etapa 9: O servidor recebe uma solicitação para amostrar mais respostas para as perguntas em um nó da Etapa 8, provavelmente para prompting de autoconsistência. Para fazer espaço para essas solicitações, vários nós são evitados.

Esse exemplo demonstra como RadixAttention lida com a alocação dinâmica e evição de nós em resposta a diferentes tipos de solicitações, garantindo a reutilização eficiente do cache KV e o gerenciamento de memória.

SGLang: Avaliação e Resultados

Resultados em Modelos de Peso Aberto

Os resultados de latência e taxa de transferência são mostrados nas figuras a seguir. SGLang melhora a taxa de transferência em até 6,4 vezes e reduz a latência em até 3,7 vezes. Essas melhorias resultam da reutilização do cache KV, da exploração do paralelismo dentro de um programa e da decodificação de saída estruturada mais rápida.

Nessas benchmarks, a taxa de acerto do cache varia de 50% a 99%. A Figura 13 (Apêndice) lista as taxas de acerto de cache alcançadas e ótimas para todos eles, mostrando que a programação orientada a cache de SGLang se aproxima de 96% da taxa de acerto ótima em média.

Resultados em Modelos Maiores com Paralelismo de Tensor

Modelos maiores, Mixtral-8x7B e Llama-70B, foram testados com paralelismo de tensor na mesma série de benchmarks, e os resultados são relatados na figura a seguir. A aceleração nos modelos maiores mostra uma tendência semelhante à observada nos modelos menores, indicando que as otimizações de SGLang se generalizam bem para modelos maiores. Guidance e LMQL foram omitidos devido à falta de implementações eficientes de paralelismo de tensor.

Resultados em Modelos Multimodais

SGLang tem suporte nativo para modelos multimodais com os primitivos de imagem e vídeo. As otimizações deste artigo são compatíveis com modelos multimodais. Para RadixAttention, o hash das imagens de entrada é computado e usado como a chave na árvore radix, permitindo a reutilização do cache KV dos tokens de imagem do mesmo imagem. LLaVA-v1.5-7B (imagem) foi executado em llava-bench-in-the-wild e LLaVA-NeXT-34B (vídeo) em ActivityNet. Como esses modelos não são bem suportados por outros sistemas de baseline, a implementação original dos autores do modelo em Hugging Face Transformers foi usada como baseline. Como mostrado na tabela a seguir, SGLang fornece uma taxa de transferência de até 6 vezes mais alta nesses benchmarks. Em llava-bench-in-the-wild, várias perguntas sobre a mesma imagem foram tratadas, e o tempo de execução de SGLang reutilizou o cache KV nesse caso.

Implantação em Produção

SGLang foi implantado na Chatbot Arena para servir modelos de peso aberto. Devido ao baixo tráfego para alguns modelos, apenas um trabalhador SGLang serve cada um. Após um mês, uma taxa de acerto de cache RadixAttention de 52,4% para LLaVA-Next-34B e 74,1% para Vicuna-33B foi observada. Os acertos de cache vieram de mensagens do sistema comuns, imagens de exemplo frequentemente reutilizadas e históricos de chat de várias rodadas. Isso reduziu a latência do primeiro token em média 1,7 vezes para Vicuna-33B.

Pensamentos Finais

Neste artigo, falamos sobre SGLang, um sistema recém-introduzido que visa fornecer execução eficiente de programas de modelos de linguagem complexos. SGLang compreende uma linguagem de frontend e um tempo de execução. O frontend simplifica a programação com primitivos para geração e controle de paralelismo, enquanto o tempo de execução acelera a execução por meio de otimizações novas, como RadixAttention para reutilização do cache KV e máquinas de estado finito comprimidas para decodificação de saída estruturada mais rápida. Experimentos demonstram que SGLang alcança até 6,4 vezes mais taxa de transferência em comparação com sistemas de inferência de ponta em vários grandes modelos de linguagem e multimodais, abordando tarefas como controle de agente, raciocínio lógico, benchmarks de aprendizado de poucos disparos, decodificação JSON, pipelines de geração aumentada por recuperação e chat de várias voltas.

Um engenheiro por profissão, um escritor por coração. Kunal é um escritor técnico com um amor e compreensão profundos de AI e ML, dedicado a simplificar conceitos complexos nestes campos por meio de sua documentação envolvente e informativa.