Reduza a latência de LLMs com roteamento por prefixo no Amazon SageMaker Inference

O problema que o roteamento por prefixo resolve

Quem já construiu uma aplicação sobre um modelo de linguagem de grande escala (LLM) sabe que os prompts geralmente têm duas partes: uma parte fixa — com instruções, documentos de referência ou histórico de conversa — e uma parte variável, que é o que o usuário digitou de fato.

Pense num bot de atendimento ao cliente. Cada requisição começa com o mesmo bloco de texto: “Você é um agente de suporte da Empresa X. Aqui estão nossas políticas…”, seguido da pergunta do cliente. Esse bloco de instruções pode ter 3.000 tokens. A pergunta do usuário, talvez 50. Isso significa que, em centenas ou milhares de requisições, o modelo está processando aquele mesmo início de 3.000 tokens repetidamente.

Frameworks de servir LLMs como vLLM e TensorRT-LLM já têm uma solução para isso: o prefix caching (cache de prefixos). Eles armazenam os pares chave-valor (KV) computados para prefixos que já foram vistos. Quando o mesmo início aparece numa nova requisição, o modelo reutiliza o cálculo em cache e processa apenas os tokens novos do final. Isso reduz significativamente o Tempo até o Primeiro Token (TTFT).

O problema surge quando você escala além de uma única instância. Com uma frota de máquinas atrás de um endpoint, as requisições são distribuídas entre todas elas. O mesmo prefixo de 3.000 tokens cai na instância A numa requisição, na instância B na seguinte, na instância C depois. Cada instância recalcula do zero, porque nenhuma delas recebe o mesmo prefixo com frequência suficiente para construir um cache confiável. O prefix caching existe, mas a camada de roteamento dilui demais as requisições para que ele funcione.

O que é o prefix-aware routing

Para resolver esse cenário, a AWS anunciou o prefix-aware routing no Amazon SageMaker Inference. É uma nova estratégia de roteamento que analisa o início de cada requisição e direciona consistentemente as requisições com o mesmo início para a mesma instância. Assim, o cache KV daquela instância se acumula e é reutilizado de verdade.

O funcionamento é direto: quando uma requisição chega ao endpoint, o SageMaker analisa o início do payload e decide qual instância deve processá-la. O mesmo início vai para a mesma instância. Inícios diferentes se distribuem entre instâncias diferentes. Se 10 requisições compartilham um prefixo, todas as 10 vão para a mesma máquina — e o cache dessa máquina fica quente para aquele prefixo.

Não é necessário rotular requisições nem gerenciar afinidade manualmente. O endpoint faz isso com base no conteúdo da requisição.

Proteções integradas

O recurso tem dois mecanismos de proteção nativos:

Proteção contra sobrecarga: Se um prefixo é muito popular e a instância-alvo já está na capacidade máxima, o endpoint redireciona a requisição para uma instância menos ocupada. Você perde o cache hit naquela requisição específica, mas evita sobrecarregar uma única máquina.
Comportamento estável durante escalonamento: Quando instâncias são adicionadas ou removidas, a maior parte das requisições continua indo para a mesma instância de antes. Apenas uma pequena fração do tráfego é redistribuída. Os caches não são invalidados a cada evento de escalonamento.

Resultados dos benchmarks

A AWS realizou benchmarks comparando o prefix-aware routing com o roteamento aleatório padrão, usando o Llama 3.1 70B Instruct em 7 instâncias ml.p5.48xlarge com vLLM (prefix caching habilitado). Foram 16 configurações de teste cobrindo endpoints de modelo único, endpoints com inference components, a API nativa Invoke e a API compatível com OpenAI. Todos os testes concluíram com 100% de sucesso.

Workloads de contexto longo

Prefixos compartilhados de 8.000 tokens, sustentados por 1 hora:

TTFT no percentil 90 (P90): redução de 33–37%
TTFT no percentil 50 (P50): redução de 71–77%
Taxa de acerto no cache KV: de aproximadamente 25% para 82%
Throughput: aumento de 15–16%

Workloads de contexto curto

Conversas no estilo ShareGPT com comprimento variável, durante 30 minutos:

TTFT no P90: redução de 24–37%
TTFT no P50: redução de 13–16%
Taxa de acerto no cache KV: de aproximadamente 30% para 80%
Throughput: aumento de 1,7–2,0%

