From 30dab36f7694103f1405264639dbdd5da232074b Mon Sep 17 00:00:00 2001
From: Xiangyu Tian <109123695+xiangyuT@users.noreply.github.com>
Date: Mon, 25 Dec 2023 14:17:06 +0800
Subject: [PATCH] [LLM] vLLM: Fix kv cache init (#9771)

Fix kv cache init
---
 python/llm/src/bigdl/llm/vllm/engine/llm_engine.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/python/llm/src/bigdl/llm/vllm/engine/llm_engine.py b/python/llm/src/bigdl/llm/vllm/engine/llm_engine.py
index 5978fc38..15a03de4 100644
--- a/python/llm/src/bigdl/llm/vllm/engine/llm_engine.py
+++ b/python/llm/src/bigdl/llm/vllm/engine/llm_engine.py
@@ -127,7 +127,7 @@ class LLMEngine:
         # self.parallel_config = parallel_config
         self.scheduler_config = scheduler_config
         self.log_stats = log_stats
-        self.kv_cache = [[dict() for _ in range(2)] for _ in range(32)]
+        self.kv_cache = [[dict(), dict()] for _ in range(model_config.hf_config.num_hidden_layers)]
         # self._verify_args()
 
         self.tokenizer = get_tokenizer(