1
0

eetpy.cpp 308 B

123456789
  1. #include <pybind11/pybind11.h>
  2. #include <torch/extension.h>
  3. #include "cutlass_kernels/fpA_intB_gemm_wrapper.h"
  4. PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
  5. {
  6. m.def("w8_a16_gemm", &w8_a16_gemm_forward_cuda, "Weight only gemm");
  7. m.def("w8_a16_gemm_", &w8_a16_gemm_forward_cuda_, "Weight only gemm inplace");
  8. }