| 1 | #!/bin/bash |
| 2 | |
| 3 | SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" |
| 4 | |
| 5 | if [[ $# -lt 1 ]]; then |
| 6 | echo "Usage: $0 <prefiller | decoder> [model]" |
| 7 | exit 1 |
| 8 | fi |
| 9 | |
| 10 | if [[ $# -eq 1 ]]; then |
| 11 | echo "Using default model: meta-llama/Llama-3.1-8B-Instruct" |
| 12 | MODEL="meta-llama/Llama-3.1-8B-Instruct" |
| 13 | else |
| 14 | echo "Using model: $2" |
| 15 | MODEL=$2 |
| 16 | fi |
| 17 | |
| 18 | # The prefillers and decoders in LMCache use the same hash seed for all chunk keys. |
| 19 | # This seed must be aligned so that decoders can identify and retrieve KV cache |
| 20 | # entries stored by prefillers. |
| 21 | # |
| 22 | # WARNING: Using a fixed hash seed is insecure and makes the application vulnerable to |
| 23 | # denial-of-service attacks. In a production environment, this should be set to a |
| 24 | # secure random value. This is set to a fixed value for demonstration purposes only. |
| 25 | export PYTHONHASHSEED=${VLLM_PYTHON_HASH_SEED:-123} |
| 26 | |
| 27 | if [[ $1 == "prefiller" ]]; then |
| 28 | # Prefiller listens on port 8100 |
| 29 | prefill_config_file=$SCRIPT_DIR/configs/lmcache-prefiller-config.yaml |
| 30 | |
| 31 | UCX_TLS=cuda_ipc,cuda_copy,tcp \ |
| 32 | LMCACHE_CONFIG_FILE=$prefill_config_file \ |
| 33 | VLLM_ENABLE_V1_MULTIPROCESSING=1 \ |
| 34 | VLLM_WORKER_MULTIPROC_METHOD=spawn \ |
| 35 | CUDA_VISIBLE_DEVICES=0 \ |
| 36 | vllm serve "$MODEL" \ |
| 37 | --port 8100 \ |
| 38 | --enforce-eager \ |
| 39 | --kv-transfer-config \ |
| 40 | '{"kv_connector":"LMCacheConnectorV1","kv_role":"kv_producer","kv_connector_extra_config": {"discard_partial_chunks": false, "lmcache_rpc_port": "producer1"}}' |
| 41 | |
| 42 | |
| 43 | elif [[ $1 == "decoder" ]]; then |
| 44 | # Decoder listens on port 8200 |
| 45 | decode_config_file=$SCRIPT_DIR/configs/lmcache-decoder-config.yaml |
| 46 | |
| 47 | UCX_TLS=cuda_ipc,cuda_copy,tcp \ |
| 48 | LMCACHE_CONFIG_FILE=$decode_config_file \ |
| 49 | VLLM_ENABLE_V1_MULTIPROCESSING=1 \ |
| 50 | VLLM_WORKER_MULTIPROC_METHOD=spawn \ |
| 51 | CUDA_VISIBLE_DEVICES=1 \ |
| 52 | vllm serve "$MODEL" \ |
| 53 | --port 8200 \ |
| 54 | --enforce-eager \ |
| 55 | --kv-transfer-config \ |
| 56 | '{"kv_connector":"LMCacheConnectorV1","kv_role":"kv_consumer","kv_connector_extra_config": {"discard_partial_chunks": false, "lmcache_rpc_port": "consumer1"}}' |
| 57 | |
| 58 | |
| 59 | else |
| 60 | echo "Invalid role: $1" |
| 61 | echo "Should be either prefiller, decoder" |
| 62 | exit 1 |
| 63 | fi |