feat: support libtorch

2025-08-26 19:17:08 -07:00
parent 239112cb4c
commit 301cbc1d75
22 changed files with 652 additions and 245 deletions
--- a/csrc/apis/attention.hpp
+++ b/csrc/apis/attention.hpp
@@ -219,22 +219,4 @@ static torch::Tensor fp8_paged_mqa_logits(const torch::Tensor& q,
    return logits;
 }

-static void register_apis(pybind11::module_& m) {
-    m.def("fp8_gemm_nt_skip_head_mid", &fp8_gemm_nt_skip_head_mid,
-          py::arg("a"), py::arg("b"), py::arg("d"), py::arg("head_splits"),
-          py::arg("recipe") = std::nullopt,
-          py::arg("compiled_dims") = "nk",
-          py::arg("disable_ue8m0_cast") = false);
-    m.def("fp8_mqa_logits", &fp8_mqa_logits,
-      py::arg("q"), py::arg("kv"), py::arg("weights"),
-      py::arg("cu_seq_len_k_start"), py::arg("cu_seq_len_k_end"),
-      py::arg("clean_logits") = true);
-    m.def("get_paged_mqa_logits_metadata", &get_paged_mqa_logits_metadata,
-          py::arg("context_lens"), py::arg("block_kv"), py::arg("num_sms"));
-    m.def("fp8_paged_mqa_logits", &fp8_paged_mqa_logits,
-          py::arg("q"), py::arg("kv_cache"), py::arg("weights"),
-          py::arg("context_lens"), py::arg("block_table"), py::arg("schedule_meta"),
-          py::arg("max_context_len"), py::arg("clean_logits") = false);
-}
-
 } // namespace deep_gemm::attention
--- a/csrc/apis/einsum.hpp
+++ b/csrc/apis/einsum.hpp
@@ -1,8 +1,5 @@
 #pragma once

-#include <pybind11/pybind11.h>
-#include <torch/python.h>
-
 #include "../utils/exception.hpp"
 #include "../utils/format.hpp"
 #include "../utils/layout.hpp"
@@ -106,10 +103,4 @@ static void einsum(const std::string& expr,
    }
 }

-static void register_apis(pybind11::module_& m) {
-    m.def("einsum", &einsum,
-          py::arg("expr"), py::arg("a"), py::arg("b"),
-          py::arg("d"), py::arg("c") = std::nullopt);
-}
-
 } // namespace deep_gemm::einsum
--- a/csrc/apis/gemm.hpp
+++ b/csrc/apis/gemm.hpp
@@ -500,84 +500,4 @@ static void cublaslt_gemm_tt(const torch::Tensor& a, const torch::Tensor& b,
    cublaslt_gemm_nt(a.transpose(0, 1), b, d, c);
 }

-static void register_apis(pybind11::module_& m) {
-    // FP8 GEMMs
-     m.def("fp8_gemm_nt", &fp8_gemm_nt,
-          py::arg("a"), py::arg("b"), py::arg("d"),
-          py::arg("c") = std::nullopt, py::arg("recipe") = std::nullopt,
-          py::arg("compiled_dims") = "nk",
-          py::arg("disable_ue8m0_cast") = false);
-    m.def("fp8_gemm_nn", &fp8_gemm_nn,
-          py::arg("a"), py::arg("b"), py::arg("d"),
-          py::arg("c") = std::nullopt, py::arg("recipe") = std::nullopt,
-          py::arg("compiled_dims") = "nk",
-          py::arg("disable_ue8m0_cast") = false);
-    m.def("fp8_gemm_tn", &fp8_gemm_tn,
-          py::arg("a"), py::arg("b"), py::arg("d"),
-          py::arg("c") = std::nullopt, py::arg("recipe") = std::nullopt,
-          py::arg("compiled_dims") = "mn",
-          py::arg("disable_ue8m0_cast") = false);
-    m.def("fp8_gemm_tt", &fp8_gemm_tt,
-          py::arg("a"), py::arg("b"), py::arg("d"),
-          py::arg("c") = std::nullopt, py::arg("recipe") = std::nullopt,
-          py::arg("compiled_dims") = "mn",
-          py::arg("disable_ue8m0_cast") = false);
-    m.def("m_grouped_fp8_gemm_nt_contiguous", &m_grouped_fp8_gemm_nt_contiguous,
-          py::arg("a"), py::arg("b"), py::arg("d"), py::arg("m_indices"),
-          py::arg("recipe") = std::nullopt, py::arg("compiled_dims") = "nk",
-          py::arg("disable_ue8m0_cast") = false);
-    m.def("m_grouped_fp8_gemm_nn_contiguous", &m_grouped_fp8_gemm_nn_contiguous,
-          py::arg("a"), py::arg("b"), py::arg("d"), py::arg("m_indices"),
-          py::arg("recipe") = std::nullopt, py::arg("compiled_dims") = "nk",
-          py::arg("disable_ue8m0_cast") = false);
-    m.def("m_grouped_fp8_gemm_nt_masked", &m_grouped_fp8_gemm_nt_masked,
-          py::arg("a"), py::arg("b"), py::arg("d"), py::arg("masked_m"),
-          py::arg("expected_m"), py::arg("recipe") = std::nullopt,
-          py::arg("compiled_dims") = "nk", py::arg("disable_ue8m0_cast") = false);
-    m.def("k_grouped_fp8_gemm_tn_contiguous", &k_grouped_fp8_gemm_tn_contiguous,
-          py::arg("a"), py::arg("b"), py::arg("d"), py::arg("ks"),
-          py::arg("ks_tensor"), py::arg("c") = std::nullopt,
-          py::arg("recipe") = std::make_tuple(1, 1, 128),
-          py::arg("compiled_dims") = "mn");
-    m.def("k_grouped_fp8_gemm_nt_contiguous", &k_grouped_fp8_gemm_nt_contiguous,
-          py::arg("a"), py::arg("b"), py::arg("d"), py::arg("ks"),
-          py::arg("ks_tensor"), py::arg("c") = std::nullopt,
-          py::arg("recipe") = std::make_tuple(1, 1, 128),
-          py::arg("compiled_dims") = "mn");
-
-    // BF16 GEMMs
-    m.def("bf16_gemm_nt", &bf16_gemm_nt,
-          py::arg("a"), py::arg("b"), py::arg("d"),
-          py::arg("c") = std::nullopt,
-          py::arg("compiled_dims") = "nk");
-    m.def("bf16_gemm_nn", &bf16_gemm_nn,
-          py::arg("a"), py::arg("b"), py::arg("d"),
-          py::arg("c") = std::nullopt,
-          py::arg("compiled_dims") = "nk");
-    m.def("bf16_gemm_tn", &bf16_gemm_tn,
-          py::arg("a"), py::arg("b"), py::arg("d"),
-          py::arg("c") = std::nullopt,
-          py::arg("compiled_dims") = "mn");
-    m.def("bf16_gemm_tt", &bf16_gemm_tt,
-          py::arg("a"), py::arg("b"), py::arg("d"),
-          py::arg("c") = std::nullopt,
-          py::arg("compiled_dims") = "mn");
-    m.def("m_grouped_bf16_gemm_nt_contiguous", &m_grouped_bf16_gemm_nt_contiguous,
-          py::arg("a"), py::arg("b"), py::arg("d"), py::arg("m_indices"),
-          py::arg("compiled_dims") = "nk");
-    m.def("m_grouped_bf16_gemm_nt_masked", &m_grouped_bf16_gemm_nt_masked,
-          py::arg("a"), py::arg("b"), py::arg("d"), py::arg("masked_m"),
-          py::arg("expected_m"), py::arg("compiled_dims") = "nk");
-
-    // cuBLASLt GEMMs
-    m.def("cublaslt_gemm_nt", &cublaslt_gemm_nt,
-          py::arg("a"), py::arg("b"), py::arg("d"), py::arg("c") = std::nullopt);
-    m.def("cublaslt_gemm_nn", &cublaslt_gemm_nn,
-          py::arg("a"), py::arg("b"), py::arg("d"), py::arg("c") = std::nullopt);
-    m.def("cublaslt_gemm_tn", &cublaslt_gemm_tn,
-          py::arg("a"), py::arg("b"), py::arg("d"), py::arg("c") = std::nullopt);
-    m.def("cublaslt_gemm_tt", &cublaslt_gemm_tt,
-          py::arg("a"), py::arg("b"), py::arg("d"), py::arg("c") = std::nullopt);
-}
-
 } // namespace deep_gemm::gemm
