feat: add kv-cache-dtype to list of environment variables

hommayushi3 · Aug 9, 2024 · 8729cfb · 8729cfb
1 parent 6d4060c
commit 8729cfb
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/endpoints-entrypoint.sh b/endpoints-entrypoint.sh
@@ -8,9 +8,11 @@ QUANTIZATION=${QUANTIZATION:-}
 DTYPE=${DTYPE:-"auto"}
 TRUST_REMOTE_CODE=${TRUST_REMOTE_CODE:-false}
 GUIDED_DECODING_BACKEND=${GUIDED_DECODING_BACKEND:-"outlines"}
+KV_CACHE_DTYPE=${KV_CACHE_DTYPE:-"auto"}
 
 # Entrypoint for the OpenAI API server
-CMD="vllm serve $MODEL_PATH --host '0.0.0.0' --port 80 --tensor-parallel-size '$NUM_SHARD' --dtype $DTYPE --guided-decoding-backend $GUIDED_DECODING_BACKEND"
+CMD="vllm serve $MODEL_PATH --host '0.0.0.0' --port 80 --tensor-parallel-size '$NUM_SHARD'"
+CMD="$CMD --dtype $DTYPE --guided-decoding-backend $GUIDED_DECODING_BACKEND --kv-cache-dtype $KV_CACHE_DTYPE"
 
 # Append --max-model-len if its value is not -1
 if [ "$MAX_MODEL_LEN" -ne -1 ]; then

diff --git a/examples/inference.py b/examples/inference.py
@@ -30,7 +30,7 @@
         ],
         max_tokens=500,
         temperature=0.0,
-        stream=STREAM
+        stream=STREAM,
     )
 
     if STREAM: