[LLM] vLLM: Fix kv cache init (#9771)

Fix kv cache init
2023-12-25 14:17:06 +08:00 · 2023-12-25 14:17:06 +08:00 · 30dab36f76
commit 30dab36f76
parent 449b387125
1 changed files with 1 additions and 1 deletions
--- 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(