SAGA: Workflow-Atomic Scheduling for AI Agent Inference on GPU Clusters

1. Declaração do Problema

Os agendadores atuais de clusters de GPU (por exemplo, vLLM, SGLang) são otimizados para inferência de tiro único, tratando cada solicitação de Modelo de Linguagem de Grande Escala (LLM) como uma unidade independente. No entanto, Agentes de IA modernos operam por meio de loops iterativos de "Pensamento-Ação-Observação", executando dezenas a centenas de chamadas de LLM encadeadas por tarefa, intercaladas com invocações de ferramentas externas (execução de código, navegação na web, consultas a bancos de dados).

Isso cria uma incompatibilidade fundamental com os sistemas de agendamento existentes, levando a três ineficiências críticas:

1. Sobrecarga de Regeneração de Cache KV: Os sistemas atuais expulsam os caches de Chave-Valor (KV) durante períodos de inatividade de chamadas de ferramentas (que podem durar de milissegundos a minutos). Quando o agente retoma, o cache deve ser regenerado, desperdiçando 38% do tempo de execução e aumentando a latência em 3–8×.
2. Utilização Fragmentada de Memória: O agendamento independente leva a uma baixa localidade de memória. A utilização da memória da GPU média apenas ~42% porque os caches estão fragmentados e expulsos prematuramente, apesar do potencial para alta reutilização.
3. Falta de Equidade no Nível de Tarefa: Os agendadores existentes otimizam para latência ou vazão por solicitação. Em ambientes multi-inquilino, isso causa "fome" para agentes com fluxos de trabalho longos e complexos, pois solicitações curtas monopolizam os recursos.

2. Metodologia: A Arquitetura SAGA

O SAGA propõe uma mudança de paradigma do agendamento no nível de solicitação para o nível de programa (fluxo de trabalho atômico). Ele trata todo o fluxo de trabalho do agente como uma única unidade agendável, em vez de uma sequência de solicitações independentes. O sistema consiste em três camadas principais:

A. Grafos de Execução de Agente (AEGs)

O SAGA modela fluxos de trabalho de agentes como grafos direcionados onde os nós representam etapas de inferência de LLM e as arestas representam dependências (incluindo chamadas de ferramentas).

• Construção: Os AEGs são derivados de dicas do framework (por exemplo, callbacks do LangChain) ou inferidos via correspondência de padrões se as dicas não estiverem disponíveis.
• Utilidade: O grafo permite que o agendador preveja etapas futuras, probabilidades de transição e a probabilidade de reutilização de cache através dos limites de chamadas de ferramentas.

B. Três Mecanismos Principais

1. Gerenciamento de Cache KV Consciente de Fluxo de Trabalho (WA-LRU):

   • Em vez do LRU padrão, o SAGA usa uma política de expulsão Consciente de Fluxo de Trabalho (WA-LRU).
   • Ele calcula uma pontuação de expulsão com base em: (1) Recência, (2) Probabilidade de Reutilização Prevista (derivada do AEG) e (3) Tamanho do Cache.
   • TTL Consciente de Chamada de Ferramenta: O sistema ajusta dinamicamente o Tempo de Vida (TTL) dos estados KV em cache com base na ferramenta específica sendo chamada e na pressão de memória atual. Isso previne a expulsão prematura durante longas esperas de ferramentas.
   • Pré-busca Especulativa: Se a próxima etapa for previsível, o SAGA busca antecipadamente o cache KV necessário durante a fase de execução da ferramenta para sobrepor I/O e computação.

2. Agrupamento com Afinidade de Sessão e Roubo de Trabalho:

   • Roteamento por Afinidade: Solicitações pertencentes à mesma sessão de agente são roteadas para o mesmo trabalhador de GPU para maximizar acertos de cache.
   • Roubo de Trabalho: Para evitar desequilíbrio de carga (atrasos), um coordenador global realiza roubo de trabalho aleatório. Se um trabalhador estiver ocioso enquanto outros estão sobrecarregados, ele rouba tarefas. Crucialmente, o sistema migra o estado do cache KV (usando mecanismos como Llumnix) para preservar a localidade, garantindo que a tarefa roubada não perca seu contexto em cache.

