fix

Author: ltd0924

fastdeploy/cache_manager/cache_messager.py

@@ -775,6 +775,7 @@ def main():
             key_cache_shape_list[2],
             key_cache_shape_list[3],
         ]
+        value_cache_shape = []
         if value_cache_shape_list:
             value_cache_shape = [
                 num_gpu_blocks,

fastdeploy/cache_manager/cache_transfer_manager.py

@@ -215,6 +215,7 @@ def _init_gpu_cache(self, args):
                 self.key_cache_shape[2],
                 self.key_cache_shape[3],
             ]
+            value_cache_shape = []
             if self.value_cache_shape:
                 value_cache_shape = [
                     num_gpu_blocks,
@@ -257,9 +258,9 @@ def _init_gpu_cache(self, args):
         logger.info(f"[rank {self.rank}/{self.n_ranks}] done init cache (full) gmem alloc : {memory_allocated()}")
 
     def _init_cpu_cache(self, args):
-        key_cache_size = args.key_cache_shape[1] * args.key_cache_shape[2] * args.key_cache_shape[3]
+        key_cache_size = self.key_cache_shape[1] * self.key_cache_shape[2] * self.key_cache_shape[3]
         if args.value_cache_shape:
-            value_cache_size = args.value_cache_shape[1] * args.value_cache_shape[2] * args.value_cache_shape[3]
+            value_cache_size = self.value_cache_shape[1] * self.value_cache_shape[2] * self.value_cache_shape[3]
         else:
             value_cache_size = 0
         if args.cache_dtype == "bfloat16":
@@ -270,7 +271,9 @@ def _init_cpu_cache(self, args):
             raise ValueError(f"Unsupported cache dtype: {args.cache_dtype}")
         key_need_to_allocate_bytes = args.num_cpu_blocks * cache_bytes * key_cache_size
         value_need_to_allocate_bytes = args.num_cpu_blocks * cache_bytes * value_cache_size
-        # logger.info(f"[rank {self.rank}/{self.n_ranks}] ..swap space size : { / 1024 ** 3:.2f}GB")
+        logger.info(
+            f"[rank {self.rank}/{self.n_ranks}] ..swap space size : {(key_need_to_allocate_bytes + value_need_to_allocate_bytes) / 1024 ** 3:.2f}GB"
+        )
         if args.num_cpu_blocks == 0:
             logger.info(f"[rank {self.rank}/{self.n_ranks}] 💡 no swap space (cpu cache) is specified.")
             self.swap_space_ready_signal.value[self.rank] = 1

fastdeploy/demo/offline_disaggregated_demo.py

@@ -20,7 +20,7 @@
 
 from fastdeploy.entrypoints.llm import LLM
 
-model_name_or_path = "/root/PaddlePaddle/ERNIE-4.5-0.3B-Paddle"
+model_name_or_path = "baidu/ERNIE-4.5-0.3B-Paddle"
 
 
 def start_decode(model_name_or_path):

fastdeploy/model_executor/layers/attention/iluvatar_attn_backend.py

@@ -194,7 +194,7 @@ def get_kv_cache_shape(
         """
         Calculate kv cache shape
         """
-        key_cache_shape = value_cache_shape = [max_num_blocks, self.kv_num_heads, self.block_size, self.head_dim]
+        key_cache_shape = value_cache_shape = [max_num_blocks, self.num_kv_heads, self.block_size, self.head_dim]
         return key_cache_shape, value_cache_shape
 
     def transpose(self, hidden_states):