View a markdown version of this page

Configurazione della cache KV su più livelli - Amazon SageMaker AI
PrerequisitiConfigura il tuo RayClusterConfigura la tua applicazione Ray ServeAbilita l'offload della CPU (opzionale)

Le traduzioni sono generate tramite traduzione automatica. In caso di conflitto tra il contenuto di una traduzione e la versione originale in Inglese, quest'ultima prevarrà.

Configurazione della cache KV su più livelli

La cache KV gestita su più livelli aggiunge un livello di cache a livello di cluster alle implementazioni Ray Serve, quindi una replica legge un prefisso di token che un'altra replica ha già calcolato. Ciò riduce il tempo necessario per inviare il primo token per documenti lunghi, conversazioni a più turni e richieste di sistema condivise. Il livello a livello di cluster viene eseguito su HyperPod Tiered Storage utilizzando LMCache con il backend di archiviazione nella documentazione di LMCache, che si integra con Ray Serve e VLLM per fornire una cache KV distribuita supportata da Tiered HyperPod Storage. HyperPod

La cache ha due livelli:

  • Un livello locale nella memoria del nodo che serve la richiesta, per il riutilizzo più rapido.

  • Un livello a livello di cluster su HyperPod Tiered Storage, un livello di memoria in pool che comprende i nodi del cluster. Una replica legge un prefisso calcolato da un'altra replica invece di ricalcolarlo.

Quando una richiesta condivide il prefisso di un token con un lavoro precedente, la distribuzione legge lo stato KV memorizzato nella cache dal livello locale, quindi dal livello a livello di cluster, prima di ricalcolare qualsiasi cosa.

Prerequisiti

Configura il tuo RayCluster

Aggiungete quanto segue alle specifiche del vostro gruppo di lavoro per esporre l'endpoint Tiered Storage e la memoria condivisa alle repliche 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

La variabile di NODE_IP ambiente fornisce l'indirizzo IP dell'host in modo che il client LMCache possa connettersi al servizio Tiered Storage sulla porta 9200. Il montaggio del volume di memoria condivisa abilita il percorso di offload della CPU per la cache KV locale.

Configura la tua applicazione Ray Serve

Usa l'LLMConfigAPI con il connettore LMCache per abilitare il caching KV nella tua distribuzione Ray Serve. L'esempio seguente serve un modello con caching KV a più livelli abilitato:

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]})

I valori di configurazione chiave sono:

  • kv_connector: impostato su per LMCacheConnectorV1 abilitare l'integrazione di LMCache con VLLM.

  • kv_role: impostato kv_both in modo che ogni replica legga e scriva nella cache condivisa.

  • LMCACHE_REMOTE_URL: l'endpoint di storage su più livelli sul nodo locale. Utilizza la variabile di NODE_IP ambiente configurata nel manifesto. RayCluster

  • LMCACHE_CHUNK_SIZE: il numero di token per blocco di cache. Valori più bassi aumentano la frequenza di accesso alla cache per i prefissi condivisi al costo di un maggior numero di voci della cache.

Abilita l'offload della CPU (opzionale)

Per aggiungere un livello di memoria CPU locale che memorizza nella cache le voci KV prima che vengano scritte su Tiered Storage, aggiungi le seguenti variabili di 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 l'offload della CPU è abilitato, le voci della cache KV vengono archiviate nella memoria della CPU sul nodo locale prima di essere scritte nello storage a più livelli a livello di cluster. Ciò fornisce letture della cache più rapide per le repliche sullo stesso nodo.