Cache KV em Camadas para LLMs no Amazon SageMaker HyperPod com Curvine

O problema: cache KV isolado por réplica custa caro

Rodar inferência de Modelos de Linguagem de Grande Escala (LLMs — Large Language Models) em escala impõe um dilema clássico: ou você paga por instâncias GPU superdimensionadas para acomodar um cache KV crescente, ou aceita um tempo até o primeiro token (TTFT — Time to First Token) alto porque prompts idênticos são reprocessados em cada requisição.

Para equipes que operam um catálogo amplo de modelos de fundação disponíveis publicamente — como Qwen, Llama e DeepSeek — distribuídos em endpoints por linha de negócio, pipelines de Geração Aumentada por Recuperação (RAG — Retrieval Augmented Generation) ou aplicações de diálogo de múltiplos turnos, esse dilema se traduz diretamente em custo de infraestrutura mais alto e experiência do usuário degradada.

A causa raiz é simples. Durante a geração, o vLLM armazena as chaves e valores de atenção de cada token já processado em um cache KV, evitando recomputação a cada passo. O prefix caching estende isso reutilizando o cache entre requisições que compartilham os mesmos tokens iniciais — como um system prompt comum. Em instâncias econômicas como a ml.g6e.4xlarge (48 GB por GPU), depois de alocar pesos do modelo e runtime, a memória restante para prefix caching é limitada. Com modelos maiores ou maior concorrência, a situação piora: taxas de acerto caem, system prompts idênticos são reprocessados em toda requisição, e réplicas vLLM em escala horizontal mantêm caches isolados entre si. Rotear para uma réplica diferente equivale a um cold start funcional.

A solução: cache KV em três camadas com Curvine

A AWS descreveu uma arquitetura de cache KV em camadas no Amazon SageMaker HyperPod que estende a hierarquia de cache além da memória GPU e CPU para um pool compartilhado e distribuído de NVMe. A solução combina duas capacidades do HyperPod — Cache KV Gerenciado em Camadas e Roteamento Inteligente — com o Curvine, um sistema de arquivos de cache distribuído leve, como camada L2 compartilhada.

Em um deployment de teste, a arquitetura alcançou até 100% de taxa de acerto de cache entre Pods, melhora de até 2,7x no TTFT e latência de leitura L2 entre nós de aproximadamente 56 ms para um prompt de cerca de 1.900 tokens.

As três camadas da hierarquia

L0 – Cache GPU (HBM). É a camada nativa de atenção paginada do vLLM, que mantém os blocos KV mais quentes com a menor latência de acesso. A capacidade é limitada ao que sobra da memória GPU depois dos pesos do modelo. Em uma GPU de 48 GB, um modelo de 7B em bf16 usa cerca de 14 GB para pesos, deixando mais de 30 GB para blocos KV — headroom confortável. Já um modelo de 32B usa cerca de 64 GB de pesos e nem cabe em uma GPU de 48 GB; mesmo com sharding, sobra muito menos memória para KV, e o cache enche rapidamente sob concorrência.

L1 – Offload para memória CPU. Quando blocos GPU são despejados, o LMCache os captura na DRAM do host antes que sejam perdidos. Esse processo roda dentro de cada Pod de inferência e é gerenciado automaticamente pelo Inference Operator do SageMaker HyperPod quando se configura enableL1Cache: true no CRD InferenceEndpointConfig. É rápido, local ao Pod e dimensionado pelo parâmetro InstanceMemoryAllocationPercentage — recomenda-se começar com 20%.

L2 – Pool NVMe distribuído e compartilhado. É aqui que acontece a reutilização entre réplicas. O Curvine agrega os drives NVMe locais das instâncias G6e/P5 em um único namespace, que um cliente FUSE (um driver em espaço de usuário que apresenta o pool como um diretório comum) monta como um PVC (PersistentVolumeClaim) ReadWriteMany em cada Pod de inferência. O LMCache lê e escreve via conector fs://, então o pool distribuído parece um diretório local. Como todos os Pods montam o mesmo namespace, um bloco KV escrito por uma réplica é imediatamente legível pelas demais.

Imagem original — fonte: Aws

