Repository navigation

Benchmarking DeepGEMM main vs migrate_pybind_to_torch_library_dev branch #9

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 13 commits into migrate_pybind_to_torch_library_dev from benchmark_branch

.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 one abi3 wheel with a CPython 3.10 floor, so one Python
+            # build covers every supported Python version in the release 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 @@

CMakeLists.txt

-Original file line number
+Diff line change
@@ Expand Up @@
     link_directories(${TORCH_INSTALL_PREFIX}/lib ${CUDA_TOOLKIT_ROOT_DIR}/lib64 ${CUDA_TOOLKIT_ROOT_DIR}/lib64/stubs)
     # The main Python API entrance
-    pybind11_add_module(_C csrc/python_api.cpp)
-    target_link_libraries(_C PRIVATE ${TORCH_LIBRARIES} torch_python)
+    pybind11_add_module(_C_extension csrc/python_api.cpp)
+    target_compile_definitions(_C_extension PRIVATE DJ_DISABLE_GIL=1)
+    target_link_libraries(_C_extension PRIVATE ${TORCH_LIBRARIES} torch_python)
     # Enable kernel code indexing with CMake-based IDEs
     cuda_add_library(deep_gemm_indexing_cuda STATIC csrc/indexing/main.cu)

csrc/apis/attention.hpp

-Original file line number
+Diff line change
@@ Expand Up / @@ -15,6 +15,8 @@ @@
     #include "../jit_kernels/impls/sm90_fp8_mqa_logits.hpp"
     #include "layout.hpp"
+    #include <torch/library.h>
+    #include "../torch_library_utils.hpp"
     #include "sm120_dispatch.hpp"
     namespace deep_gemm::attention {
@@ Expand Down Expand Up @@
                                         context_lens, block_table, schedule_meta,
                                         max_context_len, clean_logits, torch::kFloat, indices);
     }
