View a markdown version of this page

Configuration du cache KV hiérarchisé - Amazon SageMaker AI
Conditions préalablesConfigurez votre RayClusterConfigurez votre application Ray ServeActiver le déchargement du processeur (facultatif)

Les traductions sont fournies par des outils de traduction automatique. En cas de conflit entre le contenu d'une traduction et celui de la version originale en anglais, la version anglaise prévaudra.

Configuration du cache KV hiérarchisé

Le cache KV hiérarchisé géré ajoute un niveau de cache à l'échelle du cluster aux déploiements Ray Serve, de sorte qu'un réplica lit un préfixe de jeton qu'un autre réplica a déjà calculé. Cela réduit le délai de création du premier jeton pour les longs documents, les conversations à plusieurs reprises et les invites système partagées. Le niveau à l'échelle du cluster s'exécute sur un stockage HyperPod hiérarchisé utilisant LMCache avec le backend de stockage décrit dans la documentation LMCache, qui s'intègre à Ray Serve et vLLm pour fournir une mise en cache KV distribuée soutenue par un HyperPod stockage hiérarchisé. HyperPod

Le cache comporte deux niveaux :

  • Un niveau local dans la mémoire du nœud qui traite la demande, pour une réutilisation rapide.

  • Un niveau à l'échelle du cluster sur le stockage HyperPod hiérarchisé, un niveau de mémoire groupée qui couvre les nœuds du cluster. Un réplica lit un préfixe calculé par un autre réplica au lieu de le recalculer.

Lorsqu'une demande partage un préfixe de jeton avec un travail antérieur, le déploiement lit l'état KV mis en cache à partir du niveau local, puis du niveau à l'échelle du cluster, avant de recalculer quoi que ce soit.

Conditions préalables

  • HyperPod Cluster orchestré par Amazon EKS, sur lequel l' KubeRay opérateur est installé.

  • HyperPod Stockage hiérarchisé activé sur le cluster. Pour les instructions de configuration, voir Configuration du point de contrôle hiérarchisé géré.

Configurez votre RayCluster

Ajoutez ce qui suit à la spécification de votre groupe de travail pour exposer le point de terminaison de stockage hiérarchisé et la mémoire partagée aux répliques 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 variable d'NODE_IPenvironnement fournit l'adresse IP de l'hôte afin que le client LMCache puisse se connecter au service de stockage hiérarchisé sur le port 9200. Le montage du volume de mémoire partagée active le chemin de déchargement du processeur pour le cache KV local.

Configurez votre application Ray Serve

Utilisez l'LLMConfigAPI avec le connecteur LMCache pour activer la mise en cache KV dans votre déploiement Ray Serve. L'exemple suivant concerne un modèle dans lequel la mise en cache KV hiérarchisée est activée :

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

Les valeurs de configuration clés sont les suivantes :

  • kv_connector: Réglez pour LMCacheConnectorV1 activer l'intégration de LMCache avec vLLM.

  • kv_role: Réglez pour kv_both que chaque réplique lise et écrit dans le cache partagé.

  • LMCACHE_REMOTE_URL: point de terminaison du stockage hiérarchisé sur le nœud local. Utilise la variable d'NODE_IPenvironnement configurée dans le RayCluster manifeste.

  • LMCACHE_CHUNK_SIZE: nombre de jetons par bloc de cache. Des valeurs plus faibles augmentent le taux d'accès au cache pour les préfixes partagés au détriment d'un plus grand nombre d'entrées de cache.

Activer le déchargement du processeur (facultatif)

Pour ajouter un niveau de mémoire du processeur local qui met en cache les entrées KV avant qu'elles ne soient écrites dans le stockage hiérarchisé, ajoutez les variables d'environnement suivantes à : 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",
}

Lorsque le déchargement du processeur est activé, les entrées du cache KV sont stockées dans la mémoire du processeur sur le nœud local avant d'être écrites dans le stockage hiérarchisé à l'échelle du cluster. Cela permet de lire plus rapidement le cache pour les répliques sur le même nœud.