View a markdown version of this page

계층형 KV 캐시 설정 - Amazon SageMaker AI
사전 조건RayCluster 구성Ray Serve 애플리케이션 구성CPU 오프로드 활성화(선택 사항)

기계 번역으로 제공되는 번역입니다. 제공된 번역과 원본 영어의 내용이 상충하는 경우에는 영어 버전이 우선합니다.

계층형 KV 캐시 설정

관리형 계층형 KV 캐시는 클러스터 전체 캐시 계층을 Ray Serve 배포에 추가하므로 한 복제본은 다른 복제본이 이미 계산한 토큰 접두사를 읽습니다. 이렇게 하면 긴 문서, 멀티턴 대화 및 공유 시스템 프롬프트에 대한 첫 번째 토큰 생성 시간이 단축됩니다. 클러스터 전체 계층은 LMCache 설명서의 HyperPod 스토리지 백엔드와 함께 LMCache를 사용하여 HyperPod 계층형 스토리지에서 실행되며, Ray Serve 및 vLLM과 통합되어 HyperPod 계층형 스토리지에서 지원하는 분산 KV 캐싱을 제공합니다. LMCache HyperPod

캐시에는 두 개의 티어가 있습니다.

  • 가장 빠른 재사용을 위해 요청을 처리하는 노드 메모리의 로컬 티어입니다.

  • 클러스터 노드에 걸쳐 있는 풀링 메모리 계층인 HyperPod 계층형 스토리지의 클러스터 전체 계층입니다. 복제본은 다시 계산하는 대신 다른 복제본에서 계산한 접두사를 읽습니다.

요청이 이전 작업과 토큰 접두사를 공유하면 배포는 캐시된 KV 상태를 로컬 계층에서 읽은 다음 클러스터 전체 계층에서 읽은 다음 다시 계산합니다.

사전 조건

RayCluster 구성

작업자 그룹 사양에 다음을 추가하여 계층형 스토리지 엔드포인트와 공유 메모리를 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

NODE_IP 환경 변수는 LMCache 클라이언트가 포트 9200의 계층형 스토리지 서비스에 연결할 수 있도록 호스트 IP 주소를 제공합니다. 공유 메모리 볼륨 마운트는 로컬 KV 캐시의 CPU 오프로드 경로를 활성화합니다.

Ray Serve 애플리케이션 구성

LLMConfig API를 LMCache 커넥터와 함께 사용하여 Ray Serve 배포에서 KV 캐싱을 활성화합니다. 다음 예제에서는 계층형 KV 캐싱이 활성화된 모델을 제공합니다.

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

키 구성 값은 다음과 같습니다.

  • kv_connector: vLLM과의 LMCache 통합을 활성화LMCacheConnectorV1하려면 로 설정합니다.

  • kv_role: 각 복제본kv_both이 공유 캐시에서 읽고 쓰도록를 로 설정합니다.

  • LMCACHE_REMOTE_URL: 로컬 노드의 계층형 스토리지 엔드포인트입니다. RayCluster 매니페스트에 구성된 NODE_IP 환경 변수를 사용합니다.

  • LMCACHE_CHUNK_SIZE: 캐시 청크당 토큰 수입니다. 값이 작을수록 공유 접두사의 캐시 적중률이 증가하며 캐시 항목이 늘어납니다.

CPU 오프로드 활성화(선택 사항)

계층형 스토리지에 기록되기 전에 KV 항목을 캐싱하는 로컬 CPU 메모리 계층을 추가하려면 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",
}

CPU 오프로드가 활성화되면 클러스터 전체 계층형 스토리지에 기록되기 전에 KV 캐시 항목이 로컬 노드의 CPU 메모리에 저장됩니다. 이렇게 하면 동일한 노드의 복제본에 대한 캐시 읽기 속도가 빨라집니다.