Skip to content

Commit 4e9de0f

Browse files
jeffli-xilinxclaude
jeffli-xilinx
and
claude
committed
change K+V fusion to Q+K fusion, remove VRAM gate
In Q4_K_M models, Q and K share the same quant type (Q4_K) while V uses a higher-precision type (Q6_K). The previous K+V fusion silently skipped all layers due to type mismatch. Switch to Q+K fusion which actually concatenates weights, and gate on GPU device type only (no VRAM check) per review feedback. Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
1 parent 65a3422 commit 4e9de0f

3 files changed

Lines changed: 54 additions & 70 deletions

File tree

‎src/llama-graph.cpp‎

Lines changed: 26 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -1502,34 +1502,30 @@ llm_graph_qkv llm_graph_context::build_qkv(
15021502
Vcur = ggml_view_3d(ctx0, qkv, n_embd_head, n_head_kv, n_tokens,
15031503
ggml_row_size(qkv->type, n_embd_head), qkv->nb[1],
15041504
ggml_row_size(qkv->type, n_embd_q + n_embd_kv));
1505-
} else if (layer.wkv_concat && loras->empty() && !layer.wk_s && !layer.wv_s) {
1506-
Qcur = build_lora_mm(layer.wq, cur, layer.wq_s);
1507-
cb(Qcur, "Qcur", il);
1508-
if (layer.wq_b) {
1509-
Qcur = ggml_add(ctx0, Qcur, layer.wq_b);
1510-
cb(Qcur, "Qcur", il);
1511-
}
1512-
if (hparams.f_clamp_kqv > 0.0f) {
1513-
Qcur = ggml_clamp(ctx0, Qcur, -hparams.f_clamp_kqv, hparams.f_clamp_kqv);
1514-
cb(Qcur, "Qcur_clamped", il);
1515-
}
1505+
} else if (layer.wqk_concat && loras->empty() && !layer.wq_s && !layer.wk_s) {
1506+
ggml_tensor * qk = ggml_mul_mat(ctx0, layer.wqk_concat, cur);
1507+
cb(qk, "qk_concat", il);
15161508

1517-
ggml_tensor * kv = ggml_mul_mat(ctx0, layer.wkv_concat, cur);
1518-
cb(kv, "kv_concat", il);
1509+
Vcur = build_lora_mm(layer.wv, cur, layer.wv_s);
1510+
cb(Vcur, "Vcur", il);
15191511

1520-
const bool has_kv_bias = layer.wk_b || layer.wv_b;
1512+
const bool has_qk_bias = layer.wq_b || layer.wk_b;
15211513
const bool has_clamp = hparams.f_clamp_kqv > 0.0f;
15221514

1523-
if (has_kv_bias || has_clamp) {
1524-
Kcur = ggml_view_2d(ctx0, kv, n_embd_kv, n_tokens, kv->nb[1], 0);
1515+
if (has_qk_bias || has_clamp) {
1516+
Qcur = ggml_view_2d(ctx0, qk, n_embd_q, n_tokens, qk->nb[1], 0);
1517+
cb(Qcur, "Qcur", il);
1518+
Kcur = ggml_view_2d(ctx0, qk, n_embd_kv, n_tokens, qk->nb[1],
1519+
ggml_row_size(qk->type, n_embd_q));
15251520
cb(Kcur, "Kcur", il);
1526-
Vcur = ggml_view_2d(ctx0, kv, n_embd_kv, n_tokens, kv->nb[1],
1527-
ggml_row_size(kv->type, n_embd_kv));
1528-
cb(Vcur, "Vcur", il);
15291521

1522+
Qcur = ggml_cont(ctx0, Qcur);
15301523
Kcur = ggml_cont(ctx0, Kcur);
1531-
Vcur = ggml_cont(ctx0, Vcur);
15321524

1525+
if (layer.wq_b) {
1526+
Qcur = ggml_add(ctx0, Qcur, layer.wq_b);
1527+
cb(Qcur, "Qcur", il);
1528+
}
15331529
if (layer.wk_b) {
15341530
Kcur = ggml_add(ctx0, Kcur, layer.wk_b);
15351531
cb(Kcur, "Kcur", il);
@@ -1539,6 +1535,8 @@ llm_graph_qkv llm_graph_context::build_qkv(
15391535
cb(Vcur, "Vcur", il);
15401536
}
15411537
if (has_clamp) {
1538+
Qcur = ggml_clamp(ctx0, Qcur, -hparams.f_clamp_kqv, hparams.f_clamp_kqv);
1539+
cb(Qcur, "Qcur_clamped", il);
15421540
Kcur = ggml_clamp(ctx0, Kcur, -hparams.f_clamp_kqv, hparams.f_clamp_kqv);
15431541
cb(Kcur, "Kcur_clamped", il);
15441542
Vcur = ggml_clamp(ctx0, Vcur, -hparams.f_clamp_kqv, hparams.f_clamp_kqv);
@@ -1549,14 +1547,15 @@ llm_graph_qkv llm_graph_context::build_qkv(
15491547
Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens);
15501548
Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens);
15511549
} else {
1552-
Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head, n_tokens);
1553-
Kcur = ggml_view_3d(ctx0, kv, n_embd_head, n_head_kv, n_tokens,
1554-
ggml_row_size(kv->type, n_embd_head), kv->nb[1], 0);
1550+
Qcur = ggml_view_3d(ctx0, qk, n_embd_head, n_head, n_tokens,
1551+
ggml_row_size(qk->type, n_embd_head), qk->nb[1], 0);
1552+
cb(Qcur, "Qcur", il);
1553+
Kcur = ggml_view_3d(ctx0, qk, n_embd_head, n_head_kv, n_tokens,
1554+
ggml_row_size(qk->type, n_embd_head), qk->nb[1],
1555+
ggml_row_size(qk->type, n_embd_q));
15551556
cb(Kcur, "Kcur", il);
1556-
Vcur = ggml_view_3d(ctx0, kv, n_embd_head, n_head_kv, n_tokens,
1557-
ggml_row_size(kv->type, n_embd_head), kv->nb[1],
1558-
ggml_row_size(kv->type, n_embd_kv));
1559-
cb(Vcur, "Vcur", il);
1557+
1558+
Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens);
15601559
}
15611560
} else {
15621561
// separate Q/K/V path

‎src/llama-model.cpp‎

Lines changed: 27 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -1015,8 +1015,8 @@ struct llama_model::impl {
10151015

10161016
bool has_tensor_overrides;
10171017

1018-
std::vector<ggml_context_ptr> wkv_concat_ctxs;
1019-
std::vector<ggml_backend_buffer_ptr> wkv_concat_bufs;
1018+
std::vector<ggml_context_ptr> wqk_concat_ctxs;
1019+
std::vector<ggml_backend_buffer_ptr> wqk_concat_bufs;
10201020

10211021
std::vector<float> tensor_split_owned;
10221022
};
@@ -1642,68 +1642,53 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) {
16421642
}
16431643

