moe.cpp 3.91 KB
Newer Older
Jiezhong Qiu's avatar
can run  
Jiezhong Qiu committed
1
2
3
4
5
6
7
8
9
10
11
12
13
#include <torch/extension.h>
#include <torch/torch.h>

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

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


Jiezhong Qiu's avatar
update  
Jiezhong Qiu committed
17
const int num_stream=16;
Jiezhong Qiu's avatar
can run  
Jiezhong Qiu committed
18
19
20
21
22
23
24
25
26
27
28
29
30

// std::vector<torch::Tensor> 
void moe_cuda_forward(
        torch::Tensor input, // [B x D_model]
        torch::Tensor gate,  // [B x N]
        torch::Tensor weight, // [N x D_model x D_ffn]
        torch::Tensor bias // [N x D_ffn]
        ) {
    const auto batch_size = input.size(0);
    const auto num_expert = gate.size(1);
    const auto d_model = weight.size(1);
    const auto d_ffn = weight.size(2);
    auto output = input.new_zeros({batch_size, num_expert, d_ffn});
Jiezhong Qiu's avatar
update  
Jiezhong Qiu committed
31
32
    std::cout << output << std::endl;
    
Jiezhong Qiu's avatar
can run  
Jiezhong Qiu committed
33
34

    cublasHandle_t handle;
Jiezhong Qiu's avatar
update  
Jiezhong Qiu committed
35
36
    checkCudaErrors(cublasCreate(&handle));
    
Jiezhong Qiu's avatar
can run  
Jiezhong Qiu committed
37
38
    cudaStream_t stream[num_stream];
    for (size_t i=0; i<num_stream; ++i) {
Jiezhong Qiu's avatar
update  
Jiezhong Qiu committed
39
        checkCudaErrors(cudaStreamCreate(&stream[i]));
Jiezhong Qiu's avatar
can run  
Jiezhong Qiu committed
40
    }
Jiezhong Qiu's avatar
update  
Jiezhong Qiu committed
41

Jiezhong Qiu's avatar
can run  
Jiezhong Qiu committed
42
43
44
45
46
47
48
    
    size_t s;
    for (size_t i=0; i<batch_size; ++i) {
        for (size_t j=0; j<num_expert; ++j) {
            s = (i * num_expert + j) % num_stream;
            printf("i=%d j=%d goes to stream %d\n", i, j, s);
            cublasSetStream(handle, stream[s]);
Jiezhong Qiu's avatar
update  
Jiezhong Qiu committed
49
50
51
52
53
54
55
            if (input.scalar_type() == torch::ScalarType::Float) {
                float alpha = 1.0;
                float beta = 0.0;
                std::cout << input[i] << std::endl;
                std::cout << weight.index(gate[i][j]) << std::endl;
                std::cout << output[i][j] << std::endl;
                cublasSgemm(handle, 
Jiezhong Qiu's avatar
can run  
Jiezhong Qiu committed
56
57
                    CUBLAS_OP_N, 
                    CUBLAS_OP_N,
Jiezhong Qiu's avatar
update  
Jiezhong Qiu committed
58
59
60
                    1, // m
                    d_ffn, // n
                    d_model, // k
Jiezhong Qiu's avatar
can run  
Jiezhong Qiu committed
61
                    &alpha,
Jiezhong Qiu's avatar
update  
Jiezhong Qiu committed
62
63
                    input.data_ptr<float>() + i * d_model,
                    // input[i].data_ptr<float>(),
Jiezhong Qiu's avatar
can run  
Jiezhong Qiu committed
64
                    1,
Jiezhong Qiu's avatar
update  
Jiezhong Qiu committed
65
                    weight.index(gate[i][j]).data_ptr<float>(),
Jiezhong Qiu's avatar
can run  
Jiezhong Qiu committed
66
67
                    d_model,
                    &beta,
Jiezhong Qiu's avatar
update  
Jiezhong Qiu committed
68
                    output.data_ptr<float>() + i * num_expert * d_ffn + j * d_ffn,
Jiezhong Qiu's avatar
can run  
Jiezhong Qiu committed
69
70
                    1);
            } else {
Jiezhong Qiu's avatar
update  
Jiezhong Qiu committed
71
                printf("only support float!!!\n");
Jiezhong Qiu's avatar
can run  
Jiezhong Qiu committed
72
73
74
            }
        }
    }
Jiezhong Qiu's avatar
update  
Jiezhong Qiu committed
75
76
77
    cudaDeviceSynchronize();
    printf("synchronized\n");
    
Jiezhong Qiu's avatar
can run  
Jiezhong Qiu committed
78

Jiezhong Qiu's avatar
update  
Jiezhong Qiu committed
79
    
Jiezhong Qiu's avatar
can run  
Jiezhong Qiu committed
80
81
82
    for (size_t i=0; i<num_stream; ++i) {
        cudaStreamDestroy(stream[i]);
    }
Jiezhong Qiu's avatar
update  
Jiezhong Qiu committed
83
84
    std::cout << output << std::endl;
    
Jiezhong Qiu's avatar
can run  
Jiezhong Qiu committed
85
86
87
88
89
90
91
92
93
94
95
96
97
    cublasDestroy(handle);
}


// 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)


int main() {
Jiezhong Qiu's avatar
update  
Jiezhong Qiu committed
98
99
100
101
102
    int device=2;
    torch::Tensor input = torch::randn({2, 4}, torch::dtype(torch::kFloat32).device(torch::kCUDA, device));
    torch::Tensor gate = torch::zeros({2, 1}, torch::dtype(torch::kInt64).device(torch::kCUDA, device));
    torch::Tensor weight = torch::randn({2, 4, 4}, torch::dtype(torch::kFloat32).device(torch::kCUDA, device));
    torch::Tensor bias = torch::randn({2, 4}, torch::dtype(torch::kFloat32).device(torch::kCUDA, device));
Jiezhong Qiu's avatar
can run  
Jiezhong Qiu committed
103
104
105
    std::cout << input << std::endl;
    moe_cuda_forward(input, gate, weight, bias);
}