gfx1151 (Strix Halo): fuse attn_k+v into single MMVQ dispatch

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;
 }
 

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);

src/llama-model.cpp

@@ -1015,6 +1015,9 @@ struct llama_model::impl {
 
     bool has_tensor_overrides;
 
+    std::vector<ggml_context_ptr>        wkv_concat_ctxs;
+    std::vector<ggml_backend_buffer_ptr> wkv_concat_bufs;
+
     std::vector<float> 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<uint8_t> 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;
 }
 

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;