OxyGen: Unified KV Cache Management for Vision-Language-Action Models under Multi-Task Parallelism

O artigo apresenta o OxyGen, um paradigma de gerenciamento unificado de cache KV para Modelos Visão-Linguagem-Ação (VLAs) que, ao tratar o cache como um recurso compartilhado entre tarefas e ao longo do tempo, elimina redundâncias computacionais e permite a execução paralela eficiente de múltiplas tarefas em dispositivos, alcançando até 3,7 vezes mais velocidade sem degradar a qualidade das ações.

O essencial

Imagine que você é um robô doméstico superinteligente. Você está na cozinha e precisa fazer três coisas ao mesmo tempo:

1. Agir: Pegar uma maçã e colocá-la na mesa (isso precisa ser feito em frações de segundo para não derrubar nada).
2. Falar: Contar para o dono o que você está fazendo ("Estou pegando a maçã...").
3. Lembrar: Guardar na memória que a maçã estava na geladeira para usar depois.

O problema é que, até agora, os robôs eram como cozinheiros desajeitados. Eles tentavam fazer uma coisa de cada vez, ou tentavam fazer tudo ao mesmo tempo, mas de um jeito bagunçado que deixava tudo lento.

O Problema: A Cozinha Bagunçada (O "Isolamento")

No mundo da Inteligência Artificial, o robô usa um "cérebro" chamado Modelo VLA (Visão-Linguagem-Ação). Para pensar, esse cérebro cria um "rascunho mental" (chamado de KV Cache) com base no que vê (a imagem da cozinha).

O sistema antigo funcionava assim:

• Para pegar a maçã, o robô lia a imagem da cozinha, fazia o rascunho mental e agia.
• Para falar, ele lê a mesma imagem da cozinha de novo, faz outro rascunho mental idêntico e começa a falar.
• Para lembrar, ele lê a imagem mais uma vez, faz mais um rascunho e guarda.

A analogia: É como se você fosse escrever um e-mail, mas para cada parágrafo, você tivesse que ler o livro inteiro de novo antes de escrever a primeira palavra. É um desperdício enorme de tempo e energia! Além disso, como o robô tem apenas um "cérebro" (uma placa de vídeo), essas tarefas brigavam por espaço, travando umas nas outras.

A Solução: O OxyGen (O "Gerente de Cozinha Unificado")

Os pesquisadores criaram o OxyGen. Pense nele como um Gerente de Cozinha Mestre que organiza a produção de forma inteligente.

O OxyGen muda duas regras principais:

1. Compartilhamento de Rascunhos (Cross-Task KV Sharing)

Em vez de ler a imagem da cozinha três vezes, o OxyGen diz: "Ei, todos vocês (Agir, Falar, Lembrar) estão olhando a mesma coisa! Vamos fazer um único rascunho mental agora e todos podem usar esse mesmo rascunho."

• Resultado: O robô não perde tempo relendo o que já viu. Ele economiza energia e tempo precioso.

2. Batching Contínuo (Cross-Frame Continuous Batching)

Aqui está a parte mais genial.

• Agir é como um maratonista de 100 metros rasos: precisa ser rápido e terminar antes que a próxima foto da cozinha apareça (dentro de 1/60 de segundo).
• Falar é como um maratonista de longa distância: pode levar mais tempo, mas precisa manter um ritmo constante.

O sistema antigo tentava forçar o "maratonista de longa distância" a correr junto com o de 100 metros, ou esperar o de 100 metros terminar tudo antes de começar a falar.

O OxyGen cria uma esteira rolante:

• Enquanto o robô executa a ação rápida (pegar a maçã) baseada na foto de hoje, ele aproveita o tempo livre para continuar escrevendo o texto baseado na foto de ontem e de anteontem.
• Ele empilha várias tarefas de fala em uma única "rodada" de processamento, como se fosse um ônibus cheio de passageiros indo para o mesmo lugar, em vez de cada um pegar um táxi separado.

