refactor the code

Author: JC-ut0

xllm/core/distributed_runtime/llm_engine.cpp

@@ -500,14 +500,8 @@ Engine::KVCacheCapacity LLMEngine::estimate_kv_cache_capacity() {
     int64_t head_v_dim = args_.linear_value_head_dim();
 
     // Parse mamba_ssm_dtype if specified
-    int64_t ssm_dtype_size = dtype_size;
-    if (!args_.mamba_ssm_dtype().empty()) {
-      auto parsed_ssm_dtype =
-          try_get_scalar_type_from_string(args_.mamba_ssm_dtype());
-      if (parsed_ssm_dtype) {
-        ssm_dtype_size = get_dtype_size(parsed_ssm_dtype.value());
-      }
-    }
+    int64_t ssm_dtype_size =
+        resolve_ssm_dtype_size(args_.mamba_ssm_dtype(), dtype_size);
 
     int64_t linear_ssm_slot_size =
         ssm_dtype_size * n_local_linear_v_heads_ * head_k_dim * head_v_dim;

xllm/core/framework/model/model_args.h

@@ -181,6 +181,7 @@ struct ModelArgs {
   PROPERTY(int32_t, linear_value_head_dim) = 0;
   PROPERTY(int64_t, linear_num_key_heads) = 0;
   PROPERTY(int32_t, linear_num_value_heads) = 0;
+  PROPERTY(std::string, mamba_ssm_dtype);
   PROPERTY(int32_t, shared_expert_intermediate_size) = 0;
   PROPERTY(float, partial_rotary_factor) = 0.0f;
   PROPERTY(std::vector<std::string>, layer_types) = {};
@@ -339,9 +340,6 @@ struct ModelArgs {
   PROPERTY(int64_t, mm_image_shortest_edge) = 0;
   PROPERTY(int64_t, mm_image_longest_edge) = 0;
 
-  // Mamba SSM dtype
-  PROPERTY(std::string, mamba_ssm_dtype);
-
   // GLM
   PROPERTY(int64_t, mm_video_shortest_edge) = 0;
   PROPERTY(int64_t, mm_video_longest_edge) = 0;

xllm/core/runtime/worker_impl.cpp

@@ -252,15 +252,9 @@ bool WorkerImpl::allocate_kv_cache(
 
       if (is_linear_layer) {
         // Linear attention layer: only allocate conv_cache and ssm_cache
-        torch::ScalarType ssm_dtype = dtype_;
         // Parse mamba_ssm_dtype if specified
-        if (!args.mamba_ssm_dtype().empty()) {
-          auto parsed_ssm_dtype =
-              try_get_scalar_type_from_string(args.mamba_ssm_dtype());
-          if (parsed_ssm_dtype) {
-            ssm_dtype = parsed_ssm_dtype.value();
-          }
-        }
+        torch::ScalarType ssm_dtype =
+            resolve_ssm_dtype(args.mamba_ssm_dtype(), dtype_);
 
 #if defined(USE_NPU)
         aclFormat npu_format_type = ACL_FORMAT_ND;

xllm/core/util/tensor_helper.h

@@ -360,4 +360,34 @@ inline int32_t get_dtype_size(torch::ScalarType dtype) {
   return static_cast<int32_t>(torch::elementSize(dtype));
 }
 
+inline torch::ScalarType resolve_ssm_dtype(
+    const std::string& mamba_ssm_dtype_str,
+    torch::ScalarType default_dtype) {
+  if (mamba_ssm_dtype_str.empty()) {
+    return default_dtype;
+  }
+  auto parsed = try_get_scalar_type_from_string(mamba_ssm_dtype_str);
+  if (parsed) {
+    return parsed.value();
+  }
+  LOG(WARNING) << "Failed to parse mamba_ssm_dtype='" << mamba_ssm_dtype_str
+               << "', falling back to default_dtype: " << default_dtype;
+  return default_dtype;
+}
+
+inline int64_t resolve_ssm_dtype_size(
+    const std::string& mamba_ssm_dtype_str,
+    int64_t default_dtype_size) {
+  if (mamba_ssm_dtype_str.empty()) {
+    return default_dtype_size;
+  }
+  auto parsed = try_get_scalar_type_from_string(mamba_ssm_dtype_str);
+  if (parsed) {
+    return get_dtype_size(parsed.value());
+  }
+  LOG(WARNING) << "Failed to parse mamba_ssm_dtype='" << mamba_ssm_dtype_str
+               << "', falling back to default dtype size";
+  return default_dtype_size;
+}
+
 }  // namespace xllm

xllm/models/llm/qwen3_5.h

@@ -146,7 +146,7 @@ TORCH_MODULE(Qwen3_5ForCausalLM);
   SET_ARG(routed_scaling_factor, 1.0f);                                        \
   SET_ARG(stop_token_ids,                                                      \
           std::unordered_set<int32_t>({args->eos_token_id()}));                \
-  LOAD_ARG_TEXT_OR_ROOT(mamba_ssm_dtype, "mamba_ssm_dtype", "bfloat16")
+  LOAD_ARG_TEXT_OR_ROOT(mamba_ssm_dtype, "mamba_ssm_dtype", "float32")
 
 #define LOAD_QWEN3_5_TYPE_AND_DTYPE(default_model_type)         \
   LOAD_ARG_OR(model_type, "model_type", default_model_type);    \