Misha GoinandClaude
d79b52b182
Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: Claude <noreply@anthropic.com>
52 lines
1.9 KiB
C++
52 lines
1.9 KiB
C++
#include <torch/all.h>
|
|
#include "core/registration.h"
|
|
#include <torch/library.h>
|
|
#include <torch/version.h>
|
|
|
|
#include <optional>
|
|
#include <vector>
|
|
|
|
#ifdef USE_ROCM
|
|
// Defined in custom_quickreduce.cu.
|
|
using fptr_t = int64_t;
|
|
fptr_t init_custom_qr(int64_t rank, int64_t world_size,
|
|
std::optional<int64_t> qr_max_size = std::nullopt);
|
|
void qr_destroy(fptr_t _fa);
|
|
torch::Tensor qr_get_handle(fptr_t _fa);
|
|
void qr_open_handles(fptr_t _fa, const std::vector<torch::Tensor>& handles);
|
|
void qr_all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out,
|
|
int64_t quant_level, bool cast_bf2half = false);
|
|
int64_t qr_max_size();
|
|
#endif
|
|
|
|
// Note on op signatures:
|
|
// The X_meta signatures are for the meta functions corresponding to op X.
|
|
// They must be kept in sync with the signature for X. Generally, only
|
|
// functions that return Tensors require a meta function.
|
|
//
|
|
// See the following links for detailed docs on op registration and function
|
|
// schemas.
|
|
// https://docs.google.com/document/d/1_W62p8WJOQQUzPsJYa7s701JXt0qf2OfLub2sbkHOaU/edit#heading=h.ptttacy8y1u9
|
|
// https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/README.md#annotations
|
|
|
|
#ifdef USE_ROCM
|
|
TORCH_LIBRARY_FRAGMENT(CONCAT(TORCH_EXTENSION_NAME, _custom_ar), custom_ar) {
|
|
// Quick Reduce all-reduce kernels (ROCm-only; stays on legacy _C).
|
|
custom_ar.def(
|
|
"qr_all_reduce(int fa, Tensor inp, Tensor out, int quant_level, bool "
|
|
"cast_bf2half) -> ()");
|
|
custom_ar.impl("qr_all_reduce", torch::kCUDA, &qr_all_reduce);
|
|
|
|
custom_ar.def("init_custom_qr", &init_custom_qr);
|
|
custom_ar.def("qr_destroy", &qr_destroy);
|
|
custom_ar.def("qr_get_handle", &qr_get_handle);
|
|
|
|
custom_ar.def("qr_open_handles(int _fa, Tensor[](b!) handles) -> ()");
|
|
custom_ar.impl("qr_open_handles", torch::kCPU, &qr_open_handles);
|
|
|
|
custom_ar.def("qr_max_size", &qr_max_size);
|
|
}
|
|
#endif
|
|
|
|
REGISTER_EXTENSION(TORCH_EXTENSION_NAME)
|