Migrate pybind11 to TORCH_LIBRARY #393

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open

cleonard530 wants to merge 28 commits into deepseek-ai:nv_dev from cleonard530:migrate_pybind_to_torch_library

.github/workflows/publish.yml

-Original file line number
+Diff line change
@@ Expand Up / @@ -41,19 +41,16 @@ jobs: @@
             # Using ubuntu-22.04 instead of 24.04 for more compatibility (glibc). Ideally we'd use the
             # manylinux docker image, but I haven't figured out how to install CUDA on manylinux.
             os: [ubuntu-22.04]
-            python-version: ["3.8", "3.9", "3.10", "3.11", "3.12", "3.13"]
+            # setup.py builds a single abi3 wheel pinned to the cp310 floor, which every torch
+            # version here (2.4-2.8) supports, so one Python build covers the whole matrix.
+            python-version: ["3.10"]
             torch-version: ["2.4.0", "2.5.1", "2.6.0", "2.7.1", "2.8.0"]
             cuda-version: ["12.9.1"]
             # We need separate wheels that either uses C++11 ABI (-D_GLIBCXX_USE_CXX11_ABI) or not.
             # Pytorch wheels currently don't use it, but nvcr images have Pytorch compiled with C++11 ABI.
             # Without this we get import error (undefined symbol: _ZN3c105ErrorC2ENS_14SourceLocationESs)
             # when building without C++11 ABI and using it on nvcr images.
             cxx11_abi: ["FALSE", "TRUE"]
-            exclude:
-              # see https://github.com/pytorch/pytorch/blob/main/RELEASE.md#release-compatibility-matrix
-              # Pytorch < 2.5 does not support Python 3.13
-              - torch-version: "2.4.0"
-                python-version: "3.13"
         uses: ./.github/workflows/_build.yml
         with:
           runs-on: ${{ matrix.os }}
@@ Expand Down @@

csrc/apis/attention.hpp

-Original file line number
+Diff line change
@@ Expand Up / @@ -15,6 +15,8 @@ @@
     #endif
     #include "layout.hpp"
+    #include <torch/library.h>
+    #include "../torch_library_utils.hpp"
     namespace deep_gemm::attention {
@@ Expand Down Expand Up @@
     }
     #endif
