diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 227487c27eda..b92c964866f7 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -4596,6 +4596,11 @@ struct ggml_backend_cuda_device_context { int op_offload_min_batch_size; }; +static int ggml_backend_cuda_get_device_cc(ggml_backend_dev_t dev) { + ggml_backend_cuda_device_context * ctx = (ggml_backend_cuda_device_context *)dev->context; + return ggml_cuda_info().devices[ctx->device].cc; +} + static const char * ggml_backend_cuda_device_get_name(ggml_backend_dev_t dev) { ggml_backend_cuda_device_context * ctx = (ggml_backend_cuda_device_context *)dev->context; return ctx->name.c_str(); @@ -5344,6 +5349,9 @@ static void * ggml_backend_cuda_reg_get_proc_address(ggml_backend_reg_t reg, con if (strcmp(name, "ggml_backend_get_features") == 0) { return (void *)ggml_backend_cuda_get_features; } + if (strcmp(name, "ggml_backend_cuda_get_device_cc") == 0) { + return (void *)ggml_backend_cuda_get_device_cc; + } return nullptr; } diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 8c9228b38f9e..2354e4e5f066 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -1502,6 +1502,62 @@ llm_graph_qkv llm_graph_context::build_qkv( Vcur = ggml_view_3d(ctx0, qkv, n_embd_head, n_head_kv, n_tokens, ggml_row_size(qkv->type, n_embd_head), qkv->nb[1], ggml_row_size(qkv->type, n_embd_q + n_embd_kv)); + } else if (layer.wkv_concat && loras->empty() && !layer.wk_s && !layer.wv_s) { + Qcur = build_lora_mm(layer.wq, cur, layer.wq_s); + cb(Qcur, "Qcur", il); + if (layer.wq_b) { + Qcur = ggml_add(ctx0, Qcur, layer.wq_b); + cb(Qcur, "Qcur", il); + } + if (hparams.f_clamp_kqv > 0.0f) { + Qcur = ggml_clamp(ctx0, Qcur, -hparams.f_clamp_kqv, hparams.f_clamp_kqv); + cb(Qcur, "Qcur_clamped", il); + } + + ggml_tensor * kv = ggml_mul_mat(ctx0, layer.wkv_concat, cur); + cb(kv, "kv_concat", il); + + const bool has_kv_bias = layer.wk_b || layer.wv_b; + const bool has_clamp = hparams.f_clamp_kqv > 0.0f; + + if (has_kv_bias || has_clamp) { + Kcur = ggml_view_2d(ctx0, kv, n_embd_kv, n_tokens, kv->nb[1], 0); + cb(Kcur, "Kcur", il); + Vcur = ggml_view_2d(ctx0, kv, n_embd_kv, n_tokens, kv->nb[1], + ggml_row_size(kv->type, n_embd_kv)); + cb(Vcur, "Vcur", il); + + Kcur = ggml_cont(ctx0, Kcur); + Vcur = ggml_cont(ctx0, Vcur); + + if (layer.wk_b) { + Kcur = ggml_add(ctx0, Kcur, layer.wk_b); + cb(Kcur, "Kcur", il); + } + if (layer.wv_b) { + Vcur = ggml_add(ctx0, Vcur, layer.wv_b); + cb(Vcur, "Vcur", il); + } + if (has_clamp) { + Kcur = ggml_clamp(ctx0, Kcur, -hparams.f_clamp_kqv, hparams.f_clamp_kqv); + cb(Kcur, "Kcur_clamped", il); + Vcur = ggml_clamp(ctx0, Vcur, -hparams.f_clamp_kqv, hparams.f_clamp_kqv); + cb(Vcur, "Vcur_clamped", il); + } + + Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head, n_tokens); + Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens); + Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens); + } else { + Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head, n_tokens); + Kcur = ggml_view_3d(ctx0, kv, n_embd_head, n_head_kv, n_tokens, + ggml_row_size(kv->type, n_embd_head), kv->nb[1], 0); + cb(Kcur, "Kcur", il); + Vcur = ggml_view_3d(ctx0, kv, n_embd_head, n_head_kv, n_tokens, + ggml_row_size(kv->type, n_embd_head), kv->nb[1], + ggml_row_size(kv->type, n_embd_kv)); + cb(Vcur, "Vcur", il); + } } else { // separate Q/K/V path Qcur = build_lora_mm(layer.wq, cur, layer.wq_s); diff --git a/src/llama-model.cpp b/src/llama-model.cpp index adacf702d055..a9e652aea26c 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -1015,6 +1015,9 @@ struct llama_model::impl { bool has_tensor_overrides; + std::vector wkv_concat_ctxs; + std::vector wkv_concat_bufs; + std::vector tensor_split_owned; }; @@ -1638,6 +1641,62 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) { } } + { + bool fuse_kv = false; + // Check first layer only — single-GPU assumption for Strix Halo iGPU. + for (auto & layer : model->layers) { + if (!layer.wk || !layer.wk->buffer) continue; + auto buft = ggml_backend_buffer_get_type(layer.wk->buffer); + auto * dev = ggml_backend_buft_get_device(buft); + if (!dev) break; + auto * reg = ggml_backend_dev_backend_reg(dev); + if (!reg) break; + auto * fn = (int (*)(ggml_backend_dev_t)) ggml_backend_reg_get_proc_address(reg, "ggml_backend_cuda_get_device_cc"); + if (fn) { + const int cc = fn(dev); + constexpr int cc_gfx1151 = 0x1000000 + 0x1151; + fuse_kv = (cc == cc_gfx1151); + } + break; + } + + if (fuse_kv) { + LLAMA_LOG_INFO("%s: fusing attn_k + attn_v weights for gfx1151 MMVQ occupancy\n", __func__); + for (size_t il = 0; il < model->layers.size(); ++il) { + auto & layer = model->layers[il]; + if (!layer.wk || !layer.wv || layer.wqkv) continue; + if (layer.wk->type != layer.wv->type) continue; + if (layer.wk->ne[0] != layer.wv->ne[0]) continue; + if (!layer.wv->buffer || layer.wv->buffer != layer.wk->buffer) continue; + + const size_t wk_bytes = ggml_nbytes(layer.wk); + const size_t wv_bytes = ggml_nbytes(layer.wv); + + ggml_init_params ctx_params = { ggml_tensor_overhead(), nullptr, true }; + auto ctx = ggml_context_ptr(ggml_init(ctx_params)); + + auto * t = ggml_new_tensor_2d(ctx.get(), layer.wk->type, + layer.wk->ne[0], + layer.wk->ne[1] + layer.wv->ne[1]); + ggml_format_name(t, "blk.%d.attn_kv_concat.weight", (int)il); + + auto buft = ggml_backend_buffer_get_type(layer.wk->buffer); + auto * buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx.get(), buft); + if (!buf) continue; + + std::vector staging(std::max(wk_bytes, wv_bytes)); + ggml_backend_tensor_get(layer.wk, staging.data(), 0, wk_bytes); + ggml_backend_tensor_set(t, staging.data(), 0, wk_bytes); + ggml_backend_tensor_get(layer.wv, staging.data(), 0, wv_bytes); + ggml_backend_tensor_set(t, staging.data(), wk_bytes, wv_bytes); + + layer.wkv_concat = t; + pimpl->wkv_concat_ctxs.push_back(std::move(ctx)); + pimpl->wkv_concat_bufs.emplace_back(buf); + } + } + } + return true; } diff --git a/src/llama-model.h b/src/llama-model.h index 45b054cedf1d..03bc780b8219 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -256,6 +256,7 @@ struct llama_layer { struct ggml_tensor * wkv_a_mqa = nullptr; struct ggml_tensor * wkv_b = nullptr; struct ggml_tensor * wkv = nullptr; + struct ggml_tensor * wkv_concat = nullptr; struct ggml_tensor * wk_b = nullptr; struct ggml_tensor * wv_b = nullptr; struct ggml_tensor * wqkv_b = nullptr;