O Curvine em si é direto de operar: um nó Primário (chamado “Master” na documentação do Curvine) cuida de metadados e journaling, persistidos no Amazon EBS para durabilidade, enquanto componentes Worker rodam em cada nó GPU e armazenam dados no NVMe local do nó. Se um Worker falha, o cache que ele mantinha é recomputado — sem preocupação com perda de dados, já que são blocos KV reproduzíveis.

Roteamento inteligente: levando requisições à réplica certa

Uma hierarquia de três camadas só entrega seu benefício máximo se as requisições chegarem às réplicas que já têm os blocos KV relevantes. O Inference Operator do HyperPod inclui um roteador integrado com três estratégias:

  • prefix-aware (padrão): ideal para diálogo de múltiplos turnos e system prompts compartilhados
  • kv-aware: indicado para processamento de documentos longos e sessões estendidas
  • round-robin: adequado para inferência em lote sem estado e testes de carga

O roteador mantém uma árvore de prefixos (modo prefix-aware) ou consulta o estado de cache de cada worker (modo kv-aware) para selecionar a réplica com maior probabilidade de produzir um cache hit. Isso ocorre de forma transparente — sem necessidade de alterações no cliente.

Implementação passo a passo

A implementação está organizada em cinco estágios. O walkthrough assume um cluster SageMaker HyperPod orquestrado pelo Amazon EKS. Para criar um, a documentação da AWS orienta seguir o guia Orquestrando clusters SageMaker HyperPod com Amazon EKS, ou usar os templates de referência no repositório AWSome Distributed AI. É necessário provisionar pelo menos dois nós GPU — um único nó não demonstra reutilização entre nós.

Estágio 1: Habilitar o Tiered Storage do HyperPod

O Tiered Storage é uma opção no nível do cluster. Uma vez ativo, o HyperPod implanta automaticamente o DaemonSet ai-toolkit em todos os nós.

aws sagemaker update-cluster \
  --cluster-name hyperpod-cluster-eks \
  --tiered-storage-config Mode=Enable,InstanceMemoryAllocationPercentage=20 \
  --node-recovery Automatic

Atenção: chamar update-cluster apenas com --tiered-storage-config retorna ValidationException. É necessário fornecer também --node-recovery ou --instance-groups. O InstanceMemoryAllocationPercentage aceita valores entre 20 e 100 — comece com 20 e aumente conforme necessário.

Estágio 2: Instalar o Inference Operator e dependências

A forma mais conveniente é o Quick Install no console do SageMaker, que provisiona o papel IAM e instala S3 CSI, FSx CSI, Metrics Server, Cert Manager e o Inference Operator em uma única ação. A alternativa via CLI:

EKS_CLUSTER_NAME=$(aws sagemaker describe-cluster --cluster-name hyperpod-cluster-eks \
  --query 'Orchestrator.Eks.ClusterArn' --output text | cut -d'/' -f2)

for addon in aws-mountpoint-s3-csi-driver aws-fsx-csi-driver metrics-server cert-manager; do
  aws eks create-addon --cluster-name $EKS_CLUSTER_NAME --addon-name $addon --region us-west-2
done

aws eks create-addon \
  --cluster-name $EKS_CLUSTER_NAME \
  --addon-name amazon-sagemaker-hyperpod-inference \
  --configuration-values file://addon-config.json \
  --region us-west-2

Estágio 3: Implantar o cache distribuído Curvine

Antes de implantar o Curvine, é necessário atualizar o plugin VPC CNI e o driver EBS CSI, conceder ao papel aws-ebs-csi-dri-role as permissões de attach de EBS e verificar que uma StorageClass EBS está disponível no cluster. Em seguida, instala-se o CSI do Curvine via Helm:

helm repo add curvine https://curvineio.github.io/helm-charts
helm repo update
helm install curvine-csi curvine/curvine-csi \
  -n curvine --create-namespace \
  --version 0.3.2-alpha \
  --set controller.sidecars.provisioner.image=registry.k8s.io/sig-storage/csi-provisioner:v3.6.0 \
  --set node.sidecars.nodeDriverRegistrar.image=registry.k8s.io/sig-storage/csi-node-driver-registrar:v2.10.0 \
  --set controller.container.securityContext.privileged=true \
  --set node.container.securityContext.privileged=true

