Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
jeffli-xilinx wants to merge 7 commits into
base: gfx11
Choose a base branch
from
+124 −0
Open

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

Show file tree
Hide file tree
Changes from 4 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
81e0ac6
fuse attn_k + attn_v into single MMVQ dispatch via weight concatenation
jeffli-xilinx Jul 17, 2026
41b77c4
gate KV weight fusion on RDNA 3.5 devices only
jeffli-xilinx Jul 17, 2026
437cc65
narrow KV fusion gate to gfx1151 (Strix Halo) only
jeffli-xilinx Jul 17, 2026
be84ef4
address review: add buffer guard, log, and single-GPU comment
jeffli-xilinx Jul 20, 2026
65a3422
replace gfx1151 CC gate with backend-agnostic VRAM check
jeffli-xilinx Jul 21, 2026
4e9de0f
change K+V fusion to Q+K fusion, remove VRAM gate
jeffli-xilinx Jul 21, 2026
5162a2b
revert Q+K fusion back to K+V fusion
jeffli-xilinx Jul 21, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions ggml/src/ggml-cuda/ggml-cuda.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down Expand Up @@ -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;
}

Expand Down
56 changes: 56 additions & 0 deletions src/llama-graph.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
59 changes: 59 additions & 0 deletions src/llama-model.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
};

Expand Down Expand Up @@ -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);
}
Comment on lines +1654 to +1659

@roberteg16 roberteg16 Jul 21, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

AFAIK, llama.cpp folks want the graph level code backend agnostic

@jeffli-xilinx jeffli-xilinx Jul 21, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good point. The graph-level code (llama-graph.cpp) is backend-agnostic — the wkv_concat branch emits standard ggml_mul_mat + ggml_view ops, no CUDA-specific calls. Any backend that handles MUL_MAT benefits from the reduced dispatch count.

The backend-specific part is limited to llama-model.cpp: the CC detection via ggml_backend_reg_get_proc_address("ggml_backend_cuda_get_device_cc") that gates whether to create wkv_concat at load time. This keeps the CUDA/HIP coupling to a runtime proc-address lookup rather than a compile-time dependency — same pattern used elsewhere in the codebase for backend capability probing.

If this moves upstream, the CC gate could be generalized (e.g., a backend hint for "supports fused KV") to avoid the CUDA-specific string.

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;
Comment thread
jeffli-xilinx marked this conversation as resolved.
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;
}

Expand Down
1 change: 1 addition & 0 deletions src/llama-model.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down