SYCL: reduce tensor allreduce sync with pinned host buffers (#29604)

This commit is contained in:
Captain-Tripps
2026-09-30 02:25:13 -04:00
committed by GitHub
parent 649dcb1036
commit 2a53ace3be
+50 -23
View File
@@ -6995,8 +6995,48 @@ struct ggml_backend_sycl_comm_context {
std::unique_ptr<ggml_sycl_pool_alloc<uint8_t>> buf0;
std::unique_ptr<ggml_sycl_pool_alloc<uint8_t>> buf1;
int64_t buf_nelem = 0;
// pinned host staging per device context, reused across calls (queues are in-order)
uint8_t * host_out[2] = { nullptr, nullptr };
uint8_t * host_in[2] = { nullptr, nullptr };
size_t host_bytes = 0;
void free_host() {
for (int i = 0; i < 2; ++i) {
const sycl::context sctx = ((ggml_backend_sycl_context *) backends[i]->context)->stream()->get_context();
if (host_out[i]) { sycl::free(host_out[i], sctx); host_out[i] = nullptr; }
if (host_in[i]) { sycl::free(host_in[i], sctx); host_in[i] = nullptr; }
}
host_bytes = 0;
}
};
// dst0 <- src1 and dst1 <- src0 through pinned host memory, with a single host wait
static void ggml_sycl_comm_exchange(ggml_backend_sycl_comm_context * comm, queue_ptr q0, queue_ptr q1,
void * dst0, const void * src1, void * dst1, const void * src0,
size_t nbytes) {
if (comm->host_bytes < nbytes) {
q0->wait();
q1->wait();
comm->free_host();
comm->host_out[0] = sycl::malloc_host<uint8_t>(nbytes, q0->get_context());
comm->host_in[0] = sycl::malloc_host<uint8_t>(nbytes, q0->get_context());
comm->host_out[1] = sycl::malloc_host<uint8_t>(nbytes, q1->get_context());
comm->host_in[1] = sycl::malloc_host<uint8_t>(nbytes, q1->get_context());
GGML_ASSERT(comm->host_out[0] && comm->host_in[0] && comm->host_out[1] && comm->host_in[1]);
comm->host_bytes = nbytes;
}
sycl::event e0 = q0->memcpy(comm->host_out[0], src0, nbytes);
sycl::event e1 = q1->memcpy(comm->host_out[1], src1, nbytes);
e0.wait();
e1.wait();
std::memcpy(comm->host_in[0], comm->host_out[1], nbytes);
std::memcpy(comm->host_in[1], comm->host_out[0], nbytes);
q0->memcpy(dst0, comm->host_in[0], nbytes);
q1->memcpy(dst1, comm->host_in[1], nbytes);
}
void * ggml_backend_sycl_comm_init(ggml_backend_t * backends, size_t n_backends) try {
for (size_t i = 0; i < n_backends; ++i) {
if (!ggml_backend_is_sycl(backends[i])) {
@@ -7035,6 +7075,7 @@ void ggml_backend_sycl_comm_free(void * comm_ctx_v) {
try {
sctx0->stream()->wait();
sctx1->stream()->wait();
comm_ctx->free_host();
} catch (...) { /* best effort during shutdown */ }
}
@@ -7090,19 +7131,15 @@ bool ggml_backend_sycl_comm_allreduce_tensor(void * comm_ctx_v, struct ggml_tens
uint8_t * buf1 = comm_ctx->buf1->get();
// F16 native path: direct 2-byte cross-device copy + add, skipping the
// F32 round-trip the meta-backend fallback would force. Cross-device copies
// go through dev2dev_memcpy because the two devices are in separate SYCL
// contexts (a raw peer-USM q->memcpy would be a silent no-op).
// F32 round-trip the meta-backend fallback would force. The devices are in separate SYCL
// contexts (a raw peer-USM q->memcpy would be a silent no-op), so copies go through host memory.
if (tensors[0]->type == GGML_TYPE_F16) {
sycl::half * f16_out0 = (sycl::half *) tensors[0]->data;
sycl::half * f16_out1 = (sycl::half *) tensors[1]->data;
sycl::half * f16_tmp0 = (sycl::half *) buf0;
sycl::half * f16_tmp1 = (sycl::half *) buf1;
q0->wait();
q1->wait();
dev2dev_memcpy(ctx0->device, *q0, ctx1->device, *q1, f16_tmp0, tensors[1]->data, nbytes);
dev2dev_memcpy(ctx1->device, *q1, ctx0->device, *q0, f16_tmp1, tensors[0]->data, nbytes);
ggml_sycl_comm_exchange(comm_ctx, q0, q1, f16_tmp0, tensors[1]->data, f16_tmp1, tensors[0]->data, nbytes);
q0->submit([&](sycl::handler & h) {
h.parallel_for(sycl::range<1>(nelem), [=](sycl::id<1> i) {
@@ -7131,14 +7168,8 @@ bool ggml_backend_sycl_comm_allreduce_tensor(void * comm_ctx_v, struct ggml_tens
float * tmp0 = (float *) buf0;
float * tmp1 = (float *) buf1;
// COMM-D2D-FIX: the two devices are in SEPARATE SYCL contexts, so a raw
// q->memcpy of a peer USM pointer is a silent no-op. Route cross-device
// copies through dev2dev_memcpy (L0 direct copy / host staging). It is
// synchronous, so wait for the local partials to be produced first.
q0->wait();
q1->wait();
dev2dev_memcpy(ctx0->device, *q0, ctx1->device, *q1, tmp0, tensors[1]->data, nbytes);
dev2dev_memcpy(ctx1->device, *q1, ctx0->device, *q0, tmp1, tensors[0]->data, nbytes);
// separate SYCL contexts: a raw peer-USM q->memcpy is a silent no-op, so stage through host memory
ggml_sycl_comm_exchange(comm_ctx, q0, q1, tmp0, tensors[1]->data, tmp1, tensors[0]->data, nbytes);
q0->submit([&](sycl::handler & h) {
h.parallel_for(sycl::range<1>(nelem), [=](sycl::id<1> i) {
@@ -7162,21 +7193,17 @@ bool ggml_backend_sycl_comm_allreduce_tensor(void * comm_ctx_v, struct ggml_tens
uint16_t * inbox1 = outbox1 + nelem;
// Phase A: compress each device's local partial in parallel.
sycl::event c0 = q0->parallel_for(sycl::range<1>(nelem), [=](sycl::id<1> i) {
q0->parallel_for(sycl::range<1>(nelem), [=](sycl::id<1> i) {
outbox0[i] = (uint16_t) (sycl::bit_cast<uint32_t>(out0[i]) >> 16);
});
sycl::event c1 = q1->parallel_for(sycl::range<1>(nelem), [=](sycl::id<1> i) {
q1->parallel_for(sycl::range<1>(nelem), [=](sycl::id<1> i) {
outbox1[i] = (uint16_t) (sycl::bit_cast<uint32_t>(out1[i]) >> 16);
});
// Phase B: COMM-D2D-FIX-BF16 cross-device copy of compressed bytes via
// dev2dev_memcpy (separate SYCL contexts; sync copy after compress).
// Phase B: exchange the compressed bytes (in-order queues keep this after the compress kernels).
const size_t bf16_bytes = nelem * sizeof(uint16_t);
c0.wait();
c1.wait();
dev2dev_memcpy(ctx0->device, *q0, ctx1->device, *q1, inbox0, outbox1, bf16_bytes);
dev2dev_memcpy(ctx1->device, *q1, ctx0->device, *q0, inbox1, outbox0, bf16_bytes);
ggml_sycl_comm_exchange(comm_ctx, q0, q1, inbox0, outbox1, inbox1, outbox0, bf16_bytes);
// Phase C: decompress + add into local FP32 partial.
q0->submit([&](sycl::handler & h) {