mirror of https://github.com/hpcaitech/ColossalAI
119 lines
3.3 KiB
C++
119 lines
3.3 KiB
C++
#include <torch/extension.h>
|
|
|
|
|
|
torch::Tensor moe_dispatch_cuda_forward(
|
|
int s, int ec, int h,
|
|
torch::Tensor batch_tokens,
|
|
torch::Tensor mask,
|
|
torch::Tensor dest_idx);
|
|
|
|
torch::Tensor moe_dispatch_cuda_backward(
|
|
int s, int ec, int h,
|
|
torch::Tensor expert_grad,
|
|
torch::Tensor mask,
|
|
torch::Tensor dest_idx);
|
|
|
|
torch::Tensor moe_combine_cuda_forward(
|
|
int s, int e, int c, int h,
|
|
torch::Tensor expert_tokens,
|
|
torch::Tensor logits,
|
|
torch::Tensor mask,
|
|
torch::Tensor dest_idx);
|
|
|
|
std::vector<torch::Tensor> moe_combine_cuda_backward(
|
|
int s, int e, int c, int h,
|
|
torch::Tensor tokens_grad,
|
|
torch::Tensor expert_tokens,
|
|
torch::Tensor logits,
|
|
torch::Tensor mask,
|
|
torch::Tensor dest_idx);
|
|
|
|
torch::Tensor cumsum_sub_one_in_dim0(torch::Tensor mask);
|
|
|
|
#define CHECK_CUDA(x) TORCH_CHECK(x.device().is_cuda(), #x " must be a CUDA tensor")
|
|
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
|
|
#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
|
|
|
|
torch::Tensor moe_dispatch_forward(
|
|
int s, int ec, int h,
|
|
torch::Tensor batch_tokens,
|
|
torch::Tensor mask,
|
|
torch::Tensor dest_idx) {
|
|
|
|
CHECK_INPUT(batch_tokens);
|
|
CHECK_CUDA(mask);
|
|
CHECK_CUDA(dest_idx);
|
|
|
|
return moe_dispatch_cuda_forward(
|
|
s, ec, h,
|
|
batch_tokens, mask, dest_idx);
|
|
}
|
|
|
|
torch::Tensor moe_dispatch_backward(
|
|
int s, int ec, int h,
|
|
torch::Tensor expert_grad,
|
|
torch::Tensor mask,
|
|
torch::Tensor dest_idx) {
|
|
|
|
CHECK_INPUT(expert_grad);
|
|
CHECK_CUDA(mask);
|
|
CHECK_CUDA(dest_idx);
|
|
|
|
return moe_dispatch_cuda_backward(
|
|
s, ec, h,
|
|
expert_grad, mask, dest_idx);
|
|
}
|
|
|
|
torch::Tensor moe_combine_forward(
|
|
int s, int e, int c, int h,
|
|
torch::Tensor expert_tokens,
|
|
torch::Tensor logits,
|
|
torch::Tensor mask,
|
|
torch::Tensor dest_idx) {
|
|
|
|
CHECK_INPUT(expert_tokens);
|
|
CHECK_INPUT(logits);
|
|
CHECK_CUDA(mask);
|
|
CHECK_CUDA(dest_idx);
|
|
|
|
return moe_combine_cuda_forward(
|
|
s, e, c, h,
|
|
expert_tokens, logits, mask, dest_idx);
|
|
}
|
|
|
|
std::vector<torch::Tensor> moe_combine_backward(
|
|
int s, int e, int c, int h,
|
|
torch::Tensor tokens_grad,
|
|
torch::Tensor expert_tokens,
|
|
torch::Tensor logits,
|
|
torch::Tensor mask,
|
|
torch::Tensor dest_idx) {
|
|
|
|
CHECK_INPUT(tokens_grad);
|
|
CHECK_INPUT(logits);
|
|
CHECK_CUDA(mask);
|
|
CHECK_CUDA(dest_idx);
|
|
|
|
return moe_combine_cuda_backward(
|
|
s, e, c, h,
|
|
tokens_grad, expert_tokens, logits, mask, dest_idx);
|
|
}
|
|
|
|
torch::Tensor moe_cumsum(torch::Tensor mask) {
|
|
CHECK_INPUT(mask);
|
|
return cumsum_sub_one_in_dim0(mask);
|
|
}
|
|
|
|
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
|
m.def("cumsum_sub_one", &moe_cumsum,
|
|
"Fast cumsum operation in dim0");
|
|
m.def("dispatch_forward", &moe_dispatch_forward,
|
|
"Forward operation in MoE dispatch function");
|
|
m.def("dispatch_backward", &moe_dispatch_backward,
|
|
"Backward operation in MoE dispatch function");
|
|
m.def("combine_forward", &moe_combine_forward,
|
|
"Combine operation in MoE combine function");
|
|
m.def("combine_backward", &moe_combine_backward,
|
|
"Combine operation in MoE combine function");
|
|
}
|