Make various updates and fixes: (#164)

- Add BF16 support for SM90 and SM100
- Refactor Python APIs
- Other fixes and code refactoring

csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp

@@ -3,6 +3,7 @@
 #include <torch/python.h>
 
 #include "../../jit/compiler.hpp"
+#include "../../jit/device_runtime.hpp"
 #include "../../jit/kernel_runtime.hpp"
 #include "../../utils/exception.hpp"
 #include "../../utils/format.hpp"
@@ -139,7 +140,7 @@ static void sm90_m_grouped_fp8_gemm_contiguous_1d2d(const torch::Tensor& a, cons
     const auto& aligned_k = align(k, 128);
     const auto& config = get_best_config<SM90ArchSpec>(
         GemmType::MGroupedContiguous, KernelType::Kernel1D2D,
-        m, n, k, num_groups, major_a, major_b,
+        m, n, k, 1, major_a, major_b,
         torch::kFloat8_e4m3fn, d.scalar_type(), false,
         device_runtime->get_num_sms());
 
@@ -185,7 +186,7 @@ static void sm90_m_grouped_fp8_gemm_contiguous_1d2d(const torch::Tensor& a, cons
     SM90FP8Gemm1D2DRuntime::launch(runtime, args);
 }
 
-static void sm90_fp8_m_grouped_gemm_masked_1d2d(const torch::Tensor& a, const torch::Tensor& sfa,
+static void sm90_m_grouped_fp8_gemm_masked_1d2d(const torch::Tensor& a, const torch::Tensor& sfa,
                                                 const torch::Tensor& b, const torch::Tensor& sfb,
                                                 const torch::Tensor& d,
                                                 const torch::Tensor& masked_m,