Lumpiasty/llama.cpp
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

llama: refactor fused ops (#24646)

Browse Source
This commit is contained in:
Aman Gupta
2026-07-08 18:18:09 +08:00
committed by GitHub
parent bbebeec4a8
commit 90e0f5cfcb
6 changed files with 122 additions and 131 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/llama-context.cpp
+89 -122
View File
@@ -17,6 +17,7 @@
#include <cstring>
#include <limits>
#include <stdexcept>
#include <string>
//
// llama_context
@@ -30,6 +31,30 @@ static llm_graph_type ctx_type_to_graph_type(llama_context_type ctx_type) {
throw std::runtime_error("Unsupported ctx type");
}
struct llm_fused_op_probe {
llm_fused_op op;
const char * name;
uint32_t n_tokens_per_seq;
};
static const llm_fused_op_probe llm_fused_op_flash_attn_probe = {
/*.op =*/ LLM_FUSED_OP_FLASH_ATTN,
/*.name =*/ "Flash Attention",
/*.n_tokens_per_seq =*/ 1,
};
static const llm_fused_op_probe llm_fused_op_gdn_ar_probe = {
/*.op =*/ LLM_FUSED_OP_GDN_AR,
/*.name =*/ "fused Gated Delta Net (autoregressive)",
/*.n_tokens_per_seq =*/ 1,
};
static const llm_fused_op_probe llm_fused_op_gdn_ch_probe = {
/*.op =*/ LLM_FUSED_OP_GDN_CH,
/*.name =*/ "fused Gated Delta Net (chunked)",
/*.n_tokens_per_seq =*/ 16,
};
llama_context::llama_context(
const llama_model & model,
llama_context_params params) :
@@ -436,6 +461,69 @@ llama_context::~llama_context() {
ggml_opt_free(opt_ctx);
}
void llama_context::resolve_fused_ops(const llama_memory_context_i * mctx, uint32_t n_seqs) {
const char * func = __func__;
auto resolve = [&](const llm_fused_op_probe & probe, bool & enabled) {
if (!enabled) {
return;
}
const uint32_t n_tokens_probe = probe.n_tokens_per_seq*n_seqs;
auto * gf = graph_reserve(n_tokens_probe, n_seqs, n_tokens_probe, mctx, true);
if (!gf) {
throw std::runtime_error(std::string("failed to reserve graph for ") + probe.name + " check");
}
bool device_mismatch = false;
for (const auto & node : get_gf_res_reserve()->get_fused_nodes()) {
if (node.op != probe.op) {
continue;
}
GGML_ASSERT(node.il >= 0);
ggml_backend_t backend_fused = ggml_backend_sched_get_tensor_backend(sched.get(), node.tensor);
ggml_backend_dev_t device_fused = backend_fused ? ggml_backend_get_device(backend_fused) : nullptr;
// TODO: make this descriptor-specific; model.dev_layer() preserves the current behavior,
// but is still wrong for cases like --no-kv-offload.
ggml_backend_dev_t device_layer = model.dev_layer(node.il);
if (device_fused != device_layer) {
LLAMA_LOG_WARN("%s: layer %d is assigned to device %s but %s "
"is assigned to device %s (usually due to missing support)\n",
func, node.il,
device_layer ? ggml_backend_dev_name(device_layer) : "none",
probe.name,
device_fused ? ggml_backend_dev_name(device_fused) : "none");
device_mismatch = true;
break;
}
}
if (device_mismatch) {
enabled = false;
LLAMA_LOG_WARN("%s: %s not supported, set to disabled\n", func, probe.name);
} else {
enabled = true;
LLAMA_LOG_INFO("%s: %s enabled\n", func, probe.name);
}
};
if (cparams.auto_fa) {
resolve(llm_fused_op_flash_attn_probe, cparams.flash_attn);
cparams.auto_fa = false;
}
if (cparams.auto_fgdn) {
LLAMA_LOG_INFO("%s: resolving fused Gated Delta Net support:\n", func);
resolve(llm_fused_op_gdn_ar_probe, cparams.fused_gdn_ar);
resolve(llm_fused_op_gdn_ch_probe, cparams.fused_gdn_ch);
cparams.auto_fgdn = false;
}
}
void llama_context::sched_reserve() {
if (!sched_need_reserve) {
return;
@@ -475,128 +563,7 @@ void llama_context::sched_reserve() {
LLAMA_LOG_DEBUG("%s: worst-case: n_tokens = %d, n_seqs = %d, n_outputs = %d\n", __func__, n_tokens, n_seqs, n_outputs);
// resolve automatic Flash Attention use
if (cparams.auto_fa) {
auto * gf = graph_reserve(1, n_seqs, n_outputs, mctx.get(), true);
if (!gf) {
throw std::runtime_error("failed to reserve graph for Flash Attention check");
}
const size_t prefix_len = strlen(LLAMA_TENSOR_NAME_FATTN) + 1;
bool fa_device_mismatch = false;
for (int i = 0; i < ggml_graph_n_nodes(gf); i++) {
ggml_tensor * n = ggml_graph_node(gf, i);
if (n->op != GGML_OP_FLASH_ATTN_EXT) {
continue;
}
ggml_backend_dev_t device_fa = ggml_backend_get_device(ggml_backend_sched_get_tensor_backend(sched.get(), n));
// TODO: instead of the tensor names, use a map to keep track of which (FA) tensors belong to which layer
GGML_ASSERT(strncmp(n->name, LLAMA_TENSOR_NAME_FATTN "-", prefix_len) == 0);
const int il = std::stoi(n->name + prefix_len);
ggml_backend_dev_t device_kv = model.dev_layer(il);
if (device_fa != device_kv) {
LLAMA_LOG_WARN("%s: layer %d is assigned to device %s but the Flash Attention tensor "
"is assigned to device %s (usually due to missing support)\n",
__func__, il, ggml_backend_dev_name(device_kv), ggml_backend_dev_name(device_fa));
// FIXME: fa_device_mismatch logic is wrong for --no-kv-offload, but this is broken anyways
fa_device_mismatch = true;
break;
}
}
if (fa_device_mismatch) {
cparams.flash_attn = false;
LLAMA_LOG_WARN("%s: Flash Attention was auto, set to disabled\n", __func__);
} else {
cparams.flash_attn = true;
LLAMA_LOG_INFO("%s: Flash Attention was auto, set to enabled\n", __func__);
}
cparams.auto_fa = false;
}
if (cparams.auto_fgdn) {
LLAMA_LOG_INFO("%s: resolving fused Gated Delta Net support:\n", __func__);
if (cparams.fused_gdn_ar) {
auto * gf = graph_reserve(1, n_seqs, n_outputs, mctx.get(), true);
if (!gf) {
throw std::runtime_error("failed to reserve graph for fused Gated Delta Net check (autoregressive)");
}
const size_t prefix_len = strlen(LLAMA_TENSOR_NAME_FGDN_AR) + 1;
bool gdn_device_mismatch = false;
for (int i = 0; i < ggml_graph_n_nodes(gf); i++) {
ggml_tensor * n = ggml_graph_node(gf, i);
if (n->op != GGML_OP_GATED_DELTA_NET) {
continue;
}
ggml_backend_dev_t device_gdn = ggml_backend_get_device(ggml_backend_sched_get_tensor_backend(sched.get(), n));
GGML_ASSERT(strncmp(n->name, LLAMA_TENSOR_NAME_FGDN_AR "-", prefix_len) == 0);
const int il = std::stoi(n->name + prefix_len);
ggml_backend_dev_t device_kv = model.dev_layer(il);
if (device_gdn != device_kv) {
LLAMA_LOG_WARN("%s: layer %d is assigned to device %s but the fused Gated Delta Net tensor "
"is assigned to device %s (usually due to missing support)\n",
__func__, il, ggml_backend_dev_name(device_kv), ggml_backend_dev_name(device_gdn));
gdn_device_mismatch = true;
break;
}
}
if (gdn_device_mismatch) {
cparams.fused_gdn_ar = false;
LLAMA_LOG_WARN("%s: fused Gated Delta Net (autoregressive) not supported, set to disabled\n", __func__);
} else {
LLAMA_LOG_INFO("%s: fused Gated Delta Net (autoregressive) enabled\n", __func__);
}
}
if (cparams.fused_gdn_ch) {
// more than one token in the batch per sequence in order to take the chunked path
// note: n_outputs must match n_tokens for embedding models with mean/rank pooling,
// because build_pooling creates inp_mean with shape [n_tokens, n_seqs] and multiplies
// it with t_embd which is reduced to [n_outputs, ...] via out_ids. if n_outputs != n_tokens,
// the ggml_mul_mat assertion fails.
const uint32_t n_tokens_ch = 16*n_seqs;
auto * gf = graph_reserve(n_tokens_ch, n_seqs, n_tokens_ch, mctx.get(), true);
if (!gf) {
throw std::runtime_error("failed to reserve graph for fused Gated Delta Net check (chunked)");
}
const size_t prefix_len = strlen(LLAMA_TENSOR_NAME_FGDN_CH) + 1;
bool gdn_device_mismatch = false;
for (int i = 0; i < ggml_graph_n_nodes(gf); i++) {
ggml_tensor * n = ggml_graph_node(gf, i);
if (n->op != GGML_OP_GATED_DELTA_NET) {
continue;
}
ggml_backend_dev_t device_gdn = ggml_backend_get_device(ggml_backend_sched_get_tensor_backend(sched.get(), n));
GGML_ASSERT(strncmp(n->name, LLAMA_TENSOR_NAME_FGDN_CH "-", prefix_len) == 0);
const int il = std::stoi(n->name + prefix_len);
ggml_backend_dev_t device_kv = model.dev_layer(il);
if (device_gdn != device_kv) {
LLAMA_LOG_WARN("%s: layer %d is assigned to device %s but the fused Gated Delta Net tensor "
"is assigned to device %s (usually due to missing support)\n",
__func__, il, ggml_backend_dev_name(device_kv), ggml_backend_dev_name(device_gdn));
gdn_device_mismatch = true;
break;
}
}
if (gdn_device_mismatch) {
cparams.fused_gdn_ch = false;
LLAMA_LOG_WARN("%s: fused Gated Delta Net (chunked) not supported, set to disabled\n", __func__);
} else {
LLAMA_LOG_INFO("%s: fused Gated Delta Net (chunked) enabled\n", __func__);
}
}
cparams.auto_fgdn = false;
}
resolve_fused_ops(mctx.get(), n_seqs);
// reserve worst-case graph
int n_splits_pp = -1;