16441644
{
1645-
bool fuse_kv = false;
1645+
bool fuse_qk = false;
16461646
for (auto & layer : model->layers) {
1647-
if (!layer.wk || !layer.wk->buffer) continue;
1648-
auto buft = ggml_backend_buffer_get_type(layer.wk->buffer);
1647+
if (!layer.wq || !layer.wq->buffer) continue;
1648+
auto buft = ggml_backend_buffer_get_type(layer.wq->buffer);
16491649
auto * dev = ggml_backend_buft_get_device(buft);
16501650
if (!dev) break;
1651-
16521651
auto dev_type = ggml_backend_dev_type(dev);
1653-
if (dev_type != GGML_BACKEND_DEVICE_TYPE_GPU &&
1654-
dev_type != GGML_BACKEND_DEVICE_TYPE_IGPU) {
1655-
break;
1656-
}
1657-
1658-
size_t free = 0, total = 0;
1659-
ggml_backend_dev_memory(dev, &free, &total);
1660-
1661-
size_t fusion_cost = 0;
1662-
for (auto & l : model->layers) {
1663-
if (l.wk && l.wv && !l.wqkv &&
1664-
l.wk->type == l.wv->type &&
1665-
l.wk->ne[0] == l.wv->ne[0]) {
1666-
fusion_cost += ggml_nbytes(l.wk) + ggml_nbytes(l.wv);
1667-
}
1652+
if (dev_type == GGML_BACKEND_DEVICE_TYPE_GPU ||
1653+
dev_type == GGML_BACKEND_DEVICE_TYPE_IGPU) {
1654+
fuse_qk = true;
16681655
}
1669-
1670-
fuse_kv = (free > fusion_cost * 2);
16711656
break;
16721657
}
16731658

1674-
if (fuse_kv) {
1675-
LLAMA_LOG_INFO("%s: fusing attn_k + attn_v weights for improved MMVQ occupancy\n", __func__);
1659+
if (fuse_qk) {
1660+
LLAMA_LOG_INFO("%s: fusing attn_q + attn_k weights for improved MMVQ occupancy\n", __func__);
16761661
for (size_t il = 0; il < model->layers.size(); ++il) {
16771662
auto & layer = model->layers[il];
1678-
if (!layer.wk || !layer.wv || layer.wqkv) continue;
1679-
if (layer.wk->type != layer.wv->type) continue;
1680-
if (layer.wk->ne[0] != layer.wv->ne[0]) continue;
1681-
if (!layer.wv->buffer || layer.wv->buffer != layer.wk->buffer) continue;
1663+
if (!layer.wq || !layer.wk || layer.wqkv) continue;
1664+
if (layer.wq->type != layer.wk->type) continue;
1665+
if (layer.wq->ne[0] != layer.wk->ne[0]) continue;
1666+
if (!layer.wk->buffer || layer.wk->buffer != layer.wq->buffer) continue;
16821667

1668+
const size_t wq_bytes = ggml_nbytes(layer.wq);
16831669
const size_t wk_bytes = ggml_nbytes(layer.wk);
1684-
const size_t wv_bytes = ggml_nbytes(layer.wv);
16851670

16861671
ggml_init_params ctx_params = { ggml_tensor_overhead(), nullptr, true };
16871672
auto ctx = ggml_context_ptr(ggml_init(ctx_params));
16881673

1689-
auto * t = ggml_new_tensor_2d(ctx.get(), layer.wk->type,
1690-
layer.wk->ne[0],
1691-
layer.wk->ne[1] + layer.wv->ne[1]);
1692-
ggml_format_name(t, "blk.%d.attn_kv_concat.weight", (int)il);
1674+
auto * t = ggml_new_tensor_2d(ctx.get(), layer.wq->type,
1675+
layer.wq->ne[0],
1676+
layer.wq->ne[1] + layer.wk->ne[1]);
1677+
ggml_format_name(t, "blk.%d.attn_qk_concat.weight", (int)il);
16931678

1694-
auto buft = ggml_backend_buffer_get_type(layer.wk->buffer);
1679+
auto buft = ggml_backend_buffer_get_type(layer.wq->buffer);
16951680
auto * buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx.get(), buft);
16961681
if (!buf) continue;
16971682

1698-
std::vector<uint8_t> staging(std::max(wk_bytes, wv_bytes));
1683+
std::vector<uint8_t> staging(std::max(wq_bytes, wk_bytes));
1684+
ggml_backend_tensor_get(layer.wq, staging.data(), 0, wq_bytes);
1685+
ggml_backend_tensor_set(t, staging.data(), 0, wq_bytes);
16991686
ggml_backend_tensor_get(layer.wk, staging.data(), 0, wk_bytes);
1700-
ggml_backend_tensor_set(t, staging.data(), 0, wk_bytes);
1701-
ggml_backend_tensor_get(layer.wv, staging.data(), 0, wv_bytes);
1702-
ggml_backend_tensor_set(t, staging.data(), wk_bytes, wv_bytes);
1687+
ggml_backend_tensor_set(t, staging.data(), wq_bytes, wk_bytes);
17031688

1704-
layer.wkv_concat = t;
1705-
pimpl->wkv_concat_ctxs.push_back(std::move(ctx));
1706-
pimpl->wkv_concat_bufs.emplace_back(buf);
1689+
layer.wqk_concat = t;
1690+
pimpl->wqk_concat_ctxs.push_back(std::move(ctx));
1691+
pimpl->wqk_concat_bufs.emplace_back(buf);
17071692
}
17081693
}
17091694
}

‎src/llama-model.h‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -256,7 +256,7 @@ struct llama_layer {
256256
struct ggml_tensor * wkv_a_mqa = nullptr;
257257
struct ggml_tensor * wkv_b = nullptr;
258258
struct ggml_tensor * wkv = nullptr;
259-
struct ggml_tensor * wkv_concat = nullptr;
259+
struct ggml_tensor * wqk_concat = nullptr;
260260
struct ggml_tensor * wk_b = nullptr;
261261
struct ggml_tensor * wv_b = nullptr;
262262
struct ggml_tensor * wqkv_b = nullptr;

0 commit comments

Comments
 (0)