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

+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;
}