Depois, cria-se a StorageClass do Curvine e implanta-se o servidor Curvine (componentes Primary + Worker). Para workloads de cache KV, o NVMe local com hostPath é o backend recomendado — instâncias G6e/P5 vêm com NVMe (~3 GB/s), o que entrega performance muito superior ao EBS gp3, sem custo adicional, e é aceitável para caches onde perda é recuperável.

Após confirmar que todos os Pods estão em estado Running, é fundamental desabilitar imediatamente as flags de formatação para evitar que uma reinicialização futura do Pod reformate e apague os metadados existentes:

helm upgrade curvine curvine/curvine -n curvine \
  --version 0.3.2-alpha --reuse-values \
  --set cluster.formatMaster=false \
  --set cluster.formatWorker=false \
  --set cluster.formatJournal=false

Por fim, cria-se um PVC (PersistentVolumeClaim) ReadWriteMany para os Pods de inferência e aguarda-se o status Bound antes de prosseguir.

Estágio 4: Implantar o endpoint vLLM com cache KV em camadas

Com o Tiered Storage habilitado, o Operator instalado e o Curvine rodando, o estágio final implanta o endpoint de inferência. Isso envolve três passos: declarar o endpoint com o CRD InferenceEndpointConfig, fazer patch no Deployment renderizado para montar o PVC do Curvine, e sobrescrever a URL de cache L2 injetada pelo Operator para apontar ao Curvine.

Dois pontos merecem atenção especial nas variáveis de ambiente do deployment:

  • LMCACHE_REMOTE_SERDE=naive: o serializador cachegen tem um bug de serialização zip no conector de sistema de arquivos do LMCache. O serializador naive é a escolha estável.
  • PYTHONHASHSEED=0: o LMCache deriva chaves de cache de hashes Python. Sem um seed fixo, cada Pod computa chaves diferentes para o mesmo prompt e o compartilhamento entre Pods silenciosamente não funciona.

Como o Operator executa um loop de reconciliação que sobrescreve patches aplicados diretamente no Deployment, a sequência correta é: pausar o Operator, escalar o Deployment para zero, aplicar todos os patches em um único comando, escalar de volta, e restaurar o Operator somente após os Pods estarem prontos.

O patch substitui o valor de LMCACHE_REMOTE_URL para fs://localhost:0/mnt/curvine/l2cache/, adiciona o volume do PVC do Curvine e monta o volume no container vLLM. O formato localhost:0 é um placeholder — o conector fs:// do LMCache ignora host e porta e usa apenas o caminho do filesystem.

Importante para produção: o patch persiste enquanto o Operator não tiver motivo para re-renderizar o Deployment. Para uso produtivo de longa duração, a AWS recomenda substituir o patch manual por um MutatingWebhook ou uma ClusterPolicy do Kyverno que injete o volume, o mount e a URL do Curvine automaticamente a cada renderização.

Estágio 5: Verificar escrita L2, acerto e compartilhamento entre Pods

Um prompt curto não dispara escrita L2 — o LMCache armazena um bloco somente quando um prompt ultrapassa o tamanho de chunk de 256 tokens. O teste de validação usa um documento longo real e envia a requisição byte-idêntica para dois Pods diferentes em dois nós distintos: o primeiro Pod computa o prefill e escreve no Curvine; o segundo lê a mesma entrada pelo mount FUSE compartilhado.

Resultados observados (duas réplicas em dois nós G6e separados):

# POD1 (nó A) -- cold: 0 hits, computa prefill, escreve 1.925 tokens no Curvine
LMCache INFO: Reqid: ..., Total tokens 1925, LMCache hit tokens: 0, need to load: 0
LMCache INFO: Stored 1925 out of total 1925 tokens. size: 0.1028 GB, cost 10.71 ms, throughput: 9.60 GB/s

# POD2 (nó B) -- lê a MESMA entrada via Curvine: 100% de acerto entre Pods, sem prefill
LMCache INFO: Reqid: ..., Total tokens 1925, LMCache hit tokens: 1925, need to load: 1924
LMCache INFO: Retrieved 1925 out of 1925 required tokens. cost 55.75 ms, throughput: 1.84 GB/s

O acerto de 1.925/1.925 no POD2 é a prova de que o L2 compartilhado funciona: o POD2 nunca executou o prefill desse prompt, mas recuperou o cache completo via Curvine.

