feat: support deepseek prefixcache.

Author: DongheJin

xllm/core/runtime/llm_engine.cpp

@@ -41,6 +41,8 @@ limitations under the License.
 
 namespace xllm {
 
+constexpr int32_t NZ_ALIGNMENT = 16;
+
 namespace {
 uint32_t determine_micro_batches_num(const std::vector<Batch>& batch) {
   bool not_all_in_decode =
@@ -258,8 +260,10 @@ Engine::KVCacheCapacity LLMEngine::estimate_kv_cache_capacity() {
   int64_t slot_size = 0;
   if (FLAGS_enable_mla) {
     if (FLAGS_enable_prefix_cache) {
-      slot_size = dtype_size * ((args_.kv_lora_rank() + 15) / 16 * 16 +
-                                (args_.qk_rope_head_dim() + 15) / 16 * 16);
+      slot_size =
+          dtype_size *
+          ((args_.kv_lora_rank() + NZ_ALIGNMENT - 1) / NZ_ALIGNMENT +
+           (args_.qk_rope_head_dim() + NZ_ALIGNMENT - 1) / NZ_ALIGNMENT);
     } else {
       slot_size =
           dtype_size * (args_.kv_lora_rank() + args_.qk_rope_head_dim());

xllm/core/runtime/worker_impl.cpp

@@ -48,6 +48,8 @@ limitations under the License.
 namespace xllm {
 
 constexpr uint64_t MBUF_SIZE = 128 * 1024 * 1024;
+constexpr int32_t FORMAT_ND = 2;
+constexpr int32_t FORMAT_NZ = 29;
 
 WorkerImpl::WorkerImpl(const ParallelArgs& parallel_args,
                        const torch::Device& device,
@@ -88,7 +90,7 @@ bool WorkerImpl::allocate_kv_cache(
     torch::Tensor key_cache, value_cache;
 #if defined(USE_NPU)
     int32_t npu_format_type =
-        FLAGS_enable_mla && FLAGS_enable_prefix_cache ? 29 : 2;
+        FLAGS_enable_mla && FLAGS_enable_prefix_cache ? FORMAT_NZ : FORMAT_ND;
     key_cache = at_npu::native::npu_format_cast(
         torch::empty(kv_cache_shape[0], torch::dtype(dtype_).device(device_)),
         npu_format_type);

xllm/core/util/utils.cpp

@@ -80,7 +80,7 @@ torch::ScalarType parse_dtype(const std::string& dtype_str,
   }
   if ((boost::iequals(dtype_str, "float") ||
        boost::iequals(dtype_str, "float32"))) {
-    return torch::kFloat16;
+    return torch::kFloat32;
   }
 
   if (dtype_str.empty() || boost::iequals(dtype_str, "auto")) {
@@ -99,7 +99,7 @@ torch::ScalarType parse_dtype(const std::string& dtype_str) {
   }
   if ((boost::iequals(dtype_str, "float") ||
        boost::iequals(dtype_str, "float32"))) {
-    return torch::kFloat16;
+    return torch::kFloat32;
   }
 
   if (dtype_str.empty() || boost::iequals(dtype_str, "auto")) {