diff --git a/ggml/src/iqk/iqk_mul_mat.cpp b/ggml/src/iqk/iqk_mul_mat.cpp index 2db82fa7..f966fb04 100644 --- a/ggml/src/iqk/iqk_mul_mat.cpp +++ b/ggml/src/iqk/iqk_mul_mat.cpp @@ -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]);