update code

Author: kevincheng2

fastdeploy/cache_manager/cache_transfer_manager.py

@@ -319,6 +319,11 @@ 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
+        if args.cache_dtype == "block_wise_fp8":
+            cache_scales = paddle.empty(shape=[], dtype=paddle.get_default_dtype())
+            cache_scales_size = self.key_cache_shape[1] * self.key_cache_shape[2]
+            scales_key_need_to_allocate_bytes = args.num_cpu_blocks * cache_scales.element_size() * cache_scales_size
+            scales_value_need_to_allocate_bytes = args.num_cpu_blocks * cache_scales.element_size() * cache_scales_size
         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"
         )
@@ -343,13 +348,13 @@ def _init_cpu_cache(self, args):
             self.cpu_cache_kvs[key_name] = cuda_host_alloc(key_need_to_allocate_bytes)
             self.k_dst_ptrs.append(self.cpu_cache_kvs[key_name])
             if args.cache_dtype == "block_wise_fp8":
-                self.cpu_cache_kvs[key_cache_scales_name] = cuda_host_alloc(key_need_to_allocate_bytes)
+                self.cpu_cache_kvs[key_cache_scales_name] = cuda_host_alloc(scales_key_need_to_allocate_bytes)
                 self.k_scales_ptrs.append(self.cpu_cache_kvs[key_cache_scales_name])
             if value_need_to_allocate_bytes > 0:
                 self.cpu_cache_kvs[val_name] = cuda_host_alloc(value_need_to_allocate_bytes)
                 self.v_dst_ptrs.append(self.cpu_cache_kvs[val_name])
                 if args.cache_dtype == "block_wise_fp8":
-                    self.cpu_cache_kvs[value_cache_scales_name] = cuda_host_alloc(value_need_to_allocate_bytes)
+                    self.cpu_cache_kvs[value_cache_scales_name] = cuda_host_alloc(scales_value_need_to_allocate_bytes)
                     self.v_scales_ptrs.append(self.cpu_cache_kvs[value_cache_scales_name])
         logger.info(f"[rank {self.rank}/{self.n_ranks}] ✅ swap space (cpu cache) is ready!")
         self.swap_space_ready_signal.value[self.rank] = 1