feat: add linear and attention scale overrides (#1964)

This commit is contained in:
leejet authored and GitHub committed 2026-09-12 01:41:28 +08:00
1 parent 5ebce93342
commit 7f410a3793
39 files changed
+272 -46

No files matched your search

+18
View File
@@ -2,6 +2,7 @@
#include <map>
#include <utility>
#include "core/ggml_extend.h"
#include "core/ggml_extend_backend.h"
#include "core/ggml_runner.h"
#include "core/ggml_tensor_utils.h"
@@ -11,6 +12,21 @@
using namespace sd;
ggml_tensor* ggml_ext_attention_ext(GGMLRunnerContext* ctx,
ggml_tensor* q,
ggml_tensor* k,
ggml_tensor* v,
int64_t n_head,
ggml_tensor* mask,
bool skip_reshape,
bool flash_attn,
float kv_scale) {
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);
}
void GGMLRunner::alloc_params_ctx() {
ggml_init_params params;
params.mem_size = static_cast<size_t>(MAX_PARAMS_TENSOR_NUM * ggml_tensor_overhead());
@@ -510,6 +526,8 @@ GGMLRunnerContext GGMLRunner::get_context() {
runner_ctx.ggml_ctx = compute_ctx;
runner_ctx.backend = runtime_backend;
runner_ctx.flash_attn_enabled = flash_attn_enabled;
runner_ctx.linear_scale = linear_scale;
runner_ctx.attn_scale = attn_scale;
runner_ctx.conv2d_direct_enabled = conv2d_direct_enabled;
runner_ctx.circular_x_enabled = circular_x_enabled;
runner_ctx.circular_y_enabled = circular_y_enabled;
+19
View File
@@ -68,6 +68,8 @@ struct GGMLRunnerContext {
ggml_backend_t backend = nullptr;
ggml_context* ggml_ctx = nullptr;
bool flash_attn_enabled = false;
float linear_scale = 0.f;
float attn_scale = 0.f;
bool conv2d_direct_enabled = false;
bool circular_x_enabled = false;
bool circular_y_enabled = false;
@@ -113,6 +115,16 @@ struct GGMLRunnerContext {
}
};
ggml_tensor* ggml_ext_attention_ext(GGMLRunnerContext* ctx,
ggml_tensor* q,
ggml_tensor* k,
ggml_tensor* v,
int64_t n_head,
ggml_tensor* mask = nullptr,
bool skip_reshape = false,
bool flash_attn = false,
float kv_scale = 1.f);
struct GGMLRunner {
private:
std::map<ggml_backend_t, size_t> logged_compute_bytes_;
@@ -163,6 +175,8 @@ protected:
const std::string final_result_name = "ggml_runner_final_result_tensor";
bool flash_attn_enabled = false;
float linear_scale = 0.f;
float attn_scale = 0.f;
bool conv2d_direct_enabled = false;
bool circular_x_enabled = false;
bool circular_y_enabled = false;
@@ -323,6 +337,11 @@ public:
flash_attn_enabled = enabled;
}
void set_scale_overrides(float linear_scale, float attn_scale) {
this->linear_scale = linear_scale;
this->attn_scale = attn_scale;
}
void set_conv2d_direct_enabled(bool enabled) {
conv2d_direct_enabled = enabled;
}