comp files
This commit is contained in:
Aryan Singh committed
1 parent
fa3efcaf07
commit
1f77d0131d
27 files changed
+1238
-798
No files matched your search
+197
-19
@@ -3,35 +3,213 @@
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
#include "../utils/cuda_utils.cuh"
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/gemm/device/gemm.h"
|
||||
#include "/work/hdd/bche/hshah6/cutlass/include/cutlass/epilogue/thread/linear_combination_gelu.h"
|
||||
|
||||
__global__ void matmul_forward_kernel(float* out, const float* inp, const float* weight,
|
||||
const float* bias,int B,int T, int C, int OC) {
|
||||
// Implement this
|
||||
int b = blockDim.x*blockIdx.x + threadIdx.x; // B
|
||||
int t = blockDim.y*blockIdx.y + threadIdx.y; // T
|
||||
int oc = blockDim.z*blockIdx.z + threadIdx.z; // OC
|
||||
//referenced: https://github.com/NVIDIA/cutlass/blob/main/examples/14_ampere_tf32_tensorop_gemm/ampere_tf32_tensorop_gemm.cu
|
||||
|
||||
if (b >= B || t >= T || oc >= OC) {
|
||||
return;
|
||||
//removing bias overhead and fusing with cutlass
|
||||
// __global__ void bias_add_filler_kernel(float* __restrict__ out, const float* __restrict__ bias, int BT, int OC) {
|
||||
// int x = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
// if (x < BT * OC) out[x] = bias[x % OC];
|
||||
// }
|
||||
|
||||
//using big cutlass for kernels for BT>=64
|
||||
using GemmDefault = cutlass::gemm::device::Gemm<
|
||||
float, cutlass::layout::RowMajor,
|
||||
float, cutlass::layout::ColumnMajor,
|
||||
float, cutlass::layout::RowMajor,
|
||||
float,
|
||||
cutlass::arch::OpClassTensorOp,
|
||||
cutlass::arch::Sm80,
|
||||
cutlass::gemm::GemmShape<128,128,32>,
|
||||
cutlass::gemm::GemmShape<64,64,32>,
|
||||
cutlass::gemm::GemmShape<16,8,8>,
|
||||
cutlass::epilogue::thread::LinearCombination<float, 128 / cutlass::sizeof_bits<float>::value, float, float>,
|
||||
cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>,
|
||||
3 // Stages
|
||||
>;
|
||||
|
||||
//small cutlass for kernels for BT<64, we can use smaller threadblocks and warps to increase occupancy
|
||||
using GemmsmallDefault = cutlass::gemm::device::Gemm<
|
||||
float, cutlass::layout::RowMajor,
|
||||
float, cutlass::layout::ColumnMajor,
|
||||
float, cutlass::layout::RowMajor,
|
||||
float,
|
||||
cutlass::arch::OpClassTensorOp,
|
||||
cutlass::arch::Sm80,
|
||||
cutlass::gemm::GemmShape<32,128,32>,
|
||||
cutlass::gemm::GemmShape<32,64,32>,
|
||||
cutlass::gemm::GemmShape<16,8,8>,
|
||||
cutlass::epilogue::thread::LinearCombination<float, 128 / cutlass::sizeof_bits<float>::value, float, float>,
|
||||
cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>,
|
||||
2
|
||||
>;
|
||||
|
||||
using GemmMediumDefault = cutlass::gemm::device::Gemm<
|
||||
float, cutlass::layout::RowMajor,
|
||||
float, cutlass::layout::ColumnMajor,
|
||||
float, cutlass::layout::RowMajor,
|
||||
float,
|
||||
cutlass::arch::OpClassTensorOp,
|
||||
cutlass::arch::Sm80,
|
||||
cutlass::gemm::GemmShape<64,128,32>,
|
||||
cutlass::gemm::GemmShape<32,64,32>,
|
||||
cutlass::gemm::GemmShape<16,8,8>,
|
||||
cutlass::epilogue::thread::LinearCombination<float, 128 / cutlass::sizeof_bits<float>::value, float, float>,
|
||||
cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>,
|
||||
3
|
||||
>;
|
||||
|
||||
inline void matmul_forward(float* __restrict__ out, const float* __restrict__ inp, const float* __restrict__ weight,
|
||||
const float* __restrict__ bias,int B,int T, int C, int OC) {
|
||||
int BT = B*T;
|
||||
|
||||
if(bias != NULL){
|
||||
if(BT <= 32){
|
||||
GemmsmallDefault gemm_op;
|
||||
typename GemmsmallDefault::Arguments args{
|
||||
{BT, OC, C}, // problem size
|
||||
{inp, C}, // A matrix
|
||||
{weight, C}, // B matrix
|
||||
{bias, 0}, // C matrix
|
||||
{out, OC}, // D matrix (output)
|
||||
{1.0f, 1.0f} // alpha, beta for the epilogue
|
||||
};
|
||||
gemm_op(args);
|
||||
}
|
||||
else if(BT <= 128){
|
||||
GemmMediumDefault gemm_op;
|
||||
typename GemmMediumDefault::Arguments args{
|
||||
{BT, OC, C}, // problem size
|
||||
{inp, C}, // A matrix
|
||||
{weight, C}, // B matrix
|
||||
{bias, 0}, // C matrix
|
||||
{out, OC}, // D matrix (output)
|
||||
{1.0f, 1.0f} // alpha, beta for the epilogue
|
||||
};
|
||||
gemm_op(args);
|
||||
}
|
||||
else {
|
||||
GemmDefault gemm_op;
|
||||
typename GemmDefault::Arguments args{
|
||||
{BT, OC, C}, // problem size
|
||||
{inp, C}, // A matrix
|
||||
{weight, C}, // B matrix
|
||||
{bias, 0}, // C matrix
|
||||
{out, OC}, // D matrix (output)
|
||||
{1.0f, 1.0f} // alpha, beta for the epilogue
|
||||
};
|
||||
gemm_op(args);
|
||||
}
|
||||
}
|
||||
else{
|
||||
if(BT <= 32){
|
||||
GemmsmallDefault gemm_op;
|
||||
typename GemmsmallDefault::Arguments args{
|
||||
{BT, OC, C}, // problem size
|
||||
{inp, C}, // A matrix
|
||||
{weight, C}, // B matrix
|
||||
{out, OC}, // C matrix (output)
|
||||
{out, OC}, // D matrix (output)
|
||||
{1.0f, 0.0f} // alpha, beta for the epilogue
|
||||
};
|
||||
gemm_op(args);
|
||||
}
|
||||
else if(BT <= 128){
|
||||
GemmDefault gemm_op;
|
||||
typename GemmDefault::Arguments args{
|
||||
{BT, OC, C}, // problem size
|
||||
{inp, C}, // A matrix
|
||||
{weight, C}, // B matrix
|
||||
{out, OC}, // C matrix (output)
|
||||
{out, OC}, // D matrix (output)
|
||||
{1.0f, 0.0f} // alpha, beta for the epilogue
|
||||
};
|
||||
gemm_op(args);
|
||||
}
|
||||
else{
|
||||
GemmDefault gemm_op;
|
||||
typename GemmDefault::Arguments args{
|
||||
{BT, OC, C}, // problem size
|
||||
{inp, C}, // A matrix
|
||||
{weight, C}, // B matrix
|
||||
{out, OC}, // C matrix (output)
|
||||
{out, OC}, // D matrix (output)
|
||||
{1.0f, 0.0f} // alpha, beta for the epilogue
|
||||
};
|
||||
gemm_op(args);
|
||||
}
|
||||
}
|
||||
|
||||
float sum = (bias != NULL) ? bias[oc] : 0.0f;
|
||||
const float* inp_token_ptr = inp + b * T * C + t * C;
|
||||
for(int i = 0; i < C; i++){
|
||||
sum += inp_token_ptr[i] * weight[oc*C + i] ;
|
||||
}
|
||||
out[b * T * OC + t * OC + oc] = sum;
|
||||
|
||||
}
|
||||
|
||||
using Gemmgelu = cutlass::gemm::device::Gemm<
|
||||
float, cutlass::layout::RowMajor,
|
||||
float, cutlass::layout::ColumnMajor,
|
||||
float, cutlass::layout::RowMajor,
|
||||
float,
|
||||
cutlass::arch::OpClassTensorOp,
|
||||
cutlass::arch::Sm80,
|
||||
cutlass::gemm::GemmShape<128, 128, 32>,
|
||||
cutlass::gemm::GemmShape<64, 64, 32>,
|
||||
cutlass::gemm::GemmShape<16, 8, 8>,
|
||||
cutlass::epilogue::thread::LinearCombinationGELU<float, 128 / cutlass::sizeof_bits<float>::value, float, float>,
|
||||
cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>,
|
||||
3>;
|
||||
|
||||
using GemmgeluSmall = cutlass::gemm::device::Gemm<
|
||||
float, cutlass::layout::RowMajor,
|
||||
float, cutlass::layout::ColumnMajor,
|
||||
float, cutlass::layout::RowMajor,
|
||||
float,
|
||||
cutlass::arch::OpClassTensorOp,
|
||||
cutlass::arch::Sm80,
|
||||
cutlass::gemm::GemmShape<32, 128, 32>,
|
||||
cutlass::gemm::GemmShape<32, 64, 32>,
|
||||
cutlass::gemm::GemmShape<16, 8, 8>,
|
||||
cutlass::epilogue::thread::LinearCombinationGELU<float, 128 / cutlass::sizeof_bits<float>::value, float, float>,
|
||||
cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>,
|
||||
2>;
|
||||
|
||||
using GemmgeluMedium = cutlass::gemm::device::Gemm<
|
||||
float, cutlass::layout::RowMajor,
|
||||
float, cutlass::layout::ColumnMajor,
|
||||
float, cutlass::layout::RowMajor,
|
||||
float,
|
||||
cutlass::arch::OpClassTensorOp,
|
||||
cutlass::arch::Sm80,
|
||||
cutlass::gemm::GemmShape<64, 128, 32>,
|
||||
cutlass::gemm::GemmShape<32, 64, 32>,
|
||||
cutlass::gemm::GemmShape<16, 8, 8>,
|
||||
cutlass::epilogue::thread::LinearCombinationGELU<float, 128 / cutlass::sizeof_bits<float>::value, float, float>,
|
||||
cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>,
|
||||
3>;
|
||||
|
||||
// Launch kernel here
|
||||
void matmul_forward(float* out, const float* inp, const float* weight, const float* bias,
|
||||
inline void matmul_forward_fused_gelu(float* out, const float* inp, const float* weight, const float* bias,
|
||||
int B, int T, int C, int OC) {
|
||||
// Implement this
|
||||
dim3 block(8,8,8);
|
||||
dim3 grid((B+block.x-1)/block.x, (T+block.y-1)/block.y, (OC+block.z-1)/block.z);
|
||||
|
||||
matmul_forward_kernel<<<grid, block>>>(out, inp, weight, bias,B,T, C, OC);
|
||||
int BT = B*T;
|
||||
|
||||
if (BT <= 32) {
|
||||
GemmgeluSmall op;
|
||||
typename GemmgeluSmall::Arguments args{
|
||||
{BT, OC, C}, {inp, C}, {weight, C}, {bias, 0}, {out, OC}, {1.0f, 1.0f}};
|
||||
op(args);
|
||||
} else if (BT <= 128) {
|
||||
GemmgeluMedium op;
|
||||
typename GemmgeluMedium::Arguments args{
|
||||
{BT, OC, C}, {inp, C}, {weight, C}, {bias, 0}, {out, OC}, {1.0f, 1.0f}};
|
||||
op(args);
|
||||
} else {
|
||||
Gemmgelu op;
|
||||
typename Gemmgelu::Arguments args{
|
||||
{BT, OC, C}, {inp, C}, {weight, C}, {bias, 0}, {out, OC}, {1.0f, 1.0f}};
|
||||
op(args);
|
||||
}
|
||||
}
|
||||
|
||||
#endif // __MATMUL_KERNEL_CUH__
|
||||
Reference in new issue
Block a user