LLM-Inference/ik_llama.cpp
1
0
Fork 0
mirror of https://github.com/ikawrakow/ik_llama.cpp.git synced 2026-08-12 22:29:39 +04:00
Code Issues Packages Projects Releases Wiki Activity

Bucket top_k (CPU): ~3% better TG at 128k context (#2225)

Browse Source
This commit is contained in:
Kawrakow
2026-08-05 07:50:41 +03:00
committed by GitHub
parent 7257cabc49
commit 1474935e31
1 changed files with 149 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/iqk/iqk_mul_mat.cpp
+149 -2
View File
@@ -1741,6 +1741,7 @@ bool iqk_fused_delta_net(int head_dim, int n_heads, int gqa_ratio, int repeat_ty
namespace {
constexpr int k_indexer_chunks = 64;
constexpr int k_n_bucket = 64;
size_t iqk_idx_topk_work_wbs_per_thread(const struct ggml_tensor * dst, int nth) {
auto k = dst->src[0];
auto q = dst->src[1];
@@ -1754,6 +1755,9 @@ size_t iqk_idx_topk_work_wbs_per_thread(const struct ggml_tensor * dst, int nth)
size += k_indexer_chunks * q->ne[1] * sizeof(float);
size += k->ne[1] * sizeof(float);
size += k->ne[1] * sizeof(int32_t);
// We will not use iqk_bucket_sort for batch processing, so
// no need to allocate the extra work buffers.
//size += (2*k->ne[1] + k_n_bucket)*sizeof(int);
size = GGML_PAD(size, 128);
return size;
}
@@ -1765,6 +1769,138 @@ inline void iqk_f16_to_f32(int n, const ggml_fp16_t * x, float * y) {
y[i] = GGML_FP16_TO_FP32(x[i]);
}
}
#ifdef __AVX2__
inline float hmin_f32_8(__m256 x) {
__m128 min4 = _mm_min_ps(_mm256_extractf128_ps(x, 1), _mm256_castps256_ps128(x));
min4 = _mm_min_ps(min4, _mm_movehl_ps(min4, min4));
min4 = _mm_min_ss(min4, _mm_movehdup_ps(min4));
return _mm_cvtss_f32(min4);
}
#endif
// Note: result is not actually sorted in decreasing order, we just get the indices of the ntop
// values stored in idx.
// In micro-benchmark testing this code outperforms std::partial_sort by a factor of 4-6
// (factors, not percentages!).
// For DS4 running CPU-only this translates into a 3% better TG at a context of 128k tokens.
void iqk_bucket_topk(int nval, int ntop, float * values, int * idx, int * idx_inf, int nbucket, int * counts,
int * idx_aux) {
#if 0
int ngood = nval;
while (ngood > 0 && values[ngood-1] == -INFINITY) --ngood;
if (ngood <= ntop) {
for (int j = 0; j < ntop; ++j) idx[j] = j;
return;
}
float max = values[0], min = values[0];
#ifdef __AVX2__
auto vmax = _mm256_loadu_ps(values);
auto vmin = vmax;
auto vidx = _mm256_set_epi32(7,6,5,4,3,2,1,0);
auto vstep = _mm256_set1_epi32(8);
_mm256_storeu_si256((__m256i *)idx, vidx);
for (int j = 1; j < ngood/8; ++j) {
auto v = _mm256_loadu_ps(values + 8*j);
vidx = _mm256_add_epi32(vidx, vstep);
_mm256_storeu_si256((__m256i *)idx + j, vidx);
vmax = _mm256_max_ps(vmax, v);
vmin = _mm256_min_ps(vmin, v);
}
max = hmax_f32_8(vmax);
min = hmin_f32_8(vmin);
for (int j = 8*(ngood/8); j < ngood; ++j) {
float v = values[j];
max = std::max(max, v);
min = std::min(min, v);
idx[j] = j;
}
#else
for (int j = 0; j < ngood; ++j) {
float v = values[j];
max = std::max(max, v);
min = std::min(min, v);
idx[j] = j;
}
#endif
#else
// If we knew that we don't have -inf values, we could do this more efficiently.
// But we don't. At least not for sure.
// Oh, I did measure using the above commented out code, which assumes that
// -inf values if present are at the end, and I saw no real performance difference.
int ngood = 0, ninf = 0;
float max = values[0], min = values[0];
for (int j = 0; j < nval; ++j) {
if (float v = values[j]; v > -INFINITY) {
values[ngood] = v;
idx[ngood++] = j;
max = std::max(max, v);
min = std::min(min, v);
} else {
idx_inf[ninf++] = j;
}
}
#endif
if (ngood <= ntop) {
for (int j = ngood; j < ntop; ++j) idx[j] = idx_inf[j-ngood];
return;
}
if (max - min < 1e-6f) return; // we got basically the same values, so it doesn't matter which we pick
float av = (nbucket - 0.75f)/(min - max);
float bv = -av*max;
#ifdef __AVX2__
auto v_av = _mm256_set1_ps(av);
auto v_bv = _mm256_set1_ps(bv);
for (int i = 0; i < nbucket; ++i) counts[i] = 0;
for (int j = 0; j < ngood/8; ++j) {
auto v = _mm256_loadu_ps(values + 8*j);
auto xv = _mm256_fmadd_ps(v_av, v, v_bv);
auto iv = _mm256_cvtps_epi32(xv);
iv = _mm256_min_epi32(iv, _mm256_set1_epi32(nbucket-1));
auto aux = idx_aux + 8*j;
_mm256_storeu_si256((__m256i *)aux, iv);
for (int k = 0; k < 8; ++k) ++counts[aux[k]];
}
for (int j = 8*(ngood/8); j < ngood; ++j) {
int i = int(av*values[j] + bv);
i = std::min(i, nbucket-1);
idx_aux[j] = i;
++counts[i];
}
#else
for (int j = 0; j < ngood; ++j) {
int i = int(av*values[j] + bv);
i = std::min(i, nbucket-1);
idx_aux[j] = i;
++counts[i];
}
#endif
int last_bucket = 0;
int sum = 0;
for (; last_bucket < nbucket-1; ++last_bucket) {
sum += counts[last_bucket];
if (sum >= ntop) break;
}
int nhave = 0, nlast = 0;
for (int j = 0; j < ngood; ++j) {
if (idx_aux[j] < last_bucket) {
idx[nhave++] = idx[j];
} else if (idx_aux[j] == last_bucket) {
idx_inf[nlast++] = idx[j];
}
}
int n_extra = ntop - nhave;
auto compare = [values] (int l, int r) {
return values[l] > values[r];
};
if (2*n_extra < nlast) {
std::partial_sort(idx_inf, idx_inf + n_extra, idx_inf + nlast, compare);
} else {
std::sort(idx_inf, idx_inf + nlast, compare);
}
for (int j = 0; j < n_extra; ++j) idx[nhave + j] = idx_inf[j];
}
}
size_t iqk_idx_topk_work_buffer_size(const struct ggml_tensor * dst, int nthread) {
@@ -1782,6 +1918,7 @@ size_t iqk_idx_topk_work_buffer_size(const struct ggml_tensor * dst, int nthread
size += k->ne[1] * q->ne[1] * sizeof(float);
size += k->ne[1] * sizeof(float);
size += k->ne[1] * sizeof(int32_t);
size += (2*k->ne[1] + k_n_bucket)*sizeof(int);
return size;
}
@@ -1834,6 +1971,9 @@ bool iqk_indexer_topk(struct ggml_tensor * dst, void * work_buffer, barrier_t ba
auto kq = (float *)(work + quantize_size);
auto score = kq + k_indexer_chunks*q->ne[1];
auto sorted = (int32_t *)(score + k->ne[1]);
//auto idx_inf = sorted + k->ne[1];
//auto idx_aux = idx_inf + k->ne[1];
//auto counts = idx_aux + k->ne[1];
for (int iq = ith; iq < q->ne[2]; iq += nth) {
auto this_q = (const char *)q->data + iq*q->nb[2];
auto this_m = (const char *)m->data + iq*m->nb[1];
@@ -1885,6 +2025,11 @@ bool iqk_indexer_topk(struct ggml_tensor * dst, void * work_buffer, barrier_t ba
kq_i += nk;
}
}
// Here iqk_bucket_topk is not faster than just using std::partial_sort, so no need to allocate the extra
// work buffers.
// iqk_bucket_topk(k->ne[1], n_top_k, score, sorted, idx_inf, k_n_bucket, counts, idx_aux);
for (int j = 0; j < int(k->ne[1]); ++j) sorted[j] = j;
std::partial_sort(sorted, sorted + n_top_k, sorted + k->ne[1], [score] (int32_t l, int32_t r) -> bool { return score[l] > score[r]; });
std::memcpy((char *)dst->data + dst->nb[1]*iq, sorted, n_top_k*sizeof(int32_t));
@@ -1929,6 +2074,9 @@ bool iqk_indexer_topk(struct ggml_tensor * dst, void * work_buffer, barrier_t ba
auto score = kq + k->ne[1]*q->ne[1];
auto score_th = score + first;
auto sorted = (int32_t *)(score + k->ne[1]);
auto idx_inf = sorted + k->ne[1];
auto idx_aux = idx_inf + k->ne[1];
auto counts = idx_aux + k->ne[1];
for (int iq = 0; iq < q->ne[2]; ++iq) {
if (n_this_thread > 0) {
auto this_q = q_data + iq*qnb2;
@@ -1953,8 +2101,7 @@ bool iqk_indexer_topk(struct ggml_tensor * dst, void * work_buffer, barrier_t ba
}
barrier(barrier_data);
if (ith == 0) {
for (int j = 0; j < int(k->ne[1]); ++j) sorted[j] = j;
std::partial_sort(sorted, sorted + n_top_k, sorted + k->ne[1], [score] (int32_t l, int32_t r) -> bool { return score[l] > score[r]; });
iqk_bucket_topk(k->ne[1], n_top_k, score, sorted, idx_inf, k_n_bucket, counts, idx_aux);
std::memcpy((char *)dst->data + dst->nb[1]*iq, sorted, n_top_k*sizeof(int32_t));
}
if (iq + 1 < q->ne[2]) {