moe.cpp 3.32 KB
Newer Older
Jiezhong Qiu's avatar
can run  
Jiezhong Qiu committed
1
2
3
4
5
6
7
#include <torch/extension.h>
#include <torch/torch.h>

#include <cstdio>
#include <iostream>
#include <vector>

Jiezhong Qiu's avatar
update  
Jiezhong Qiu committed
8
9
// CUDA runtime
#include <cuda.h>                                                                                             
Jiezhong Qiu's avatar
can run  
Jiezhong Qiu committed
10
11
12
13
14
#include <cuda_runtime.h>                                                                                                 
#include <cublas_v2.h>                                                                                                    
                                                                                                                            
// CUDA and CUBLAS functions                                                                                              
//#include <helper_functions.h>                                                                                             
Jiezhong Qiu's avatar
update  
Jiezhong Qiu committed
15
#include <helper_cuda.h> 
Jiezhong Qiu's avatar
can run  
Jiezhong Qiu committed
16

Jiezhong Qiu's avatar
update  
Jiezhong Qiu committed
17
18
19
20
21
22
23
24
25
26
template <typename scalar_t>
void moe_first_linear_cuda_forward(
        const scalar_t* input,
        const size_t* gate,
        const scalar_t* weight,
        scalar_t* output,
        const size_t batch_size,
        const size_t top_k,
        const size_t in_feat,
        const size_t out_feat);
Jiezhong Qiu's avatar
can run  
Jiezhong Qiu committed
27

Jiezhong Qiu's avatar
update  
Jiezhong Qiu committed
28
29
30
31
32
33
34
35
36
// C++ interface

// NOTE: AT_ASSERT has become AT_CHECK on master after 0.4.
#define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " must be contiguous")
#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)


std::vector<torch::Tensor> moe_first_linear_forward(
Jiezhong Qiu's avatar
update  
Jiezhong Qiu committed
37
38
        torch::Tensor input, // [B x D_model]
        torch::Tensor gate,  // [B x K]
Jiezhong Qiu's avatar
update  
Jiezhong Qiu committed
39
        torch::Tensor weight // [N x D_ffn x D_model]
Jiezhong Qiu's avatar
update  
Jiezhong Qiu committed
40
        ) {
Jiezhong Qiu's avatar
update  
Jiezhong Qiu committed
41
42
43
    CHECK_INPUT(input);
    CHECK_INPUT(gate);
    CHECK_INPUT(weight);
Jiezhong Qiu's avatar
update  
Jiezhong Qiu committed
44
45
46
47
48
49
50
51
52
53
54
    /*
        The bias term should have been merged into weight. Note the following fact that 
        Wx+b = [W b] [x]
                     [1]  
    */
    const auto batch_size = input.size(0);
    const auto top_k = gate.size(1);
    const auto num_expert = weight.size(0);
    const auto out_feat = weight.size(1);
    const auto in_feat = weight.size(2);
    
Jiezhong Qiu's avatar
update  
Jiezhong Qiu committed
55
    printf("b=%ld, expert=%ld, in_feat (d_model)=%ld, out_feat (d_ffn)=%ld, topk=%ld\n", batch_size, num_expert, in_feat, out_feat, top_k);
Jiezhong Qiu's avatar
update  
Jiezhong Qiu committed
56
57
    auto output = input.new_zeros({batch_size, top_k, out_feat});

Jiezhong Qiu's avatar
update  
Jiezhong Qiu committed
58
59
    AT_DISPATCH_FLOATING_TYPES(input.scalar_type(), "moe_first_linear_forward", ([&] {
        moe_first_linear_cuda_forward<scalar_t>(
Jiezhong Qiu's avatar
update  
Jiezhong Qiu committed
60
61
62
63
64
65
66
67
68
69
            input.data_ptr<scalar_t>(),
            gate.data_ptr<size_t>(),
            weight.data_ptr<scalar_t>(),
            output.data_ptr<scalar_t>(),
            batch_size,
            top_k,
            in_feat,
            out_feat
        );
    }));
Jiezhong Qiu's avatar
update  
Jiezhong Qiu committed
70
    return {output, };
Jiezhong Qiu's avatar
can run  
Jiezhong Qiu committed
71
72
73
}


Jiezhong Qiu's avatar
update  
Jiezhong Qiu committed
74
/*
Jiezhong Qiu's avatar
can run  
Jiezhong Qiu committed
75
int main() {
Jiezhong Qiu's avatar
update  
Jiezhong Qiu committed
76
    int device=2;
Jiezhong Qiu's avatar
update  
Jiezhong Qiu committed
77
    torch::Tensor input = torch::randn({2048, 512}, torch::dtype(torch::kFloat32).device(torch::kCUDA, device));
Jiezhong Qiu's avatar
update  
Jiezhong Qiu committed
78
    torch::Tensor gate = torch::zeros({2048, 2}, torch::dtype(torch::kInt64));
Jiezhong Qiu's avatar
update  
Jiezhong Qiu committed
79
80
    torch::Tensor weight = torch::randn({2, 512, 2048}, torch::dtype(torch::kFloat32).device(torch::kCUDA, device));
    checkCudaErrors(cudaSetDevice(device));
Jiezhong Qiu's avatar
update  
Jiezhong Qiu committed
81
    moe_cuda_forward(input, gate, weight);
Jiezhong Qiu's avatar
update  
Jiezhong Qiu committed
82
83
84
85
86
87
}
*/

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("forward", &moe_first_linear_forward, "MoE first linear forward (CUDA)");
  // m.def("backward", &lltm_backward, "LLTM backward (CUDA)");
Jiezhong Qiu's avatar
can run  
Jiezhong Qiu committed
88
}