mirror of
https://github.com/ggml-org/llama.cpp.git
synced 2026-08-12 22:31:11 +04:00
Multi-stream MoE optimization
This commit is contained in:
@@ -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,
|
||||
|
||||
@@ -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) {
|
||||
|
||||
Reference in New Issue
Block a user