--- a/csrc/apis/layout.hpp
+++ b/csrc/apis/layout.hpp
@@ -69,17 +69,4 @@ static torch::Tensor transform_k_grouped_sf_into_required_layout(const torch::Te
    DG_HOST_UNREACHABLE("Unknown cases");
 }

-static void register_apis(pybind11::module_& m) {
-    m.def("transform_sf_into_required_layout", &transform_sf_into_required_layout,
-      py::arg("sf"), py::arg("mn"), py::arg("k"), py::arg("recipe"),
-      py::arg("num_groups") = std::nullopt, py::arg("is_sfa") = false,
-      py::arg("disable_ue8m0_cast") = false);
-
-    m.def("get_tma_aligned_size", &get_tma_aligned_size);
-    m.def("get_mk_alignment_for_contiguous_layout", &get_mk_alignment_for_contiguous_layout);
-    m.def("get_mn_major_tma_aligned_tensor", &get_mn_major_tma_aligned_tensor);
-    m.def("get_mn_major_tma_aligned_packed_ue8m0_tensor", &get_mn_major_tma_aligned_packed_ue8m0_tensor);
-    m.def("get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor", &get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor);
-}
-
 } // namespace deep_gemm::layout
--- a/csrc/apis/runtime.hpp
+++ b/csrc/apis/runtime.hpp
@@ -5,24 +5,6 @@

 namespace deep_gemm::runtime {

-static void register_apis(pybind11::module_& m) {
-    m.def("set_num_sms", [&](const int& new_num_sms) {
-        device_runtime->set_num_sms(new_num_sms);
-    });
-    m.def("get_num_sms", [&]() {
-       return device_runtime->get_num_sms();
-    });
-    m.def("set_tc_util", [&](const int& new_tc_util) {
-        device_runtime->set_tc_util(new_tc_util);
-    });
-    m.def("get_tc_util", [&]() {
-        return device_runtime->get_tc_util();
-    });
-
-    m.def("init", [&](const std::string& library_root_path, const std::string& cuda_home_path_by_python) {
-        Compiler::prepare_init(library_root_path, cuda_home_path_by_python);
-        KernelRuntime::prepare_init(cuda_home_path_by_python);
-    });
-}
+// The init and other functions are now exposed via TORCH_LIBRARY in python_api.cpp

 } // namespace deep_gemm::runtime
