SGLang-TorchTPU quase dobra a velocidade de geração de tokens do Muse Spark Open

Um agente de programação leu o repositório, executou uma ferramenta e fez uma pausa antes da próxima resposta. É fácil atribuir essa pausa ao tamanho do modelo. Mas o servidor pode estar gastando tempo carregando pesos, trocando dados entre aceleradores e reprocessando um contexto já conhecido. Para a equipe de produto, o problema é bastante concreto: o usuário espera enquanto a infraestrutura pela qual se paga realiza trabalho preparatório.

A equipe do SGLang-TorchTPU descreveu a execução do Muse Spark Open, com 816 bilhões de parâmetros, em um único host com quatro TPU v7. O principal resultado com uma única requisição: o tempo por token de saída caiu de 3,78 para 1,95 ms. As medições e os detalhes da implementação foram publicados no blog do PyTorch.

O acelerador precisa receber os dados na hora certa

Imagine uma cozinha em que o cozinheiro para antes de cada tarefa para esperar a próxima cesta de ingredientes. Comprar um fogão mais potente não elimina essa espera. Outro caminho é organizar as entregas para que a próxima cesta chegue enquanto o prato atual está sendo preparado.

A equipe aplicou uma ideia semelhante: fundiu a execução de 62 camadas em um único programa e sobrepôs a transferência de dados à computação. Isso permite preparar com antecedência o que será necessário na próxima etapa do trabalho.

A aplicação prática dessa abordagem começa pela análise do perfil de execução. Suponha que você ofereça um assistente interno para uma equipe pequena. Em geral, várias pessoas conversam com ele ao mesmo tempo, e todas querem uma resposta gerada de forma fluida. Nesse cenário, vale examinar quanto tempo leva cada etapa de geração e qual parcela desse tempo é gasta esperando dados.

Um serviço com muitos usuários enfrenta outro desafio: atender à fila sem um aumento brusco da latência. Por isso, um número impressionante de velocidade ainda não responde quantos usuários o sistema suporta. É preciso fazer uma série de medições com diferentes cargas e as mesmas condições.

Uma pausa do agente é motivo para preservar o contexto

O segundo resultado diz respeito à memória. Um cache adicional na RAM do host aumentou a vazão em cargas de trabalho de agentes em até 77%.

Pense em um assistente que procura um bug em um projeto. Ele lê arquivos, executa testes, recebe a saída e volta ao raciocínio. Uma parte significativa da conversa permanece igual entre essas ações. Se o servidor preservou os dados do contexto já processados, o próximo turno pode começar com menos trabalho repetido.

Para a equipe de produto, isso traz uma pergunta útil: o que acontece com o estado da sessão enquanto o usuário ou a ferramenta fica em silêncio? Uma pausa, por si só, não significa que a conversa acabou. Especialmente quando o agente está esperando uma compilação, uma busca externa ou o resultado de um comando demorado.

Uma maneira prática de verificar isso é criar um cenário específico. Dê ao agente um documento grande, faça vários turnos curtos com pausas e compare o tempo até o início da resposta. Depois, repita o experimento com a memória cheia. Assim, ficará mais claro quanto o comportamento do serviço depende da preservação do contexto e o que acontece quando sessões antigas são removidas do cache.

O que aplicar no seu serviço

Uma aceleração de quase duas vezes não se traduz automaticamente em economia de orçamento. A vantagem do programa fundido diminui conforme aumenta o número de requisições simultâneas. Além disso, os autores ainda planejam incorporar parte do trabalho à branch principal do projeto.

Quem desenvolve deve começar com três medições: tempo até o primeiro token, intervalos entre os tokens seguintes e número de requisições concluídas com uma determinada latência. Quem cuida do produto deve acrescentar uma jornada de usuário realista, com ferramentas e pausas. Quem fundou a empresa deve calcular o custo de uma tarefa concluída com sucesso e com tempo de resposta aceitável.

Também há um cenário específico para testar a fila: um usuário envia um documento enorme, enquanto outro faz uma pergunta curta. Observe quanto tempo o segundo espera. Esse experimento mostra rapidamente se a distribuição de recursos entrega a qualidade de serviço prometida.

O valor da publicação está no conjunto de soluções de engenharia que podem ser testadas. Antes de trocar o hardware, vale descobrir exatamente onde o servidor fica ocioso, qual trabalho repete e quais requisições deixa ocupar a fila por tempo demais.