Benchmarks: quando o L2 compartilhado compensa

A AWS realizou dois benchmarks complementares usando Qwen2-7B-Instruct (fp16, paralelismo tensorial 1), servido por duas réplicas vLLM em dois nós GPU separados.

O primeiro benchmark variou prompts de 500 a 3.000 tokens em instâncias ml.g5.4xlarge. Para prompts acima de 1.000 tokens, o Pod B atingiu 100% de acerto dos tokens escritos pelo Pod A e pulou o prefill completamente: o speedup cresceu de 1,7x em 1.000 tokens até 2,7x em 2.500 tokens, onde uma requisição fria de 774 ms foi concluída em 287 ms. Abaixo de ~1.000 tokens, o round-trip do L2 custa tanto quanto recomputar o prefill (0,99x em 500 tokens) — para prompts curtos, as camadas L0/L1 são o lugar certo para acertar, que é exatamente o que o roteamento prefix-aware incentiva.

O segundo benchmark simulou uma conversa de quatro turnos com histórico acumulado (530 a 2.114 tokens de prompt) em instâncias ml.g6e.4xlarge e ml.g6.16xlarge. Todos os turnos atingiram 100% no L2, reduzindo a latência total da conversa de 4,21 s para 3,25 s (1,30x), com speedup por turno crescendo conforme o histórico acumulava — de 1,22x em 530 tokens a 1,34x em 2.114 tokens.

Duas conclusões práticas dos números: escritas custam pouco (caem no Worker NVMe local a velocidade próxima do disco); leituras são o custo a gerenciar, pois cruzam a rede do Pod — por isso o roteamento prefix-aware importa tanto. Além disso, a economia melhora com o comprimento do contexto até a zona de conforto da rede (~2.500 tokens nesse ambiente), e workloads dominados por prompts abaixo de ~1.000 tokens devem depender de L0/L1 em vez de L2.

Benefícios e operação

A arquitetura entrega benefícios mensuráveis em múltiplas dimensões. Instâncias econômicas como ml.g6e.4xlarge (48 GB por GPU) passam a ser suficientes para hospedar modelos como Qwen2-7B e Llama-3-8B com altas taxas de acerto de cache KV em prompts longos. O cache KV não fica mais limitado pela memória GPU: os dados mais quentes ficam na GPU, os mornos na CPU e os frios são descarregados para o NVMe local agregado, com capacidade L2 escalando linearmente (100 Gi de NVMe por Worker × N nós). Workloads multi-tenant de prompts longos que antes exigiam ml.g6e.12xlarge ou ml.p5.24xlarge podem ser migrados para instâncias menores.

Imagem original — fonte: Aws

Do ponto de vista operacional, o SageMaker HyperPod implanta o DaemonSet ai-toolkit automaticamente em todos os nós, com novos nós se auto-integrando sem intervenção. Com o Curvine, os Workers usam hostPath + NVMe de instância sem custo adicional de armazenamento, enquanto o componente Master usa apenas um pequeno volume EBS (10–50 Gi) para metadados. O Inference Operator gerencia o ciclo de vida dos Pods vLLM, roteamento inteligente, Cert Manager e métricas — as equipes de workload mantêm apenas um único CRD.

Para quais workloads faz sentido?

Essa arquitetura é adequada para workloads com alta sobreposição de prompts: pipelines RAG que reutilizam contexto recuperado, diálogo de múltiplos turnos onde o histórico de sessão se acumula, e deployments multi-tenant que compartilham um template de system prompt. Se o tráfego tem esse perfil e a equipe quer menor TTFT em instâncias menores, vale avaliar.

A mesma camada L2 também sustenta o serving com prefill/decode desagregados (PD) no Inference Operator do HyperPod, que roteia a troca KV entre pools de prefill e decode pelo mesmo stack LMCache — o Curvine pode ser estendido para modelos grandes multi-nó via o override padrão de LMCACHE_REMOTE_URL.

Para começar

Fonte

Tiered KV cache for large LLMs on Amazon SageMaker HyperPod with Curvine (https://aws.amazon.com/blogs/machine-learning/tiered-kv-cache-for-large-llms-on-amazon-sagemaker-hyperpod-with-curvine/)

Comments

Leave a Reply

Your email address will not be published. Required fields are marked *