--- a/csrc/jit/device_runtime.hpp
+++ b/csrc/jit/device_runtime.hpp
@@ -17,6 +17,11 @@ class DeviceRuntime {
    cublasLtHandle_t cublaslt_handle{};
    std::shared_ptr<torch::Tensor> cublaslt_workspace;

+    // cuBLASLt utils
+    static constexpr size_t kCublasLtWorkspaceSize = 32 * 1024 * 1024;
+    cublasLtHandle_t cublaslt_handle{};
+    std::shared_ptr<torch::Tensor> cublaslt_workspace;
+
 public:
    explicit DeviceRuntime() {
        cublaslt_workspace = std::make_shared<torch::Tensor>(torch::empty({kCublasLtWorkspaceSize}, dtype(torch::kByte).device(at::kCUDA)));
--- a/csrc/jit_kernels/impls/runtime_utils.hpp
+++ b/csrc/jit_kernels/impls/runtime_utils.hpp
@@ -1,7 +1,6 @@
 #pragma once

 #include <cuda.h>
-#include <torch/python.h>

 #include "../../utils/math.hpp"
 #include "../heuristics/sm90.hpp"
@@ -75,10 +74,6 @@ static CUtensorMapSwizzle mode_into_tensor_map_swizzle(const int& mode, const in
    }
 #endif

-    DG_HOST_ASSERT(base == 0);
-    switch (mode) {
-        case   0:
-        case  16: return CU_TENSOR_MAP_SWIZZLE_NONE;
        case  32: return CU_TENSOR_MAP_SWIZZLE_32B;
        case  64: return CU_TENSOR_MAP_SWIZZLE_64B;
        case 128: return CU_TENSOR_MAP_SWIZZLE_128B;
--- a/csrc/jit_kernels/impls/sm100_bf16_gemm.hpp
+++ b/csrc/jit_kernels/impls/sm100_bf16_gemm.hpp
@@ -1,7 +1,5 @@
 #pragma once

-#include <torch/python.h>
-
 #include "../../jit/compiler.hpp"
 #include "../../jit/device_runtime.hpp"
 #include "../../jit/kernel_runtime.hpp"
--- a/csrc/jit_kernels/impls/sm100_bmk_bnk_mn.hpp
+++ b/csrc/jit_kernels/impls/sm100_bmk_bnk_mn.hpp
@@ -1,7 +1,5 @@
 #pragma once

-#include <torch/python.h>
-
 #include "../../jit/compiler.hpp"
 #include "../../jit/device_runtime.hpp"
 #include "../../jit/kernel_runtime.hpp"
@@ -134,4 +132,4 @@ static void sm100_bmn_bnk_mn_gemm(const torch::Tensor &a,
    SM100BmkBnkMnRuntime::launch(runtime, args);
 }

-} // namespace deep_gemm
+} // namespace deep_gemm
--- a/csrc/jit_kernels/impls/sm100_fp8_gemm_1d1d.hpp
+++ b/csrc/jit_kernels/impls/sm100_fp8_gemm_1d1d.hpp
@@ -1,7 +1,5 @@
 #pragma once

-#include <torch/python.h>
-
 #include "../../jit/compiler.hpp"
 #include "../../jit/device_runtime.hpp"
 #include "../../jit/kernel_runtime.hpp"
--- a/csrc/jit_kernels/impls/sm100_fp8_gemm_1d2d.hpp
+++ b/csrc/jit_kernels/impls/sm100_fp8_gemm_1d2d.hpp
@@ -1,7 +1,5 @@
 #pragma once

-#include <torch/python.h>
-
 #include "../../jit/compiler.hpp"
 #include "../../jit/device_runtime.hpp"
 #include "../../jit/kernel_runtime.hpp"
--- a/csrc/jit_kernels/impls/sm90_bf16_gemm.hpp
+++ b/csrc/jit_kernels/impls/sm90_bf16_gemm.hpp
@@ -1,7 +1,5 @@
 #pragma once

-#include <torch/python.h>
-
 #include "../../jit/compiler.hpp"
 #include "../../jit/kernel_runtime.hpp"
 #include "../../utils/exception.hpp"
--- a/csrc/jit_kernels/impls/sm90_bmk_bnk_mn.hpp
+++ b/csrc/jit_kernels/impls/sm90_bmk_bnk_mn.hpp
@@ -1,7 +1,5 @@
 #pragma once

-#include <torch/python.h>
-
 #include "../../jit/compiler.hpp"
 #include "../../jit/device_runtime.hpp"
 #include "../../jit/kernel_runtime.hpp"
--- a/csrc/jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp
+++ b/csrc/jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp
@@ -1,7 +1,5 @@
 #pragma once

-#include <torch/python.h>
-
 #include "../../jit/compiler.hpp"
 #include "../../jit/device_runtime.hpp"
 #include "../../jit/kernel_runtime.hpp"
@@ -133,7 +131,7 @@ static void sm90_fp8_gemm_1d1d(const torch::Tensor& a, const torch::Tensor& sfa,
    const auto& code = SM90FP8Gemm1D1DRuntime::generate(args);
    const auto& runtime = compiler->build("sm90_fp8_gemm_1d1d", code);

-    SM90FP8Gemm1D1DRuntime::launch(runtime, args);
+    MAYBE_LAUNCH(SM90FP8Gemm1D1DRuntime::launch(runtime, args));
 }

 static void sm90_fp8_k_grouped_gemm_1d1d(const torch::Tensor& a, const torch::Tensor& sfa,
@@ -208,7 +206,7 @@ static void sm90_fp8_k_grouped_gemm_1d1d(const torch::Tensor& a, const torch::Te
    const auto& code = SM90FP8Gemm1D1DRuntime::generate(args);
    const auto& runtime = compiler->build("sm90_fp8_gemm_1d1d", code);

-    SM90FP8Gemm1D1DRuntime::launch(runtime, args);
+    MAYBE_LAUNCH(SM90FP8Gemm1D1DRuntime::launch(runtime, args));
 }

 } // namespace deep_gemm
--- a/csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp
+++ b/csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp
@@ -1,7 +1,5 @@
 #pragma once

-#include <torch/python.h>
-
 #include "../../jit/compiler.hpp"
 #include "../../jit/device_runtime.hpp"
 #include "../../jit/kernel_runtime.hpp"
--- a/csrc/jit_kernels/impls/smxx_layout.hpp
+++ b/csrc/jit_kernels/impls/smxx_layout.hpp
@@ -1,7 +1,5 @@
 #pragma once

-#include <torch/python.h>
-
 #include "../../jit/kernel_runtime.hpp"
 #include "../../utils/exception.hpp"
 #include "../../utils/format.hpp"
--- a/csrc/python_api.cpp
+++ b/csrc/python_api.cpp
@@ -1,5 +1,11 @@
-#include <pybind11/pybind11.h>
-#include <torch/python.h>
+#include <torch/library.h>
+#include <torch/types.h>
+#include <vector>
+#include <string>
+#include <optional>
+#include <tuple>
+#include <numeric>
+#include <Python.h>

 #include "apis/attention.hpp"
 #include "apis/einsum.hpp"
@@ -7,17 +13,510 @@
 #include "apis/layout.hpp"
 #include "apis/runtime.hpp"

+#include "jit/compiler.hpp"
+#include "jit/device_runtime.hpp"
+#include "jit/kernel_runtime.hpp"
+
 #ifndef TORCH_EXTENSION_NAME
-#define TORCH_EXTENSION_NAME deep_gemm_cpp
+#define TORCH_EXTENSION_NAME deep_gemm
 #endif

-// ReSharper disable once CppParameterMayBeConstPtrOrRef
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-    m.doc() = "DeepGEMM C++ library";

-    deep_gemm::attention::register_apis(m);
-    deep_gemm::einsum::register_apis(m);
-    deep_gemm::gemm::register_apis(m);
-    deep_gemm::layout::register_apis(m);
-    deep_gemm::runtime::register_apis(m);
+#define _CONCAT(A, B) A##B
+#define CONCAT(A, B) _CONCAT(A, B)
+
+#define _STRINGIFY(A) #A
+#define STRINGIFY(A) _STRINGIFY(A)
+
+#define REGISTER_EXTENSION(NAME)                                                                      \
+  PyMODINIT_FUNC CONCAT(PyInit_, NAME)() {                                                            \
+    static struct PyModuleDef module = {PyModuleDef_HEAD_INIT, STRINGIFY(NAME), nullptr, 0, nullptr}; \
+    return PyModule_Create(&module);                                                                  \
+  }
+
+namespace {
+
+std::optional<std::tuple<int, int, int>> to_recipe_tuple(const c10::optional<c10::IntArrayRef>& recipe_opt) {
+    if (!recipe_opt.has_value()) {
+        return std::nullopt;
+    }
+    auto recipe_ref = recipe_opt.value();
+    TORCH_CHECK(recipe_ref.size() == 3, "Recipe must be a list/tuple of 3 integers.");
+    return std::make_tuple(static_cast<int>(recipe_ref[0]), static_cast<int>(recipe_ref[1]), static_cast<int>(recipe_ref[2]));
 }
+
+std::tuple<int, int, int> to_recipe_tuple_default(c10::IntArrayRef recipe_ref) {
+    TORCH_CHECK(recipe_ref.size() == 3, "Recipe must be a list/tuple of 3 integers.");
+    return std::make_tuple(static_cast<int>(recipe_ref[0]), static_cast<int>(recipe_ref[1]), static_cast<int>(recipe_ref[2]));
+}
+
+// Accept Tensor, (Tensor, Tensor) tuple, or [Tensor, Tensor] list; return (tensor, scale)
+std::pair<at::Tensor, at::Tensor> parse_tensor_or_tuple(const c10::IValue& input) {
+    if (input.isTuple()) {
+        auto tuple = input.toTuple();
+        TORCH_CHECK(tuple->elements().size() >= 2, "Expected (Tensor, Tensor) tuple");
+        return {tuple->elements()[0].toTensor(), tuple->elements()[1].toTensor()};
+    } else if (input.isList()) {
+        auto list = input.toList();
+        TORCH_CHECK(list.size() >= 2, "Expected [Tensor, Tensor] list");
+        return {list.get(0).toTensor(), list.get(1).toTensor()};
+    } else if (input.isTensor()) {
+        auto tensor = input.toTensor();
+        auto scale = at::ones({1}, tensor.options().dtype(at::kFloat));
+        return {tensor, scale};
+    }
+    TORCH_CHECK(false, "Expected Tensor, (Tensor, Tensor) tuple, or [Tensor, Tensor] list");
+}
+
+} // anonymous namespace
+
+namespace deep_gemm_wrappers {
+
+// Runtime wrappers
+void set_num_sms_wrapper(int64_t new_num_sms) {
+    deep_gemm::device_runtime->set_num_sms(new_num_sms);
+}
+
+int64_t get_num_sms_wrapper() {
+    return deep_gemm::device_runtime->get_num_sms();
+}
+
+void set_tc_util_wrapper(int64_t new_tc_util) {
+    deep_gemm::device_runtime->set_tc_util(new_tc_util);
+}
+
+int64_t get_tc_util_wrapper() {
+    return deep_gemm::device_runtime->get_tc_util();
+}
+
+void init_wrapper(const std::string& library_root_path, const std::string& cuda_home_path_by_python) {
+    deep_gemm::Compiler::prepare_init(library_root_path, cuda_home_path_by_python);
+    deep_gemm::KernelRuntime::prepare_init(cuda_home_path_by_python);
+}
+
+// Scalar layout utility wrappers (int64_t signatures for PyTorch registration)
+int64_t get_tma_aligned_size_wrapper(int64_t x, int64_t element_size);
+int64_t get_mk_alignment_for_contiguous_layout_wrapper();
+
+// Layout wrappers
+torch::Tensor transform_sf_into_required_layout_wrapper(const torch::Tensor& sf, int64_t mn, int64_t k, c10::IntArrayRef recipe, const c10::optional<int64_t>& num_groups, bool is_sfa, bool disable_ue8m0_cast) {
+    return deep_gemm::layout::transform_sf_into_required_layout(sf, mn, k, to_recipe_tuple_default(recipe), num_groups, is_sfa, disable_ue8m0_cast);
+}
+
+torch::Tensor get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor_wrapper(const torch::Tensor& sf, const torch::Tensor& ks_tensor, c10::List<int64_t> ks) {
+    std::vector<int> ks_vec;
+    ks_vec.reserve(ks.size());
+    for (const auto& k_val : ks) {
+        ks_vec.push_back(k_val);
+    }
+    return deep_gemm::get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor(sf, ks_tensor, ks_vec);
+}
+
+// GEMM wrappers
+void fp8_gemm_nt_wrapper(const torch::Tensor& a_val, const torch::Tensor& a_scale, const torch::Tensor& b_val, const torch::Tensor& b_scale, const torch::Tensor& d, const c10::optional<torch::Tensor>& c, const c10::optional<c10::IntArrayRef>& recipe, const std::string& compiled_dims, bool disable_ue8m0_cast) {
+    deep_gemm::gemm::fp8_gemm_nt({a_val, a_scale}, {b_val, b_scale}, d, c, to_recipe_tuple(recipe), compiled_dims, disable_ue8m0_cast);
+}
+
+void fp8_gemm_nn_wrapper(const torch::Tensor& a_val, const torch::Tensor& a_scale, const torch::Tensor& b_val, const torch::Tensor& b_scale, const torch::Tensor& d, const c10::optional<torch::Tensor>& c, const c10::optional<c10::IntArrayRef>& recipe, const std::string& compiled_dims, bool disable_ue8m0_cast) {
+    deep_gemm::gemm::fp8_gemm_nn({a_val, a_scale}, {b_val, b_scale}, d, c, to_recipe_tuple(recipe), compiled_dims, disable_ue8m0_cast);
+}
+
+void fp8_gemm_tn_wrapper(const torch::Tensor& a_val, const torch::Tensor& a_scale, const torch::Tensor& b_val, const torch::Tensor& b_scale, const torch::Tensor& d, const c10::optional<torch::Tensor>& c, const c10::optional<c10::IntArrayRef>& recipe, const std::string& compiled_dims, bool disable_ue8m0_cast) {
+    deep_gemm::gemm::fp8_gemm_tn({a_val, a_scale}, {b_val, b_scale}, d, c, to_recipe_tuple(recipe), compiled_dims, disable_ue8m0_cast);
+}
+
+void fp8_gemm_tt_wrapper(const torch::Tensor& a_val, const torch::Tensor& a_scale, const torch::Tensor& b_val, const torch::Tensor& b_scale, const torch::Tensor& d, const c10::optional<torch::Tensor>& c, const c10::optional<c10::IntArrayRef>& recipe, const std::string& compiled_dims, bool disable_ue8m0_cast) {
+    deep_gemm::gemm::fp8_gemm_tt({a_val, a_scale}, {b_val, b_scale}, d, c, to_recipe_tuple(recipe), compiled_dims, disable_ue8m0_cast);
+}
+
+void m_grouped_fp8_gemm_nt_contiguous_wrapper(const torch::Tensor& a_val, const torch::Tensor& a_scale, const torch::Tensor& b_val, const torch::Tensor& b_scale, const torch::Tensor& d, const torch::Tensor& m_indices, const c10::optional<c10::IntArrayRef>& recipe, const std::string& compiled_dims, bool disable_ue8m0_cast) {
+    deep_gemm::gemm::m_grouped_fp8_gemm_nt_contiguous({a_val, a_scale}, {b_val, b_scale}, d, m_indices, to_recipe_tuple(recipe), compiled_dims, disable_ue8m0_cast);
+}
+
+void m_grouped_fp8_gemm_nn_contiguous_wrapper(const torch::Tensor& a_val, const torch::Tensor& a_scale, const torch::Tensor& b_val, const torch::Tensor& b_scale, const torch::Tensor& d, const torch::Tensor& m_indices, const c10::optional<c10::IntArrayRef>& recipe, const std::string& compiled_dims, bool disable_ue8m0_cast) {
+    deep_gemm::gemm::m_grouped_fp8_gemm_nn_contiguous({a_val, a_scale}, {b_val, b_scale}, d, m_indices, to_recipe_tuple(recipe), compiled_dims, disable_ue8m0_cast);
+}
+
+void m_grouped_fp8_gemm_nt_masked_wrapper(const torch::Tensor& a_val, const torch::Tensor& a_scale, const torch::Tensor& b_val, const torch::Tensor& b_scale, const torch::Tensor& d, const torch::Tensor& masked_m, int64_t expected_m, const c10::optional<c10::IntArrayRef>& recipe, const std::string& compiled_dims, bool disable_ue8m0_cast) {
+    deep_gemm::gemm::m_grouped_fp8_gemm_nt_masked({a_val, a_scale}, {b_val, b_scale}, d, masked_m, expected_m, to_recipe_tuple(recipe), compiled_dims, disable_ue8m0_cast);
+}
+
+void k_grouped_fp8_gemm_nt_contiguous_wrapper(const torch::Tensor& a_val, const torch::Tensor& a_scale, const torch::Tensor& b_val, const torch::Tensor& b_scale, const torch::Tensor& d, c10::List<int64_t> ks, const torch::Tensor& ks_tensor, const c10::optional<torch::Tensor>& c, c10::IntArrayRef recipe, const std::string& compiled_dims) {
+    std::vector<int> ks_vec;
+    ks_vec.reserve(ks.size());
+    for(const auto i : ks) {
+        ks_vec.push_back(i);
+    }
+    deep_gemm::gemm::k_grouped_fp8_gemm_nt_contiguous({a_val, a_scale}, {b_val, b_scale}, d, ks_vec, ks_tensor, c, to_recipe_tuple_default(recipe), compiled_dims);
+}
+
+void k_grouped_fp8_gemm_tn_contiguous_wrapper(const torch::Tensor& a_val, const torch::Tensor& a_scale, const torch::Tensor& b_val, const torch::Tensor& b_scale, const torch::Tensor& d, c10::List<int64_t> ks, const torch::Tensor& ks_tensor, const c10::optional<torch::Tensor>& c, c10::IntArrayRef recipe, const std::string& compiled_dims) {
+    std::vector<int> ks_vec;
+    ks_vec.reserve(ks.size());
+    for(const auto i : ks) {
+        ks_vec.push_back(i);
+    }
+    deep_gemm::gemm::k_grouped_fp8_gemm_tn_contiguous({a_val, a_scale}, {b_val, b_scale}, d, ks_vec, ks_tensor, c, to_recipe_tuple_default(recipe), compiled_dims);
+}
+
+void bf16_gemm_nt_wrapper(const torch::Tensor& a, const torch::Tensor& b, const torch::Tensor& d, const c10::optional<torch::Tensor>& c, const std::string& compiled_dims) {
+    deep_gemm::gemm::bf16_gemm_nt(a, b, d, c, compiled_dims);
+}
+
+void bf16_gemm_nn_wrapper(const torch::Tensor& a, const torch::Tensor& b, const torch::Tensor& d, const c10::optional<torch::Tensor>& c, const std::string& compiled_dims) {
+    deep_gemm::gemm::bf16_gemm_nn(a, b, d, c, compiled_dims);
+}
+
+void bf16_gemm_tn_wrapper(const torch::Tensor& a, const torch::Tensor& b, const torch::Tensor& d, const c10::optional<torch::Tensor>& c, const std::string& compiled_dims) {
+    deep_gemm::gemm::bf16_gemm_tn(a, b, d, c, compiled_dims);
+}
+
+void bf16_gemm_tt_wrapper(const torch::Tensor& a, const torch::Tensor& b, const torch::Tensor& d, const c10::optional<torch::Tensor>& c, const std::string& compiled_dims) {
+    deep_gemm::gemm::bf16_gemm_tt(a, b, d, c, compiled_dims);
+}
+
+void m_grouped_bf16_gemm_nt_contiguous_wrapper(const torch::Tensor& a, const torch::Tensor& b, const torch::Tensor& d, const torch::Tensor& m_indices, const std::string& compiled_dims) {
+    deep_gemm::gemm::m_grouped_bf16_gemm_nt_contiguous(a, b, d, m_indices, compiled_dims);
+}
+
+void m_grouped_bf16_gemm_nt_masked_wrapper(const torch::Tensor& a, const torch::Tensor& b, const torch::Tensor& d, const torch::Tensor& masked_m, int64_t expected_m, const std::string& compiled_dims) {
+    deep_gemm::gemm::m_grouped_bf16_gemm_nt_masked(a, b, d, masked_m, expected_m, compiled_dims);
+}
+
+// Attention wrappers
+void fp8_gemm_nt_skip_head_mid_wrapper(const torch::Tensor& a_val, const torch::Tensor& a_scale, const torch::Tensor& b_val, const torch::Tensor b_scale, const torch::Tensor& d, const c10::IntArrayRef& head_splits, const c10::optional<c10::IntArrayRef>& recipe, const std::string& compiled_dims, bool disable_ue8m0_cast) {
+    deep_gemm::attention::fp8_gemm_nt_skip_head_mid({a_val, a_scale}, {b_val, b_scale}, d, to_recipe_tuple_default(head_splits), to_recipe_tuple(recipe), compiled_dims, disable_ue8m0_cast);
+}
+
+torch::Tensor fp8_mqa_logits_wrapper(const torch::Tensor& q, const torch::Tensor& k, const torch::Tensor& v, const torch::Tensor& weight, const torch::Tensor& cu_seq_len_k_start, const torch::Tensor& cu_seq_len_k_end, bool clean_logits) {
+    return deep_gemm::attention::fp8_mqa_logits(q, {k, v}, weight, cu_seq_len_k_start, cu_seq_len_k_end, clean_logits);
+}
+
+torch::Tensor get_paged_mqa_logits_metadata_wrapper(const torch::Tensor& context_lens, int64_t block_kv, int64_t num_sms) {
+    return deep_gemm::attention::get_paged_mqa_logits_metadata(context_lens, block_kv, num_sms);
+}
+
+torch::Tensor fp8_paged_mqa_logits_wrapper(const torch::Tensor& q, const torch::Tensor& fused_kv_cache, const torch::Tensor& weight, const torch::Tensor& context_lens, const torch::Tensor& block_table, const torch::Tensor& schedule_meta, const int64_t max_context_len, bool clean_logits) {
+    return deep_gemm::attention::fp8_paged_mqa_logits(q, fused_kv_cache, weight, context_lens, block_table, schedule_meta, max_context_len, clean_logits); 
+}
+
+} // namespace deep_gemm_wrappers
+
+TORCH_LIBRARY(deep_gemm, m) {
+    // runtime APIs (explicit schema + impl for stable type behavior)
+    m.def("set_num_sms(int new_num_sms) -> ()");
+    m.impl("set_num_sms", [](int64_t new_num_sms) {
+        deep_gemm::device_runtime->set_num_sms(static_cast<int>(new_num_sms));
+    });
+
+    m.def("get_num_sms() -> int");
+    m.impl("get_num_sms", []() -> int64_t {
+        return static_cast<int64_t>(deep_gemm::device_runtime->get_num_sms());
+    });
+
+    m.def("set_tc_util(int new_tc_util) -> ()");
+    m.impl("set_tc_util", [](int64_t new_tc_util) {
+        deep_gemm::device_runtime->set_tc_util(static_cast<int>(new_tc_util));
+    });
+
+    m.def("get_tc_util() -> int");
+    m.impl("get_tc_util", []() -> int64_t {
+        return static_cast<int64_t>(deep_gemm::device_runtime->get_tc_util());
+    });
+
+    m.def("init(str library_root_path, str cuda_home_path_by_torch) -> ()");
+    m.impl("init", [](const std::string& library_root_path, const std::string& cuda_home_path_by_torch) {
+        deep_gemm_wrappers::init_wrapper(library_root_path, cuda_home_path_by_torch);
+    });
+
+    // layout APIs
+    m.def("transform_sf_into_required_layout(Tensor sf, int mn, int k, int[] recipe, int? num_groups=None, bool is_sfa=False, bool disable_ue8m0_cast=False) -> Tensor", deep_gemm_wrappers::transform_sf_into_required_layout_wrapper);
+
+    m.def("get_tma_aligned_size(int size, int element_size) -> int");
+    m.impl("get_tma_aligned_size", [](int64_t size, int64_t element_size) -> int64_t {
+        return deep_gemm_wrappers::get_tma_aligned_size_wrapper(size, element_size);
+    });
+
+    m.def("get_mk_alignment_for_contiguous_layout() -> int");
+    m.impl("get_mk_alignment_for_contiguous_layout", []() -> int64_t {
+        return deep_gemm_wrappers::get_mk_alignment_for_contiguous_layout_wrapper();
+    });
+    m.def("get_mn_major_tma_aligned_tensor(Tensor a) -> Tensor");
+    m.impl("get_mn_major_tma_aligned_tensor", [](const torch::Tensor& a) -> torch::Tensor {
+        return deep_gemm::get_mn_major_tma_aligned_tensor(a);
+    });
+
+    m.def("get_mn_major_tma_aligned_packed_ue8m0_tensor(Tensor a) -> Tensor");
+    m.impl("get_mn_major_tma_aligned_packed_ue8m0_tensor", [](const torch::Tensor& a) -> torch::Tensor {
+        return deep_gemm::get_mn_major_tma_aligned_packed_ue8m0_tensor(a);
+    });
+
+    m.def("get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor(Tensor a, Tensor ks_tensor, int[] ks) -> Tensor");
+    m.impl("get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor", [](const torch::Tensor& a,
+                                                                        const torch::Tensor& ks_tensor,
+                                                                        at::IntArrayRef ks_ref) -> torch::Tensor {
+        std::vector<int> ks_vec(ks_ref.begin(), ks_ref.end());
+        return deep_gemm::get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor(a, ks_tensor, ks_vec);
+    });
+
+    // gemm APIs (explicit schema + impl)
+    m.def(R"(fp8_gemm_nt(Any a, Any b, Tensor d, Tensor? c=None, int[]? recipe=None, str compiled_dims="nk", bool disable_ue8m0_cast=False) -> ())");
+    m.impl("fp8_gemm_nt", torch::kCUDA, [](const c10::IValue& a_input, const c10::IValue& b_input,
+                                            const torch::Tensor& d,
+                                            const c10::optional<torch::Tensor>& c,
+                                            const c10::optional<c10::IntArrayRef>& recipe,
+                                            const std::string& compiled_dims,
+                                            bool disable_ue8m0_cast) {
+        auto [a_val, a_scale] = parse_tensor_or_tuple(a_input);
+        auto [b_val, b_scale] = parse_tensor_or_tuple(b_input);
+        deep_gemm_wrappers::fp8_gemm_nt_wrapper(a_val, a_scale, b_val, b_scale, d, c, recipe, compiled_dims, disable_ue8m0_cast);
+    });
+
+    m.def(R"(fp8_gemm_nn(Any a, Any b, Tensor d, Tensor? c=None, int[]? recipe=None, str compiled_dims="nk", bool disable_ue8m0_cast=False) -> ())");
+    m.impl("fp8_gemm_nn", torch::kCUDA, [](const c10::IValue& a_input, const c10::IValue& b_input,
+                                            const torch::Tensor& d,
+                                            const c10::optional<torch::Tensor>& c,
+                                            const c10::optional<c10::IntArrayRef>& recipe,
+                                            const std::string& compiled_dims,
+                                            bool disable_ue8m0_cast) {
+        auto [a_val, a_scale] = parse_tensor_or_tuple(a_input);
+        auto [b_val, b_scale] = parse_tensor_or_tuple(b_input);
+        deep_gemm_wrappers::fp8_gemm_nn_wrapper(a_val, a_scale, b_val, b_scale, d, c, recipe, compiled_dims, disable_ue8m0_cast);
+    });
+
+    m.def(R"(fp8_gemm_tn(Any a, Any b, Tensor d, Tensor? c=None, int[]? recipe=None, str compiled_dims="mn", bool disable_ue8m0_cast=False) -> ())");
+    m.impl("fp8_gemm_tn", torch::kCUDA, [](const c10::IValue& a_input, const c10::IValue& b_input,
+                                            const torch::Tensor& d,
+                                            const c10::optional<torch::Tensor>& c,
+                                            const c10::optional<c10::IntArrayRef>& recipe,
+                                            const std::string& compiled_dims,
+                                            bool disable_ue8m0_cast) {
+        auto [a_val, a_scale] = parse_tensor_or_tuple(a_input);
+        auto [b_val, b_scale] = parse_tensor_or_tuple(b_input);
+        deep_gemm_wrappers::fp8_gemm_tn_wrapper(a_val, a_scale, b_val, b_scale, d, c, recipe, compiled_dims, disable_ue8m0_cast);
+    });
+
+    m.def(R"(fp8_gemm_tt(Any a, Any b, Tensor d, Tensor? c=None, int[]? recipe=None, str compiled_dims="mn", bool disable_ue8m0_cast=False) -> ())");
+    m.impl("fp8_gemm_tt", torch::kCUDA, [](const c10::IValue& a_input, const c10::IValue& b_input,
+                                            const torch::Tensor& d,
+                                            const c10::optional<torch::Tensor>& c,
+                                            const c10::optional<c10::IntArrayRef>& recipe,
+                                            const std::string& compiled_dims,
+                                            bool disable_ue8m0_cast) {
+        auto [a_val, a_scale] = parse_tensor_or_tuple(a_input);
+        auto [b_val, b_scale] = parse_tensor_or_tuple(b_input);
+        deep_gemm_wrappers::fp8_gemm_tt_wrapper(a_val, a_scale, b_val, b_scale, d, c, recipe, compiled_dims, disable_ue8m0_cast);
+    });
+
+    m.def(R"(m_grouped_fp8_gemm_nt_contiguous(Any a, Any b, Tensor d, Tensor m_indices, int[]? recipe=None, str compiled_dims="nk", bool disable_ue8m0_cast=False) -> ())");
+    m.impl("m_grouped_fp8_gemm_nt_contiguous", torch::kCUDA, [](const c10::IValue& a_input, const c10::IValue& b_input,
+                                                                 const torch::Tensor& d,
+                                                                 const torch::Tensor& m_indices,
+                                                                 const c10::optional<c10::IntArrayRef>& recipe,
+                                                                 const std::string& compiled_dims,
+                                                                 bool disable_ue8m0_cast) {
+        auto [a_val, a_scale] = parse_tensor_or_tuple(a_input);
+        auto [b_val, b_scale] = parse_tensor_or_tuple(b_input);
+        deep_gemm_wrappers::m_grouped_fp8_gemm_nt_contiguous_wrapper(a_val, a_scale, b_val, b_scale, d, m_indices, recipe, compiled_dims, disable_ue8m0_cast);
+    });
+
+    m.def(R"(m_grouped_fp8_gemm_nn_contiguous(Any a, Any b, Tensor d, Tensor m_indices, int[]? recipe=None, str compiled_dims="nk", bool disable_ue8m0_cast=False) -> ())");
+    m.impl("m_grouped_fp8_gemm_nn_contiguous", torch::kCUDA, [](const c10::IValue& a_input, const c10::IValue& b_input,
+                                                                 const torch::Tensor& d,
+                                                                 const torch::Tensor& m_indices,
+                                                                 const c10::optional<c10::IntArrayRef>& recipe,
+                                                                 const std::string& compiled_dims,
+                                                                 bool disable_ue8m0_cast) {
+        auto [a_val, a_scale] = parse_tensor_or_tuple(a_input);
+        auto [b_val, b_scale] = parse_tensor_or_tuple(b_input);
+        deep_gemm_wrappers::m_grouped_fp8_gemm_nn_contiguous_wrapper(a_val, a_scale, b_val, b_scale, d, m_indices, recipe, compiled_dims, disable_ue8m0_cast);
+    });
+
+    m.def(R"(m_grouped_fp8_gemm_nt_masked(Any a, Any b, Tensor d, Tensor masked_m, int expected_m, int[]? recipe=None, str compiled_dims="nk", bool disable_ue8m0_cast=False) -> ())");
+    m.impl("m_grouped_fp8_gemm_nt_masked", torch::kCUDA, [](const c10::IValue& a_input, const c10::IValue& b_input,
+                                                             const torch::Tensor& d,
+                                                             const torch::Tensor& masked_m,
+                                                             int64_t expected_m,
+                                                             const c10::optional<c10::IntArrayRef>& recipe,
+                                                             const std::string& compiled_dims,
+                                                             bool disable_ue8m0_cast) {
+        auto [a_val, a_scale] = parse_tensor_or_tuple(a_input);
+        auto [b_val, b_scale] = parse_tensor_or_tuple(b_input);
+        deep_gemm_wrappers::m_grouped_fp8_gemm_nt_masked_wrapper(a_val, a_scale, b_val, b_scale, d, masked_m, expected_m, recipe, compiled_dims, disable_ue8m0_cast);
+    });
+
+    m.def(R"(k_grouped_fp8_gemm_nt_contiguous(Any a, Any b, Tensor d, int[] ks, Tensor ks_tensor, Tensor? c=None, int[] recipe=[1, 1, 128], str compiled_dims="mn") -> ())");
+    m.impl("k_grouped_fp8_gemm_nt_contiguous", torch::kCUDA, [](const c10::IValue& a_input, const c10::IValue& b_input,
+                                                                 const torch::Tensor& d,
+                                                                 at::IntArrayRef ks,
+                                                                 const torch::Tensor& ks_tensor,
+                                                                 const c10::optional<torch::Tensor>& c,
+                                                                 c10::IntArrayRef recipe,
+                                                                 const std::string& compiled_dims) {
+        auto [a_val, a_scale] = parse_tensor_or_tuple(a_input);
+        auto [b_val, b_scale] = parse_tensor_or_tuple(b_input);
+        std::vector<int64_t> ks64(ks.begin(), ks.end());
+        c10::List<int64_t> ks_list(ks64);
+        deep_gemm_wrappers::k_grouped_fp8_gemm_nt_contiguous_wrapper(a_val, a_scale, b_val, b_scale, d, ks_list, ks_tensor, c, recipe, compiled_dims);
+    });
+
+    m.def(R"(k_grouped_fp8_gemm_tn_contiguous(Any a, Any b, Tensor d, int[] ks, Tensor ks_tensor, Tensor? c=None, int[] recipe=[1, 1, 128], str compiled_dims="mn") -> ())");
+    m.impl("k_grouped_fp8_gemm_tn_contiguous", torch::kCUDA, [](const c10::IValue& a_input, const c10::IValue& b_input,
+                                                                 const torch::Tensor& d,
+                                                                 at::IntArrayRef ks,
+                                                                 const torch::Tensor& ks_tensor,
+                                                                 const c10::optional<torch::Tensor>& c,
+                                                                 c10::IntArrayRef recipe,
+                                                                 const std::string& compiled_dims) {
+        auto [a_val, a_scale] = parse_tensor_or_tuple(a_input);
+        auto [b_val, b_scale] = parse_tensor_or_tuple(b_input);
+        std::vector<int64_t> ks64(ks.begin(), ks.end());
+        c10::List<int64_t> ks_list(ks64);
+        deep_gemm_wrappers::k_grouped_fp8_gemm_tn_contiguous_wrapper(a_val, a_scale, b_val, b_scale, d, ks_list, ks_tensor, c, recipe, compiled_dims);
+    });
+
+    /*
+     * BF16 GEMM
+     */
+
+    m.def(R"(bf16_gemm_nt(Tensor a, Tensor b, Tensor d, Tensor? c=None, str compiled_dims="") -> ())");
+    m.impl("bf16_gemm_nt", torch::kCUDA, [](const torch::Tensor& a, const torch::Tensor& b, const torch::Tensor& d,
+                                             const c10::optional<torch::Tensor>& c,
+                                             const std::string& compiled_dims) {
+        deep_gemm_wrappers::bf16_gemm_nt_wrapper(a, b, d, c, compiled_dims);
+    });
+
+    m.def(R"(bf16_gemm_nn(Tensor a, Tensor b, Tensor d, Tensor? c=None, str compiled_dims="") -> ())");
+    m.impl("bf16_gemm_nn", torch::kCUDA, [](const torch::Tensor& a, const torch::Tensor& b, const torch::Tensor& d,
+                                             const c10::optional<torch::Tensor>& c,
+                                             const std::string& compiled_dims) {
+        deep_gemm_wrappers::bf16_gemm_nn_wrapper(a, b, d, c, compiled_dims);
+    });
+
+    m.def(R"(bf16_gemm_tn(Tensor a, Tensor b, Tensor d, Tensor? c=None, str compiled_dims="") -> ())");
+    m.impl("bf16_gemm_tn", torch::kCUDA, [](const torch::Tensor& a, const torch::Tensor& b, const torch::Tensor& d,
+                                             const c10::optional<torch::Tensor>& c,
+                                             const std::string& compiled_dims) {
+        deep_gemm_wrappers::bf16_gemm_tn_wrapper(a, b, d, c, compiled_dims);
+    });
+
+    m.def(R"(bf16_gemm_tt(Tensor a, Tensor b, Tensor d, Tensor? c=None, str compiled_dims="") -> ())");
+    m.impl("bf16_gemm_tt", torch::kCUDA, [](const torch::Tensor& a, const torch::Tensor& b, const torch::Tensor& d,
+                                             const c10::optional<torch::Tensor>& c,
+                                             const std::string& compiled_dims) {
+        deep_gemm_wrappers::bf16_gemm_tt_wrapper(a, b, d, c, compiled_dims);
+    });
+
+    m.def(R"(m_grouped_bf16_gemm_nt_contiguous(Tensor a, Tensor b, Tensor d, Tensor m_indices, str compiled_dims="") -> ())");
+    m.impl("m_grouped_bf16_gemm_nt_contiguous", torch::kCUDA, [](const torch::Tensor& a, const torch::Tensor& b,
+                                                                  const torch::Tensor& d, const torch::Tensor& m_indices,
+                                                                  const std::string& compiled_dims) {
+        deep_gemm_wrappers::m_grouped_bf16_gemm_nt_contiguous_wrapper(a, b, d, m_indices, compiled_dims);
+    });
+
+    m.def(R"(m_grouped_bf16_gemm_nt_masked(Tensor a, Tensor b, Tensor d, Tensor masked_m, int expected_m, str compiled_dims="") -> ())");
+    m.impl("m_grouped_bf16_gemm_nt_masked", torch::kCUDA, [](const torch::Tensor& a, const torch::Tensor& b, const torch::Tensor& d,
+                                                              const torch::Tensor& masked_m, int64_t expected_m,
+                                                              const std::string& compiled_dims) {
+        deep_gemm_wrappers::m_grouped_bf16_gemm_nt_masked_wrapper(a, b, d, masked_m, expected_m, compiled_dims);
+    });
+
+    /*
+     * cublas gemm
+     */
+    // cuBLASLt GEMMs
+    m.def(R"(cublaslt_gemm_nt(Tensor a, Tensor b, Tensor d, Tensor? c) -> ())");
+    m.impl("cublaslt_gemm_nt", torch::kCUDA, [](const torch::Tensor& a, const torch::Tensor& b, const torch::Tensor& d,
+                                             const c10::optional<torch::Tensor>& c) {
+        deep_gemm::gemm::cublaslt_gemm_nt(a, b, d, c);
+    });
+
+    m.def(R"(cublaslt_gemm_nn(Tensor a, Tensor b, Tensor d, Tensor? c) -> ())");
+    m.impl("cublaslt_gemm_nn", torch::kCUDA, [](const torch::Tensor& a, const torch::Tensor& b, const torch::Tensor& d,
+                                             const c10::optional<torch::Tensor>& c) {
+        deep_gemm::gemm::cublaslt_gemm_nn(a, b, d, c);
+    });
+
+    m.def(R"(cublaslt_gemm_tn(Tensor a, Tensor b, Tensor d, Tensor? c) -> ())");
+    m.impl("cublaslt_gemm_tn", torch::kCUDA, [](const torch::Tensor& a, const torch::Tensor& b, const torch::Tensor& d,
+                                             const c10::optional<torch::Tensor>& c) {
+        deep_gemm::gemm::cublaslt_gemm_tn(a, b, d, c);
+    });
+
+    m.def(R"(cublaslt_gemm_tt(Tensor a, Tensor b, Tensor d, Tensor? c) -> ())");
+    m.impl("cublaslt_gemm_tt", torch::kCUDA, [](const torch::Tensor& a, const torch::Tensor& b, const torch::Tensor& d,
+                                             const c10::optional<torch::Tensor>& c) {
+        deep_gemm::gemm::cublaslt_gemm_tt(a, b, d, c);
+    });
+
+    /*
+     * Attention
+     */
+    m.def(R"(fp8_gemm_nt_skip_head_mid(Any a, Any b, Tensor d, int[] head_splits, int[]? recipe=None, str compiled_dims="nk", bool disable_ue8m0_cast=False) -> ())");
+    m.impl("fp8_gemm_nt_skip_head_mid", torch::kCUDA, [](const c10::IValue& a_input, const c10::IValue& b_input,
+                                                        const torch::Tensor& d, 
+                                                        const c10::IntArrayRef& head_splits, 
+                                                        const c10::optional<c10::IntArrayRef>& recipe, 
+                                                        const std::string& compiled_dims, 
+                                                        bool disable_ue8m0_cast) {
+        auto [a_val, a_scale] = parse_tensor_or_tuple(a_input);
+        auto [b_val, b_scale] = parse_tensor_or_tuple(b_input);
+        deep_gemm_wrappers::fp8_gemm_nt_skip_head_mid_wrapper(a_val, a_scale, b_val, b_scale, d, head_splits, recipe, compiled_dims, disable_ue8m0_cast);
+    });
+    
+    m.def(R"(fp8_mqa_logits(Tensor q, Any kv, Tensor weights, Tensor cu_seq_len_k_start, Tensor cu_seq_len_k_end, bool clean_logits=True) -> Tensor)");
+    m.impl("fp8_mqa_logits", torch::kCUDA, [](
+        const torch::Tensor& q, 
+        const c10::IValue& kv, 
+        const torch::Tensor& weights, 
+        const torch::Tensor& cu_seq_len_k_start, 
+        const torch::Tensor& cu_seq_len_k_end, 
+        bool clean_logits
+    ) -> torch::Tensor {
+        auto [k, v] = parse_tensor_or_tuple(kv);
+        return deep_gemm_wrappers::fp8_mqa_logits_wrapper(q, k, v, weights, cu_seq_len_k_start, cu_seq_len_k_end, clean_logits);
+    });
+
+    m.def(R"(get_paged_mqa_logits_metadata(Tensor context_lens, int block_kv, int num_sms) -> Tensor)");
+    m.impl("get_paged_mqa_logits_metadata", torch::kCUDA, [](
+        const torch::Tensor& context_lens, 
+        int64_t block_kv, 
+        int64_t num_sms
+    ) -> torch::Tensor {
+        return deep_gemm_wrappers::get_paged_mqa_logits_metadata_wrapper(context_lens, block_kv, num_sms);
+    });
+
+    m.def(R"(fp8_paged_mqa_logits(Tensor q, Tensor fused_kv_cache, Tensor weights, Tensor context_lens, Tensor block_table, Tensor schedule_meta, int max_context_len, bool clean_logits) -> Tensor)");
+    m.impl("fp8_paged_mqa_logits", torch::kCUDA, [](
+        const torch::Tensor& q, 
+        const torch::Tensor& fused_kv_cache, 
+        const torch::Tensor& weights, 
+        const torch::Tensor& context_lens, 
+        const torch::Tensor& block_table, 
+        const torch::Tensor& schedule_meta, 
+        int64_t max_context_len, 
+        bool clean_logits
+    ) -> torch::Tensor {
+        return deep_gemm_wrappers::fp8_paged_mqa_logits_wrapper(q, fused_kv_cache, weights, context_lens, block_table, schedule_meta, max_context_len, clean_logits);
+    });
+
+    /*
+     * einsum
+     */
+    m.def(R"(einsum(str expr, Tensor a, Tensor b, Tensor d, Tensor? c=None) -> ())");
+    m.impl("einsum", torch::kCUDA, [](const std::string& expr, const torch::Tensor& a, const torch::Tensor& b, const torch::Tensor& d, const c10::optional<torch::Tensor>& c) {
+        deep_gemm::einsum::einsum(expr, a, b, d, c);
+    });
+
+}
+
+// Provide single definitions for the declared wrappers
+int64_t deep_gemm_wrappers::get_tma_aligned_size_wrapper(int64_t x, int64_t element_size) {
+    return static_cast<int64_t>(deep_gemm::get_tma_aligned_size(static_cast<int>(x), static_cast<int>(element_size)));
+}
+
+int64_t deep_gemm_wrappers::get_mk_alignment_for_contiguous_layout_wrapper() {
+    return static_cast<int64_t>(deep_gemm::get_mk_alignment_for_contiguous_layout());
+}
+
+REGISTER_EXTENSION(deep_gemm_cpp)
--- a/csrc/utils/layout.hpp
+++ b/csrc/utils/layout.hpp
@@ -1,7 +1,6 @@
 #pragma once

 #include <cute/arch/mma_sm100_umma.hpp>
-#include <torch/python.h>

 #include "math.hpp"
 #include "exception.hpp"
--- a/csrc/utils/math.hpp
+++ b/csrc/utils/math.hpp
@@ -1,6 +1,6 @@
 #pragma once

-#include <torch/python.h>
+#include <cstdint>

 #include "exception.hpp"

--- a/deep_gemm/init.py
+++ b/deep_gemm/init.py
@@ -1,5 +1,7 @@
 import os
 import subprocess
+import torch
+import torch.utils.cpp_extension

 # Set some default environment provided at setup
 try:
@@ -11,62 +13,12 @@ try:
 except ImportError:
    pass

-# Configs
-import deep_gemm_cpp
-from deep_gemm_cpp import (
-    set_num_sms,
-    get_num_sms,
-    set_tc_util,
-    get_tc_util,
-)
-
-# Kernels
-from deep_gemm_cpp import (
-    # FP8 GEMMs
-    fp8_gemm_nt, fp8_gemm_nn,
-    fp8_gemm_tn, fp8_gemm_tt,
-    fp8_gemm_nt_skip_head_mid,
-    m_grouped_fp8_gemm_nt_contiguous,
-    m_grouped_fp8_gemm_nn_contiguous,
-    m_grouped_fp8_gemm_nt_masked,
-    k_grouped_fp8_gemm_nt_contiguous,
-    k_grouped_fp8_gemm_tn_contiguous,
-    # BF16 GEMMs
-    bf16_gemm_nt, bf16_gemm_nn,
-    bf16_gemm_tn, bf16_gemm_tt,
-    m_grouped_bf16_gemm_nt_contiguous,
-    m_grouped_bf16_gemm_nt_masked,
-    # cuBLASLt GEMMs
-    cublaslt_gemm_nt, cublaslt_gemm_nn,
-    cublaslt_gemm_tn, cublaslt_gemm_tt,
-    # Einsum kernels
-    einsum,
-    # Attention kernels
-    fp8_mqa_logits,
-    get_paged_mqa_logits_metadata,
-    fp8_paged_mqa_logits,
-    # Layout kernels
-    transform_sf_into_required_layout
-)
-
-# Some alias for legacy supports
-# TODO: remove these later
-fp8_m_grouped_gemm_nt_masked = m_grouped_fp8_gemm_nt_masked
-bf16_m_grouped_gemm_nt_masked = m_grouped_bf16_gemm_nt_masked
-
-# Some utils
-from . import testing
-from . import utils
-from .utils import *
+from . import deep_gemm_cpp  # noqa: F401  # Registers ops into torch.ops without touching CUDA


-# Initialize CPP modules
 def _find_cuda_home() -> str:
-    # TODO: reuse PyTorch API later
-    # For some PyTorch versions, the original `_find_cuda_home` will initialize CUDA, which is incompatible with process forks
    cuda_home = os.environ.get('CUDA_HOME') or os.environ.get('CUDA_PATH')
    if cuda_home is None:
-        # noinspection PyBroadException
        try:
            with open(os.devnull, 'w') as devnull:
                nvcc = subprocess.check_output(['which', 'nvcc'], stderr=devnull).decode().rstrip('\r\n')
@@ -79,7 +31,106 @@ def _find_cuda_home() -> str:
    return cuda_home


-deep_gemm_cpp.init(
-    os.path.dirname(os.path.abspath(__file__)), # Library root directory path
-    _find_cuda_home()                           # CUDA home
-)
+# Lazy runtime init to be fork-safe on Linux (avoid initializing CUDA before fork)
+_dg_initialized = False
+
+def _ensure_initialized() -> None:
+    global _dg_initialized
+    if _dg_initialized:
+        return
+    library_root = os.path.dirname(os.path.abspath(__file__))
+    torch.ops.deep_gemm.init(library_root, _find_cuda_home())
+    _dg_initialized = True
+
+
+def _wrap_op(name: str):
+    def _fn(*args, **kwargs):
+        _ensure_initialized()
+        return getattr(torch.ops.deep_gemm, name)(*args, **kwargs)
+    return _fn
+
+set_num_sms = _wrap_op('set_num_sms')
+get_num_sms = _wrap_op('get_num_sms')
+set_tc_util = _wrap_op('set_tc_util')
+get_tc_util = _wrap_op('get_tc_util')
+
+fp8_gemm_nt = _wrap_op('fp8_gemm_nt')
+fp8_gemm_nn = _wrap_op('fp8_gemm_nn')
+fp8_gemm_tn = _wrap_op('fp8_gemm_tn')
+fp8_gemm_tt = _wrap_op('fp8_gemm_tt')
+m_grouped_fp8_gemm_nt_contiguous = _wrap_op('m_grouped_fp8_gemm_nt_contiguous')
+m_grouped_fp8_gemm_nn_contiguous = _wrap_op('m_grouped_fp8_gemm_nn_contiguous')
+# Export both canonical name and backward-compat alias
+m_grouped_fp8_gemm_nt_masked = _wrap_op('m_grouped_fp8_gemm_nt_masked')
+fp8_m_grouped_gemm_nt_masked = m_grouped_fp8_gemm_nt_masked
+k_grouped_fp8_gemm_nt_contiguous = _wrap_op('k_grouped_fp8_gemm_nt_contiguous')
+k_grouped_fp8_gemm_tn_contiguous = _wrap_op('k_grouped_fp8_gemm_tn_contiguous')
+
+# BF16 GEMMs
+bf16_gemm_nt = _wrap_op('bf16_gemm_nt')
+bf16_gemm_nn = _wrap_op('bf16_gemm_nn')
+bf16_gemm_tn = _wrap_op('bf16_gemm_tn')
+bf16_gemm_tt = _wrap_op('bf16_gemm_tt')
+m_grouped_bf16_gemm_nt_contiguous = _wrap_op('m_grouped_bf16_gemm_nt_contiguous')
+m_grouped_bf16_gemm_nt_masked = _wrap_op('m_grouped_bf16_gemm_nt_masked')
+
+# cuBLASLt GEMMs
+cublaslt_gemm_nt = _wrap_op('cublaslt_gemm_nt')
+cublaslt_gemm_nn = _wrap_op('cublaslt_gemm_nn')
+cublaslt_gemm_tn = _wrap_op('cublaslt_gemm_tn')
+cublaslt_gemm_tt = _wrap_op('cublaslt_gemm_tt')
+
+# Attention kernel
+fp8_gemm_nt_skip_head_mid = _wrap_op('fp8_gemm_nt_skip_head_mid')
+fp8_mqa_logits = _wrap_op('fp8_mqa_logits')
+get_paged_mqa_logits_metadata = _wrap_op('get_paged_mqa_logits_metadata')
+fp8_paged_mqa_logits = _wrap_op('fp8_paged_mqa_logits')
+
+# Einsum kernel
+einsum = _wrap_op('einsum')
+
+# Layout kernels
+transform_sf_into_required_layout = _wrap_op('transform_sf_into_required_layout')
+
+# Utility functions
+get_tma_aligned_size = _wrap_op('get_tma_aligned_size')
+get_mk_alignment_for_contiguous_layout = _wrap_op('get_mk_alignment_for_contiguous_layout')
+get_mn_major_tma_aligned_tensor = _wrap_op('get_mn_major_tma_aligned_tensor')
+get_mn_major_tma_aligned_packed_ue8m0_tensor = _wrap_op('get_mn_major_tma_aligned_packed_ue8m0_tensor')
+get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor = _wrap_op('get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor')
+
+# Some utils
+from . import testing
+from . import utils
+from .utils import *
+
+def _verify_ops_loaded():
+    expected_ops = [
+        'init', 'set_num_sms', 'get_num_sms', 'set_tc_util', 'get_tc_util',
+        'fp8_gemm_nt', 'fp8_gemm_nn', 'fp8_gemm_tn', 'fp8_gemm_tt',
+        'm_grouped_fp8_gemm_nt_contiguous', 'm_grouped_fp8_gemm_nn_contiguous',
+        'm_grouped_fp8_gemm_nt_masked', 'k_grouped_fp8_gemm_nt_contiguous',
+        'k_grouped_fp8_gemm_tn_contiguous',
+        'transform_sf_into_required_layout', 'get_tma_aligned_size',
+        'get_mk_alignment_for_contiguous_layout', 'get_mn_major_tma_aligned_tensor',
+        'get_mn_major_tma_aligned_packed_ue8m0_tensor',
+        'get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor',
+        'fp8_gemm_nt_skip_head_mid', 'fp8_mqa_logits',
+        'get_paged_mqa_logits_metadata', 'fp8_paged_mqa_logits',
+        'einsum',
+        'cublaslt_gemm_nt', 'cublaslt_gemm_nn',
+        'cublaslt_gemm_tn', 'cublaslt_gemm_tt',
+    ]
+    
+    available_ops = list(torch.ops.deep_gemm.__dict__.keys())
+    missing_ops = [op for op in expected_ops if op not in available_ops]
+    
+    if missing_ops:
+        print(f"Warning: Missing operations: {missing_ops}")
+
+
+_ensure_initialized()
+
+
+if __debug__:
+    _verify_ops_loaded()
--- a/deep_gemm/utils/layout.py
+++ b/deep_gemm/utils/layout.py
@@ -1,11 +1,13 @@
-from deep_gemm_cpp import (
-    get_tma_aligned_size,
-    get_mk_alignment_for_contiguous_layout,
-    get_mn_major_tma_aligned_tensor,
-    get_mn_major_tma_aligned_packed_ue8m0_tensor,
-    get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor
-)
+import torch
+from .. import _ensure_initialized

-# Some alias
-get_m_alignment_for_contiguous_layout = get_mk_alignment_for_contiguous_layout
-get_k_alignment_for_contiguous_layout = get_mk_alignment_for_contiguous_layout
+_ensure_initialized()
+
+get_tma_aligned_size = torch.ops.deep_gemm.get_tma_aligned_size
+get_mk_alignment_for_contiguous_layout = torch.ops.deep_gemm.get_mk_alignment_for_contiguous_layout
+get_mn_major_tma_aligned_tensor = torch.ops.deep_gemm.get_mn_major_tma_aligned_tensor
+get_mn_major_tma_aligned_packed_ue8m0_tensor = torch.ops.deep_gemm.get_mn_major_tma_aligned_packed_ue8m0_tensor
+get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor = torch.ops.deep_gemm.get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor
+
+get_m_alignment_for_contiguous_layout = torch.ops.deep_gemm.get_mk_alignment_for_contiguous_layout
+get_k_alignment_for_contiguous_layout = torch.ops.deep_gemm.get_mk_alignment_for_contiguous_layout
--- a/tests/test_einsum.py
+++ b/tests/test_einsum.py
@@ -8,6 +8,17 @@ from deep_gemm.testing import (
 )


+def get_cuda_version():
+    if torch.version.cuda:
+        return tuple(map(int, torch.version.cuda.split(".")))
+    return (0, 0)
+
+def nvjet_accessable():
+    cuda_version = get_cuda_version()
+    if cuda_version[0] > 12 or (cuda_version[0] == 12 and cuda_version[1] >= 6):
+        return True
+    return False
+
 def test_bmk_bnk_mn() -> None:
    print('Testing "bmk, bnk -> mn":')
    for s in (129, 4096, 8192):
@@ -81,5 +92,6 @@ if __name__ == '__main__':
    print(f' > {deep_gemm.__path__}\n')

    test_bmk_bnk_mn()
-    test_bhr_hdr_bhd()
-    test_bhd_hdr_bhr()
+    if nvjet_accessable():
+        test_bhr_hdr_bhd()
+        test_bhd_hdr_bhr()