default to get data_type startargs from config.json and fix deepseekv…

lightllm/models/deepseek2/triton_kernel/context_flashattention_nopad.py

@@ -177,6 +177,9 @@ def context_attention_fwd(
     else:
         BLOCK = 128 if not TESLA else 64
 
+    if q_nope.dtype == torch.float32:
+        BLOCK = BLOCK // 4
+
     sm_scale = softmax_scale
     batch, head = b_seq_len.shape[0], q_nope.shape[1]
     kv_group_num = q_nope.shape[1]  # deepseekv2 的 group 就是q的head数量，类似于MQA
@@ -370,6 +373,9 @@ def context_attention_fwd_no_prompt_cache(
     else:
         BLOCK = 128 if not TESLA else 64
 
+    if q_nope.dtype == torch.float32:
+        BLOCK = BLOCK // 4
+
     sm_scale = softmax_scale
     batch, head = b_seq_len.shape[0], q_nope.shape[1]
     kv_group_num = q_nope.shape[1]

lightllm/server/api_server.py

@@ -449,7 +449,7 @@ def make_argument_parser() -> argparse.ArgumentParser:
         "--data_type",
         type=str,
         choices=["fp16", "float16", "bf16", "bfloat16", "fp32", "float32"],
-        default="float16",
+        default=None,
         help="the data type of the model weight",
     )
     parser.add_argument("--return_all_prompt_logprobs", action="store_true", help="return all prompt tokens logprobs")
@@ -562,6 +562,12 @@ def main():
 
         args.eos_id = get_eos_token_ids(args.model_dir)
 
+    if args.data_type is None:
+        from lightllm.utils.config_utils import get_dtype
+
+        args.data_type = get_dtype(args.model_dir)
+        assert args.data_type in ["fp16", "float16", "bf16", "bfloat16", "fp32", "float32"]
+
     logger.info(f"all start args:{args}")
 
     can_use_ports = alloc_can_use_network_port(num=6 + args.tp, used_nccl_port=args.nccl_port)

lightllm/utils/config_utils.py

@@ -1,5 +1,8 @@
 import json
 import os
+from lightllm.utils.log_utils import init_logger
+
+logger = init_logger(__name__)
 
 
 def get_config_json(model_path: str):
@@ -16,3 +19,13 @@ def get_eos_token_ids(model_path: str):
     if isinstance(eos_token_id, list):
         return eos_token_id
     assert False, "error eos_token_id format in config.json"
+
+
+def get_dtype(model_path: str):
+    config_json = get_config_json(model_path)
+    try:
+        torch_dtype = config_json["torch_dtype"]
+        return torch_dtype
+    except:
+        logger.warning("torch_dtype not in config.json, use float16 as default")
+        return "float16"