A lógica é simples: quanto mais longo o prefixo compartilhado, maior o ganho. Workloads de contexto longo se beneficiam mais porque há mais computação a ser pulada em cada cache hit. Workloads de contexto curto ainda se beneficiam, mas os prefixos compartilhados são menores, então a economia por requisição é proporcionalmente menor.

Overhead de roteamento

A lógica do prefix-aware routing adiciona 1,3–1,9 milissegundos por requisição. Nos testes, o TTFT do modelo variou entre 63 e 280 milissegundos — o custo de roteamento é desprezível. A distribuição de tráfego também se manteve equilibrada: cada uma das 7 instâncias recebeu entre 13,3% e 15,4% das requisições, dentro de 1% de uma divisão perfeitamente uniforme. Sem pontos quentes.

Estratégias de roteamento disponíveis no SageMaker Inference

Com este lançamento, o Amazon SageMaker Inference passa a oferecer três estratégias de roteamento para endpoints em tempo real:

RANDOM (padrão): Distribui as requisições uniformemente entre as instâncias. Recomendado para workloads de uso geral, modelos que não são LLMs, ou situações em que as requisições são intercambiáveis.
LEAST_OUTSTANDING_REQUESTS: Envia cada requisição para a instância com menos requisições em andamento. Recomendado quando os tempos de processamento variam e você quer manter todas as instâncias igualmente ocupadas.
PREFIX_AWARE (novo): Envia requisições com o mesmo prefixo de prompt para a mesma instância. Recomendado para workloads de LLM em que muitas requisições compartilham texto em comum no início e o framework de servir tem prefix caching habilitado.

A estratégia é configurada por variante de produção na configuração do endpoint. É possível alternar entre elas atualizando a configuração do endpoint sem reimplantar o modelo.

Quando usar o prefix-aware routing

O recurso entrega valor quando as requisições compartilham texto no início. Os padrões com maior impacto são:

Aplicações RAG (Geração Aumentada por Recuperação): Você recupera um documento e o coloca antes da pergunta do usuário. Quando vários usuários fazem perguntas sobre o mesmo documento, todos compartilham aquele documento como prefixo. O roteamento os direciona para a mesma instância, onde o cache KV daquele documento já está quente.
Conversas multi-turno: Cada turno inclui o histórico completo dos turnos anteriores. Conforme a conversa cresce, esse histórico compartilhado se torna um prefixo mais longo e mais custoso. Rotear por esse prefixo mantém o cache da conversa em uma única instância ao longo dos turnos.
Bots e assistentes com templates: Bots com instruções longas e estruturadas (políticas, regras de formatação, definições de persona) enviam essas mesmas instruções com cada requisição. Apenas a mensagem do usuário no final muda. Com o prefix-aware routing, aquele bloco de instruções caro é processado uma vez, não milhares de vezes.
Completação de código: Assistentes de código incluem conteúdo de arquivos como contexto. Enquanto um desenvolvedor trabalha no mesmo arquivo, cada requisição de completação compartilha aquele conteúdo como prefixo.

Como habilitar o prefix-aware routing

O roteamento por prefixo é configurado ao criar a configuração do endpoint. Dois parâmetros controlam o comportamento:

PrefixLength (1024–65536): Quanto do início da requisição usar para roteamento. Para a API nativa do SageMaker Invoke, são bytes do início do corpo da requisição. Para a API compatível com OpenAI, são caracteres do texto de mensagem extraído. Configure para cobrir o prefixo compartilhado mais conteúdo único suficiente para distribuir workloads diferentes entre instâncias.
ConcurrencyThreshold (1–1024): O número máximo de requisições em andamento na instância-alvo antes de o overflow ser acionado. Se a instância-alvo estiver nesse limite, a requisição vai para uma instância menos carregada.

Exemplo de criação de configuração de endpoint:

aws sagemaker create-endpoint-config \
  --endpoint-config-name example-llm-config \
  --production-variants '[{
    "VariantName": "AllTraffic",
    "ModelName": "example-llm-model",
    "InitialInstanceCount": 3,
    "InstanceType": "ml.p5.48xlarge",
    "RoutingConfig": {
      "RoutingStrategy": "PREFIX_AWARE",
      "PrefixAwareRoutingConfig": {
        "PrefixLength": 4096,
        "ConcurrencyThreshold": 10
      }
    }
  }]'

