[LLM] vLLM: Delete last_kv_cache before prefilling (#9619)

Remove last_kv_cache before prefilling to reduce peak memory usage.
2023-12-07 11:32:33 +08:00 · 2023-12-07 11:32:33 +08:00 · deee65785c
commit deee65785c
parent 48b85593b3
1 changed files with 2 additions and 0 deletions
--- a/python/llm/src/bigdl/llm/vllm/model_executor/models/bigdl_llama.py
+++ b/python/llm/src/bigdl/llm/vllm/model_executor/models/bigdl_llama.py
@ -186,6 +186,8 @@ class BigDLLlamaForCausalLM(BigDLModelForCausalLM):
                "use_cache": True,
                # "return_dict": True,
            }
+            if self.last_kv_cache:
+                del self.last_kv_cache
        # pdb.set_trace()

        if self.device.type == 'xpu':