mirror of
https://github.com/ikawrakow/ik_llama.cpp.git
synced 2026-02-25 23:54:10 +00:00
Experimenting with flash attention on Zen4
This version outperforms no-FA at 8k tokens, but then performance becomes the same at 16k and worse at 32k.
This commit is contained in:
Iwan Kawrakow
@@ -6770,6 +6770,19 @@ void iqk_flash_helper_3(int ne00,
|
||||
const void * mask, // mask. If not null, assumed to be fp16. nk elements
|
||||
float scale,
|
||||
float * qkv) {
|
||||
constexpr int q_step = 8;
|
||||
constexpr int k_step = 32; //16;
|
||||
//if (nq1%q_step != 0 || nk1%k_step != 0) printf("Oops(%s): nq1 = %d, nk1 = %d\n", __func__, nq1, nk1);
|
||||
if (nq1%q_step != 0 || nk1%k_step != 0) {
|
||||
for (int iq1 = 0; iq1 < nq1; ++iq1) {
|
||||
iqk_flash_helper_2(false, ne00, nk1, stride_k, stride_v,
|
||||
q, k, v, (const void *)((const char *)mask + iq1*stride_m),
|
||||
scale, 1.0f, nullptr, qkv);
|
||||
q += stride_q;
|
||||
qkv += stride_qkv;
|
||||
}
|
||||
return;
|
||||
}
|
||||
stride_q /= sizeof(float);
|
||||
// The following works
|
||||
//for (int iq1 = 0; iq1 < nq1; ++iq1) {
|
||||
@@ -6789,71 +6802,140 @@ void iqk_flash_helper_3(int ne00,
|
||||
// q += stride_q;
|
||||
// qkv += stride_qkv;
|
||||
//}
|
||||
float cache[256];
|
||||
float S[16], M[16];
|
||||
const ggml_half h_inf = GGML_FP32_TO_FP16(-INFINITY);
|
||||
float cache[q_step*k_step];
|
||||
float S[q_step], M[q_step];
|
||||
__m512 vk[8];
|
||||
for (int i1 = 0; i1 < nq1/16; ++i1) {
|
||||
for (int j = 0; j < 16; ++j) {
|
||||
auto R = qkv + (16*i1 + j)*stride_qkv;
|
||||
std::memset(R, 0, 128*sizeof(float));
|
||||
__m512 vms[q_step];
|
||||
bool need_scaling[q_step];
|
||||
auto vscale = _mm512_set1_ps(scale);
|
||||
for (int i1 = 0; i1 < nq1/q_step; ++i1) {
|
||||
for (int j = 0; j < q_step; ++j) {
|
||||
//auto R = qkv + (q_step*i1 + j)*stride_qkv;
|
||||
//std::memset(R, 0, 128*sizeof(float));
|
||||
S[j] = 0; M[j] = -INFINITY;
|
||||
}
|
||||
for (int k1 = 0; k1 < nk1/16; ++k1) {
|
||||
for (int l1 = 0; l1 < 16; ++l1) {
|
||||
auto kr = (const ggml_half *)((const char *)k + (16*k1 + l1)*stride_k);
|
||||
for (int k1 = 0; k1 < nk1/k_step; ++k1) {
|
||||
///////////////////////////////////////////////////////////////////////////////////
|
||||
//const ggml_half * kr = (const ggml_half *)((const char *)k + stride_k*16*k1);
|
||||
//DataInfo info{cache, (const char *)(q + 16*i1*stride_q), 16*sizeof(float), size_t(stride_q)*sizeof(float), 0, 0, nullptr, 0};
|
||||
//mul_mat_fX_fY_T<4, ggml_half, float>(ne00, (const void *)kr, stride_k, info, 16);
|
||||
///////////////////////////////////////////////////////////////////////////////////
|
||||
//for (int l1 = 0; l1 < 16; ++l1) {
|
||||
// for (int m1 = 0; m1 < 16; ++m1) {
|
||||
// const ggml_half * mp = (const ggml_half *)((const char *)mask + stride_m*(16*i1 + m1)) + 16*k1;
|
||||
// cache[16*m1 + l1] = GGML_FP16_TO_FP32(mp[l1]) == -INFINITY ? -INFINITY : scale*cache[16*m1 + l1];
|
||||
// }
|
||||
//}
|
||||
for (int l1 = 0; l1 < k_step; ++l1) {
|
||||
auto kr = (const ggml_half *)((const char *)k + (k_step*k1 + l1)*stride_k);
|
||||
for (int i = 0; i < 8; ++i) vk[i] = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)kr + i));
|
||||
for (int m1 = 0; m1 < 16; ++m1) {
|
||||
// q index is 16*i1 + m1
|
||||
// k index is 16*k1 + l1
|
||||
const ggml_half * mp = (const ggml_half *)((const char *)mask + stride_m*(16*i1 + m1)) + 16*k1;
|
||||
if (GGML_FP16_TO_FP32(mp[l1]) == -INFINITY) {
|
||||
cache[16*m1 + l1] = -INFINITY;
|
||||
for (int m1 = 0; m1 < q_step; ++m1) {
|
||||
// q index is q_step*i1 + m1
|
||||
// k index is k_step*k1 + l1
|
||||
const ggml_half * mp = (const ggml_half *)((const char *)mask + stride_m*(q_step*i1 + m1)) + k_step*k1;
|
||||
if (mp[l1] == h_inf) {
|
||||
cache[k_step*m1 + l1] = -INFINITY;
|
||||
continue;
|
||||
}
|
||||
auto qr = q + (16*i1 + m1)*stride_q;
|
||||
auto qr = q + (q_step*i1 + m1)*stride_q;
|
||||
auto vsum = _mm512_mul_ps(vk[0], _mm512_loadu_ps(qr));
|
||||
for (int i = 1; i < 8; ++i) vsum = _mm512_fmadd_ps(vk[i], _mm512_loadu_ps(qr + 16*i), vsum);
|
||||
cache[16*m1 + l1] = scale*_mm512_reduce_add_ps(vsum);
|
||||
cache[k_step*m1 + l1] = _mm512_reduce_add_ps(vsum);
|
||||
}
|
||||
}
|
||||
for (int j = 0; j < 16; ++j) {
|
||||
auto R = qkv + (16*i1 + j)*stride_qkv;
|
||||
auto val = _mm512_loadu_ps(cache + 16*j);
|
||||
auto smax = _mm512_reduce_max_ps(val);
|
||||
for (int i = 0; i < 8; ++i) vk[i] = _mm512_loadu_ps(R + 16*i);
|
||||
// This variant is much slower than the one below
|
||||
for (int j = 0; j < q_step; ++j) {
|
||||
auto R = qkv + (q_step*i1 + j)*stride_qkv;
|
||||
auto val1 = _mm512_mul_ps(vscale, _mm512_loadu_ps(cache + k_step*j));
|
||||
auto val2 = _mm512_mul_ps(vscale, _mm512_loadu_ps(cache + k_step*j + 16));
|
||||
auto smax = _mm512_reduce_max_ps(_mm512_max_ps(val1, val2));
|
||||
//auto val = _mm512_mul_ps(vscale, _mm512_loadu_ps(cache + k_step*j));
|
||||
//auto smax = _mm512_reduce_max_ps(val);
|
||||
need_scaling[j] = false;
|
||||
if (smax > M[j]) {
|
||||
if (M[j] > -INFINITY) {
|
||||
float m = expf(M[j] - smax);
|
||||
auto vm = _mm512_set1_ps(m);
|
||||
for (int i = 0; i < 8; ++i) {
|
||||
vk[i] = _mm512_mul_ps(vm, vk[i]);
|
||||
//auto r = _mm512_loadu_ps(R + 16*i);
|
||||
//_mm512_storeu_ps(R + 16*i, _mm512_mul_ps(vm, r));
|
||||
}
|
||||
vms[j] = _mm512_set1_ps(m);
|
||||
need_scaling[j] = true;
|
||||
S[j] *= m;
|
||||
} else {
|
||||
for (int i = 0; i < 8; ++i) vk[i] = _mm512_setzero_ps();
|
||||
//std::memset(R, 0, 128*sizeof(float));
|
||||
std::memset(R, 0, 128*sizeof(float));
|
||||
S[j] = 0;
|
||||
}
|
||||
M[j] = smax;
|
||||
}
|
||||
val = v_expf(_mm512_sub_ps(val, _mm512_set1_ps(M[j])));
|
||||
S[j] += _mm512_reduce_add_ps(val);
|
||||
_mm512_storeu_ps(cache + 16*j, val);
|
||||
for (int l1 = 0; l1 < 16; ++l1) {
|
||||
if (cache[16*j + l1] < -20.0f) continue;
|
||||
auto vr = (const ggml_half *)((const char *)v + (16*k1 + l1)*stride_v);
|
||||
auto vs = _mm512_set1_ps(cache[16*j + l1]);
|
||||
for (int i = 0; i < 8; ++i) {
|
||||
vk[i] = _mm512_fmadd_ps(vs, _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i*)vr+i)), vk[i]);
|
||||
val1 = v_expf(_mm512_sub_ps(val1, _mm512_set1_ps(M[j])));
|
||||
val2 = v_expf(_mm512_sub_ps(val2, _mm512_set1_ps(M[j])));
|
||||
S[j] += _mm512_reduce_add_ps(_mm512_add_ps(val1, val2));
|
||||
_mm512_storeu_ps(cache + k_step*j, val1);
|
||||
_mm512_storeu_ps(cache + k_step*j + 16, val2);
|
||||
//val = v_expf(_mm512_sub_ps(val, _mm512_set1_ps(M[j])));
|
||||
//S[j] += _mm512_reduce_add_ps(val);
|
||||
//_mm512_storeu_ps(cache + k_step*j, val);
|
||||
}
|
||||
for (int i = 0; i < 8; ++i) {
|
||||
for (int j = 0; j < q_step; ++j) {
|
||||
auto R = qkv + (q_step*i1 + j)*stride_qkv;
|
||||
vk[j] = _mm512_loadu_ps(R + 16*i);
|
||||
if (need_scaling[j]) vk[j] = _mm512_mul_ps(vk[j], vms[j]);
|
||||
}
|
||||
for (int l1 = 0; l1 < k_step; ++l1) {
|
||||
auto vr = (const ggml_half *)((const char *)v + (k_step*k1 + l1)*stride_v);
|
||||
auto v = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i*)vr+i));
|
||||
for (int j = 0; j < q_step; ++j) {
|
||||
vk[j] = _mm512_fmadd_ps(v, _mm512_set1_ps(cache[k_step*j + l1]), vk[j]);
|
||||
}
|
||||
}
|
||||
for (int i = 0; i < 8; ++i) _mm512_storeu_ps(R + 16*i, vk[i]);
|
||||
for (int j = 0; j < q_step; ++j) {
|
||||
auto R = qkv + (q_step*i1 + j)*stride_qkv;
|
||||
_mm512_storeu_ps(R + 16*i, vk[j]);
|
||||
}
|
||||
}
|
||||
//for (int j = 0; j < q_step; ++j) {
|
||||
// auto R = qkv + (q_step*i1 + j)*stride_qkv;
|
||||
// //auto val1 = _mm512_mul_ps(vscale, _mm512_loadu_ps(cache + k_step*j));
|
||||
// //auto val2 = _mm512_mul_ps(vscale, _mm512_loadu_ps(cache + k_step*j + 16));
|
||||
// //auto smax = _mm512_reduce_max_ps(_mm512_max_ps(val1, val2));
|
||||
// auto val = _mm512_mul_ps(vscale, _mm512_loadu_ps(cache + k_step*j));
|
||||
// auto smax = _mm512_reduce_max_ps(val);
|
||||
// for (int i = 0; i < 8; ++i) vk[i] = _mm512_loadu_ps(R + 16*i);
|
||||
// if (smax > M[j]) {
|
||||
// if (M[j] > -INFINITY) {
|
||||
// float m = expf(M[j] - smax);
|
||||
// auto vm = _mm512_set1_ps(m);
|
||||
// for (int i = 0; i < 8; ++i) {
|
||||
// vk[i] = _mm512_mul_ps(vm, vk[i]);
|
||||
// }
|
||||
// S[j] *= m;
|
||||
// } else {
|
||||
// for (int i = 0; i < 8; ++i) vk[i] = _mm512_setzero_ps();
|
||||
// S[j] = 0;
|
||||
// }
|
||||
// M[j] = smax;
|
||||
// }
|
||||
// //auto vm = _mm512_set1_ps(M[j]);
|
||||
// //val1 = v_expf(_mm512_sub_ps(val1, vm));
|
||||
// //val2 = v_expf(_mm512_sub_ps(val2, vm));
|
||||
// //S[j] += _mm512_reduce_add_ps(_mm512_add_ps(val1, val2));
|
||||
// //_mm512_storeu_ps(cache + k_step*j, val1);
|
||||
// //_mm512_storeu_ps(cache + k_step*j + 16, val2);
|
||||
// val = v_expf(_mm512_sub_ps(val, _mm512_set1_ps(M[j])));
|
||||
// S[j] += _mm512_reduce_add_ps(val);
|
||||
// _mm512_storeu_ps(cache + k_step*j, val);
|
||||
// for (int l1 = 0; l1 < k_step; ++l1) {
|
||||
// if (cache[k_step*j + l1] < -20.0f) continue;
|
||||
// auto vr = (const ggml_half *)((const char *)v + (k_step*k1 + l1)*stride_v);
|
||||
// auto vs = _mm512_set1_ps(cache[k_step*j + l1]);
|
||||
// for (int i = 0; i < 8; ++i) {
|
||||
// vk[i] = _mm512_fmadd_ps(vs, _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i*)vr+i)), vk[i]);
|
||||
// }
|
||||
// }
|
||||
// for (int i = 0; i < 8; ++i) _mm512_storeu_ps(R + 16*i, vk[i]);
|
||||
//}
|
||||
}
|
||||
for (int j = 0; j < 16; ++j) {
|
||||
auto R = qkv + (16*i1 + j)*stride_qkv;
|
||||
for (int j = 0; j < q_step; ++j) {
|
||||
auto R = qkv + (q_step*i1 + j)*stride_qkv;
|
||||
GGML_ASSERT(S[j] > 0);
|
||||
if (S[j] > 0) {
|
||||
auto norm = _mm512_set1_ps(1/S[j]);
|
||||
|
||||
Reference in New Issue
Block a user