Commit 77c1f3992 for llama.cpp
commit 77c1f39929d0cae9488d846fff0925cd0f1738c8
Author: lhez <lih@qti.qualcomm.com>
Date: Thu Oct 15 00:11:51 2026 -0700
opencl: apply flash attn softcap before mask (#30546)
diff --git a/ggml/src/ggml-opencl/kernels/flash_attn_f16.cl b/ggml/src/ggml-opencl/kernels/flash_attn_f16.cl
index f9797d346..cb6fd8930 100644
--- a/ggml/src/ggml-opencl/kernels/flash_attn_f16.cl
+++ b/ggml/src/ggml-opencl/kernels/flash_attn_f16.cl
@@ -175,6 +175,12 @@ __kernel void flash_attn_f16(
ACC_TYPE s1 = (dot_acc1.s0 + dot_acc1.s1 + dot_acc1.s2 + dot_acc1.s3) * scale;
ACC_TYPE s2 = (dot_acc2.s0 + dot_acc2.s1 + dot_acc2.s2 + dot_acc2.s3) * scale;
ACC_TYPE s3 = (dot_acc3.s0 + dot_acc3.s1 + dot_acc3.s2 + dot_acc3.s3) * scale;
+ if (logit_softcap > 0.0f) {
+ s0 = logit_softcap * tanh(s0 / logit_softcap);
+ s1 = logit_softcap * tanh(s1 / logit_softcap);
+ s2 = logit_softcap * tanh(s2 / logit_softcap);
+ s3 = logit_softcap * tanh(s3 / logit_softcap);
+ }
if (is_causal) {
const int causal_limit = n_kv - n_q + my_query_row;
@@ -196,13 +202,6 @@ __kernel void flash_attn_f16(
if (k_row3 < n_kv) s3 += slope * (ACC_TYPE)mask_ptr[k_row3];
}
- if (logit_softcap > 0.0f) {
- s0 = logit_softcap * tanh(s0 / logit_softcap);
- s1 = logit_softcap * tanh(s1 / logit_softcap);
- s2 = logit_softcap * tanh(s2 / logit_softcap);
- s3 = logit_softcap * tanh(s3 / logit_softcap);
- }
-
const ACC_TYPE m_new = max(m_i, max(max(s0, s1), max(s2, s3)));
const ACC_TYPE scale_prev = native_exp(m_i - m_new);
const ACC_TYPE p0 = native_exp(s0 - m_new);
@@ -331,13 +330,13 @@ __kernel void flash_attn_f16_q1(
dot_acc = mad(q_priv[k], CONVERT_ACC4(k_ptr[k]), dot_acc);
}
ACC_TYPE score = (dot_acc.s0 + dot_acc.s1 + dot_acc.s2 + dot_acc.s3) * scale;
+ if (logit_softcap > 0.0f) {
+ score = logit_softcap * tanh(score / logit_softcap);
+ }
if (mask_base != NULL) {
const global DATA_TYPE* mask_ptr = (const global DATA_TYPE*)(mask_base);
score += slope * (ACC_TYPE)mask_ptr[k_idx];
}
- if (logit_softcap > 0.0f) {
- score = logit_softcap * tanh(score / logit_softcap);
- }
m_i = max(m_i, score);
}
@@ -367,13 +366,13 @@ __kernel void flash_attn_f16_q1(
dot_acc = mad(q_priv[k], CONVERT_ACC4(k_ptr[k]), dot_acc);
}
ACC_TYPE score = (dot_acc.s0 + dot_acc.s1 + dot_acc.s2 + dot_acc.s3) * scale;
+ if (logit_softcap > 0.0f) {
+ score = logit_softcap * tanh(score / logit_softcap);
+ }
if (mask_base != NULL) {
const global DATA_TYPE* mask_ptr = (const global DATA_TYPE*)(mask_base);
score += slope * (ACC_TYPE)mask_ptr[k_idx];
}
- if (logit_softcap > 0.0f) {
- score = logit_softcap * tanh(score / logit_softcap);
- }
const ACC_TYPE p = exp(score - m_final);
l_i += p;
FA_UNROLL
diff --git a/ggml/src/ggml-opencl/kernels/flash_attn_f32.cl b/ggml/src/ggml-opencl/kernels/flash_attn_f32.cl
index 5911524e1..15b9cf30a 100644
--- a/ggml/src/ggml-opencl/kernels/flash_attn_f32.cl
+++ b/ggml/src/ggml-opencl/kernels/flash_attn_f32.cl
@@ -176,6 +176,12 @@ __kernel void flash_attn_f32(
ACC_TYPE s1 = (dot_acc1.s0 + dot_acc1.s1 + dot_acc1.s2 + dot_acc1.s3) * scale;
ACC_TYPE s2 = (dot_acc2.s0 + dot_acc2.s1 + dot_acc2.s2 + dot_acc2.s3) * scale;
ACC_TYPE s3 = (dot_acc3.s0 + dot_acc3.s1 + dot_acc3.s2 + dot_acc3.s3) * scale;
+ if (logit_softcap > 0.0f) {
+ s0 = logit_softcap * tanh(s0 / logit_softcap);
+ s1 = logit_softcap * tanh(s1 / logit_softcap);
+ s2 = logit_softcap * tanh(s2 / logit_softcap);
+ s3 = logit_softcap * tanh(s3 / logit_softcap);
+ }
if (is_causal) {
const int causal_limit = n_kv - n_q + my_query_row;
@@ -197,13 +203,6 @@ __kernel void flash_attn_f32(
if (k_row3 < n_kv) s3 += slope * (ACC_TYPE)mask_ptr[k_row3];
}
- if (logit_softcap > 0.0f) {
- s0 = logit_softcap * tanh(s0 / logit_softcap);
- s1 = logit_softcap * tanh(s1 / logit_softcap);
- s2 = logit_softcap * tanh(s2 / logit_softcap);
- s3 = logit_softcap * tanh(s3 / logit_softcap);
- }
-
const ACC_TYPE m_new = max(m_i, max(max(s0, s1), max(s2, s3)));
const ACC_TYPE scale_prev = native_exp(m_i - m_new);
const ACC_TYPE p0 = native_exp(s0 - m_new);
@@ -332,13 +331,13 @@ __kernel void flash_attn_f32_q1(
dot_acc = mad(q_priv[k], CONVERT_ACC4(k_ptr[k]), dot_acc);
}
ACC_TYPE score = (dot_acc.s0 + dot_acc.s1 + dot_acc.s2 + dot_acc.s3) * scale;
+ if (logit_softcap > 0.0f) {
+ score = logit_softcap * tanh(score / logit_softcap);
+ }
if (mask_base != NULL) {
const global MASK_DATA_TYPE* mask_ptr = (const global MASK_DATA_TYPE*)(mask_base);
score += slope * (ACC_TYPE)mask_ptr[k_idx];
}
- if (logit_softcap > 0.0f) {
- score = logit_softcap * tanh(score / logit_softcap);
- }
m_i = max(m_i, score);
}
@@ -368,13 +367,13 @@ __kernel void flash_attn_f32_q1(
dot_acc = mad(q_priv[k], CONVERT_ACC4(k_ptr[k]), dot_acc);
}
ACC_TYPE score = (dot_acc.s0 + dot_acc.s1 + dot_acc.s2 + dot_acc.s3) * scale;
+ if (logit_softcap > 0.0f) {
+ score = logit_softcap * tanh(score / logit_softcap);
+ }
if (mask_base != NULL) {
const global MASK_DATA_TYPE* mask_ptr = (const global MASK_DATA_TYPE*)(mask_base);
score += slope * (ACC_TYPE)mask_ptr[k_idx];
}
- if (logit_softcap > 0.0f) {
- score = logit_softcap * tanh(score / logit_softcap);
- }
const ACC_TYPE p = exp(score - m_final);
l_i += p;
FA_UNROLL
diff --git a/ggml/src/ggml-opencl/kernels/flash_attn_f32_f16.cl b/ggml/src/ggml-opencl/kernels/flash_attn_f32_f16.cl
index 27553c729..d1cc14801 100644
--- a/ggml/src/ggml-opencl/kernels/flash_attn_f32_f16.cl
+++ b/ggml/src/ggml-opencl/kernels/flash_attn_f32_f16.cl
@@ -336,6 +336,10 @@ __kernel void FA_TILE_NAME(
ACC_TYPE score0 = partial0 * scale;
ACC_TYPE score1 = partial1 * scale;
+ if (logit_softcap > 0.0f) {
+ score0 = logit_softcap * tanh(score0 / logit_softcap);
+ score1 = logit_softcap * tanh(score1 / logit_softcap);
+ }
if (!query_valid) { score0 = FA_M_INIT; score1 = FA_M_INIT; }
if (is_causal) {
@@ -359,11 +363,6 @@ __kernel void FA_TILE_NAME(
}
}
- if (logit_softcap > 0.0f) {
- score0 = logit_softcap * tanh(score0 / logit_softcap);
- score1 = logit_softcap * tanh(score1 / logit_softcap);
- }
-
const ACC_TYPE m_new = max(m_i, max(score0, score1));
// Whole tile masked (m_new == FA_M_INIT): force the exp() args
// far negative so the tile contributes 0, not exp(0)=1.
@@ -408,6 +407,9 @@ __kernel void FA_TILE_NAME(
score += local_partial[j][q_lane * N_SPLIT + s];
}
score *= scale;
+ if (logit_softcap > 0.0f) {
+ score = logit_softcap * tanh(score / logit_softcap);
+ }
if (is_causal && k_row > (n_kv - n_q + my_query_row)) score = FA_M_INIT;
if (k_row >= n_kv) score = FA_M_INIT;
@@ -424,10 +426,6 @@ __kernel void FA_TILE_NAME(
}
}
- if (logit_softcap > 0.0f) {
- score = logit_softcap * tanh(score / logit_softcap);
- }
-
m_new = max(m_new, score);
local_p[q_lane][j] = score;
}
@@ -502,6 +500,12 @@ __kernel void FA_TILE_NAME(
ACC_TYPE s1 = (dot_acc1.s0 + dot_acc1.s1 + dot_acc1.s2 + dot_acc1.s3) * scale;
ACC_TYPE s2 = (dot_acc2.s0 + dot_acc2.s1 + dot_acc2.s2 + dot_acc2.s3) * scale;
ACC_TYPE s3 = (dot_acc3.s0 + dot_acc3.s1 + dot_acc3.s2 + dot_acc3.s3) * scale;
+ if (logit_softcap > 0.0f) {
+ s0 = logit_softcap * tanh(s0 / logit_softcap);
+ s1 = logit_softcap * tanh(s1 / logit_softcap);
+ s2 = logit_softcap * tanh(s2 / logit_softcap);
+ s3 = logit_softcap * tanh(s3 / logit_softcap);
+ }
if (is_causal) {
const int causal_limit = n_kv - n_q + my_query_row;
@@ -531,13 +535,6 @@ __kernel void FA_TILE_NAME(
}
}
- if (logit_softcap > 0.0f) {
- s0 = logit_softcap * tanh(s0 / logit_softcap);
- s1 = logit_softcap * tanh(s1 / logit_softcap);
- s2 = logit_softcap * tanh(s2 / logit_softcap);
- s3 = logit_softcap * tanh(s3 / logit_softcap);
- }
-
const ACC_TYPE m_new = max(m_i, max(max(s0, s1), max(s2, s3)));
// Whole tile masked (m_new == FA_M_INIT): force the exp() args
// far negative so the tile contributes 0, not exp(0)=1.
@@ -748,13 +745,13 @@ __kernel void flash_attn_f32_f16_q1(
dot_acc = mad(q_shared[k], CONVERT_KV_ACC4(k_ptr[k]), dot_acc);
}
ACC_TYPE score = (dot_acc.s0 + dot_acc.s1 + dot_acc.s2 + dot_acc.s3) * scale;
+ if (logit_softcap > 0.0f) {
+ score = logit_softcap * tanh(score / logit_softcap);
+ }
if (mask_base != NULL) {
const global MASK_DATA_TYPE* mask_ptr = (const global MASK_DATA_TYPE*)(mask_base);
score += slope * (ACC_TYPE)mask_ptr[k_idx];
}
- if (logit_softcap > 0.0f) {
- score = logit_softcap * tanh(score / logit_softcap);
- }
m_i = max(m_i, score);
}
@@ -776,13 +773,13 @@ __kernel void flash_attn_f32_f16_q1(
dot_acc = mad(q_shared[k], CONVERT_KV_ACC4(k_ptr[k]), dot_acc);
}
ACC_TYPE score = (dot_acc.s0 + dot_acc.s1 + dot_acc.s2 + dot_acc.s3) * scale;
+ if (logit_softcap > 0.0f) {
+ score = logit_softcap * tanh(score / logit_softcap);
+ }
if (mask_base != NULL) {
const global MASK_DATA_TYPE* mask_ptr = (const global MASK_DATA_TYPE*)(mask_base);
score += slope * (ACC_TYPE)mask_ptr[k_idx];
}
- if (logit_softcap > 0.0f) {
- score = logit_softcap * tanh(score / logit_softcap);
- }
const ACC_TYPE p = exp(score - m_final);
l_i += p;
FA_UNROLL
@@ -1122,13 +1119,13 @@ __kernel void flash_attn_f32_f16_q1_local_tile(
if (tid == 0) {
float s = red[0] * scale;
+ if (logit_softcap > 0.0f) {
+ s = logit_softcap * tanh(s / logit_softcap);
+ }
if (mask_base != NULL) {
const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base;
s += slope * (float) mask_ptr[kv_idx];
}
- if (logit_softcap > 0.0f) {
- s = logit_softcap * tanh(s / logit_softcap);
- }
score_shared = s;
}
barrier(CLK_LOCAL_MEM_FENCE);
@@ -1323,13 +1320,13 @@ __kernel void flash_attn_f32_f16_q1_local_mq_split(
contrib += q_shared[h * DK + d] * (float) k_tile[j * DK + d];
}
float s = sub_group_reduce_add(contrib) * scale;
+ if (logit_softcap > 0.0f) {
+ s = logit_softcap * tanh(s / logit_softcap);
+ }
if (mask_base[h] != NULL) {
const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base[h];
s += slope[h] * (float) mask_ptr[kv_idx];
}
- if (logit_softcap > 0.0f) {
- s = logit_softcap * tanh(s / logit_softcap);
- }
score[h] = s;
}
@@ -1510,13 +1507,13 @@ __kernel void flash_attn_f32_f16_q1_vec_mq(
for (int h = 0; h < MQ_GQA; ++h) {
const ACC_TYPE dot_partial = dot4[h].s0 + dot4[h].s1 + dot4[h].s2 + dot4[h].s3;
ACC_TYPE s = sub_group_reduce_add(dot_partial) * scale;
+ if (logit_softcap > 0.0f) {
+ s = logit_softcap * tanh(s / logit_softcap);
+ }
if (mask_base[h] != NULL) {
const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base[h];
s += slope[h] * (ACC_TYPE) mask_ptr[k_idx];
}
- if (logit_softcap > 0.0f) {
- s = logit_softcap * tanh(s / logit_softcap);
- }
score[h] = s;
}
@@ -2448,13 +2445,13 @@ __kernel void flash_attn_f32_f16_q1_vec_mq_split_k_img(
for (int h = 0; h < MQ_GQA; ++h) {
const ACC_TYPE dot_partial = dot4[h].s0 + dot4[h].s1 + dot4[h].s2 + dot4[h].s3;
ACC_TYPE s = sub_group_reduce_add(dot_partial) * scale;
+ if (logit_softcap > 0.0f) {
+ s = logit_softcap * tanh(s / logit_softcap);
+ }
if (mask_base[h] != NULL) {
const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base[h];
s += slope[h] * (ACC_TYPE) mask_ptr[k_idx];
}
- if (logit_softcap > 0.0f) {
- s = logit_softcap * tanh(s / logit_softcap);
- }
score[h] = s;
}
@@ -2633,13 +2630,13 @@ __kernel void flash_attn_f32_f16_q1_split(
dot_acc = mad(q_shared[k], CONVERT_KV_ACC4(k_ptr[k]), dot_acc);
}
ACC_TYPE score = (dot_acc.s0 + dot_acc.s1 + dot_acc.s2 + dot_acc.s3) * scale;
+ if (logit_softcap > 0.0f) {
+ score = logit_softcap * tanh(score / logit_softcap);
+ }
if (mask_base != NULL) {
const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) (mask_base);
score += slope * (ACC_TYPE) mask_ptr[k_idx];
}
- if (logit_softcap > 0.0f) {
- score = logit_softcap * tanh(score / logit_softcap);
- }
m_i = max(m_i, score);
}
@@ -2662,13 +2659,13 @@ __kernel void flash_attn_f32_f16_q1_split(
dot_acc = mad(q_shared[k], CONVERT_KV_ACC4(k_ptr[k]), dot_acc);
}
ACC_TYPE score = (dot_acc.s0 + dot_acc.s1 + dot_acc.s2 + dot_acc.s3) * scale;
+ if (logit_softcap > 0.0f) {
+ score = logit_softcap * tanh(score / logit_softcap);
+ }
if (mask_base != NULL) {
const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) (mask_base);
score += slope * (ACC_TYPE) mask_ptr[k_idx];
}
- if (logit_softcap > 0.0f) {
- score = logit_softcap * tanh(score / logit_softcap);
- }
const ACC_TYPE p = exp(score - m_c);
l_i += p;
#pragma unroll
diff --git a/ggml/src/ggml-opencl/kernels/flash_attn_f32_q4_0.cl b/ggml/src/ggml-opencl/kernels/flash_attn_f32_q4_0.cl
index 48adba4f7..3e75dc942 100644
--- a/ggml/src/ggml-opencl/kernels/flash_attn_f32_q4_0.cl
+++ b/ggml/src/ggml-opencl/kernels/flash_attn_f32_q4_0.cl
@@ -287,14 +287,14 @@ __kernel void flash_attn_f32_q4_0_q1(
#endif
}
score *= scale;
+ if (logit_softcap > 0.0f) {
+ score = logit_softcap * tanh(score / logit_softcap);
+ }
if (mask_base != NULL) {
const global MASK_DATA_TYPE* mask_ptr = (const global MASK_DATA_TYPE*)(mask_base);
score += slope * (ACC_TYPE)mask_ptr[k_idx];
}
- if (logit_softcap > 0.0f) {
- score = logit_softcap * tanh(score / logit_softcap);
- }
// Online softmax step.
const ACC_TYPE m_new = max(m_i, score);
@@ -554,14 +554,14 @@ __kernel void flash_attn_f32_q4_0_q1_vec(
ACC_TYPE dot_partial = dot4.s0 + dot4.s1 + dot4.s2 + dot4.s3;
ACC_TYPE score = sub_group_reduce_add(dot_partial) * scale;
#endif
+ if (logit_softcap > 0.0f) {
+ score = logit_softcap * tanh(score / logit_softcap);
+ }
if (mask_base != NULL) {
const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base;
score += slope * (ACC_TYPE) mask_ptr[k_idx];
}
- if (logit_softcap > 0.0f) {
- score = logit_softcap * tanh(score / logit_softcap);
- }
const ACC_TYPE m_new = max(m_i, score);
const ACC_TYPE scale_prev = native_exp(m_i - m_new);
@@ -745,13 +745,13 @@ __kernel void flash_attn_f32_q4_0_q1_split(
#endif
}
score *= scale;
+ if (logit_softcap > 0.0f) {
+ score = logit_softcap * tanh(score / logit_softcap);
+ }
if (mask_base != NULL) {
const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) (mask_base);
score += slope * (ACC_TYPE) mask_ptr[k_idx];
}
- if (logit_softcap > 0.0f) {
- score = logit_softcap * tanh(score / logit_softcap);
- }
// Online softmax step.
const ACC_TYPE m_new = max(m_i, score);
@@ -1037,13 +1037,13 @@ __kernel void flash_attn_f32_q4_0_q1_vec_mq_split(
#pragma unroll
for (int h = 0; h < MQ_GQA; ++h) {
ACC_TYPE s = sub_group_reduce_add(lane_contrib[h]) * scale;
+ if (logit_softcap > 0.0f) {
+ s = logit_softcap * tanh(s / logit_softcap);
+ }
if (mask_base[h] != NULL) {
const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base[h];
s += slope[h] * (ACC_TYPE) mask_ptr[k_idx];
}
- if (logit_softcap > 0.0f) {
- s = logit_softcap * tanh(s / logit_softcap);
- }
score[h] = s;
}
#else
@@ -1067,13 +1067,13 @@ __kernel void flash_attn_f32_q4_0_q1_vec_mq_split(
for (int h = 0; h < MQ_GQA; ++h) {
const ACC_TYPE dot_partial = dot4[h].s0 + dot4[h].s1 + dot4[h].s2 + dot4[h].s3;
ACC_TYPE s = sub_group_reduce_add(dot_partial) * scale;
+ if (logit_softcap > 0.0f) {
+ s = logit_softcap * tanh(s / logit_softcap);
+ }
if (mask_base[h] != NULL) {
const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base[h];
s += slope[h] * (ACC_TYPE) mask_ptr[k_idx];
}
- if (logit_softcap > 0.0f) {
- s = logit_softcap * tanh(s / logit_softcap);
- }
score[h] = s;
}
#endif
@@ -1385,13 +1385,13 @@ __kernel void flash_attn_f32_q4_0_q1_vec_mq_split_c8(
s += sub_group_shuffle_xor(s, step);
}
s *= scale;
+ if (logit_softcap > 0.0f) {
+ s = logit_softcap * tanh(s / logit_softcap);
+ }
if (mask_base[h] != NULL) {
const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base[h];
s += slope[h] * (ACC_TYPE) mask_ptr[k_safe];
}
- if (logit_softcap > 0.0f) {
- s = logit_softcap * tanh(s / logit_softcap);
- }
score[h] = valid ? s : FA_M_INIT;
}
@@ -1852,9 +1852,15 @@ __kernel void flash_attn_f32_q4_0(
s2 = sub_group_shuffle(s2, tri_base + 0) + sub_group_shuffle(s2, tri_base + 1) + sub_group_shuffle(s2, tri_base + 2);
s3 = sub_group_shuffle(s3, tri_base + 0) + sub_group_shuffle(s3, tri_base + 1) + sub_group_shuffle(s3, tri_base + 2);
#endif
- if (!query_valid) { s0 = FA_M_INIT; s1 = FA_M_INIT; s2 = FA_M_INIT; s3 = FA_M_INIT; }
#endif
s0 *= scale; s1 *= scale; s2 *= scale; s3 *= scale;
+ if (logit_softcap > 0.0f) {
+ s0 = logit_softcap * tanh(s0 / logit_softcap);
+ s1 = logit_softcap * tanh(s1 / logit_softcap);
+ s2 = logit_softcap * tanh(s2 / logit_softcap);
+ s3 = logit_softcap * tanh(s3 / logit_softcap);
+ }
+ if (!query_valid) { s0 = FA_M_INIT; s1 = FA_M_INIT; s2 = FA_M_INIT; s3 = FA_M_INIT; }
if (is_causal) {
const int causal_limit = n_kv - n_q + my_query_row;
@@ -1876,12 +1882,6 @@ __kernel void flash_attn_f32_q4_0(
if (k_row2 < n_kv) s2 += slope * (ACC_TYPE) mask_ptr[k_row2];
if (k_row3 < n_kv) s3 += slope * (ACC_TYPE) mask_ptr[k_row3];
}
- if (logit_softcap > 0.0f) {
- s0 = logit_softcap * tanh(s0 / logit_softcap);
- s1 = logit_softcap * tanh(s1 / logit_softcap);
- s2 = logit_softcap * tanh(s2 / logit_softcap);
- s3 = logit_softcap * tanh(s3 / logit_softcap);
- }
const ACC_TYPE m_new = max(m_i, max(max(s0, s1), max(s2, s3)));
// Whole tile masked (m_new == FA_M_INIT): force the exp() args
diff --git a/ggml/src/ggml-opencl/kernels/flash_attn_f32_q8_0.cl b/ggml/src/ggml-opencl/kernels/flash_attn_f32_q8_0.cl
index f50912d21..75d0d9b0a 100644
--- a/ggml/src/ggml-opencl/kernels/flash_attn_f32_q8_0.cl
+++ b/ggml/src/ggml-opencl/kernels/flash_attn_f32_q8_0.cl
@@ -232,14 +232,14 @@ __kernel void flash_attn_f32_q8_0_q1(
#endif
}
score *= scale;
+ if (logit_softcap > 0.0f) {
+ score = logit_softcap * tanh(score / logit_softcap);
+ }
if (mask_base != NULL) {
const global MASK_DATA_TYPE* mask_ptr = (const global MASK_DATA_TYPE*)(mask_base);
score += slope * (ACC_TYPE)mask_ptr[k_idx];
}
- if (logit_softcap > 0.0f) {
- score = logit_softcap * tanh(score / logit_softcap);
- }
// Online softmax step.
const ACC_TYPE m_new = max(m_i, score);
@@ -437,14 +437,14 @@ __kernel void flash_attn_f32_q8_0_q1_vec(
}
ACC_TYPE dot_partial = dot4.s0 + dot4.s1 + dot4.s2 + dot4.s3;
ACC_TYPE score = sub_group_reduce_add(dot_partial) * scale;
+ if (logit_softcap > 0.0f) {
+ score = logit_softcap * tanh(score / logit_softcap);
+ }
if (mask_base != NULL) {
const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base;
score += slope * (ACC_TYPE) mask_ptr[k_idx];
}
- if (logit_softcap > 0.0f) {
- score = logit_softcap * tanh(score / logit_softcap);
- }
const ACC_TYPE m_new = max(m_i, score);
const ACC_TYPE scale_prev = native_exp(m_i - m_new);
@@ -627,13 +627,13 @@ __kernel void flash_attn_f32_q8_0_q1_split(
#endif
}
score *= scale;
+ if (logit_softcap > 0.0f) {
+ score = logit_softcap * tanh(score / logit_softcap);
+ }
if (mask_base != NULL) {
const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) (mask_base);
score += slope * (ACC_TYPE) mask_ptr[k_idx];
}
- if (logit_softcap > 0.0f) {
- score = logit_softcap * tanh(score / logit_softcap);
- }
// Online softmax step.
const ACC_TYPE m_new = max(m_i, score);
@@ -882,13 +882,13 @@ __kernel void flash_attn_f32_q8_0_q1_vec_mq_split(
for (int h = 0; h < MQ_GQA; ++h) {
const ACC_TYPE dot_partial = dot4[h].s0 + dot4[h].s1 + dot4[h].s2 + dot4[h].s3;
ACC_TYPE s = sub_group_reduce_add(dot_partial) * scale;
+ if (logit_softcap > 0.0f) {
+ s = logit_softcap * tanh(s / logit_softcap);
+ }
if (mask_base[h] != NULL) {
const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base[h];
s += slope[h] * (ACC_TYPE) mask_ptr[k_idx];
}
- if (logit_softcap > 0.0f) {
- s = logit_softcap * tanh(s / logit_softcap);
- }
score[h] = s;
}
@@ -1151,13 +1151,13 @@ __kernel void flash_attn_f32_q8_0_q1_vec_mq_split_c8(
s += sub_group_shuffle_xor(s, step);
}
s *= scale;
+ if (logit_softcap > 0.0f) {
+ s = logit_softcap * tanh(s / logit_softcap);
+ }
if (mask_base[h] != NULL) {
const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base[h];
s += slope[h] * (ACC_TYPE) mask_ptr[k_safe];
}
- if (logit_softcap > 0.0f) {
- s = logit_softcap * tanh(s / logit_softcap);
- }
score[h] = valid ? s : FA_M_INIT;
}
@@ -1652,9 +1652,15 @@ __kernel void flash_attn_f32_q8_0(
s2 = sub_group_shuffle(s2, tri_base + 0) + sub_group_shuffle(s2, tri_base + 1) + sub_group_shuffle(s2, tri_base + 2);
s3 = sub_group_shuffle(s3, tri_base + 0) + sub_group_shuffle(s3, tri_base + 1) + sub_group_shuffle(s3, tri_base + 2);
#endif
- if (!query_valid) { s0 = FA_M_INIT; s1 = FA_M_INIT; s2 = FA_M_INIT; s3 = FA_M_INIT; }
#endif
s0 *= scale; s1 *= scale; s2 *= scale; s3 *= scale;
+ if (logit_softcap > 0.0f) {
+ s0 = logit_softcap * tanh(s0 / logit_softcap);
+ s1 = logit_softcap * tanh(s1 / logit_softcap);
+ s2 = logit_softcap * tanh(s2 / logit_softcap);
+ s3 = logit_softcap * tanh(s3 / logit_softcap);
+ }
+ if (!query_valid) { s0 = FA_M_INIT; s1 = FA_M_INIT; s2 = FA_M_INIT; s3 = FA_M_INIT; }
if (is_causal) {
const int causal_limit = n_kv - n_q + my_query_row;
@@ -1676,12 +1682,6 @@ __kernel void flash_attn_f32_q8_0(
if (k_row2 < n_kv) s2 += slope * (ACC_TYPE) mask_ptr[k_row2];
if (k_row3 < n_kv) s3 += slope * (ACC_TYPE) mask_ptr[k_row3];
}
- if (logit_softcap > 0.0f) {
- s0 = logit_softcap * tanh(s0 / logit_softcap);
- s1 = logit_softcap * tanh(s1 / logit_softcap);
- s2 = logit_softcap * tanh(s2 / logit_softcap);
- s3 = logit_softcap * tanh(s3 / logit_softcap);
- }
const ACC_TYPE m_new = max(m_i, max(max(s0, s1), max(s2, s3)));
// Whole tile masked (m_new == FA_M_INIT): force the exp() args