Multi-stream MoE optimization

This commit is contained in:
Alexey Kopytko
2026-05-13 09:24:21 +09:00
parent 85a9b5be7f
commit 938929c3f1
2 changed files with 95 additions and 10 deletions
+54 -10
View File
@@ -322,7 +322,16 @@ struct ggml_backend_sycl_context {
std::string name;
optimize_feature opt_feature;
// Active stream slot used by stream() / stream_dnnl() / pool() (no-arg defaults).
// Multi-stream consumers set this via ggml_sycl_stream_override_guard so that
// existing call sites that use the default-arg form route to a non-zero slot
// without having to thread a parameter through every layer of dispatch.
int curr_stream_no = 0;
queue_ptr qptrs[GGML_SYCL_MAX_DEVICES][GGML_SYCL_MAX_STREAMS] = { { nullptr } };
// Owns per-slot in-order queues for slot > 0. Slot 0 aliases dpct's default_queue
// and is not stored here. unique_ptr gives stable addresses across vector growth.
std::vector<std::unique_ptr<sycl::queue>> owned_queues;
explicit ggml_backend_sycl_context(int device) :
device(device),
@@ -332,13 +341,20 @@ struct ggml_backend_sycl_context {
queue_ptr stream(int device, int stream) {
if (qptrs[device][stream] == nullptr) {
qptrs[device][stream] = &(dpct::get_device(device).default_queue());
if (stream == 0) {
qptrs[device][stream] = &(dpct::get_device(device).default_queue());
} else {
ggml_sycl_set_device(device);
owned_queues.emplace_back(std::make_unique<sycl::queue>(
dpct::get_device(device).create_in_order_queue()));
qptrs[device][stream] = owned_queues.back().get();
}
}
return qptrs[device][stream];
}
queue_ptr stream() {
return stream(device, 0);
return stream(device, curr_stream_no);
}
#if GGML_SYCL_DNNL
@@ -383,7 +399,7 @@ struct ggml_backend_sycl_context {
}
}
dnnl::stream stream_dnnl() {
return stream_dnnl(device, 0);
return stream_dnnl(device, curr_stream_no);
}
dnnl::memory get_scratchpad_mem(const dnnl::memory::desc & scratchpad_md,
const dnnl::engine & eng, const queue_ptr q) {
@@ -405,8 +421,12 @@ struct ggml_backend_sycl_context {
}
#endif
// pool
std::unique_ptr<ggml_sycl_pool> pools[GGML_SYCL_MAX_DEVICES];
// pool — one per (device, stream slot). Per-slot pools give multi-stream consumers
// race-free concurrent allocation: each in-order queue's pool reuses buffers in
// submission order, so a buffer freed and re-alloced on the same slot is naturally
// serialized. Slot 0 is the default; the per-slot pools are lazily created on
// first use, so single-GPU/single-stream paths only ever instantiate slot 0.
std::unique_ptr<ggml_sycl_pool> pools[GGML_SYCL_MAX_DEVICES][GGML_SYCL_MAX_STREAMS];
std::unordered_map<sycl::queue *, std::unique_ptr<ggml_sycl_pool_alloc<uint8_t>>> scratchpad_map;
std::unique_ptr<ggml_sycl_fattn_kv_buffers> fattn_bufs[GGML_SYCL_MAX_DEVICES];
@@ -417,17 +437,21 @@ struct ggml_backend_sycl_context {
static std::unique_ptr<ggml_sycl_pool> new_pool_for_host(queue_ptr qptr, int device);
ggml_sycl_pool & pool(int device, int stream) {
if (pools[device][stream] == nullptr) {
pools[device][stream] = new_pool_for_device(this->stream(device, stream), device);
}
return *pools[device][stream];
}
static std::unique_ptr<ggml_sycl_fattn_kv_buffers> new_fattn_kv_buffers(queue_ptr qptr, int device);
ggml_sycl_pool & pool(int device) {
if (pools[device] == nullptr) {
pools[device] = new_pool_for_device(stream(device,0), device);
}
return *pools[device];
return pool(device, 0);
}
ggml_sycl_pool & pool() {
return pool(device);
return pool(device, curr_stream_no);
}
ggml_sycl_fattn_kv_buffers & fattn_buffers(int device) {
@@ -455,6 +479,26 @@ struct ggml_backend_sycl_context {
ggml_sycl_pool & host_pool() { return host_pool(device); }
};
// RAII guard: scopes an override of ggml_backend_sycl_context::curr_stream_no.
// Used by multi-stream consumers (e.g. per-expert MoE loop) to redirect the
// default stream() / stream_dnnl() / pool() lookups to a specific slot for
// the duration of a code region, then restore the prior slot on exit.
struct ggml_sycl_stream_override_guard {
ggml_backend_sycl_context & ctx;
int saved;
ggml_sycl_stream_override_guard(ggml_backend_sycl_context & ctx_, int new_slot) :
ctx(ctx_), saved(ctx_.curr_stream_no) {
ctx.curr_stream_no = new_slot;
}
~ggml_sycl_stream_override_guard() {
ctx.curr_stream_no = saved;
}
ggml_sycl_stream_override_guard(const ggml_sycl_stream_override_guard &) = delete;
ggml_sycl_stream_override_guard & operator=(const ggml_sycl_stream_override_guard &) = delete;
};
// common device functions
static __dpct_inline__ float warp_reduce_sum(float x,
+41
View File
@@ -3075,6 +3075,27 @@ static void ggml_sycl_op_mul_mat(ggml_backend_sycl_context & ctx, const ggml_ten
int64_t is_max = (ne11 + MUL_MAT_SRC1_COL_STRIDE - 1) / MUL_MAT_SRC1_COL_STRIDE;
is_max = is_max <= GGML_SYCL_MAX_STREAMS ? is_max : GGML_SYCL_MAX_STREAMS;
GGML_SYCL_DEBUG("[SYCL][TENSOR-SPLIT] mul_mat ne11=%ld stride=%d slots=%ld devices=%d\n",
(long) ne11, (int) MUL_MAT_SRC1_COL_STRIDE, (long) is_max,
ggml_sycl_info().device_count);
// each non-main device's slot 0 waits on its own slots 1..is_max-1, so
// buffers allocated from its slot-0 pool (and used across slots) are
// safe to reuse on slot 0 after this op returns. before commit
// 7277d77b1 every slot aliased slot 0 and this dep was implicit; now
// slots are distinct queues and must be joined explicitly.
for (int i = 0; i < ggml_sycl_info().device_count; ++i) {
if (i == ctx.device || dev[i].row_low == dev[i].row_high) {
continue;
}
ggml_sycl_set_device(i);
for (int64_t is = 1; is < is_max; ++is) {
SYCL_CHECK(CHECK_TRY_ERROR(
ctx.stream(i, 0)->ext_oneapi_submit_barrier(
{*src0_extra->events[i][is]})));
}
}
ggml_sycl_set_device(ctx.device);
for (int i = 0; i < ggml_sycl_info().device_count; ++i) {
if (dev[i].row_low == dev[i].row_high) {
@@ -4184,6 +4205,12 @@ static void ggml_sycl_mul_mat_id(ggml_backend_sycl_context & ctx,
src1_row.data = src1_contiguous.get();
dst_row.data = dst_contiguous.get();
// Multi-stream the per-expert iterations across N stream slots so that
// independent experts run concurrently. Each iteration's allocations come
// from that slot's pool (per-stream isolation prevents the cross-stream
// alias hazard the shared pool would otherwise create).
const int n_streams = (int) std::min<int64_t>(GGML_SYCL_MAX_STREAMS, n_as);
for (int64_t i02 = 0; i02 < n_as; i02++) {
int64_t num_src1_rows = 0;
for (int64_t iid1 = 0; iid1 < ids->ne[1]; iid1++) {
@@ -4204,6 +4231,8 @@ static void ggml_sycl_mul_mat_id(ggml_backend_sycl_context & ctx,
continue;
}
ggml_sycl_stream_override_guard slot_guard(ctx, (int)(i02 % n_streams));
const queue_ptr stream = ctx.stream(); // per-iter, follows the override
ggml_sycl_pool_alloc<int> dev_cur_src1_row(ctx.pool(), 1);
ggml_sycl_pool_alloc<mmid_row_mapping> dev_row_mapping(ctx.pool(), num_src1_rows);
@@ -4278,6 +4307,18 @@ static void ggml_sycl_mul_mat_id(ggml_backend_sycl_context & ctx,
});
}
}
// Barrier-and-join: ensure all per-stream work completes before slot 0
// proceeds. The next op in the graph runs on slot 0 by default and may
// read tensors written by experts dispatched to slots 1..n_streams-1.
if (n_streams > 1) {
std::vector<sycl::event> events;
events.reserve(n_streams - 1);
for (int s = 1; s < n_streams; ++s) {
events.push_back(ctx.stream(ctx.device, s)->ext_oneapi_submit_barrier());
}
ctx.stream(ctx.device, 0)->ext_oneapi_submit_barrier(events);
}
}
}
catch (sycl::exception const &exc) {