30 lines
1014 B
Plaintext
30 lines
1014 B
Plaintext
#pragma once
|
|
|
|
#include "common.cuh"
|
|
#include "ggml-backend-impl.h"
|
|
|
|
#include <cstddef>
|
|
|
|
// Opaque pipeline context -- owns all pinned buffers, streams, and events.
|
|
struct ggml_cuda_ar_pipeline;
|
|
|
|
// Allocate a pipeline for n_devices GPUs.
|
|
// devices[] holds the CUDA device IDs in rank order.
|
|
// Returns nullptr on allocation failure.
|
|
ggml_cuda_ar_pipeline * ggml_cuda_ar_pipeline_init(
|
|
const int * devices, size_t n_devices);
|
|
|
|
// Release all resources owned by the pipeline.
|
|
void ggml_cuda_ar_pipeline_free(ggml_cuda_ar_pipeline * pipeline);
|
|
|
|
// Execute an in-place AllReduce (sum) across tensors[0..n_devices-1].
|
|
// tensors[i] must live on the device managed by backends[i] and be
|
|
// contiguous F32, F16, or BF16.
|
|
// Preconditions are checked by the CUDA comm dispatcher before calling this.
|
|
// Returns true once the reduction work has been enqueued successfully.
|
|
bool ggml_cuda_ar_allreduce(
|
|
ggml_cuda_ar_pipeline * pipeline,
|
|
ggml_backend_t * backends,
|
|
ggml_tensor ** tensors);
|
|
|