-    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_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,
-              py::arg("schedule_meta") = std::nullopt);
-        m.def("get_mqa_logits_metadata", &get_mqa_logits_metadata,
-              py::arg("cu_seq_len_k_start"), py::arg("cu_seq_len_k_end"),
-              py::arg("num_kv_tokens"), py::arg("num_heads"));
-        m.def("get_sparse_mqa_logits_metadata", &get_sparse_mqa_logits_metadata,
-              py::arg("cu_seq_len_k_start"), py::arg("cu_seq_len_k_end"), py::arg("num_kv_tokens"),
-              py::arg("sparse_kv_block_indices"), py::arg("qk_dtype"), py::arg("sparse_block_kv"),
-              py::arg("use_unaligned_ks") = false);
-        m.def("get_paged_sparse_mqa_logits_metadata", &get_paged_sparse_mqa_logits_metadata,
-              py::arg("context_lens"), py::arg("block_table"), py::arg("indices"), py::arg("page_kv"),
-              py::arg("sparse_kv_block_indices"), py::arg("qk_dtype"), py::arg("sparse_block_kv"));
-        m.def("fp8_fp4_sparse_mqa_logits", &fp8_fp4_sparse_mqa_logits,
-              py::arg("q"), py::arg("kv"), py::arg("weights"), py::arg("metadata"),
-              py::arg("num_max_sparse_blocks"), py::arg("sparse_block_kv"),
-              py::arg("use_unaligned_ks") = false);
-        m.def("fp8_fp4_paged_sparse_mqa_logits", &fp8_fp4_paged_sparse_mqa_logits,
-              py::arg("q"), py::arg("kv_cache"), py::arg("weights"), py::arg("metadata"),
-              py::arg("num_max_sparse_blocks"), py::arg("sparse_block_kv"));
-        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);
+    }  // namespace deep_gemm::attention
+    namespace deep_gemm::torch_registration {
+    using namespace deep_gemm::torch_utils;
+    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,
+        const std::optional<torch::Tensor>& schedule_meta) {
+        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, schedule_meta);
+    }
+    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);
+    }
+    static torch::Tensor get_mqa_logits_metadata(
+        const torch::Tensor& cu_seq_len_k_start,
+        const torch::Tensor& cu_seq_len_k_end,
+        const int64_t& num_kv_tokens,
+        const int64_t& num_heads) {
+        return attention::get_mqa_logits_metadata(
+            cu_seq_len_k_start, cu_seq_len_k_end, num_kv_tokens, num_heads);
+    }
+    static torch::Tensor get_sparse_mqa_logits_metadata(
+        const torch::Tensor& cu_seq_len_k_start,
+        const torch::Tensor& cu_seq_len_k_end,
+        const int64_t& num_kv_tokens,
+        const torch::Tensor& sparse_kv_block_indices,
+        const at::ScalarType& qk_dtype,
+        const int64_t& sparse_block_kv,
+        const bool& use_unaligned_ks) {
+        return attention::get_sparse_mqa_logits_metadata(
+            cu_seq_len_k_start, cu_seq_len_k_end, num_kv_tokens, sparse_kv_block_indices, qk_dtype, sparse_block_kv, use_unaligned_ks);
+    }
+    static torch::Tensor get_paged_sparse_mqa_logits_metadata(
+        const torch::Tensor& context_lens,
+        const torch::Tensor& block_table,
+        const torch::Tensor& indices,
+        const int64_t& page_kv,
+        const torch::Tensor& sparse_kv_block_indices,
+        const at::ScalarType& qk_dtype,
+        const int64_t& sparse_block_kv) {
+        return attention::get_paged_sparse_mqa_logits_metadata(
+            context_lens, block_table, indices, page_kv, sparse_kv_block_indices, qk_dtype, sparse_block_kv);
+    }
+    static torch::Tensor fp8_fp4_sparse_mqa_logits(
+        const torch::Tensor& q,
+        const std::optional<torch::Tensor>& q_sf,
+        const torch::Tensor& kv,
+        const torch::Tensor& kv_sf,
+        const torch::Tensor& weights,
+        const torch::Tensor& metadata,
+        const int64_t& num_max_sparse_blocks,
+        const int64_t& sparse_block_kv,
+        const bool& use_unaligned_ks) {
+        return attention::fp8_fp4_sparse_mqa_logits(
+            {q, q_sf}, {kv, kv_sf}, weights, metadata, num_max_sparse_blocks, sparse_block_kv, use_unaligned_ks);
+    }
+    static torch::Tensor fp8_fp4_paged_sparse_mqa_logits(
+        const torch::Tensor& q,
+        const std::optional<torch::Tensor>& q_sf,
+        const torch::Tensor& fused_kv_cache,
+        const torch::Tensor& weights,
+        const torch::Tensor& metadata,
+        const int64_t& num_max_sparse_blocks,
+        const int64_t& sparse_block_kv) {
+        return attention::fp8_fp4_paged_sparse_mqa_logits(
+            {q, q_sf}, fused_kv_cache, weights, metadata, num_max_sparse_blocks, sparse_block_kv);
+    }
+    } // namespace deep_gemm::torch_registration
+    TORCH_LIBRARY_FRAGMENT(deep_gemm, m) {
+        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? schedule_meta=None) -> Tensor");
+        m.def("get_mqa_logits_metadata(Tensor cu_seq_len_k_start, Tensor cu_seq_len_k_end, int num_kv_tokens, int num_heads) -> Tensor");
+        m.def("get_sparse_mqa_logits_metadata(Tensor cu_seq_len_k_start, Tensor cu_seq_len_k_end, int num_kv_tokens, Tensor sparse_kv_block_indices, ScalarType qk_dtype, int sparse_block_kv, bool use_unaligned_ks=False) -> Tensor");
+        m.def("get_paged_sparse_mqa_logits_metadata(Tensor context_lens, Tensor block_table, Tensor indices, int page_kv, Tensor sparse_kv_block_indices, ScalarType qk_dtype, int sparse_block_kv) -> Tensor");
+        m.def("fp8_fp4_sparse_mqa_logits(Tensor q, Tensor? q_sf, Tensor kv, Tensor kv_sf, Tensor weights, Tensor metadata, int num_max_sparse_blocks, int sparse_block_kv, bool use_unaligned_ks=False) -> Tensor");
+        m.def("fp8_fp4_paged_sparse_mqa_logits(Tensor q, Tensor? q_sf, Tensor kv_cache, Tensor weights, Tensor metadata, int num_max_sparse_blocks, int sparse_block_kv) -> 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");
         // 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_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");
     }