Os Resultados: Robôs Mais Rápidos e Eficientes

Com o OxyGen, o robô consegue:

• Agir 3,7 vezes mais rápido: Ele consegue pegar objetos com uma fluidez incrível (70 vezes por segundo!), sem travar.
• Falar muito mais: Ele consegue descrever o que está fazendo em alta velocidade (mais de 200 palavras por segundo).
• Não perder qualidade: Ele não fica "bobo" ou desajeitado. A qualidade da ação continua perfeita, como se ele tivesse feito tudo do jeito antigo, mas muito mais rápido.
• Economizar energia: Como ele não faz cálculos repetidos, gasta menos bateria.

Resumo em uma Frase

O OxyGen é como transformar uma cozinha onde cada cozinheiro tinha que ler o livro de receitas do zero para cada prato, em uma cozinha onde há um único livro aberto na mesa, e todos os cozinheiros (ação, fala, memória) trabalham juntos, sem brigar pelo espaço, fazendo tudo ao mesmo tempo de forma super rápida e eficiente.

Isso permite que os robôs do futuro não apenas "façam" coisas, mas conversem e aprendam enquanto trabalham, exatamente como os humanos fazem.

Resumo técnico

Título: OxyGen: Gerenciamento Unificado de KV Cache para Modelos Visão-Linguagem-Ação sob Paralelismo de Múltiplas Tarefas

1. O Problema

A área de IA Embutida (Embodied AI) avança para agentes capazes de executar múltiplas tarefas simultaneamente a partir de observações compartilhadas, como manipulação robótica, conversação e construção de memória. Modelos recentes, conhecidos como Modelos Visão-Linguagem-Ação (VLAs) baseados em Mistura de Transformers (MoT), suportam arquiteturalmente a geração de saídas heterogêneas (ações e linguagem) a partir de um único modelo.

No entanto, os sistemas de inferência existentes falham em alcançar um paralelismo eficiente para implantação em dispositivos (on-device) devido a duas ineficiências principais no gerenciamento de recursos:

1. Computação Redundante: Quando múltiplas tarefas (ex: gerar uma ação e gerar um texto) compartilham a mesma observação de entrada, os sistemas atuais codificam essa observação repetidamente, gerando entradas de KV Cache (Chave-Valor) idênticas para cada tarefa.
2. Contenção de Recursos: Mesmo que o KV Cache fosse compartilhado, tarefas com restrições de tempo diferentes (ex: ações robóticas com prazos rígidos por quadro vs. geração de linguagem com prazos flexíveis) competem pelo mesmo hardware (geralmente uma única GPU), bloqueando-se mutuamente.

O artigo identifica que a raiz do problema é o gerenciamento isolado do KV Cache, onde cada tarefa trata seu cache como um recurso privado, perdendo oportunidades de compartilhamento e agendamento coordenado.

2. Metodologia: OxyGen

Os autores propõem o OxyGen, um sistema de inferência eficiente para VLAs MoT que introduz um novo paradigma: o Gerenciamento Unificado de KV Cache. Neste paradigma, o KV Cache é tratado como um recurso de primeira classe, compartilhado entre tarefas e ao longo do tempo.

O sistema implementa duas otimizações principais habilitadas por essa abstração:

A. Compartilhamento de KV entre Tarefas (Cross-Task KV Sharing)

• Mecanismo: Quando múltiplas tarefas operam na mesma observação (no mesmo quadro/tempo $t$), a observação é codificada apenas uma vez pelo backbone do VLM.
• Benefício: O KV Cache resultante ($K_t$) é reutilizado por todos os especialistas (o especialista de ação e o especialista de linguagem) simultaneamente. Isso elimina a necessidade de repetir o estágio de prefill (preenchimento inicial) para cada tarefa, reduzindo drasticamente a latência inicial.

B. Batching Contínuo entre Quadros (Cross-Frame Continuous Batching)