-    static void register_apis(pybind11::module_& m) {
+    }  // namespace deep_gemm::attention
+    namespace deep_gemm::torch_registration {
+    using namespace deep_gemm::torch_utils;
+    #if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE
+    static void fp8_gemm_nt_skip_head_mid(
+        const torch::Tensor& a, const torch::Tensor& sfa,
+        const torch::Tensor& b, const torch::Tensor& sfb,
+        const torch::Tensor& d,
+        const std::vector<int64_t>& head_splits,
+        const c10::optional<std::vector<int64_t>>& recipe,
+        const std::string& compiled_dims,
+        const bool& disable_ue8m0_cast) {
+        attention::fp8_gemm_nt_skip_head_mid(
+            {a, sfa}, {b, sfb}, d,
+            list_to_tuple3(head_splits),
+            list_to_recipe3(recipe),
+            compiled_dims, disable_ue8m0_cast);
+    }
+    static torch::Tensor fp8_fp4_mqa_logits(
+        const torch::Tensor& q, const c10::optional<torch::Tensor>& q_sf,
+        const torch::Tensor& kv, const torch::Tensor& kv_sf,
+        const torch::Tensor& weights,
+        const torch::Tensor& cu_seq_len_k_start,
+        const torch::Tensor& cu_seq_len_k_end,
+        const bool& clean_logits,
+        const int64_t& max_seqlen_k,
+        at::ScalarType logits_dtype) {
+        return attention::fp8_fp4_mqa_logits(
+            std::make_tuple(q, q_sf),
+            std::make_tuple(kv, kv_sf),
+            weights, cu_seq_len_k_start, cu_seq_len_k_end,
+            clean_logits, static_cast<int>(max_seqlen_k),
+            logits_dtype);
+    }
+    static torch::Tensor get_paged_mqa_logits_metadata(
+        const torch::Tensor& context_lens, const int64_t& block_kv,
+        const int64_t& num_sms, const c10::optional<torch::Tensor>& indices) {
+        return attention::get_paged_mqa_logits_metadata(
+            context_lens, static_cast<int>(block_kv),
+            static_cast<int>(num_sms), indices);
+    }
+    static torch::Tensor fp8_fp4_paged_mqa_logits(
+        const torch::Tensor& q, const c10::optional<torch::Tensor>& q_sf,
+        const torch::Tensor& kv_cache,
+        const torch::Tensor& weights,
+        const torch::Tensor& context_lens,
+        const torch::Tensor& block_table,
+        const torch::Tensor& schedule_meta,
+        const int64_t& max_context_len,
+        const bool& clean_logits,
+        at::ScalarType logits_dtype,
+        const c10::optional<torch::Tensor>& indices) {
+        return attention::fp8_fp4_paged_mqa_logits(
+            std::make_tuple(q, q_sf),
+            kv_cache, weights, context_lens, block_table, schedule_meta,
+            static_cast<int>(max_context_len), clean_logits,
+            logits_dtype, indices);
+    }
+    static torch::Tensor fp8_mqa_logits(
+        const torch::Tensor& q,
+        const torch::Tensor& kv, const torch::Tensor& kv_sf,
+        const torch::Tensor& weights,
+        const torch::Tensor& cu_seq_len_k_start,
+        const torch::Tensor& cu_seq_len_k_end,
+        const bool& clean_logits,
+        const int64_t& max_seqlen_k) {
+        return attention::fp8_mqa_logits(
+            q, std::make_tuple(kv, kv_sf), weights,
+            cu_seq_len_k_start, cu_seq_len_k_end,
+            clean_logits, static_cast<int>(max_seqlen_k));
+    }
+    static torch::Tensor fp8_paged_mqa_logits(
+        const torch::Tensor& q,
+        const torch::Tensor& kv_cache,
+        const torch::Tensor& weights,
+        const torch::Tensor& context_lens,
+        const torch::Tensor& block_table,
+        const torch::Tensor& schedule_meta,
+        const int64_t& max_context_len,
+        const bool& clean_logits,
+        const c10::optional<torch::Tensor>& indices) {
+        return attention::fp8_paged_mqa_logits(
+            q, kv_cache, weights,
+            context_lens, block_table, schedule_meta,
+            static_cast<int>(max_context_len), clean_logits, indices);
+    }
+    #endif
+    }  // namespace deep_gemm::torch_registration
+    TORCH_LIBRARY_FRAGMENT(deep_gemm, m) {
     #if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE
-        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_fp4_mqa_logits", &fp8_fp4_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,
-              py::arg("max_seqlen_k") = 0,
-              py::arg("logits_dtype") = torch::kFloat32);
-        m.def("get_paged_mqa_logits_metadata", &get_paged_mqa_logits_metadata,
-              py::arg("context_lens"), py::arg("block_kv"), py::arg("num_sms"),
-              py::arg("indices") = std::nullopt);
-        m.def("fp8_fp4_paged_mqa_logits", &fp8_fp4_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,
-              py::arg("logits_dtype") = torch::kFloat32,
-              py::arg("indices") = std::nullopt);
-        // Legacy API
-        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,
-              py::arg("max_seqlen_k") = 0);
-        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,
-              py::arg("indices") = std::nullopt);
+        m.def(
+            "fp8_gemm_nt_skip_head_mid(Tensor a, Tensor sfa, Tensor b, Tensor sfb, Tensor(d!) d, int[3] head_splits, int[3]? recipe=None, str compiled_dims='nk', bool disable_ue8m0_cast=False) -> ()");
+        m.def(
+            "fp8_fp4_mqa_logits(Tensor q, Tensor? q_sf, Tensor kv, Tensor kv_sf, Tensor weights, Tensor cu_seq_len_k_start, Tensor cu_seq_len_k_end, bool clean_logits=True, int max_seqlen_k=0, ScalarType logits_dtype=float) -> Tensor");
+        m.def(
+            "get_paged_mqa_logits_metadata(Tensor context_lens, int block_kv, int num_sms, Tensor? indices=None) -> Tensor");
+        m.def(
+            "fp8_fp4_paged_mqa_logits(Tensor q, Tensor? q_sf, Tensor kv_cache, Tensor weights, Tensor context_lens, Tensor block_table, Tensor schedule_meta, int max_context_len, bool clean_logits=False, ScalarType logits_dtype=float, Tensor? indices=None) -> Tensor");
+        m.def(
+            "fp8_mqa_logits(Tensor q, Tensor kv, Tensor kv_sf, Tensor weights, Tensor cu_seq_len_k_start, Tensor cu_seq_len_k_end, bool clean_logits=True, int max_seqlen_k=0) -> Tensor");
+        m.def(
+            "fp8_paged_mqa_logits(Tensor q, Tensor kv_cache, Tensor weights, Tensor context_lens, Tensor block_table, Tensor schedule_meta, int max_context_len, bool clean_logits=False, Tensor? indices=None) -> Tensor");
     #endif
     }