3. Partilha Justa de Agente (AFS):

   • Uma nova métrica de equidade baseada no Tempo de Conclusão da Tarefa em vez de contagem de solicitações.
   • Usa análise de deriva de Lyapunov para garantir que inquilinos com tarefas urgentes e de longa duração recebam recursos proporcionais, prevenindo a fome. O sistema ajusta dinamicamente as prioridades de agendamento com base na "urgência" (trabalho restante vs. prazo) de cada tarefa do agente.

3. Principais Contribuições

• Abstração Atômica de Fluxo de Trabalho: O primeiro sistema a tratar todo o programa do agente como a unidade de agendamento de primeira classe, expondo explicitamente a estrutura do fluxo de trabalho (DAGs) ao agendador.
• Limites Teóricos:
  • Provou que o WA-LRU alcança uma razão competitiva dentro de 1,31× da política offline ótima de Bélády (o melhor possível teoricamente para expulsão de cache) em traços de produção.
  • Fornecia uma prova formal para a Partilha Justa de Agente (AFS) usando deriva de Lyapunov, garantindo desvio limitado nos tempos de conclusão de tarefas sob heterogeneidade de demanda limitada.
• Design de Sistema Integrado: Combina agendamento distribuído, cache preditivo e garantias de equidade em um framework unificado, abordando lacunas deixadas por trabalhos anteriores (por exemplo, cache de prefixo do vLLM, expulsão do KVFlow, migração do Llumnix).

4. Resultados Experimentais

Avaliado em um cluster de 64 GPUs (8 nós × 8 NVIDIA A100) usando cargas de trabalho SWE-bench (agentes de codificação) e WebArena (agentes de navegador).

• Redução de Latência:
  • O SAGA reduz o Tempo de Conclusão da Tarefa (TCT) em 1,64× (média geométrica) comparado ao estado da arte vLLM v0.15.1 com Cache de Prefixo Automático (APC).
  • Comparado ao vLLM puro, a melhoria é de 3,01×.
  • Especificamente, no SWE-bench, o TCT caiu de ~612s (vLLM) para 203s (SAGA).
• Eficiência de Memória:
  • A utilização da memória da GPU aumentou de ~42% (vLLM) para 71,3% (SAGA).
  • O tempo de regeneração do cache KV foi reduzido de 38% do tempo total para apenas 8%.
• Equidade:
  • Alcançou 99,2% de alcance de SLO (tarefas atendendo aos prazos) sob interferência multi-inquilino, comparado a 67,3% para o vLLM.
  • Eliminou a latência de "cauda longa" para inquilinos leves, garantindo desempenho consistente entre cargas de trabalho pesadas e leves.
• Compromisso:
  • O SAGA otimiza para latência e conclusão de tarefas, resultando em uma redução de ~30% na vazão de pico comparado ao agendamento de lotes otimizado para vazão. Os autores argumentam que isso é um compromisso aceitável para implantações de agentes interativos sensíveis à latência.

5. Significado

• Ponte da Lacuna: O SAGA demonstra que as ineficiências em sistemas de IA compostos não são inerentes aos modelos, mas são artefatos de abstrações no nível de solicitação. Ao expor a estrutura do fluxo de trabalho, os agendadores podem se aproximar da eficiência offline ótima.
• Validação Empírica da Teoria: O artigo fornece o primeiro limite de razão competitiva empírica (1,31×) para gerenciamento de cache online em fluxos de trabalho de agentes, validando que a consciência do fluxo de trabalho é suficiente para quase igualar políticas offline ótimas.
• Pronto para Produção: O sistema é implementado como uma extensão do vLLM e suporta principais frameworks de agentes (LangChain, AutoGen) sem exigir alterações de código nos próprios agentes.
• Direção Futura: Estabelece uma nova linha de base para "Sistemas de IA Compostos", sugerindo que a infraestrutura futura deve evoluir de servir solicitações estáticas para orquestrar fluxos de trabalho dinâmicos e com estado.

Em resumo, o SAGA resolve os gargalos críticos de latência e eficiência de agentes de IA ao mudar a unidade de agendamento da solicitação individual para todo o fluxo de trabalho, aproveitando o cache preditivo e a alocação de recursos consciente de equidade para alcançar desempenho quase ótimo em clusters de GPU.