View a markdown version of this page

Configurando o cache KV em camadas - SageMaker IA da Amazon
Pré-requisitosConfigure seu RayClusterConfigure seu aplicativo Ray ServeHabilitar o descarregamento da CPU (opcional)

As traduções são geradas por tradução automática. Em caso de conflito entre o conteúdo da tradução e da versão original em inglês, a versão em inglês prevalecerá.

Configurando o cache KV em camadas

O cache KV gerenciado em camadas adiciona uma camada de cache em todo o cluster às implantações do Ray Serve, então uma réplica lê um prefixo de token que outra réplica já calculou. Isso reduz o tempo até o primeiro token para documentos longos, conversas em vários turnos e solicitações compartilhadas do sistema. O nível de todo o cluster é executado no armazenamento HyperPod hierárquico usando o LMCache com o back-end de armazenamento na documentação do LMCache, que se integra ao Ray Serve e ao vLLM para fornecer cache KV distribuído apoiado pelo HyperPod armazenamento hierárquico. HyperPod

O cache tem duas camadas:

  • Uma camada local na memória do nó que atende à solicitação, para a reutilização mais rápida.

  • Uma camada de armazenamento HyperPod hierárquico que abrange todo o cluster, uma camada de memória agrupada que abrange os nós do cluster. Uma réplica lê um prefixo calculado por outra réplica em vez de recalculá-lo.

Quando uma solicitação compartilha um prefixo de token com um trabalho anterior, a implantação lê o estado do KV em cache da camada local e, em seguida, da camada de todo o cluster, antes de recalcular qualquer coisa.

Pré-requisitos

  • Um HyperPod cluster orquestrado pelo Amazon EKS, com o KubeRay operador instalado.

  • HyperPod Armazenamento hierárquico ativado no cluster. Para obter instruções de configuração, consulte Configurando pontos de verificação hierárquicos gerenciados.

Configure seu RayCluster

Adicione o seguinte à especificação do seu grupo de trabalho para expor o endpoint de armazenamento hierárquico e a memória compartilhada às réplicas do Ray Serve.

workerGroupSpecs:
- template:
    spec:
      volumes:
        - name: host-shm
          hostPath:
            path: /dev/shm
      containers:
      - name: ray-worker
        env:
          - name: NODE_IP
            valueFrom:
              fieldRef:
                fieldPath: status.hostIP
        volumeMounts:
          - name: host-shm
            mountPath: /dev/shm/ai_toolkit_cache
            subPath: ai_toolkit_cache

A variável de NODE_IP ambiente fornece o endereço IP do host para que o cliente LMCache possa se conectar ao serviço de armazenamento hierárquico na porta 9200. A montagem do volume de memória compartilhada permite o caminho de descarga da CPU para o cache KV local.

Configure seu aplicativo Ray Serve

Use a LLMConfig API com o conector LMCache para habilitar o cache KV em sua implantação do Ray Serve. O exemplo a seguir serve um modelo com o cache KV em camadas ativado:

from ray.serve.llm import LLMConfig, build_openai_app

llm_config = LLMConfig(
    model_loading_config={
        "model_id": "my-llm",
        "model_source": "Qwen/Qwen-7B-Chat",
    },
    engine_kwargs={
        "trust_remote_code": True,
        "kv_transfer_config": {
            "kv_connector": "LMCacheConnectorV1",
            "kv_role": "kv_both",
        },
    },
    runtime_env={
        "env_vars": {
            "LMCACHE_REMOTE_URL": "sagemaker-hyperpod://$(NODE_IP):9200",
            "LMCACHE_EXTRA_CONFIG": '{"sagemaker_hyperpod_bucket": "lmcache", "sagemaker_hyperpod_shared_memory_name": "ai_toolkit_cache"}',
            "LMCACHE_CHUNK_SIZE": "256",
        },
    },
    accelerator_type="A10G",
)

app = build_openai_app({"llm_configs": [llm_config]})

Os principais valores de configuração são:

  • kv_connector: defina como LMCacheConnectorV1 para ativar a integração do LMCache com o vLLM.

  • kv_role: defina kv_both para que cada réplica leia e grave no cache compartilhado.

  • LMCACHE_REMOTE_URL: O endpoint de armazenamento hierárquico no nó local. Usa a variável de NODE_IP ambiente configurada no RayCluster manifesto.

  • LMCACHE_CHUNK_SIZE: o número de tokens por bloco de cache. Valores menores aumentam a taxa de acerto do cache para prefixos compartilhados ao custo de mais entradas de cache.

Habilitar o descarregamento da CPU (opcional)

Para adicionar uma camada de memória de CPU local que armazena em cache as entradas KV antes de serem gravadas no armazenamento hierárquico, adicione as seguintes variáveis de ambiente a: runtime_env

"env_vars": {
    "LMCACHE_REMOTE_URL": "sagemaker-hyperpod://$(NODE_IP):9200",
    "LMCACHE_EXTRA_CONFIG": '{"sagemaker_hyperpod_bucket": "lmcache", "sagemaker_hyperpod_shared_memory_name": "ai_toolkit_cache"}',
    "LMCACHE_CHUNK_SIZE": "256",
    "LMCACHE_LOCAL_CPU": "True",
    "LMCACHE_MAX_LOCAL_CPU_SIZE": "100",
}

Quando o descarregamento da CPU está ativado, as entradas de cache KV são armazenadas na memória da CPU no nó local antes de serem gravadas no armazenamento hierárquico de todo o cluster. Isso fornece leituras de cache mais rápidas para réplicas no mesmo nó.