-    } // namespace deep_gemm::attention
+    TORCH_LIBRARY_IMPL(deep_gemm, CUDA, m) {
+        using namespace deep_gemm::torch_registration;
+    #if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE
+        m.impl("fp8_gemm_nt_skip_head_mid", TORCH_FN(fp8_gemm_nt_skip_head_mid));
+        m.impl("fp8_fp4_mqa_logits", TORCH_FN(fp8_fp4_mqa_logits));
+        m.impl("get_paged_mqa_logits_metadata", TORCH_FN(get_paged_mqa_logits_metadata));
+        m.impl("fp8_fp4_paged_mqa_logits", TORCH_FN(fp8_fp4_paged_mqa_logits));
+        m.impl("fp8_mqa_logits", TORCH_FN(fp8_mqa_logits));
+        m.impl("fp8_paged_mqa_logits", TORCH_FN(fp8_paged_mqa_logits));
+    #endif
+    }

csrc/apis/einsum.hpp

-Original file line number
+Diff line change
@@ -1,7 +1,6 @@
     #pragma once
-    #include <pybind11/pybind11.h>
-    #include <torch/python.h>
+    #include <torch/all.h>
     #include "../utils/exception.hpp"
     #include "../utils/format.hpp"
@@ Expand All / @@ -19,6 +18,8 @@ @@
     #include "../jit_kernels/impls/sm120_fp8_fp4_gemm_1d1d.hpp"
     #include "../jit_kernels/impls/smxx_cublaslt.hpp"
     #endif
+    #include <torch/library.h>
+    #include "../torch_library_utils.hpp"
     namespace deep_gemm::einsum {
@@ Expand Down Expand Up / @@ -268,17 +269,45 @@ static void fp8_einsum(const std::string& expr, @@
     }
     #endif
-    static void register_apis(pybind11::module_& m) {
+    }  // namespace deep_gemm::einsum
+    namespace deep_gemm::torch_registration {
+    using namespace deep_gemm::torch_utils;
     #if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE
-        m.def("einsum", &einsum,
-              py::arg("expr"), py::arg("a"), py::arg("b"),
-              py::arg("d"), py::arg("c") = std::nullopt,
-              py::arg("use_cublaslt") = false);
-        m.def("fp8_einsum", &fp8_einsum,
-              py::arg("expr"), py::arg("a"), py::arg("b"),
-              py::arg("d"),  py::arg("c") = std::nullopt,
-              py::arg("recipe") = std::make_tuple(1, 128, 128));
+    static void einsum(const std::string& expr,
+                       const torch::Tensor& a, const torch::Tensor& b,
+                       const torch::Tensor& d, const c10::optional<torch::Tensor>& c,
+                       const bool& use_cublaslt) {
+        einsum::einsum(expr, a, b, d, c, use_cublaslt);
+    }
+    static void fp8_einsum(const std::string& expr,
+                           const torch::Tensor& a, const torch::Tensor& sfa,
+                           const torch::Tensor& b, const torch::Tensor& sfb,
+                           const torch::Tensor& d, const c10::optional<torch::Tensor>& c,
+                           const std::vector<int64_t>& recipe) {
+        einsum::fp8_einsum(expr, {a, sfa}, {b, sfb}, d, c, list_to_tuple3(recipe));
+    }
+    #endif
+    }  // namespace deep_gemm::torch_registration
+    TORCH_LIBRARY_FRAGMENT(deep_gemm, m) {
+    #if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE
+        m.def(
+            "einsum(str expr, Tensor a, Tensor b, Tensor(d!) d, Tensor? c=None, bool use_cublaslt=False) -> ()");
+        m.def(
+            "fp8_einsum(str expr, Tensor a, Tensor sfa, Tensor b, Tensor sfb, Tensor(d!) d, Tensor? c=None, int[3] recipe) -> ()");
     #endif
     }
-    } // namespace deep_gemm::einsum
+    TORCH_LIBRARY_IMPL(deep_gemm, CUDA, m) {
+        using namespace deep_gemm::torch_registration;
+    #if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE
+        m.impl("einsum", TORCH_FN(einsum));
+        m.impl("fp8_einsum", TORCH_FN(fp8_einsum));
+    #endif
+    }

Navigation Menu

Uh oh!

Migrate pybind11 to TORCH_LIBRARY #393

Uh oh!

Diff view

Diff view

There are no files selected for viewing

Uh oh!

Uh oh!

Migrate pybind11 to TORCH_LIBRARY #393

Are you sure you want to change the base?

Uh oh!

Migrate pybind11 to TORCH_LIBRARY #393

Uh oh!

Uh oh!

Diff view

Diff view

There are no files selected for viewing

Uh oh!