-    } // namespace deep_gemm::attention
+    TORCH_LIBRARY_IMPL(deep_gemm, CUDA, m) {
+        using namespace deep_gemm::torch_registration;
+        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_mqa_logits_metadata", TORCH_FN(deep_gemm::torch_registration::get_mqa_logits_metadata));
+        m.impl("get_sparse_mqa_logits_metadata", TORCH_FN(deep_gemm::torch_registration::get_sparse_mqa_logits_metadata));
+        m.impl("get_paged_sparse_mqa_logits_metadata", TORCH_FN(deep_gemm::torch_registration::get_paged_sparse_mqa_logits_metadata));
+        m.impl("fp8_fp4_sparse_mqa_logits", TORCH_FN(deep_gemm::torch_registration::fp8_fp4_sparse_mqa_logits));
+        m.impl("fp8_fp4_paged_sparse_mqa_logits", TORCH_FN(deep_gemm::torch_registration::fp8_fp4_paged_sparse_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));
+        // Legacy API
+        m.impl("fp8_mqa_logits", TORCH_FN(fp8_mqa_logits));
+        m.impl("fp8_paged_mqa_logits", TORCH_FN(fp8_paged_mqa_logits));
+    }

csrc/apis/config.hpp

            
                      Original file line number
                      Diff line number
                      Diff line change
                  
    @@ -1,47 +1,41 @@
  
    #pragma once

    #include <torch/library.h>

    #include "../jit_kernels/heuristics/runtime.hpp"

    #include "../runtime/runtime.hpp"

    namespace deep_gemm::config {

    static void register_apis(pybind11::module_& m) {

        m.def("init", [&](const std::string& library_root_path) {

    TORCH_LIBRARY_FRAGMENT(deep_gemm, m) {

        using namespace deep_gemm;

        m.def("init(str library_root_path) -> ()", [&](const std::string& library_root_path) {

            init_jit(library_root_path);

        });

        m.def("set_num_sms", [&](const int& new_num_sms) {

        m.def("set_num_sms(int new_num_sms) -> ()", [&](const int64_t& new_num_sms) {

            runtime->set_num_sms(new_num_sms);

        });

        m.def("get_num_sms", [&]() {

           return runtime->get_num_sms();

        m.def("get_num_sms() -> int", [&]() {

           return int64_t(runtime->get_num_sms());

        });

        m.def("set_tc_util", [&](const int& new_tc_util) {

        m.def("set_tc_util(int new_tc_util) -> ()", [&](const int64_t& new_tc_util) {

            runtime->set_tc_util(new_tc_util);

        });

        m.def("get_tc_util", [&]() {

            return runtime->get_tc_util();

        m.def("get_tc_util() -> int", [&]() {

            return int64_t(runtime->get_tc_util());

        });

        m.def("set_pdl", [](const bool& new_enable_pdl) {

        m.def("set_pdl(bool new_enable_pdl) -> ()", [](const bool& new_enable_pdl) {

            jit->default_launch_options.enable_pdl = new_enable_pdl;

        });

        m.def("get_pdl", []() {

        m.def("get_pdl() -> bool", []() {

            return *jit->default_launch_options.enable_pdl;

        });

        m.def("use_deterministic_algorithms", [&](const bool enabled) {

        m.def("use_deterministic_algorithms(bool enabled) -> ()", [&](const bool enabled) {

            heuristics_runtime->use_deterministic_algorithms(enabled);

        });

        m.def("set_ignore_compile_dims", [&](const bool& new_value) {

        m.def("set_ignore_compile_dims(bool new_value) -> ()", [&](const bool& new_value) {

            heuristics_runtime->set_ignore_compile_dims(new_value);

        });

        m.def("set_block_size_multiple_of", [&](const std::variant<int, std::tuple<int, int>>& new_value) {

            if (std::holds_alternative<int>(new_value)) {

                auto x = std::get<int>(new_value);

                heuristics_runtime->set_block_size_multiple_of(x, x);

            } else {

                auto [x, y] = std::get<std::tuple<int, int>>(new_value);

                heuristics_runtime->set_block_size_multiple_of(x, y);

            }

        m.def("set_block_size_multiple_of(int[] new_value) -> ()", [&](const std::vector<int64_t>& new_value) {

            DG_HOST_ASSERT(new_value.size() == 1 or new_value.size() == 2);

            heuristics_runtime->set_block_size_multiple_of(new_value[0], new_value.size() == 1 ? new_value[0] : new_value[1]);

        });

    }

    } // namespace deep_gemm::config

Navigation Menu

Repository navigation

Benchmarking DeepGEMM main vs migrate_pybind_to_torch_library_dev branch #9

Uh oh!

Diff view

Diff view

There are no files selected for viewing

Uh oh!

Benchmarking DeepGEMM main vs migrate_pybind_to_torch_library_dev branch #9

Are you sure you want to change the base?

Uh oh!

Benchmarking DeepGEMM main vs migrate_pybind_to_torch_library_dev branch #9

Uh oh!

Uh oh!

Diff view

Diff view

There are no files selected for viewing

Uh oh!