mirror of
https://github.com/leejet/stable-diffusion.cpp.git
synced 2026-11-06 08:53:07 +08:00
feat: add configurable Qwen cache types and early cache scheduling (#2045)
This commit is contained in:
@@ -21,11 +21,12 @@ ggml_tensor* ggml_ext_attention_ext(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* mask,
|
||||
bool skip_reshape,
|
||||
bool flash_attn,
|
||||
float kv_scale) {
|
||||
float kv_scale,
|
||||
bool* used_flash_attn) {
|
||||
if (ctx->attn_scale > 0.f) {
|
||||
kv_scale = ctx->attn_scale;
|
||||
}
|
||||
return ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, n_head, mask, skip_reshape, flash_attn, kv_scale, ctx->sage_attn_enabled);
|
||||
return ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, n_head, mask, skip_reshape, flash_attn, kv_scale, ctx->sage_attn_enabled, used_flash_attn);
|
||||
}
|
||||
|
||||
void GGMLRunner::alloc_params_ctx() {
|
||||
@@ -515,9 +516,10 @@ GGMLRunner::~GGMLRunner() {
|
||||
free_params_ctx();
|
||||
}
|
||||
|
||||
GGMLRunnerContext GGMLRunner::get_context() {
|
||||
GGMLRunnerContext GGMLRunner::get_context(ggml_cgraph* graph) {
|
||||
GGMLRunnerContext runner_ctx;
|
||||
runner_ctx.ggml_ctx = compute_ctx;
|
||||
runner_ctx.graph = graph;
|
||||
runner_ctx.backend = runtime_backend;
|
||||
runner_ctx.flash_attn_enabled = flash_attn_enabled;
|
||||
runner_ctx.sage_attn_enabled = sage_attn_enabled;
|
||||
@@ -532,8 +534,8 @@ GGMLRunnerContext GGMLRunner::get_context() {
|
||||
runner_ctx.get_cache_tensor = [this](const std::string& name) {
|
||||
return this->get_cache_tensor_by_name(name);
|
||||
};
|
||||
runner_ctx.cache_tensor = [this](const std::string& name, ggml_tensor* tensor) {
|
||||
this->cache(name, tensor);
|
||||
runner_ctx.cache_tensor = [this, graph](const std::string& name, ggml_tensor* tensor) {
|
||||
this->cache(name, tensor, graph);
|
||||
};
|
||||
runner_ctx.set_backend_tensor_data = [this](ggml_tensor* tensor, const void* data) {
|
||||
this->set_backend_tensor_data(tensor, data);
|
||||
@@ -575,7 +577,7 @@ ggml_tensor* GGMLRunner::to_backend(ggml_tensor* tensor) {
|
||||
}
|
||||
}
|
||||
|
||||
void GGMLRunner::cache(const std::string name, ggml_tensor* tensor) {
|
||||
void GGMLRunner::cache(const std::string name, ggml_tensor* tensor, ggml_cgraph* graph) {
|
||||
if (tensor != nullptr && tensor->view_src != nullptr) {
|
||||
tensor = ggml_cont(compute_ctx, tensor);
|
||||
}
|
||||
@@ -583,6 +585,10 @@ void GGMLRunner::cache(const std::string name, ggml_tensor* tensor) {
|
||||
ggml_set_output(tensor);
|
||||
}
|
||||
cache_.stage(name, tensor);
|
||||
if (graph != nullptr && tensor != nullptr) {
|
||||
// Schedule the cache output here so its source can be reused before graph end.
|
||||
ggml_build_forward_expand(graph, tensor);
|
||||
}
|
||||
}
|
||||
|
||||
std::optional<sd::Tensor<float>> GGMLRunner::compute(get_graph_cb_t get_graph,
|
||||
|
||||
Reference in New Issue
Block a user