diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 3148f27816..3607bacd6e 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -603,8 +603,20 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str }; auto get_split_segments = [&](int axis, uint32_t il) -> std::vector> { + // TODO: clarify why this is necessary specifically for these models + // TODO: deduplicate condition [TAG_SPLIT_QGATE_QWEN] if (ud->model->arch == LLM_ARCH_QWEN3NEXT || ud->model->arch == LLM_ARCH_QWEN35 || ud->model->arch == LLM_ARCH_QWEN35MOE || ud->model->arch == LLM_ARCH_QWEN4EXP) { + + // fused full attention layers with Q gate tensors that need n_embd doubled: + if (!hparams.is_recr(il) && (std::regex_match(tensor_name, pattern_qkv_weight) || std::regex_match(tensor_name, pattern_qkv_bias))) { + const int64_t n_embd = hparams.n_head(il) * hparams.n_embd_head_k(il) * 2; + const int64_t n_embd_gqa = hparams.n_embd_v_gqa(il); + GGML_ASSERT(hparams.n_embd_k_gqa(il) == n_embd_gqa); + GGML_ASSERT(tensor->ne[axis] == n_embd + 2*n_embd_gqa); + return {{n_embd, 1}, {n_embd_gqa, 2}}; + } + const int64_t head_k_dim = hparams.ssm_d_state; const int64_t head_v_dim = hparams.ssm_d_state; const int64_t n_k_heads = hparams.ssm_n_group; @@ -654,9 +666,9 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str } if (std::regex_match(tensor_name, pattern_qkv_weight) || std::regex_match(tensor_name, pattern_qkv_bias)) { - const int64_t n_embd = hparams.n_embd; + const int64_t n_embd = hparams.n_head(il) * hparams.n_embd_head_k(il); const int64_t n_embd_gqa = hparams.n_embd_v_gqa(il); - GGML_ASSERT(hparams.n_embd_k_gqa() == n_embd_gqa); + GGML_ASSERT(hparams.n_embd_k_gqa(il) == n_embd_gqa); GGML_ASSERT(tensor->ne[axis] == n_embd + 2*n_embd_gqa); return {{n_embd, 1}, {n_embd_gqa, 2}}; } @@ -742,6 +754,7 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str if (std::regex_match(tensor_name, pattern_q_weight) || std::regex_match(tensor_name, pattern_q_bias)) { GGML_ASSERT(segments.size() == 1); // some models have Q gate tensors, for those cases the granularity needs to be doubled: + // TODO: deduplicate condition [TAG_SPLIT_QGATE_QWEN] if (ud->model->arch == LLM_ARCH_QWEN3NEXT || ud->model->arch == LLM_ARCH_QWEN35 || ud->model->arch == LLM_ARCH_QWEN35MOE || ud->model->arch == LLM_ARCH_QWEN4EXP) { return {std::lcm(2*n_embd_q, blck_size_perf)}; @@ -769,6 +782,12 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str } if (std::regex_match(tensor_name, pattern_qkv_weight) || std::regex_match(tensor_name, pattern_qkv_bias)) { GGML_ASSERT(segments.size() == 2); + // fused full attention layers need Q gate tensors handled like above: + // TODO: deduplicate condition [TAG_SPLIT_QGATE_QWEN] + if (ud->model->arch == LLM_ARCH_QWEN3NEXT || ud->model->arch == LLM_ARCH_QWEN35 || ud->model->arch == LLM_ARCH_QWEN35MOE || + ud->model->arch == LLM_ARCH_QWEN4EXP) { + return {std::lcm(2*n_embd_q, blck_size_perf), granularity_kv}; + } return {granularity_q, granularity_kv}; } }