gemm_cuda.h 188 Bytes
Newer Older
Ji Lin's avatar
Ji Lin committed
1
2
3
4
#include <torch/extension.h>

torch::Tensor gemm_forward_cuda(torch::Tensor _in_feats, torch::Tensor _kernel,
    torch::Tensor _scaling_factors, torch::Tensor _zeros, int split_k_iters);