Lumpiasty/llama.cpp
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

sycl: split long rows in TOP_K instead of one work-group per row (#27847)

Browse Source
This commit is contained in:
Titaniumtown
2026-08-30 08:57:08 +03:00
committed by GitHub
parent d882575cc8
commit 370cb12e8b
1 changed files with 219 additions and 79 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-sycl/ggml-sycl.cpp
+219 -79
View File
@@ -2402,7 +2402,138 @@ static void argsort_f32_i32_sycl(const float *x, int *dst, const int ncols,
}
}
// Scan and block merge, shared by every launch shape below so a partitioned row uses the
// same insertion order as an unpartitioned one.
//
// src_map != nullptr: report src_map[col] instead of col, so a merge pass can carry the
// original column index through.
// out_vals != nullptr: also emit the k winning values, for a later merge pass.
// swap01: emit in the output order the single-pass path uses.
static void top_k_scan_merge_f32(
const float * src_vals,
const int32_t * src_map,
const int begin,
const int end,
const int k,
const int block_size,
float * shared_vals,
int * shared_idx,
float * out_vals,
int32_t * out_idx,
const bool swap01,
const sycl::nd_item<1> & item_ct1
) {
const int tid = item_ct1.get_local_id(0);
// The running top-k lives in SLM (shared local memory) rather than a private array:
// an array indexed by a runtime position cannot be register-allocated, so a private
// one lands in scratch, i.e. device memory, and insertion is this kernel's dominant
// cost.
//
// Lane-strided (lv[i * block_size]) rather than lane-blocked (lv[i]) so a given i is
// contiguous across lanes; a k-strided layout would put every lane of a shift step in
// the same SLM bank.
float * lv = shared_vals + tid;
int * li = shared_idx + tid;
for (int i = 0; i < k; i++) {
lv[i * block_size] = -FLT_MAX;
li[i * block_size] = -1;
}
// The k-th best, cached in a register. The reject test is taken for the large
// majority of elements scanned, and in that case touches no memory.
float kth = -FLT_MAX;
for (int col = begin + tid; col < end; col += block_size) {
float val = src_vals[col];
if (val > kth) {
int pos = k - 1;
while (pos > 0 && val > lv[(pos - 1) * block_size]) {
pos--;
}
for (int i = k - 1; i > pos; i--) {
lv[i * block_size] = lv[(i - 1) * block_size];
li[i * block_size] = li[(i - 1) * block_size];
}
lv[pos * block_size] = val;
li[pos * block_size] = src_map ? src_map[col] : col;
kth = lv[(k - 1) * block_size];
}
}
item_ct1.barrier(sycl::access::fence_space::local_space);
if (tid != 0) {
return;
}
// Same treatment for the merge accumulator, past the per-lane region.
float * fv = shared_vals + (size_t) k * block_size;
int * fi = shared_idx + (size_t) k * block_size;
for (int i = 0; i < k; i++) {
fv[i] = -FLT_MAX;
fi[i] = -1;
}
float fkth = -FLT_MAX;
// Candidates are visited in the same (t, i) order as before, so tie-breaking is
// unchanged.
for (int t = 0; t < block_size; t++) {
for (int i = 0; i < k; i++) {
float val = shared_vals[i * block_size + t];
if (val <= fkth) {
// Lane t's list is sorted descending, so once one of its entries loses
// to the k-th best, every later entry loses too. fkth only rises, so
// that stays true for the rest of the merge. This turns the merge from
// block_size*k steps into roughly block_size plus the candidates
// accepted.
break;
}
int idx = shared_idx[i * block_size + t];
int pos = k - 1;
while (pos > 0 && val > fv[pos - 1]) {
pos--;
}
for (int j = k - 1; j > pos; j--) {
fv[j] = fv[j - 1];
fi[j] = fi[j - 1];
}
fv[pos] = val;
fi[pos] = idx;
fkth = fv[k - 1];
}
}
if (out_vals) {
for (int i = 0; i < k; i++) {
out_vals[i] = fv[i];
}
}
for (int i = 0; i < k; i++) {
out_idx[i] = fi[i];
}
if (swap01 && k > 1) {
int32_t temp = out_idx[0];
out_idx[0] = out_idx[1];
out_idx[1] = temp;
}
}
static void top_k_f32_sycl(
ggml_backend_sycl_context & ctx,
const float * src,
int32_t * dst_indices,
const int64_t ncols,
@@ -2410,98 +2541,107 @@ static void top_k_f32_sycl(
const int k,
dpct::queue_ptr main_stream
) {
// A row is scanned by exactly one work-group, so a vocabulary-sized row leaves the
// rest of the device idle. What the scan is short of is memory requests in flight,
// not bandwidth or per-request latency, so lanes in flight is the lever: split the
// row across independent work-groups, have each emit its partition's top-k, and
// merge those nsplit*k candidates in a second launch.
//
// split_block trades parallelism against SLM residency. Its cost is
// (split_block + 1) * k * 8 bytes of SLM per group, so at the k <= 32 ceiling 128
// lanes need about 33 KB, which leaves a single resident group per Xe-core. Revisit
// if the supported k ever grows.
constexpr int split_block = 128;
constexpr int max_splits = 128;
constexpr int min_cols = 8192;
int nsplit = 1;
if (ncols >= min_cols) {
// A partition is then always >= split_block = 128 columns, hence always more than
// the k <= 32 ceiling, so no pass is ever padded with -FLT_MAX sentinels.
const int64_t want = ncols / split_block;
nsplit = (int) (want > max_splits ? max_splits : want);
}
if (nsplit > 1) {
const int nchunk = (int) ((ncols + nsplit - 1) / nsplit);
const size_t ncand = (size_t) nrows * nsplit * k;
ggml_sycl_pool_alloc<float> part_vals(ctx.pool(), ncand);
ggml_sycl_pool_alloc<int32_t> part_idx(ctx.pool(), ncand);
float * pv = part_vals.get();
int32_t * pi = part_idx.get();
const sycl::range<1> block_dims(split_block);
main_stream->submit([&](sycl::handler &cgh) {
sycl::local_accessor<float, 1> shared_vals(sycl::range<1>((split_block + 1) * k), cgh);
sycl::local_accessor<int, 1> shared_idx(sycl::range<1>((split_block + 1) * k), cgh);
cgh.parallel_for(
sycl::nd_range<1>(sycl::range<1>(nrows * nsplit) * block_dims, block_dims),
[=](sycl::nd_item<1> item_ct1) {
const int grp = item_ct1.get_group(0);
const int row = grp / nsplit;
const int part = grp % nsplit;
const int begin = part * nchunk;
int end = begin + nchunk;
if (end > (int) ncols) {
end = (int) ncols;
}
top_k_scan_merge_f32(
src + (int64_t) row * ncols, nullptr, begin, end, k, split_block,
shared_vals.get_multi_ptr<sycl::access::decorated::no>().get(),
shared_idx.get_multi_ptr<sycl::access::decorated::no>().get(),
pv + (size_t) grp * k, pi + (size_t) grp * k, false, item_ct1);
});
});
main_stream->submit([&](sycl::handler &cgh) {
sycl::local_accessor<float, 1> shared_vals(sycl::range<1>((split_block + 1) * k), cgh);
sycl::local_accessor<int, 1> shared_idx(sycl::range<1>((split_block + 1) * k), cgh);
cgh.parallel_for(
sycl::nd_range<1>(sycl::range<1>(nrows) * block_dims, block_dims),
[=](sycl::nd_item<1> item_ct1) {
const int row = item_ct1.get_group(0);
const size_t off = (size_t) row * nsplit * k;
top_k_scan_merge_f32(
pv + off, pi + off, 0, nsplit * k, k, split_block,
shared_vals.get_multi_ptr<sycl::access::decorated::no>().get(),
shared_idx.get_multi_ptr<sycl::access::decorated::no>().get(),
nullptr, dst_indices + (int64_t) row * k, true, item_ct1);
});
});
return;
}
const int block_size = 128;
const sycl::range<1> block_dims(block_size);
const sycl::range<1> grid_dims(nrows);
main_stream->submit([&](sycl::handler &cgh) {
sycl::local_accessor<float, 1> shared_vals(sycl::range<1>(block_size * k), cgh);
sycl::local_accessor<int, 1> shared_idx(sycl::range<1>(block_size * k), cgh);
sycl::local_accessor<float, 1> shared_vals(sycl::range<1>((block_size + 1) * k), cgh);
sycl::local_accessor<int, 1> shared_idx(sycl::range<1>((block_size + 1) * k), cgh);
cgh.parallel_for(
sycl::nd_range<1>(grid_dims * block_dims, block_dims),
[=](sycl::nd_item<1> item_ct1) {
const int row = item_ct1.get_group(0);
const int tid = item_ct1.get_local_id(0);
if (row >= nrows) return;
const float * src_row = src + row * ncols;
int32_t * dst_idx_row = dst_indices + row * k;
float local_vals[32];
int local_idx[32];
for (int i = 0; i < k; i++) {
local_vals[i] = -FLT_MAX;
local_idx[i] = -1;
}
for (int col = tid; col < ncols; col += block_size) {
float val = src_row[col];
if (val > local_vals[k-1]) {
int pos = k - 1;
while (pos > 0 && val > local_vals[pos - 1]) {
pos--;
}
for (int i = k - 1; i > pos; i--) {
local_vals[i] = local_vals[i - 1];
local_idx[i] = local_idx[i - 1];
}
local_vals[pos] = val;
local_idx[pos] = col;
}
}
for (int i = 0; i < k; i++) {
shared_vals[tid * k + i] = local_vals[i];
shared_idx[tid * k + i] = local_idx[i];
}
item_ct1.barrier(sycl::access::fence_space::local_space);
if (tid == 0) {
float final_vals[32];
int final_idx[32];
for (int i = 0; i < k; i++) {
final_vals[i] = -FLT_MAX;
final_idx[i] = -1;
}
for (int t = 0; t < block_size; t++) {
for (int i = 0; i < k; i++) {
float val = shared_vals[t * k + i];
int idx = shared_idx[t * k + i];
if (val > final_vals[k-1]) {
int pos = k - 1;
while (pos > 0 && val > final_vals[pos - 1]) {
pos--;
}
for (int j = k - 1; j > pos; j--) {
final_vals[j] = final_vals[j - 1];
final_idx[j] = final_idx[j - 1];
}
final_vals[pos] = val;
final_idx[pos] = idx;
}
}
}
for (int i = 0; i < k; i++) {
dst_idx_row[i] = final_idx[i];
}
if (k > 1) {
int32_t temp = dst_idx_row[0];
dst_idx_row[0] = dst_idx_row[1];
dst_idx_row[1] = temp;
}
}
top_k_scan_merge_f32(
src + (int64_t) row * ncols, nullptr, 0, (int) ncols, k, block_size,
shared_vals.get_multi_ptr<sycl::access::decorated::no>().get(),
shared_idx.get_multi_ptr<sycl::access::decorated::no>().get(),
nullptr, dst_indices + (int64_t) row * k, true, item_ct1);
});
});
}
@@ -2902,7 +3042,7 @@ static void ggml_sycl_op_top_k(ggml_backend_sycl_context & ctx, ggml_tensor * ds
GGML_ASSERT(k > 0 && k <= 32);
GGML_ASSERT(k <= ncols);
top_k_f32_sycl(src0_dd, dst_dd, ncols, nrows, k, main_stream);
top_k_f32_sycl(ctx, src0_dd, dst_dd, ncols, nrows, k, main_stream);
}
inline void ggml_sycl_op_argmax(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {