fix deepseek

ModelTC · Oct 24, 2024 · 0aebaec · 0aebaec
1 parent 07a7da5
commit 0aebaec
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/lightllm/common/deepseek2_mem_manager.py b/lightllm/common/deepseek2_mem_manager.py
@@ -4,5 +4,9 @@
 
 
 class Deepseek2MemoryManager(MemoryManager):
+
+    def get_cell_size(self):
+        return self.head_num * self.head_dim * self.layer_num * torch._utils._element_size(self.dtype)
+
     def _init_buffers(self, size, dtype, head_num, head_dim, layer_num):
         self.kv_buffer = [torch.empty((size, head_num, head_dim), dtype=dtype, device="cuda") for _ in range(layer_num)]