From 3dca28d23b1fac0c38c746fd3e15851572a54460 Mon Sep 17 00:00:00 2001 From: liam Date: Thu, 6 Feb 2025 22:39:16 +0800 Subject: [PATCH] :zap: fix moe.cpp int overflow problem --- ktransformers/ktransformers_ext/operators/llamafile/moe.cpp | 4 ++-- ktransformers/local_chat.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/ktransformers/ktransformers_ext/operators/llamafile/moe.cpp b/ktransformers/ktransformers_ext/operators/llamafile/moe.cpp index 0fcf9df..a131b1f 100644 --- a/ktransformers/ktransformers_ext/operators/llamafile/moe.cpp +++ b/ktransformers/ktransformers_ext/operators/llamafile/moe.cpp @@ -224,7 +224,7 @@ void MOE::forward_many(int qlen, int k, const uint64_t* expert_ids, const float* int stride = QK_K; int nth = config_.intermediate_size / stride; backend->do_work_stealing_job(nth * config_.expert_num, nullptr, [&](int task_id) { - int expert_idx = task_id / nth; + uint64_t expert_idx = task_id / nth; int ith = task_id % nth; void* gate_input_ptr = m_local_gate_input_ptr_[expert_idx]; void* gate_proj_ptr = (uint8_t*)gate_proj_ + (expert_idx * config_.intermediate_size + ith * stride) * config_.hidden_size * ggml_type_size(config_.gate_type) / ggml_blck_size(config_.gate_type); @@ -246,7 +246,7 @@ void MOE::forward_many(int qlen, int k, const uint64_t* expert_ids, const float* stride = QK_K; nth = config_.hidden_size / stride; backend->do_work_stealing_job(nth * config_.expert_num, nullptr, [&](int task_id) { - int expert_idx = task_id / nth; + uint64_t expert_idx = task_id / nth; int ith = task_id % nth; void* down_input_ptr = m_local_down_input_ptr_[expert_idx]; void* down_proj_ptr = (uint8_t*)down_proj_ + (expert_idx * config_.hidden_size + ith * stride) * config_.intermediate_size * ggml_type_size(config_.down_type) / ggml_blck_size(config_.down_type); diff --git a/ktransformers/local_chat.py b/ktransformers/local_chat.py index 513f480..93f0ab8 100644 --- a/ktransformers/local_chat.py +++ b/ktransformers/local_chat.py @@ -32,7 +32,7 @@ custom_models = { ktransformer_rules_dir = os.path.dirname(os.path.abspath(__file__)) + "/optimize/optimize_rules/" default_optimize_rules = { "DeepseekV2ForCausalLM": ktransformer_rules_dir + "DeepSeek-V2-Chat.yaml", - "DeepseekV3ForCausalLM": ktransformer_rules_dir + "DeepSeek-V3-Chat.yaml", + "DeepseekV3ForCausalLM": ktransformer_rules_dir + "DeepSeek-V3-Chat-multi-gpu.yaml", "Qwen2MoeForCausalLM": ktransformer_rules_dir + "Qwen2-57B-A14B-Instruct.yaml", "LlamaForCausalLM": ktransformer_rules_dir + "Internlm2_5-7b-Chat-1m.yaml", "MixtralForCausalLM": ktransformer_rules_dir + "Mixtral.yaml",