• Mecanismo: O sistema desacopla o fluxo de inferência do loop de controle por quadro tradicional.
  • Tarefas em Tempo Real (Ações): Completam-se dentro do quadro atual para atender a prazos rígidos (ex: 50-70 Hz).
  • Tarefas de Streaming (Linguagem): São agrupadas (batched) continuamente através de múltiplos quadros. O gerenciador mantém o estado de geração incremental de várias solicitações de linguagem pendentes e as decodifica em paralelo.
• Benefício: Permite que a decodificação de linguagem aproveite a capacidade de processamento paralelo de hardware (GPUs) ao longo do tempo, aumentando a taxa de transferência (throughput) de tokens sem comprometer a frequência de controle das ações.

Arquitetura do Gerenciador

O gerenciador unificado mantém estados de geração resumíveis ($\sigma_t$) para cada solicitação, contendo o KV Cache, o buffer de tokens e um sinal de término. Ele oferece operações para armazenar, recuperar, atualizar e agrupar (batch) esses estados, permitindo que a inferência seja interrompida e retomada entre quadros sem recálculo.

3. Principais Contribuições

1. Formulação do Problema: Identificaram o gerenciamento isolado de KV Cache como a causa raiz da ineficiência em sistemas de inferência MoT VLA para paralelismo de múltiplas tarefas.
2. Novo Paradigma: Propuseram o gerenciamento unificado de KV Cache, que permite otimizações como compartilhamento entre tarefas e batching contínuo entre quadros.
3. Implementação e Avaliação: Implementaram o OxyGen no modelo $\pi0.5$ (o VLA MoT mais popular, com base no sistema openpi) e avaliaram em configurações robóticas representativas, demonstrando ganhos significativos de desempenho.

4. Resultados Experimentais

Os experimentos foram realizados em uma GPU NVIDIA RTX 4090 (plataforma representativa para inferência on-device) utilizando três benchmarks: LIBERO, DROID e ALOHA.

• Aceleração Geral: O OxyGen alcançou um speedup de até 3.7x em comparação com a execução isolada (sequencial ou paralela ingênua).
• Desempenho Simultâneo: O sistema conseguiu atingir simultaneamente:
  • > 200 tokens/s de throughput na decodificação de linguagem.
  • 70 Hz de frequência de geração de ações (crítico para controle robótico suave).
• Qualidade da Ação: A taxa de sucesso nas tarefas (avaliada no benchmark LIBERO) permaneceu inalterada, confirmando que a otimização não degrada a qualidade das ações geradas.
• Eficiência Energética:
  • Redução de até 78% no consumo de energia por solicitação.
  • Redução de até 47% no consumo médio de potência.
  • O custo de memória adicional foi modesto (~15%), enquanto a paralelização ingênua quase dobrou o uso de memória.

Análise de Ablação:

• O Compartilhamento de KV entre Tarefas forneceu o ganho inicial (1.4x) ao eliminar a redundância de prefill.
• O Batching Contínuo entre Quadros foi crucial para manter a frequência de ações alta (estável em ~60-70 Hz) mesmo com longos processos de decodificação de linguagem, onde sistemas sem batching caíam drasticamente (de 49.9 Hz para 19.1 Hz).

5. Significado e Impacto

O trabalho OxyGen é fundamental para a viabilidade de agentes robóticos autônomos complexos que precisam interagir naturalmente com humanos e ambientes dinâmicos. Ao resolver o gargalo de inferência em modelos multimodais unificados, o sistema permite que robôs:

• Planejem e conversem enquanto executam manipulações físicas precisas.
• Operem em hardware limitado (como GPUs de borda) sem sacrificar a fluidez do controle ou a velocidade de resposta.
• Utilizem a capacidade computacional de forma mais eficiente, reduzindo custos energéticos e térmicos.

Em resumo, o OxyGen demonstra que o gerenciamento inteligente de recursos de inferência (especificamente o KV Cache) é tão crítico quanto o avanço dos próprios modelos para a próxima geração de IA Embutida.