furyhawk/llama.cpp
1
0
Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-07-21 10:15:53 +00:00
Code Issues 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;
src/llama-context.h
+4
View File
@@ -262,6 +262,10 @@ private:
llm_graph_cb graph_get_cb() const;
// disable auto fused ops (Flash Attention, Gated Delta Net) whose op lands on a device
// that differs from the layer it belongs to (usually due to missing backend support)
void resolve_fused_ops(const llama_memory_context_i * mctx, uint32_t n_seqs);
// TODO: read/write lora adapters and cvec
size_t state_write_data(llama_io_write_i & io);
size_t state_read_data (llama_io_read_i & io);
src/llama-graph.cpp
+8 -1
View File
@@ -1192,6 +1192,7 @@ void llm_graph_result::reset() {
params = {};
inputs.clear();
fused_nodes.clear();
buf_compute_meta.resize(ggml_tensor_overhead()*max_nodes + ggml_graph_overhead_custom(max_nodes, false));
@@ -1293,6 +1294,10 @@ llm_graph_input_i * llm_graph_result::add_input(llm_graph_input_ptr input) {
return inputs.back().get();
}
void llm_graph_result::add_fused_node(llm_graph_fused_node result) {
fused_nodes.push_back(result);
}
void llm_graph_result::set_params(const llm_graph_params & params) {
this->params = params;
}
@@ -1352,6 +1357,8 @@ void llm_graph_context::cb(ggml_tensor * cur, const char * name, int il) const {
}
}
ggml_tensor * llm_graph_context::build_cvec(
ggml_tensor * cur,
int il) const {
@@ -2402,7 +2409,7 @@ ggml_tensor * llm_graph_context::build_attn_mha(
cur = ggml_flash_attn_ext(ctx0, q, k, v, kq_mask, kq_scale, hparams.f_max_alibi_bias,
hparams.attn_soft_cap ? hparams.f_attn_logit_softcapping : 0.0f);
cb(cur, LLAMA_TENSOR_NAME_FATTN, il);
res->add_fused_node({LLM_FUSED_OP_FLASH_ATTN, cur, il});
ggml_flash_attn_ext_add_sinks(cur, sinks);
ggml_flash_attn_ext_set_prec (cur, GGML_PREC_F32);
src/llama-graph.h
+17
View File
@@ -38,6 +38,12 @@ enum llm_graph_type {
LLM_GRAPH_TYPE_DECODER_MTP,
};
enum llm_fused_op {
LLM_FUSED_OP_FLASH_ATTN,
LLM_FUSED_OP_GDN_AR,
LLM_FUSED_OP_GDN_CH,
};
enum llm_ffn_op_type : int {
LLM_FFN_NONE = 0, // sentinel: unset; archs must assign before use
LLM_FFN_SILU,
@@ -775,6 +781,12 @@ struct llm_graph_params {
}
};
struct llm_graph_fused_node {
llm_fused_op op;
ggml_tensor * tensor;
int il;
};
class llm_graph_result {
public:
llm_graph_result(int64_t max_nodes);
@@ -808,6 +820,10 @@ public:
llm_graph_input_i * add_input(llm_graph_input_ptr input);
void add_fused_node(llm_graph_fused_node result);
const std::vector<llm_graph_fused_node> & get_fused_nodes() const { return fused_nodes; }
void set_params(const llm_graph_params & params);
// important graph nodes
@@ -826,6 +842,7 @@ public:
std::map<llama_seq_id, ggml_tensor *> t_sampled_probs;
std::vector<llm_graph_input_ptr> inputs;
std::vector<llm_graph_fused_node> fused_nodes;
ggml_context_ptr ctx_compute;
src/llama-impl.h
-4
View File
@@ -103,7 +103,3 @@ std::string llama_format_tensor_shape(const std::vector<int64_t> & ne);
std::string llama_format_tensor_shape(const struct ggml_tensor * t);
std::string gguf_kv_to_str(const struct gguf_context * ctx_gguf, int i);
#define LLAMA_TENSOR_NAME_FATTN "__fattn__"
#define LLAMA_TENSOR_NAME_FGDN_AR "__fgdn_ar__"
#define LLAMA_TENSOR_NAME_FGDN_CH "__fgdn_ch__"
src/models/delta-net-base.cpp
+4 -4
View File
@@ -401,9 +401,9 @@ std::pair<ggml_tensor *, ggml_tensor *> llm_build_delta_net_base::build_delta_ne
// K=1: output carries the final state only. state s is 4D [S_v, S_v, H_v, n_seqs].
ggml_tensor * result = ggml_gated_delta_net(ctx0, q, k, v, g, b, s, /*K=*/1);
if (n_tokens == 1) {
cb(result, LLAMA_TENSOR_NAME_FGDN_AR, il);
res->add_fused_node({LLM_FUSED_OP_GDN_AR, result, il});
} else {
cb(result, LLAMA_TENSOR_NAME_FGDN_CH, il);
res->add_fused_node({LLM_FUSED_OP_GDN_CH, result, il});
}
ggml_tensor * output = ggml_view_4d(ctx0, result,
@@ -566,9 +566,9 @@ ggml_tensor * llm_build_delta_net_base::build_recurrent_attn(
// state s is 4D [S_v, S_v, H_v, n_seqs]; K snapshot slots are written into the output.
ggml_tensor * gdn_out = ggml_gated_delta_net(ctx0, q, k, v, g, b, s, K);
if (n_seq_tokens > 1) {
cb(gdn_out, LLAMA_TENSOR_NAME_FGDN_CH, il);
res->add_fused_node({LLM_FUSED_OP_GDN_CH, gdn_out, il});
} else {
cb(gdn_out, LLAMA_TENSOR_NAME_FGDN_AR, il);
res->add_fused_node({LLM_FUSED_OP_GDN_AR, gdn_out, il});
}
const int64_t attn_score_elems = S_v * H_v * n_seq_tokens * n_seqs;