Em seguida, crie o endpoint normalmente:

aws sagemaker create-endpoint \
  --endpoint-name example-llm-endpoint \
  --endpoint-config-name example-llm-config

Nenhuma alteração no contêiner do modelo ou no framework de servir é necessária. O prefix-aware routing opera inteiramente na camada de roteamento do endpoint.

Como invocar o endpoint

Nada muda na forma de chamar o endpoint. As mesmas APIs InvokeEndpoint e InvokeEndpointWithResponseStream funcionam exatamente como antes:

aws sagemaker-runtime invoke-endpoint \
  --endpoint-name example-llm-endpoint \
  --content-type application/json \
  --body fileb://request.json \
  output.json

O mesmo vale para a API de Chat Completion compatível com OpenAI:

from openai import OpenAI
from sagemaker.core.token_generator import generate_token

client = OpenAI(
    base_url=f"https://runtime.sagemaker.us-west-2.amazonaws.com"
             f"/endpoints/example-llm-endpoint/openai/v1",
    api_key=generate_token(region="us-west-2")
)

response = client.chat.completions.create(
    model="example-model",
    messages=[
        {"role": "user", "content": "What is your return policy?"},
    ],
)

Isolamento de prefixo multi-tenant

Se diferentes tenants compartilham as mesmas instruções de prompt, mas você quer que sejam roteados separadamente (para manter contextos de cache independentes), é possível passar um ID opcional:

API nativa Invoke: defina o header X-Amzn-SageMaker-Prefix-Aware-Id (até 64 caracteres ASCII).
API OpenAI: inclua o campo prompt_cache_key no corpo da requisição.

Esse ID se combina com o prefixo para que requisições com prefixos idênticos, mas IDs diferentes, sejam direcionadas para instâncias diferentes.

Inference components e adaptadores LoRA

O prefix-aware routing funciona com endpoints de inference components e adaptadores de Adaptação de Baixa Classificação Dinâmica (LoRA). Para inference components, o comportamento é o mesmo de endpoints de modelo único. Para adaptadores LoRA, opera dentro do conjunto de instâncias sticky do adaptador, usando seleção baseada em prefixo entre as instâncias que já têm o adaptador carregado.

Orientações práticas

Habilite o prefix caching no seu framework de servir. O prefix-aware routing garante que prefixos repetidos cheguem à mesma instância, mas o contêiner precisa ter o prefix caching ativado para de fato armazenar e reutilizar os pares KV. No vLLM, isso está habilitado por padrão nas versões recentes. Outros frameworks podem exigir configuração explícita.
Mantenha a serialização das requisições consistente. Para a API nativa Invoke, o PrefixLength opera sobre bytes brutos. Espaços em branco no JSON, ordenação de chaves e formatação afetam o roteamento. Se o mesmo prompt for serializado de formas diferentes entre requisições, elas podem acabar em instâncias diferentes.
Dimensione o PrefixLength com cuidado. Muito curto e todas as requisições com o mesmo prefixo curto serão canalizadas para uma instância, acionando overflow. Muito longo e pequenas diferenças no payload (como valores de temperatura) dispersam requisições que deveriam ficar juntas. Comece com o comprimento do seu prefixo compartilhado mais uma margem modesta.
Você precisa de pelo menos duas instâncias. Com uma instância, todas as requisições vão para o mesmo lugar independentemente da estratégia.
Monitore as taxas de acerto no cache. Habilite a observabilidade detalhada do SageMaker para acompanhar as taxas de acerto no cache KV no nível do modelo. Isso confirma se o prefix-aware routing está funcionando para o seu workload específico.

Considerações finais

O prefix-aware routing está disponível hoje nos endpoints de inferência em tempo real do SageMaker. Para acessar os novos parâmetros RoutingStrategy e PrefixAwareRoutingConfig, basta atualizar o AWS SDK ou a CLI para a versão mais recente.

Para workloads de LLM com prefixos compartilhados — RAG, conversas multi-turno, bots com templates longos ou assistentes de código — esse recurso pode representar uma redução expressiva de latência e custo computacional sem nenhuma alteração no modelo ou no framework de servir.

Fonte

Reduce LLM latency with prefix-aware routing on Amazon SageMaker Inference (https://aws.amazon.com/blogs/machine-learning/reduce-llm-latency-with-prefix-aware-routing-on-amazon-sagemaker-inference/)