diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index a7b3e6b880..2903785635 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -41,7 +41,9 @@ 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. @@ -49,11 +51,6 @@ jobs: # 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 }} diff --git a/CMakeLists.txt b/CMakeLists.txt index 70797aa793..9880a258f6 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -27,8 +27,9 @@ include_directories(${CUDA_TOOLKIT_ROOT_DIR}/include/cccl ${TORCH_INCLUDE_DIRS} 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) diff --git a/csrc/apis/attention.hpp b/csrc/apis/attention.hpp index 010710cff8..191bb98e20 100644 --- a/csrc/apis/attention.hpp +++ b/csrc/apis/attention.hpp @@ -15,6 +15,8 @@ #include "../jit_kernels/impls/sm90_fp8_mqa_logits.hpp" #include "layout.hpp" +#include +#include "../torch_library_utils.hpp" #include "sm120_dispatch.hpp" namespace deep_gemm::attention { @@ -613,57 +615,194 @@ static torch::Tensor fp8_paged_mqa_logits(const torch::Tensor& q, 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& head_splits, + const c10::optional>& 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& 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& 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(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& indices) { + return attention::get_paged_mqa_logits_metadata( + context_lens, static_cast(block_kv), + static_cast(num_sms), indices); +} + +static torch::Tensor fp8_fp4_paged_mqa_logits( + const torch::Tensor& q, const c10::optional& 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& 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(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(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& indices) { + return attention::fp8_paged_mqa_logits( + q, kv_cache, weights, + context_lens, block_table, schedule_meta, + static_cast(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& 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& 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)); +} diff --git a/csrc/apis/config.hpp b/csrc/apis/config.hpp index a0a6dd7366..5654502f42 100644 --- a/csrc/apis/config.hpp +++ b/csrc/apis/config.hpp @@ -1,47 +1,41 @@ #pragma once +#include + #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>& new_value) { - if (std::holds_alternative(new_value)) { - auto x = std::get(new_value); - heuristics_runtime->set_block_size_multiple_of(x, x); - } else { - auto [x, y] = std::get>(new_value); - heuristics_runtime->set_block_size_multiple_of(x, y); - } + m.def("set_block_size_multiple_of(int[] new_value) -> ()", [&](const std::vector& 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 diff --git a/csrc/apis/einsum.hpp b/csrc/apis/einsum.hpp index 40cfe02ef6..611fe7b96f 100644 --- a/csrc/apis/einsum.hpp +++ b/csrc/apis/einsum.hpp @@ -1,10 +1,12 @@ #pragma once +#include +#include "../torch_library_utils.hpp" + #include #include -#include -#include +#include #include "../utils/exception.hpp" #include "../utils/layout.hpp" @@ -297,14 +299,33 @@ static void fp8_einsum(const std::string& expr, DG_HOST_UNREACHABLE(std::format("Unsupported einsum expression: {}", 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); - 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)); +} // namespace deep_gemm::einsum + +namespace deep_gemm::torch_registration { +using namespace deep_gemm::torch_utils; + +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& c, + const std::vector& recipe, + const std::optional& sfd) { + std::variant> output = d; + if (sfd.has_value()) output = std::make_pair(d, *sfd); + einsum::fp8_einsum(expr, {a, sfa}, {b, sfb}, output, c, list_to_tuple3(recipe)); +} +} // namespace deep_gemm::torch_registration + +TORCH_LIBRARY_FRAGMENT(deep_gemm, m) { + m.def( + "einsum(str expr, Tensor a, Tensor b, Tensor(d!) d, Tensor? c=None) -> ()"); + m.def( + "fp8_einsum(str expr, Tensor a, Tensor sfa, Tensor b, Tensor sfb, Tensor(d!) d, Tensor? c=None, int[3] recipe, Tensor(sfd!)? sfd=None) -> ()"); } -} // namespace deep_gemm::einsum +TORCH_LIBRARY_IMPL(deep_gemm, CUDA, m) { + using namespace deep_gemm::torch_registration; + + m.impl("einsum", TORCH_FN(deep_gemm::einsum::einsum)); + m.impl("fp8_einsum", TORCH_FN(fp8_einsum)); +} diff --git a/csrc/apis/gemm.hpp b/csrc/apis/gemm.hpp index 58fa1b0c8f..2ef98e2b95 100644 --- a/csrc/apis/gemm.hpp +++ b/csrc/apis/gemm.hpp @@ -1,5 +1,8 @@ #pragma once +#include +#include "../torch_library_utils.hpp" + #include "../utils/compatibility.hpp" #include "../jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp" @@ -849,148 +852,313 @@ static void batched_symm(const torch::Tensor& a, const torch::Tensor& b, cublaslt_batched_gemm(a_3d, b_transposed, d_3d, m, k, m, num_batches, major_a, major_b); } -static void register_apis(pybind11::module_& m) { - - // FP8 FP4 GEMMs - m.def("fp8_fp4_gemm_nt", &fp8_fp4_gemm_nt, - py::arg("a"), py::arg("b"), py::arg("d"), - py::arg("c") = std::nullopt, py::arg("recipe") = std::nullopt, - py::arg("recipe_a") = std::nullopt, py::arg("recipe_b") = std::nullopt, - py::arg("compiled_dims") = "nk", - py::arg("disable_ue8m0_cast") = false, - py::arg("alpha") = std::nullopt); - m.def("fp8_fp4_gemm_nn", &fp8_fp4_gemm_nn, - py::arg("a"), py::arg("b"), py::arg("d"), - py::arg("c") = std::nullopt, py::arg("recipe") = std::nullopt, - py::arg("recipe_a") = std::nullopt, py::arg("recipe_b") = std::nullopt, - py::arg("compiled_dims") = "nk", - py::arg("disable_ue8m0_cast") = false, - py::arg("alpha") = std::nullopt); - m.def("fp8_fp4_gemm_tn", &fp8_fp4_gemm_tn, - py::arg("a"), py::arg("b"), py::arg("d"), - py::arg("c") = std::nullopt, py::arg("recipe") = std::nullopt, - py::arg("recipe_a") = std::nullopt, py::arg("recipe_b") = std::nullopt, - py::arg("compiled_dims") = "mn", - py::arg("disable_ue8m0_cast") = false, - py::arg("alpha") = std::nullopt); - m.def("fp8_fp4_gemm_tt", &fp8_fp4_gemm_tt, - py::arg("a"), py::arg("b"), py::arg("d"), - py::arg("c") = std::nullopt, py::arg("recipe") = std::nullopt, - py::arg("recipe_a") = std::nullopt, py::arg("recipe_b") = std::nullopt, - py::arg("compiled_dims") = "mn", - py::arg("disable_ue8m0_cast") = false, - py::arg("alpha") = std::nullopt); - m.def("m_grouped_fp8_fp4_gemm_nt_contiguous", &m_grouped_fp8_fp4_gemm_nt_contiguous, - py::arg("a"), py::arg("b"), py::arg("d"), py::arg("grouped_layout"), - py::arg("recipe") = std::nullopt, - py::arg("recipe_a") = std::nullopt, py::arg("recipe_b") = std::nullopt, - py::arg("compiled_dims") = "nk", - py::arg("disable_ue8m0_cast") = false, - py::arg("use_psum_layout") = false, - py::arg("ensure_zero_padding") = true, - py::arg("expected_m_for_psum_layout") = std::nullopt); - m.def("m_grouped_fp8_fp4_gemm_nn_contiguous", &m_grouped_fp8_fp4_gemm_nn_contiguous, - py::arg("a"), py::arg("b"), py::arg("d"), py::arg("grouped_layout"), - py::arg("recipe") = std::nullopt, - py::arg("recipe_a") = std::nullopt, py::arg("recipe_b") = std::nullopt, - py::arg("compiled_dims") = "nk", - py::arg("disable_ue8m0_cast") = false, - py::arg("use_psum_layout") = false, - py::arg("ensure_zero_padding") = true); - m.def("m_grouped_fp8_fp4_gemm_nt_masked", &m_grouped_fp8_fp4_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("recipe_a") = std::nullopt, py::arg("recipe_b") = 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_cpu"), py::arg("grouped_layout"), - py::arg("c") = std::nullopt, - py::arg("recipe") = std::make_tuple(1, 1, 128), - py::arg("compiled_dims") = "mn", - py::arg("use_psum_layout") = false); - 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_cpu"), py::arg("grouped_layout"), - py::arg("c") = std::nullopt, - py::arg("recipe") = std::make_tuple(1, 1, 128), - py::arg("compiled_dims") = "mn", - py::arg("use_psum_layout") = false); - m.def("k_grouped_fp4_gemm_nt_contiguous", &k_grouped_fp4_gemm_nt_contiguous, - py::arg("a"), py::arg("b"), py::arg("d"), - py::arg("ks_cpu"), py::arg("grouped_layout"), - py::arg("c") = std::nullopt, - py::arg("recipe") = std::make_tuple(1, 1, 32), - py::arg("compiled_dims") = "mn", - py::arg("use_psum_layout") = false); - - // FP4 and FP8 GEMM alias names - m.attr("fp4_gemm_nt") = m.attr("fp8_fp4_gemm_nt"); - m.attr("fp8_gemm_nt") = m.attr("fp8_fp4_gemm_nt"); - m.attr("fp8_gemm_nn") = m.attr("fp8_fp4_gemm_nn"); - m.attr("fp8_gemm_tn") = m.attr("fp8_fp4_gemm_tn"); - m.attr("fp8_gemm_tt") = m.attr("fp8_fp4_gemm_tt"); - m.attr("m_grouped_fp4_gemm_nt_contiguous") = m.attr("m_grouped_fp8_fp4_gemm_nt_contiguous"); - m.attr("m_grouped_fp8_gemm_nt_contiguous") = m.attr("m_grouped_fp8_fp4_gemm_nt_contiguous"); - m.attr("m_grouped_fp8_gemm_nn_contiguous") = m.attr("m_grouped_fp8_fp4_gemm_nn_contiguous"); - m.attr("m_grouped_fp4_gemm_nt_masked") = m.attr("m_grouped_fp8_fp4_gemm_nt_masked"); - m.attr("m_grouped_fp8_gemm_nt_masked") = m.attr("m_grouped_fp8_fp4_gemm_nt_masked"); - // 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", - py::arg("alpha") = std::nullopt); - 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", - py::arg("alpha") = std::nullopt); - 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", - py::arg("alpha") = std::nullopt); - 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", - py::arg("alpha") = std::nullopt); - 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("grouped_layout"), - py::arg("compiled_dims") = "nk", - py::arg("use_psum_layout") = false, - py::arg("ensure_zero_padding") = true, - py::arg("expected_m_for_psum_layout") = std::nullopt); - m.def("m_grouped_bf16_gemm_nn_contiguous", &m_grouped_bf16_gemm_nn_contiguous, - py::arg("a"), py::arg("b"), py::arg("d"), py::arg("grouped_layout"), - py::arg("compiled_dims") = "nk", - py::arg("use_psum_layout") = false, - py::arg("ensure_zero_padding") = true); - 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"); - m.def("k_grouped_bf16_gemm_tn_contiguous", &k_grouped_bf16_gemm_tn_contiguous, - py::arg("a"), py::arg("b"), py::arg("d"), - py::arg("ks_cpu"), py::arg("grouped_layout"), - py::arg("c") = std::nullopt, - py::arg("compiled_dims") = "mn", - py::arg("use_psum_layout") = false); - // 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); - m.def("cublaslt_nvfp4_gemm_nt", &cublaslt_nvfp4_gemm_nt, - py::arg("a"), py::arg("b"), py::arg("d"), py::arg("c") = std::nullopt); - m.def("batched_syrk", &batched_syrk, - py::arg("a"), py::arg("d")); - m.def("batched_symm", &batched_symm, - py::arg("a"), py::arg("b"), py::arg("d")); -} - -} // namespace deep_gemm::gemm +} // namespace deep_gemm::gemm + +namespace deep_gemm::torch_registration { +using namespace deep_gemm::torch_utils; + +static void fp8_fp4_gemm_nt( + const torch::Tensor& a, const torch::Tensor& sfa, + const torch::Tensor& b, const torch::Tensor& sfb, + const torch::Tensor& d, const c10::optional& c, + const c10::optional>& recipe, + const c10::optional>& recipe_a, + const c10::optional>& recipe_b, + const std::string& compiled_dims, const bool& disable_ue8m0_cast, + const std::optional& alpha) { + gemm::fp8_fp4_gemm_nt({a, sfa}, {b, sfb}, d, c, + list_to_recipe3(recipe), list_to_recipe2(recipe_a), list_to_recipe2(recipe_b), + compiled_dims, disable_ue8m0_cast, alpha); +} + +static void fp8_fp4_gemm_nn( + const torch::Tensor& a, const torch::Tensor& sfa, + const torch::Tensor& b, const torch::Tensor& sfb, + const torch::Tensor& d, const c10::optional& c, + const c10::optional>& recipe, + const c10::optional>& recipe_a, + const c10::optional>& recipe_b, + const std::string& compiled_dims, const bool& disable_ue8m0_cast, + const std::optional& alpha) { + gemm::fp8_fp4_gemm_nn({a, sfa}, {b, sfb}, d, c, + list_to_recipe3(recipe), list_to_recipe2(recipe_a), list_to_recipe2(recipe_b), + compiled_dims, disable_ue8m0_cast, alpha); +} + +static void fp8_fp4_gemm_tn( + const torch::Tensor& a, const torch::Tensor& sfa, + const torch::Tensor& b, const torch::Tensor& sfb, + const torch::Tensor& d, const c10::optional& c, + const c10::optional>& recipe, + const c10::optional>& recipe_a, + const c10::optional>& recipe_b, + const std::string& compiled_dims, const bool& disable_ue8m0_cast, + const std::optional& alpha) { + gemm::fp8_fp4_gemm_tn({a, sfa}, {b, sfb}, d, c, + list_to_recipe3(recipe), list_to_recipe2(recipe_a), list_to_recipe2(recipe_b), + compiled_dims, disable_ue8m0_cast, alpha); +} + +static void fp8_fp4_gemm_tt( + const torch::Tensor& a, const torch::Tensor& sfa, + const torch::Tensor& b, const torch::Tensor& sfb, + const torch::Tensor& d, const c10::optional& c, + const c10::optional>& recipe, + const c10::optional>& recipe_a, + const c10::optional>& recipe_b, + const std::string& compiled_dims, const bool& disable_ue8m0_cast, + const std::optional& alpha) { + gemm::fp8_fp4_gemm_tt({a, sfa}, {b, sfb}, d, c, + list_to_recipe3(recipe), list_to_recipe2(recipe_a), list_to_recipe2(recipe_b), + compiled_dims, disable_ue8m0_cast, alpha); +} + +static void m_grouped_fp8_fp4_gemm_nt_contiguous( + const torch::Tensor& a, const torch::Tensor& sfa, + const torch::Tensor& b, const torch::Tensor& sfb, + const torch::Tensor& d, const torch::Tensor& grouped_layout, + const c10::optional>& recipe, + const c10::optional>& recipe_a, + const c10::optional>& recipe_b, + const std::string& compiled_dims, const bool& disable_ue8m0_cast, + const bool& use_psum_layout, const bool& ensure_zero_padding, + const c10::optional& expected_m_for_psum_layout) { + gemm::m_grouped_fp8_fp4_gemm_nt_contiguous( + {a, sfa}, {b, sfb}, d, grouped_layout, + list_to_recipe3(recipe), list_to_recipe2(recipe_a), list_to_recipe2(recipe_b), + compiled_dims, disable_ue8m0_cast, use_psum_layout, ensure_zero_padding, + expected_m_for_psum_layout.has_value() + ? std::make_optional(static_cast(expected_m_for_psum_layout.value())) + : std::nullopt); +} + +static void m_grouped_fp8_fp4_gemm_nn_contiguous( + const torch::Tensor& a, const torch::Tensor& sfa, + const torch::Tensor& b, const torch::Tensor& sfb, + const torch::Tensor& d, const torch::Tensor& grouped_layout, + const c10::optional>& recipe, + const c10::optional>& recipe_a, + const c10::optional>& recipe_b, + const std::string& compiled_dims, const bool& disable_ue8m0_cast, + const bool& use_psum_layout, const bool& ensure_zero_padding) { + gemm::m_grouped_fp8_fp4_gemm_nn_contiguous( + {a, sfa}, {b, sfb}, d, grouped_layout, + list_to_recipe3(recipe), list_to_recipe2(recipe_a), list_to_recipe2(recipe_b), + compiled_dims, disable_ue8m0_cast, use_psum_layout, ensure_zero_padding); +} + +static void m_grouped_fp8_fp4_gemm_nt_masked( + 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, + const int64_t& expected_m, + const c10::optional>& recipe, + const c10::optional>& recipe_a, + const c10::optional>& recipe_b, + const std::string& compiled_dims, const bool& disable_ue8m0_cast) { + gemm::m_grouped_fp8_fp4_gemm_nt_masked( + {a, sfa}, {b, sfb}, d, masked_m, static_cast(expected_m), + list_to_recipe3(recipe), list_to_recipe2(recipe_a), list_to_recipe2(recipe_b), + compiled_dims, disable_ue8m0_cast); +} + +static void k_grouped_fp8_gemm_tn_contiguous( + const torch::Tensor& a, const torch::Tensor& sfa, + const torch::Tensor& b, const torch::Tensor& sfb, + const torch::Tensor& d, + const c10::optional>& ks_cpu, + const torch::Tensor& grouped_layout, + const c10::optional& c, + const std::vector& recipe, + const std::string& compiled_dims, const bool& use_psum_layout) { + gemm::k_grouped_fp8_gemm_tn_contiguous( + {a, sfa}, {b, sfb}, d, + list_to_optional_vector_int(ks_cpu), grouped_layout, c, + list_to_tuple3(recipe), compiled_dims, use_psum_layout); +} + +static void k_grouped_fp8_gemm_nt_contiguous( + const torch::Tensor& a, const torch::Tensor& sfa, + const torch::Tensor& b, const torch::Tensor& sfb, + const torch::Tensor& d, + const c10::optional>& ks_cpu, + const torch::Tensor& grouped_layout, + const c10::optional& c, + const std::vector& recipe, + const std::string& compiled_dims, const bool& use_psum_layout) { + gemm::k_grouped_fp8_gemm_nt_contiguous( + {a, sfa}, {b, sfb}, d, + list_to_optional_vector_int(ks_cpu), grouped_layout, c, + list_to_tuple3(recipe), compiled_dims, use_psum_layout); +} + +static void k_grouped_fp4_gemm_nt_contiguous( + const torch::Tensor& a, + const torch::Tensor& sfa, + const torch::Tensor& b, + const torch::Tensor& sfb, + const torch::Tensor& d, + const std::optional>& ks_cpu, + const torch::Tensor& grouped_layout, + const std::optional& c, + const std::vector& recipe, + const std::string& compiled_dims, + const bool& use_psum_layout) { + gemm::k_grouped_fp4_gemm_nt_contiguous( + {a, sfa}, {b, sfb}, d, list_to_optional_vector_int(ks_cpu), grouped_layout, c, list_to_tuple3(recipe), compiled_dims, use_psum_layout); +} + +static void bf16_gemm_nt( + const torch::Tensor& a, const torch::Tensor& b, const torch::Tensor& d, + const c10::optional& c, const std::string& compiled_dims, + const std::optional& alpha) { + gemm::bf16_gemm_nt(a, b, d, c, compiled_dims, alpha); +} + +static void bf16_gemm_nn( + const torch::Tensor& a, const torch::Tensor& b, const torch::Tensor& d, + const c10::optional& c, const std::string& compiled_dims, + const std::optional& alpha) { + gemm::bf16_gemm_nn(a, b, d, c, compiled_dims, alpha); +} + +static void bf16_gemm_tn( + const torch::Tensor& a, const torch::Tensor& b, const torch::Tensor& d, + const c10::optional& c, const std::string& compiled_dims, + const std::optional& alpha) { + gemm::bf16_gemm_tn(a, b, d, c, compiled_dims, alpha); +} + +static void bf16_gemm_tt( + const torch::Tensor& a, const torch::Tensor& b, const torch::Tensor& d, + const c10::optional& c, const std::string& compiled_dims, + const std::optional& alpha) { + gemm::bf16_gemm_tt(a, b, d, c, compiled_dims, alpha); +} + +static void m_grouped_bf16_gemm_nt_contiguous( + const torch::Tensor& a, const torch::Tensor& b, const torch::Tensor& d, + const torch::Tensor& grouped_layout, const std::string& compiled_dims, + const bool& use_psum_layout, const bool& ensure_zero_padding, + const c10::optional& expected_m_for_psum_layout) { + gemm::m_grouped_bf16_gemm_nt_contiguous( + a, b, d, grouped_layout, compiled_dims, + use_psum_layout, ensure_zero_padding, + expected_m_for_psum_layout.has_value() + ? std::make_optional(static_cast(expected_m_for_psum_layout.value())) + : std::nullopt); +} + +static void m_grouped_bf16_gemm_nt_masked( + const torch::Tensor& a, const torch::Tensor& b, const torch::Tensor& d, + const torch::Tensor& masked_m, const int64_t& expected_m, + const std::string& compiled_dims) { + gemm::m_grouped_bf16_gemm_nt_masked(a, b, d, masked_m, static_cast(expected_m), compiled_dims); +} + +static void k_grouped_bf16_gemm_tn_contiguous( + const torch::Tensor& a, const torch::Tensor& b, const torch::Tensor& d, + const c10::optional>& ks_cpu, + const torch::Tensor& grouped_layout, + const c10::optional& c, + const std::string& compiled_dims, const bool& use_psum_layout) { + gemm::k_grouped_bf16_gemm_tn_contiguous( + a, b, d, list_to_optional_vector_int(ks_cpu), + grouped_layout, c, compiled_dims, use_psum_layout); +} + +static void cublaslt_nvfp4_gemm_nt( + const torch::Tensor& a, + const torch::Tensor& sfa, + const torch::Tensor& b, + const torch::Tensor& sfb, + const torch::Tensor& d, + const std::optional& c) { + gemm::cublaslt_nvfp4_gemm_nt( + {a, sfa}, {b, sfb}, d, c); +} +} // namespace deep_gemm::torch_registration + +TORCH_LIBRARY_FRAGMENT(deep_gemm, m) { + // GEMM — FP8/FP4 + m.def( + "fp8_fp4_gemm_nt(Tensor a, Tensor sfa, Tensor b, Tensor sfb, Tensor(d!) d, Tensor? c=None, int[3]? recipe=None, int[2]? recipe_a=None, int[2]? recipe_b=None, str compiled_dims='nk', bool disable_ue8m0_cast=False, float? alpha=None) -> ()"); + m.def( + "fp8_fp4_gemm_nn(Tensor a, Tensor sfa, Tensor b, Tensor sfb, Tensor(d!) d, Tensor? c=None, int[3]? recipe=None, int[2]? recipe_a=None, int[2]? recipe_b=None, str compiled_dims='nk', bool disable_ue8m0_cast=False, float? alpha=None) -> ()"); + m.def( + "fp8_fp4_gemm_tn(Tensor a, Tensor sfa, Tensor b, Tensor sfb, Tensor(d!) d, Tensor? c=None, int[3]? recipe=None, int[2]? recipe_a=None, int[2]? recipe_b=None, str compiled_dims='mn', bool disable_ue8m0_cast=False, float? alpha=None) -> ()"); + m.def( + "fp8_fp4_gemm_tt(Tensor a, Tensor sfa, Tensor b, Tensor sfb, Tensor(d!) d, Tensor? c=None, int[3]? recipe=None, int[2]? recipe_a=None, int[2]? recipe_b=None, str compiled_dims='mn', bool disable_ue8m0_cast=False, float? alpha=None) -> ()"); + m.def( + "m_grouped_fp8_fp4_gemm_nt_contiguous(Tensor a, Tensor sfa, Tensor b, Tensor sfb, Tensor(d!) d, Tensor grouped_layout, int[3]? recipe=None, int[2]? recipe_a=None, int[2]? recipe_b=None, str compiled_dims='nk', bool disable_ue8m0_cast=False, bool use_psum_layout=False, bool ensure_zero_padding=True, int? expected_m_for_psum_layout=None) -> ()"); + m.def( + "m_grouped_fp8_fp4_gemm_nn_contiguous(Tensor a, Tensor sfa, Tensor b, Tensor sfb, Tensor(d!) d, Tensor grouped_layout, int[3]? recipe=None, int[2]? recipe_a=None, int[2]? recipe_b=None, str compiled_dims='nk', bool disable_ue8m0_cast=False, bool use_psum_layout=False, bool ensure_zero_padding=True) -> ()"); + m.def( + "m_grouped_fp8_fp4_gemm_nt_masked(Tensor a, Tensor sfa, Tensor b, Tensor sfb, Tensor(d!) d, Tensor masked_m, int expected_m, int[3]? recipe=None, int[2]? recipe_a=None, int[2]? recipe_b=None, str compiled_dims='nk', bool disable_ue8m0_cast=False) -> ()"); + m.def( + "k_grouped_fp8_gemm_tn_contiguous(Tensor a, Tensor sfa, Tensor b, Tensor sfb, Tensor(d!) d, int[]? ks_cpu, Tensor grouped_layout, Tensor? c=None, int[3] recipe, str compiled_dims='mn', bool use_psum_layout=False) -> ()"); + m.def( + "k_grouped_fp8_gemm_nt_contiguous(Tensor a, Tensor sfa, Tensor b, Tensor sfb, Tensor(d!) d, int[]? ks_cpu, Tensor grouped_layout, Tensor? c=None, int[3] recipe, str compiled_dims='mn', bool use_psum_layout=False) -> ()"); + m.def("k_grouped_fp4_gemm_nt_contiguous(Tensor a, Tensor sfa, Tensor b, Tensor sfb, Tensor(a!) d, int[]? ks_cpu, Tensor grouped_layout, Tensor? c=None, int[3] recipe, str compiled_dims='mn', bool use_psum_layout=False) -> ()"); + + // GEMM — BF16 + m.def( + "bf16_gemm_nt(Tensor a, Tensor b, Tensor(d!) d, Tensor? c=None, str compiled_dims='nk', float? alpha=None) -> ()"); + m.def( + "bf16_gemm_nn(Tensor a, Tensor b, Tensor(d!) d, Tensor? c=None, str compiled_dims='nk', float? alpha=None) -> ()"); + m.def( + "bf16_gemm_tn(Tensor a, Tensor b, Tensor(d!) d, Tensor? c=None, str compiled_dims='mn', float? alpha=None) -> ()"); + m.def( + "bf16_gemm_tt(Tensor a, Tensor b, Tensor(d!) d, Tensor? c=None, str compiled_dims='mn', float? alpha=None) -> ()"); + m.def( + "m_grouped_bf16_gemm_nt_contiguous(Tensor a, Tensor b, Tensor(d!) d, Tensor grouped_layout, str compiled_dims='nk', bool use_psum_layout=False, bool ensure_zero_padding=True, int? expected_m_for_psum_layout=None) -> ()"); + m.def( + "m_grouped_bf16_gemm_nn_contiguous(Tensor a, Tensor b, Tensor(d!) d, Tensor grouped_layout, str compiled_dims='nk', bool use_psum_layout=False, bool ensure_zero_padding=True) -> ()"); + m.def( + "m_grouped_bf16_gemm_nt_masked(Tensor a, Tensor b, Tensor(d!) d, Tensor masked_m, int expected_m, str compiled_dims='nk') -> ()"); + m.def( + "k_grouped_bf16_gemm_tn_contiguous(Tensor a, Tensor b, Tensor(d!) d, int[]? ks_cpu, Tensor grouped_layout, Tensor? c=None, str compiled_dims='mn', bool use_psum_layout=False) -> ()"); + + // GEMM — cuBLASLt + m.def("cublaslt_gemm_nt(Tensor a, Tensor b, Tensor(d!) d, Tensor? c=None) -> ()"); + m.def("cublaslt_gemm_nn(Tensor a, Tensor b, Tensor(d!) d, Tensor? c=None) -> ()"); + m.def("cublaslt_gemm_tn(Tensor a, Tensor b, Tensor(d!) d, Tensor? c=None) -> ()"); + m.def("cublaslt_gemm_tt(Tensor a, Tensor b, Tensor(d!) d, Tensor? c=None) -> ()"); + + m.def("cublaslt_nvfp4_gemm_nt(Tensor a, Tensor sfa, Tensor b, Tensor sfb, Tensor(a!) d, Tensor? c=None) -> ()"); + m.def("batched_syrk(Tensor a, Tensor(a!) d) -> ()"); + m.def("batched_symm(Tensor a, Tensor b, Tensor(a!) d) -> ()"); +} + +TORCH_LIBRARY_IMPL(deep_gemm, CUDA, m) { + using namespace deep_gemm::torch_registration; + + m.impl("fp8_fp4_gemm_nt", TORCH_FN(fp8_fp4_gemm_nt)); + m.impl("fp8_fp4_gemm_nn", TORCH_FN(fp8_fp4_gemm_nn)); + m.impl("fp8_fp4_gemm_tn", TORCH_FN(fp8_fp4_gemm_tn)); + m.impl("fp8_fp4_gemm_tt", TORCH_FN(fp8_fp4_gemm_tt)); + m.impl("m_grouped_fp8_fp4_gemm_nt_contiguous", TORCH_FN(m_grouped_fp8_fp4_gemm_nt_contiguous)); + m.impl("m_grouped_fp8_fp4_gemm_nn_contiguous", TORCH_FN(m_grouped_fp8_fp4_gemm_nn_contiguous)); + m.impl("m_grouped_fp8_fp4_gemm_nt_masked", TORCH_FN(m_grouped_fp8_fp4_gemm_nt_masked)); + m.impl("k_grouped_fp8_gemm_tn_contiguous", TORCH_FN(k_grouped_fp8_gemm_tn_contiguous)); + m.impl("k_grouped_fp8_gemm_nt_contiguous", TORCH_FN(k_grouped_fp8_gemm_nt_contiguous)); + m.impl("k_grouped_fp4_gemm_nt_contiguous", TORCH_FN(deep_gemm::torch_registration::k_grouped_fp4_gemm_nt_contiguous)); + + m.impl("bf16_gemm_nt", TORCH_FN(bf16_gemm_nt)); + m.impl("bf16_gemm_nn", TORCH_FN(bf16_gemm_nn)); + m.impl("bf16_gemm_tn", TORCH_FN(bf16_gemm_tn)); + m.impl("bf16_gemm_tt", TORCH_FN(bf16_gemm_tt)); + m.impl("m_grouped_bf16_gemm_nt_contiguous", TORCH_FN(m_grouped_bf16_gemm_nt_contiguous)); + m.impl("m_grouped_bf16_gemm_nn_contiguous", TORCH_FN(deep_gemm::gemm::m_grouped_bf16_gemm_nn_contiguous)); + m.impl("m_grouped_bf16_gemm_nt_masked", TORCH_FN(m_grouped_bf16_gemm_nt_masked)); + m.impl("k_grouped_bf16_gemm_tn_contiguous", TORCH_FN(k_grouped_bf16_gemm_tn_contiguous)); + + m.impl("cublaslt_gemm_nt", TORCH_FN(deep_gemm::gemm::cublaslt_gemm_nt)); + m.impl("cublaslt_gemm_nn", TORCH_FN(deep_gemm::gemm::cublaslt_gemm_nn)); + m.impl("cublaslt_gemm_tn", TORCH_FN(deep_gemm::gemm::cublaslt_gemm_tn)); + m.impl("cublaslt_gemm_tt", TORCH_FN(deep_gemm::gemm::cublaslt_gemm_tt)); + + m.impl("cublaslt_nvfp4_gemm_nt", TORCH_FN(deep_gemm::torch_registration::cublaslt_nvfp4_gemm_nt)); + m.impl("batched_syrk", TORCH_FN(deep_gemm::gemm::batched_syrk)); + m.impl("batched_symm", TORCH_FN(deep_gemm::gemm::batched_symm)); +} diff --git a/csrc/apis/hyperconnection.hpp b/csrc/apis/hyperconnection.hpp index 6a99c43ce6..acd4a2c940 100644 --- a/csrc/apis/hyperconnection.hpp +++ b/csrc/apis/hyperconnection.hpp @@ -1,5 +1,7 @@ #pragma once +#include + #include "../utils/compatibility.hpp" #include "../jit_kernels/impls/sm90_tf32_hc_prenorm_gemm.hpp" @@ -58,10 +60,27 @@ static void tf32_hc_prenorm_gemm(const torch::Tensor& a, } -static void register_apis(pybind11::module_& m) { - m.def("tf32_hc_prenorm_gemm", &tf32_hc_prenorm_gemm, - py::arg("a"), py::arg("b"), py::arg("d"), py::arg("sqr_sum"), - py::arg("num_splits") = std::nullopt); +} // namespace deep_gemm::hyperconnection + +namespace deep_gemm::torch_registration { +static void tf32_hc_prenorm_gemm(const torch::Tensor& a, const torch::Tensor& b, + const torch::Tensor& d, const torch::Tensor& sqr_sum, + const c10::optional& num_splits) { + hyperconnection::tf32_hc_prenorm_gemm( + a, b, d, sqr_sum, + num_splits.has_value() + ? std::make_optional(static_cast(num_splits.value())) + : std::nullopt); +} +} // namespace deep_gemm::torch_registration + +TORCH_LIBRARY_FRAGMENT(deep_gemm, m) { + m.def( + "tf32_hc_prenorm_gemm(Tensor a, Tensor b, Tensor(d!) d, Tensor(sqr_sum!) sqr_sum, int? num_splits=None) -> ()"); } -} // namespace deep_gemm::hyperconnection +TORCH_LIBRARY_IMPL(deep_gemm, CUDA, m) { + using namespace deep_gemm::torch_registration; + + m.impl("tf32_hc_prenorm_gemm", TORCH_FN(tf32_hc_prenorm_gemm)); +} diff --git a/csrc/apis/layout.hpp b/csrc/apis/layout.hpp index e37496777a..0ccb52905f 100644 --- a/csrc/apis/layout.hpp +++ b/csrc/apis/layout.hpp @@ -1,5 +1,8 @@ #pragma once +#include +#include "../torch_library_utils.hpp" + #include "../jit_kernels/heuristics/runtime.hpp" #include "../utils/layout.hpp" #include "../utils/compatibility.hpp" @@ -132,30 +135,85 @@ 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") = std::nullopt, - py::arg("disable_ue8m0_cast") = false, - py::arg("psum_layout") = std::nullopt); - - m.def("get_tma_aligned_size", &get_tma_aligned_size); - 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, - py::arg("sf"), py::arg("psum_layout") = std::nullopt); - m.def("get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor", &get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor, - py::arg("sf"), py::arg("grouped_layout"), py::arg("ks_cpu"), py::arg("gran_k"), py::arg("k_alignment"), - py::arg("use_psum_layout") = false); - m.def("set_mk_alignment_for_contiguous_layout", [&](const int& new_value) { - heuristics_runtime->set_mk_alignment_for_contiguous_layout(new_value); - }); - m.def("get_mk_alignment_for_contiguous_layout", [&]() { - return heuristics_runtime->get_mk_alignment_for_contiguous_layout(); - }); - m.def("get_theoretical_mk_alignment_for_contiguous_layout", [&](const std::optional& expected_m) { - return heuristics_runtime->get_theoretical_mk_alignment_for_contiguous_layout(expected_m); - }, py::arg("expected_m") = std::nullopt); +} // namespace deep_gemm::layout + +namespace deep_gemm::torch_registration { +using namespace deep_gemm::torch_utils; + +static torch::Tensor transform_sf_into_required_layout( + const torch::Tensor& sf, const int64_t& mn, const int64_t& k, + const std::vector& recipe, + const c10::optional& num_groups, + const c10::optional& is_sfa, + const bool& disable_ue8m0_cast, + const c10::optional& psum_layout) { + return layout::transform_sf_into_required_layout( + sf, static_cast(mn), static_cast(k), + list_to_recipe_variant(recipe), + num_groups.has_value() ? std::make_optional(static_cast(num_groups.value())) : std::nullopt, + is_sfa, + disable_ue8m0_cast, + psum_layout); +} + +static torch::Tensor get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor( + const torch::Tensor& sf, const torch::Tensor& grouped_layout, + const c10::optional>& ks_cpu, + const int64_t& gran_k, const int64_t& k_alignment, + const bool& use_psum_layout) { + return ::deep_gemm::get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor( + sf, grouped_layout, + list_to_optional_vector_int(ks_cpu), + static_cast(gran_k), static_cast(k_alignment), + use_psum_layout); +} + +static int64_t get_tma_aligned_size(const int64_t& x, const int64_t& element_size) { + return ::deep_gemm::get_tma_aligned_size(static_cast(x), static_cast(element_size)); +} + +static void set_mk_alignment_for_contiguous_layout(const int64_t& new_value) { + heuristics_runtime->set_mk_alignment_for_contiguous_layout(static_cast(new_value)); } -} // namespace deep_gemm::layout +static int64_t get_mk_alignment_for_contiguous_layout() { + return heuristics_runtime->get_mk_alignment_for_contiguous_layout(); +} + +static int64_t get_theoretical_mk_alignment_for_contiguous_layout( + const c10::optional& expected_m) { + return HeuristicsRuntime::get_theoretical_mk_alignment_for_contiguous_layout( + expected_m.has_value() ? std::make_optional(static_cast(expected_m.value())) : std::nullopt); +} +} // namespace deep_gemm::torch_registration + +TORCH_LIBRARY_FRAGMENT(deep_gemm, m) { + m.def( + "transform_sf_into_required_layout(Tensor(a) sf, int mn, int k, int[] recipe, int? num_groups=None, bool? is_sfa=None, bool disable_ue8m0_cast=False, Tensor? psum_layout=None) -> Tensor(a)"); + m.def("get_tma_aligned_size(int x, int element_size) -> int", TORCH_FN(deep_gemm::torch_registration::get_tma_aligned_size)); + m.def("get_mn_major_tma_aligned_tensor(Tensor(a) sf) -> Tensor(a)"); + m.def( + "get_mn_major_tma_aligned_packed_ue8m0_tensor(Tensor sf, Tensor? psum_layout=None) -> Tensor"); + m.def( + "get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor(Tensor sf, Tensor grouped_layout, int[]? ks_cpu, int gran_k, int k_alignment, bool use_psum_layout=False) -> Tensor"); + + m.def("set_mk_alignment_for_contiguous_layout(int new_value) -> ()", + TORCH_FN(deep_gemm::torch_registration::set_mk_alignment_for_contiguous_layout)); + m.def("get_mk_alignment_for_contiguous_layout() -> int", + TORCH_FN(deep_gemm::torch_registration::get_mk_alignment_for_contiguous_layout)); + m.def("get_theoretical_mk_alignment_for_contiguous_layout(int? expected_m=None) -> int", + TORCH_FN(deep_gemm::torch_registration::get_theoretical_mk_alignment_for_contiguous_layout)); +} + +TORCH_LIBRARY_IMPL(deep_gemm, CUDA, m) { + using namespace deep_gemm::torch_registration; + + m.impl("transform_sf_into_required_layout", + TORCH_FN(transform_sf_into_required_layout)); + m.impl("get_mn_major_tma_aligned_tensor", + TORCH_FN(deep_gemm::get_mn_major_tma_aligned_tensor)); + m.impl("get_mn_major_tma_aligned_packed_ue8m0_tensor", + TORCH_FN(deep_gemm::get_mn_major_tma_aligned_packed_ue8m0_tensor)); + m.impl("get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor", + TORCH_FN(get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor)); +} diff --git a/csrc/apis/mega_gate.hpp b/csrc/apis/mega_gate.hpp index 617a43fb81..29cc79128c 100644 --- a/csrc/apis/mega_gate.hpp +++ b/csrc/apis/mega_gate.hpp @@ -1,5 +1,8 @@ #pragma once +#include +#include "../torch_library_utils.hpp" + #include #include #include @@ -7,7 +10,7 @@ #include #include -#include +#include #include "../runtime/runtime.hpp" #include "../jit_kernels/impls/sm100_bf16_mega_gate.hpp" @@ -191,7 +194,7 @@ bf16_mega_gate(const torch::Tensor& x, return {topk_idx, topk_weights}; } -static pybind11::dict get_bf16_mega_gate_config(const int& num_tokens, const int& hidden, +static c10::Dict get_bf16_mega_gate_config(const int& num_tokens, const int& hidden, const int& num_routed_experts, const int& num_topk) { DG_HOST_ASSERT(num_tokens > 0 and num_tokens <= static_cast(layout::mega_gate::kNumMaxTokens)); DG_HOST_ASSERT(hidden > 0 and hidden % 256 == 0); @@ -200,35 +203,64 @@ static pybind11::dict get_bf16_mega_gate_config(const int& num_tokens, const int const auto num_device_sms = runtime->get_num_sms(); const auto config = get_sm100_bf16_mega_gate_config(num_tokens, hidden, num_routed_experts, num_device_sms, true, true, true); - pybind11::dict result; - result["block_tokens"] = config.block_tokens; - result["num_mma_ctas"] = config.num_mma_ctas; - result["num_split_k"] = config.num_split_k; - result["num_expert_groups"] = config.num_expert_groups; - result["num_gate_warpgroups"] = config.num_gate_warpgroups; - result["num_sms"] = config.num_launch_sms; + c10::Dict result; + result.insert("block_tokens", config.block_tokens); + result.insert("num_mma_ctas", config.num_mma_ctas); + result.insert("num_split_k", config.num_split_k); + result.insert("num_expert_groups", config.num_expert_groups); + result.insert("num_gate_warpgroups", config.num_gate_warpgroups); + result.insert("num_sms", config.num_launch_sms); return result; } -static void register_apis(pybind11::module_& m) { - m.def("get_bf16_mega_gate_config", &get_bf16_mega_gate_config, - py::arg("num_tokens"), py::arg("hidden"), - py::arg("num_routed_experts"), py::arg("num_topk")); - m.def("bf16_mega_gate", &bf16_mega_gate, - py::arg("x"), py::arg("weight"), py::arg("num_topk"), - py::arg("use_shared_as_routed"), py::arg("num_shared_experts"), - py::arg("routed_scaling_factor"), py::arg("ep_rank"), - py::arg("scoring_func") = "identity", - py::arg("mask") = std::nullopt, - py::arg("bias") = std::nullopt, - py::arg("image_bias") = std::nullopt, - py::arg("image_token_mask") = std::nullopt, - py::arg("fix_routing_mask") = std::nullopt, - py::arg("to_physical_map") = std::nullopt, - py::arg("logical_count") = std::nullopt, - py::arg("unmapped_topk_idx") = std::nullopt, - py::arg("force_random") = std::nullopt, - py::arg("out") = std::nullopt); +} // namespace deep_gemm::mega_gate + +namespace deep_gemm::torch_registration { +using namespace deep_gemm::torch_utils; + +static c10::Dict get_bf16_mega_gate_config( + const int64_t& num_tokens, + const int64_t& hidden, + const int64_t& num_routed_experts, + const int64_t& num_topk) { + return mega_gate::get_bf16_mega_gate_config( + num_tokens, hidden, num_routed_experts, num_topk); } -} // namespace deep_gemm::mega_gate +static void bf16_mega_gate(const torch::Tensor& x, + const torch::Tensor& weight, + const int64_t& num_topk, + const bool& use_shared_as_routed, + const int64_t& num_shared_experts, + const double& routed_scaling_factor, + const int64_t& ep_rank, + const std::string& scoring_func, + const std::optional& mask, + const std::optional& bias, + const std::optional& image_bias, + const std::optional& image_token_mask, + const std::optional& fix_routing_mask, + const std::optional& to_physical_map, + const std::optional& logical_count, + const std::optional& unmapped_topk_idx, + const std::optional& force_random, + const torch::Tensor& topk_idx, const torch::Tensor& topk_weights) { + mega_gate::bf16_mega_gate(x, weight, num_topk, use_shared_as_routed, num_shared_experts, routed_scaling_factor, + ep_rank, scoring_func, mask, bias, image_bias, image_token_mask, fix_routing_mask, to_physical_map, + logical_count, unmapped_topk_idx, force_random, std::make_tuple(topk_idx, topk_weights)); +} +} // namespace deep_gemm::torch_registration + +TORCH_LIBRARY_FRAGMENT(deep_gemm, m) { + m.def("get_bf16_mega_gate_config(int num_tokens, int hidden, int num_routed_experts, int num_topk) -> Dict(str, int)"); + + m.def("bf16_mega_gate(Tensor x, Tensor weight, int num_topk, bool use_shared_as_routed, int num_shared_experts, float routed_scaling_factor, int ep_rank, str scoring_func, Tensor? mask, Tensor? bias, Tensor? image_bias, Tensor? image_token_mask, Tensor? fix_routing_mask, Tensor? to_physical_map, Tensor? logical_count, Tensor(c!)? unmapped_topk_idx, Tensor? force_random, Tensor(a!) topk_idx, Tensor(b!) topk_weights) -> ()"); +} + +TORCH_LIBRARY_IMPL(deep_gemm, CatchAll, m) { + m.impl("get_bf16_mega_gate_config", TORCH_FN(deep_gemm::torch_registration::get_bf16_mega_gate_config)); +} + +TORCH_LIBRARY_IMPL(deep_gemm, CUDA, m) { + m.impl("bf16_mega_gate", TORCH_FN(deep_gemm::torch_registration::bf16_mega_gate)); +} diff --git a/csrc/apis/mega_mhc.hpp b/csrc/apis/mega_mhc.hpp index 1cbc0c8230..3a962b58cd 100644 --- a/csrc/apis/mega_mhc.hpp +++ b/csrc/apis/mega_mhc.hpp @@ -1,5 +1,8 @@ #pragma once +#include +#include "../torch_library_utils.hpp" + #include #include @@ -251,37 +254,48 @@ static void mega_mhc(const torch::Tensor& x, shared_sf_block_m, rmsnorm_scale, num_tokens, hidden, num_splits, num_sms); } -static void register_apis(pybind11::module_& m) { - m.def( - "mega_mhc", - &mega_mhc, - py::arg("x"), - py::arg("residual"), - py::arg("shifted_prev_mix"), - py::arg("post_mix"), - py::arg("comb_res_mix"), - py::arg("fn"), - py::arg("mix_scales"), - py::arg("mix_bases"), - py::arg("hc_mult"), - py::arg("hc_norm_eps"), - py::arg("hc_pre_eps"), - py::arg("hc_post_scale"), - py::arg("sinkhorn_eps"), - py::arg("num_sinkhorn_iters"), - py::arg("rmsnorm_weight"), - py::arg("rmsnorm_eps"), - py::arg("rmsnorm_scale"), - py::arg("new_residual"), - py::arg("new_prev_mix"), - py::arg("new_post_mix"), - py::arg("new_comb_res_mix"), - py::arg("y_bf16") = std::nullopt, - py::arg("y_fp8") = std::nullopt, - py::arg("y_gemm_sf") = std::nullopt, - py::arg("y_routed_sf") = std::nullopt, - py::arg("y_shared_sf") = std::nullopt, - py::arg("shared_sf_block_m") = 0); +} // namespace deep_gemm::mega_mhc + +namespace deep_gemm::torch_registration { +using namespace deep_gemm::torch_utils; + +static void mega_mhc( + const torch::Tensor& x, + const torch::Tensor& residual, + const std::optional& shifted_prev_mix, + const torch::Tensor& post_mix, + const torch::Tensor& comb_res_mix, + const torch::Tensor& fn, + const torch::Tensor& mix_scales, + const torch::Tensor& mix_bases, + const int64_t& hc_mult, + const double& hc_norm_eps, + const double& hc_pre_eps, + const double& hc_post_scale, + const double& sinkhorn_eps, + const int64_t& num_sinkhorn_iters, + const torch::Tensor& rmsnorm_weight, + const double& rmsnorm_eps, + const double& rmsnorm_scale, + const torch::Tensor& new_residual, + const std::optional& new_prev_mix, + const torch::Tensor& new_post_mix, + const torch::Tensor& new_comb_res_mix, + const std::optional& y_bf16, + const std::optional& y_fp8, + const std::optional& y_gemm_sf, + const std::optional& y_routed_sf, + const std::optional& y_shared_sf, + const int64_t& shared_sf_block_m) { + mega_mhc::mega_mhc( + x, residual, shifted_prev_mix, post_mix, comb_res_mix, fn, mix_scales, mix_bases, hc_mult, hc_norm_eps, hc_pre_eps, hc_post_scale, sinkhorn_eps, num_sinkhorn_iters, rmsnorm_weight, rmsnorm_eps, rmsnorm_scale, new_residual, new_prev_mix, new_post_mix, new_comb_res_mix, y_bf16, y_fp8, y_gemm_sf, y_routed_sf, y_shared_sf, shared_sf_block_m); +} +} // namespace deep_gemm::torch_registration + +TORCH_LIBRARY_FRAGMENT(deep_gemm, m) { + m.def("mega_mhc(Tensor x, Tensor residual, Tensor? shifted_prev_mix, Tensor post_mix, Tensor comb_res_mix, Tensor fn, Tensor mix_scales, Tensor mix_bases, int hc_mult, float hc_norm_eps, float hc_pre_eps, float hc_post_scale, float sinkhorn_eps, int num_sinkhorn_iters, Tensor rmsnorm_weight, float rmsnorm_eps, float rmsnorm_scale, Tensor(a!) new_residual, Tensor(b!)? new_prev_mix, Tensor(c!) new_post_mix, Tensor(d!) new_comb_res_mix, Tensor(e!)? y_bf16=None, Tensor(f!)? y_fp8=None, Tensor(g!)? y_gemm_sf=None, Tensor(h!)? y_routed_sf=None, Tensor(i!)? y_shared_sf=None, int shared_sf_block_m=0) -> ()"); } -} // namespace deep_gemm::mega_mhc +TORCH_LIBRARY_IMPL(deep_gemm, CUDA, m) { + m.impl("mega_mhc", TORCH_FN(deep_gemm::torch_registration::mega_mhc)); +} diff --git a/csrc/apis/mega_moe.hpp b/csrc/apis/mega_moe.hpp index f1dfe95495..c99682417e 100644 --- a/csrc/apis/mega_moe.hpp +++ b/csrc/apis/mega_moe.hpp @@ -1,9 +1,12 @@ #pragma once +#include +#include "../torch_library_utils.hpp" + #include -#include +#include +#include #include -#include #include #include @@ -15,6 +18,13 @@ namespace deep_gemm::mega { +static int checked_int(const int64_t value) { + // Preserve pybind behavior by rejecting values that do not fit exactly in a C++ int. + DG_HOST_ASSERT(value >= std::numeric_limits::min()); + DG_HOST_ASSERT(value <= std::numeric_limits::max()); + return static_cast(value); +} + static int get_token_alignment_for_mega_moe() { return layout::kLCMCandidateBlockM; } @@ -30,10 +40,72 @@ static int get_block_m_for_mega_moe( return block_m; } -static std::tuple(const torch::Tensor&)>> -get_symm_buffer_size_for_mega_moe( +struct SymmBufferLayoutInfo { + int64_t num_bytes = 0; + int64_t input_token_base = 0; + int64_t input_sf_base = 0; + int64_t input_topk_idx_base = 0; + int64_t input_topk_weights_base = 0; + int64_t shared_l1_sf_base = 0; + int64_t shared_l2_token_base = 0; + int64_t shared_l2_sf_base = 0; + int64_t l1_token_base = 0; + int64_t l1_sf_base = 0; + int64_t l2_token_base = 0; + int64_t l2_sf_base = 0; + bool with_sf = false; + int num_max_tokens_per_rank = 0; + int num_topk = 0; + int hidden = 0; + int intermediate_hidden = 0; + int num_shared_experts = 0; + int shared_intermediate_hidden = 0; + int num_ring_tokens = 0; + int num_sf_ring_tokens = 0; + + // Flatten into a plain `int[]` so it can cross the TORCH_LIBRARY boundary + std::vector to_int_list() const { + return { + num_bytes, input_token_base, input_sf_base, input_topk_idx_base, input_topk_weights_base, + shared_l1_sf_base, shared_l2_token_base, shared_l2_sf_base, + l1_token_base, l1_sf_base, l2_token_base, l2_sf_base, + static_cast(with_sf), + num_max_tokens_per_rank, num_topk, + hidden, intermediate_hidden, num_shared_experts, shared_intermediate_hidden, + num_ring_tokens, num_sf_ring_tokens, + }; + } + + static SymmBufferLayoutInfo from_int_list(const std::vector& values) { + DG_HOST_ASSERT(static_cast(values.size()) == 21); + SymmBufferLayoutInfo info; + info.num_bytes = values[0]; + info.input_token_base = values[1]; + info.input_sf_base = values[2]; + info.input_topk_idx_base = values[3]; + info.input_topk_weights_base = values[4]; + info.shared_l1_sf_base = values[5]; + info.shared_l2_token_base = values[6]; + info.shared_l2_sf_base = values[7]; + info.l1_token_base = values[8]; + info.l1_sf_base = values[9]; + info.l2_token_base = values[10]; + info.l2_sf_base = values[11]; + // `with_sf` is a bool, encoded as 0/1 since the list is all `int64_t`. + info.with_sf = values[12] != 0; + info.num_max_tokens_per_rank = checked_int(values[13]); + info.num_topk = checked_int(values[14]); + info.hidden = checked_int(values[15]); + info.intermediate_hidden = checked_int(values[16]); + info.num_shared_experts = checked_int(values[17]); + info.shared_intermediate_hidden = checked_int(values[18]); + info.num_ring_tokens = checked_int(values[19]); + info.num_sf_ring_tokens = checked_int(values[20]); + return info; + } +}; + +static SymmBufferLayoutInfo build_symm_buffer_layout( const int& num_ranks, const int& num_experts, const int& num_max_tokens_per_rank, const int& num_topk, const int& hidden, const int& intermediate_hidden, @@ -92,65 +164,103 @@ get_symm_buffer_size_for_mega_moe( DG_HOST_ASSERT(num_sf_ring_tokens % 4 == 0); } - // Slice function: creates tensor views from the raw buffer. - // NOTES: `x_sf` is K-major, while `l1_acts_sf` and `l2_acts_sf` are M-major - auto slice_input_buffers = [=](const torch::Tensor& buffer) { + SymmBufferLayoutInfo layout_info; + layout_info.num_bytes = mega_buffer.get_num_bytes(); + layout_info.input_token_base = reinterpret_cast(mega_buffer.input_token_buffer.base); + layout_info.input_sf_base = reinterpret_cast(mega_buffer.input_sf_buffer.base); + layout_info.input_topk_idx_base = reinterpret_cast(mega_buffer.input_topk_idx_buffer.base); + layout_info.input_topk_weights_base = reinterpret_cast(mega_buffer.input_topk_weights_buffer.base); + layout_info.shared_l1_sf_base = reinterpret_cast(mega_buffer.shared_l1_sf_buffer.base); + layout_info.shared_l2_token_base = reinterpret_cast(mega_buffer.shared_l2_token_buffer.base); + layout_info.shared_l2_sf_base = reinterpret_cast(mega_buffer.shared_l2_sf_buffer.base); + layout_info.l1_token_base = reinterpret_cast(mega_buffer.l1_token_buffer.base); + layout_info.l1_sf_base = reinterpret_cast(mega_buffer.l1_sf_buffer.base); + layout_info.l2_token_base = reinterpret_cast(mega_buffer.l2_token_buffer.base); + layout_info.l2_sf_base = reinterpret_cast(mega_buffer.l2_sf_buffer.base); + layout_info.with_sf = with_sf; + layout_info.num_max_tokens_per_rank = num_max_tokens_per_rank; + layout_info.num_topk = num_topk; + layout_info.hidden = hidden; + layout_info.intermediate_hidden = intermediate_hidden; + layout_info.num_shared_experts = num_shared_experts; + layout_info.shared_intermediate_hidden = shared_intermediate_hidden; + layout_info.num_ring_tokens = num_ring_tokens; + layout_info.num_sf_ring_tokens = num_sf_ring_tokens; + return layout_info; +} + +static std::tuple> get_symm_buffer_size_for_mega_moe( + const int& num_ranks, const int& num_experts, + const int& num_max_tokens_per_rank, const int& num_topk, + const int& hidden, const int& intermediate_hidden, + const std::string& mma_type, const std::string& activation, + const int& num_shared_experts = 0) { + const auto layout_info = build_symm_buffer_layout( + num_ranks, num_experts, num_max_tokens_per_rank, num_topk, + hidden, intermediate_hidden, mma_type, activation, num_shared_experts); + return std::make_tuple(layout_info.num_bytes, layout_info.to_int_list()); +} + +using SymmBufferSlice = std::tuple; + +static SymmBufferSlice slice_symm_buffer_from_layout( + const torch::Tensor& buffer, const SymmBufferLayoutInfo& layout_info) { auto x = torch::from_blob( - math::advance_ptr(buffer.data_ptr(), reinterpret_cast(mega_buffer.input_token_buffer.base)), - {num_max_tokens_per_rank, hidden}, - torch::TensorOptions().dtype(with_sf ? torch::kFloat8_e4m3fn : torch::kBFloat16).device(buffer.device())); - auto x_sf = with_sf ? torch::from_blob( - math::advance_ptr(buffer.data_ptr(), reinterpret_cast(mega_buffer.input_sf_buffer.base)), - {num_max_tokens_per_rank, hidden / 128}, + math::advance_ptr(buffer.data_ptr(), layout_info.input_token_base), + {layout_info.num_max_tokens_per_rank, layout_info.hidden}, + torch::TensorOptions().dtype(layout_info.with_sf ? torch::kFloat8_e4m3fn : torch::kBFloat16).device(buffer.device())); + auto x_sf = layout_info.with_sf ? torch::from_blob( + math::advance_ptr(buffer.data_ptr(), layout_info.input_sf_base), + {layout_info.num_max_tokens_per_rank, layout_info.hidden / 128}, torch::TensorOptions().dtype(torch::kInt).device(buffer.device())) : torch::Tensor(); auto topk_idx = torch::from_blob( - math::advance_ptr(buffer.data_ptr(), reinterpret_cast(mega_buffer.input_topk_idx_buffer.base)), - {num_max_tokens_per_rank, num_topk}, + math::advance_ptr(buffer.data_ptr(), layout_info.input_topk_idx_base), + {layout_info.num_max_tokens_per_rank, layout_info.num_topk}, torch::TensorOptions().dtype(torch::kInt64).device(buffer.device())); auto topk_weights = torch::from_blob( - math::advance_ptr(buffer.data_ptr(), reinterpret_cast(mega_buffer.input_topk_weights_buffer.base)), - {num_max_tokens_per_rank, num_topk}, + math::advance_ptr(buffer.data_ptr(), layout_info.input_topk_weights_base), + {layout_info.num_max_tokens_per_rank, layout_info.num_topk}, torch::TensorOptions().dtype(torch::kFloat32).device(buffer.device())); auto shared_l1_acts = x; - auto shared_l1_acts_sf = (with_sf and num_shared_experts > 0) ? torch::from_blob( - math::advance_ptr(buffer.data_ptr(), reinterpret_cast(mega_buffer.shared_l1_sf_buffer.base)), - {layout::get_num_max_shared_sf_tokens(num_max_tokens_per_rank), hidden / 128}, - {1, layout::get_num_max_shared_sf_tokens(num_max_tokens_per_rank)}, + auto shared_l1_acts_sf = (layout_info.with_sf and layout_info.num_shared_experts > 0) ? torch::from_blob( + math::advance_ptr(buffer.data_ptr(), layout_info.shared_l1_sf_base), + {layout::get_num_max_shared_sf_tokens(layout_info.num_max_tokens_per_rank), layout_info.hidden / 128}, + {1, layout::get_num_max_shared_sf_tokens(layout_info.num_max_tokens_per_rank)}, torch::TensorOptions().dtype(torch::kInt).device(buffer.device())) : torch::Tensor(); - auto shared_l2_acts = num_shared_experts > 0 ? torch::from_blob( - math::advance_ptr(buffer.data_ptr(), reinterpret_cast(mega_buffer.shared_l2_token_buffer.base)), - {num_max_tokens_per_rank, shared_intermediate_hidden}, - torch::TensorOptions().dtype(with_sf ? torch::kFloat8_e4m3fn : torch::kBFloat16).device(buffer.device())) : torch::Tensor(); - auto shared_l2_acts_sf = (with_sf and num_shared_experts > 0) ? torch::from_blob( - math::advance_ptr(buffer.data_ptr(), reinterpret_cast(mega_buffer.shared_l2_sf_buffer.base)), - {layout::get_num_max_shared_sf_tokens(num_max_tokens_per_rank), shared_intermediate_hidden / 128}, - {1, layout::get_num_max_shared_sf_tokens(num_max_tokens_per_rank)}, + auto shared_l2_acts = layout_info.num_shared_experts > 0 ? torch::from_blob( + math::advance_ptr(buffer.data_ptr(), layout_info.shared_l2_token_base), + {layout_info.num_max_tokens_per_rank, layout_info.shared_intermediate_hidden}, + torch::TensorOptions().dtype(layout_info.with_sf ? torch::kFloat8_e4m3fn : torch::kBFloat16).device(buffer.device())) : torch::Tensor(); + auto shared_l2_acts_sf = (layout_info.with_sf and layout_info.num_shared_experts > 0) ? torch::from_blob( + math::advance_ptr(buffer.data_ptr(), layout_info.shared_l2_sf_base), + {layout::get_num_max_shared_sf_tokens(layout_info.num_max_tokens_per_rank), layout_info.shared_intermediate_hidden / 128}, + {1, layout::get_num_max_shared_sf_tokens(layout_info.num_max_tokens_per_rank)}, torch::TensorOptions().dtype(torch::kInt).device(buffer.device())) : torch::Tensor(); auto l1_acts = torch::from_blob( - math::advance_ptr(buffer.data_ptr(), reinterpret_cast(mega_buffer.l1_token_buffer.base)), - {num_ring_tokens, hidden}, - torch::TensorOptions().dtype(with_sf ? torch::kFloat8_e4m3fn : torch::kBFloat16).device(buffer.device())); - auto l1_acts_sf = with_sf ? torch::from_blob( - math::advance_ptr(buffer.data_ptr(), reinterpret_cast(mega_buffer.l1_sf_buffer.base)), - {num_sf_ring_tokens, hidden / 128}, - {1, num_sf_ring_tokens}, + math::advance_ptr(buffer.data_ptr(), layout_info.l1_token_base), + {layout_info.num_ring_tokens, layout_info.hidden}, + torch::TensorOptions().dtype(layout_info.with_sf ? torch::kFloat8_e4m3fn : torch::kBFloat16).device(buffer.device())); + auto l1_acts_sf = layout_info.with_sf ? torch::from_blob( + math::advance_ptr(buffer.data_ptr(), layout_info.l1_sf_base), + {layout_info.num_sf_ring_tokens, layout_info.hidden / 128}, + {1, layout_info.num_sf_ring_tokens}, torch::TensorOptions().dtype(torch::kInt).device(buffer.device())) : torch::Tensor(); auto l2_acts = torch::from_blob( - math::advance_ptr(buffer.data_ptr(), reinterpret_cast(mega_buffer.l2_token_buffer.base)), - {num_ring_tokens, intermediate_hidden}, - torch::TensorOptions().dtype(with_sf ? torch::kFloat8_e4m3fn : torch::kBFloat16).device(buffer.device())); - auto l2_acts_sf = with_sf ? torch::from_blob( - math::advance_ptr(buffer.data_ptr(), reinterpret_cast(mega_buffer.l2_sf_buffer.base)), - {num_sf_ring_tokens, intermediate_hidden / 128}, - {1, num_sf_ring_tokens}, + math::advance_ptr(buffer.data_ptr(), layout_info.l2_token_base), + {layout_info.num_ring_tokens, layout_info.intermediate_hidden}, + torch::TensorOptions().dtype(layout_info.with_sf ? torch::kFloat8_e4m3fn : torch::kBFloat16).device(buffer.device())); + auto l2_acts_sf = layout_info.with_sf ? torch::from_blob( + math::advance_ptr(buffer.data_ptr(), layout_info.l2_sf_base), + {layout_info.num_sf_ring_tokens, layout_info.intermediate_hidden / 128}, + {1, layout_info.num_sf_ring_tokens}, torch::TensorOptions().dtype(torch::kInt).device(buffer.device())) : torch::Tensor(); return std::make_tuple(x, x_sf, topk_idx, topk_weights, shared_l1_acts, shared_l1_acts_sf, shared_l2_acts, shared_l2_acts_sf, l1_acts, l1_acts_sf, l2_acts, l2_acts_sf); - }; - return {mega_buffer.get_num_bytes(), slice_input_buffers}; } static void fp8_fp4_mega_moe( @@ -251,20 +361,21 @@ static void fp8_fp4_mega_moe( // Check buffer bytes const auto num_ranks = static_cast(sym_buffer_ptrs.size()); const auto num_experts_ = num_experts_per_rank * num_ranks; - const auto [num_required_bytes, slice] = get_symm_buffer_size_for_mega_moe( + const auto layout_info = build_symm_buffer_layout( num_ranks, num_experts, num_max_tokens_per_rank, num_topk, hidden, intermediate_hidden, weight_dtype == torch::kFloat8_e4m3fn ? "fp8xfp8" : "fp8xfp4", activation, num_shared_experts ); - DG_HOST_ASSERT(sym_buffer.nbytes() >= static_cast(num_required_bytes)); + DG_HOST_ASSERT(sym_buffer.nbytes() >= static_cast(layout_info.num_bytes)); DG_HOST_ASSERT(num_experts == num_experts_); // Already registered tensors const auto [x, x_sf, topk_idx, topk_weights, shared_l1_acts, shared_l1_acts_sf, shared_l2_acts, shared_l2_acts_sf, - l1_acts, l1_acts_sf, l2_acts, l2_acts_sf] = slice(sym_buffer); + l1_acts, l1_acts_sf, l2_acts, l2_acts_sf] = + slice_symm_buffer_from_layout(sym_buffer, layout_info); // Dispatch into different architectures if (arch_major == 10) { @@ -377,19 +488,20 @@ static void bf16_mega_moe( // Check buffer bytes const auto num_ranks = static_cast(sym_buffer_ptrs.size()); const auto num_experts_ = num_experts_per_rank * num_ranks; - const auto [num_required_bytes, slice] = get_symm_buffer_size_for_mega_moe( + const auto layout_info = build_symm_buffer_layout( num_ranks, num_experts, num_max_tokens_per_rank, num_topk, hidden, intermediate_hidden, "bf16xbf16", activation, num_shared_experts ); - DG_HOST_ASSERT(sym_buffer.nbytes() >= static_cast(num_required_bytes)); + DG_HOST_ASSERT(sym_buffer.nbytes() >= static_cast(layout_info.num_bytes)); DG_HOST_ASSERT(num_experts == num_experts_); // Already registered tensors const auto [x, _x_sf, topk_idx, topk_weights, shared_l1_acts, _shared_l1_acts_sf, shared_l2_acts, _shared_l2_acts_sf, - l1_acts, _l1_acts_sf, l2_acts, _l2_acts_sf] = slice(sym_buffer); + l1_acts, _l1_acts_sf, l2_acts, _l2_acts_sf] = + slice_symm_buffer_from_layout(sym_buffer, layout_info); // Dispatch into different architectures if (arch_major == 10) { @@ -417,12 +529,141 @@ static void bf16_mega_moe( sym_buffer.zero_(); } -static void register_apis(pybind11::module_& m) { - m.def("get_token_alignment_for_mega_moe", &get_token_alignment_for_mega_moe); - m.def("get_block_m_for_mega_moe", &get_block_m_for_mega_moe); - m.def("get_symm_buffer_size_for_mega_moe", &get_symm_buffer_size_for_mega_moe); - m.def("fp8_fp4_mega_moe", &fp8_fp4_mega_moe); - m.def("bf16_mega_moe", &bf16_mega_moe); +} // namespace deep_gemm::mega + +namespace deep_gemm::torch_registration { + +static std::tuple checked_recipe(const std::vector& recipe) { + DG_HOST_ASSERT(recipe.size() == 3); + return {mega::checked_int(recipe[0]), mega::checked_int(recipe[1]), mega::checked_int(recipe[2])}; +} + +static std::optional> make_optional_tensor_pair( + const std::optional& value, + const std::optional& scale_factor) { + // Preserve pybind behavior: an optional tensor tuple must contain both tensors or neither. + DG_HOST_ASSERT(value.has_value() == scale_factor.has_value()); + if (not value.has_value()) + return std::nullopt; + return std::make_tuple(*value, *scale_factor); } -} // namespace deep_gemm::mega +static int64_t get_token_alignment_for_mega_moe() { + return mega::get_token_alignment_for_mega_moe(); +} + +static int64_t get_block_m_for_mega_moe( + const int64_t& num_ranks, + const int64_t& num_experts, + const int64_t& num_max_tokens_per_rank, + const int64_t& num_tokens, + const int64_t& num_topk, + const std::string& mma_type) { + return mega::get_block_m_for_mega_moe( + mega::checked_int(num_ranks), mega::checked_int(num_experts), + mega::checked_int(num_max_tokens_per_rank), mega::checked_int(num_tokens), + mega::checked_int(num_topk), mma_type); +} + +static void fp8_fp4_mega_moe( + const torch::Tensor& y, + const torch::Tensor& l1_weights_tuple, + const torch::Tensor& l1_weights_tuple_sf, + const torch::Tensor& l2_weights_tuple, + const torch::Tensor& l2_weights_tuple_sf, + const std::optional& shared_l1_weights_tuple_opt, + const std::optional& shared_l1_weights_tuple_opt_sf, + const std::optional& shared_l2_weights_tuple_opt, + const std::optional& shared_l2_weights_tuple_opt_sf, + const std::optional& cumulative_local_expert_recv_stats, + const torch::Tensor& sym_buffer, + const std::vector& sym_buffer_ptrs, + const int64_t& rank_idx, + const int64_t& num_max_tokens_per_rank, + const int64_t& num_experts, + const int64_t& num_topk, + const std::vector& recipe, + const std::string& activation, + const std::optional& activation_clamp_opt, + const bool& fast_math, + const double& activation_alpha, + const double& activation_beta) { + mega::fp8_fp4_mega_moe( + y, {l1_weights_tuple, l1_weights_tuple_sf}, {l2_weights_tuple, l2_weights_tuple_sf}, + make_optional_tensor_pair(shared_l1_weights_tuple_opt, shared_l1_weights_tuple_opt_sf), + make_optional_tensor_pair(shared_l2_weights_tuple_opt, shared_l2_weights_tuple_opt_sf), + cumulative_local_expert_recv_stats, sym_buffer, sym_buffer_ptrs, + mega::checked_int(rank_idx), mega::checked_int(num_max_tokens_per_rank), + mega::checked_int(num_experts), mega::checked_int(num_topk), checked_recipe(recipe), + activation, activation_clamp_opt, fast_math, + static_cast(activation_alpha), static_cast(activation_beta)); +} + +static void bf16_mega_moe( + const torch::Tensor& y, + const torch::Tensor& l1_weights, + const torch::Tensor& l2_weights, + const std::optional& shared_l1_weights_opt, + const std::optional& shared_l2_weights_opt, + const std::optional& cumulative_local_expert_recv_stats, + const torch::Tensor& sym_buffer, + const std::vector& sym_buffer_ptrs, + const int64_t& rank_idx, + const int64_t& num_max_tokens_per_rank, + const int64_t& num_experts, + const int64_t& num_topk, + const std::string& activation, + const std::optional& activation_clamp_opt, + const bool& fast_math, + const double& activation_alpha, + const double& activation_beta) { + mega::bf16_mega_moe( + y, l1_weights, l2_weights, shared_l1_weights_opt, shared_l2_weights_opt, + cumulative_local_expert_recv_stats, sym_buffer, sym_buffer_ptrs, + mega::checked_int(rank_idx), mega::checked_int(num_max_tokens_per_rank), + mega::checked_int(num_experts), mega::checked_int(num_topk), activation, + activation_clamp_opt, fast_math, + static_cast(activation_alpha), static_cast(activation_beta)); +} + +static std::tuple> get_symm_buffer_size_for_mega_moe( + int64_t num_ranks, int64_t num_experts, int64_t num_max_tokens_per_rank, + int64_t num_topk, int64_t hidden, int64_t intermediate_hidden, + const std::string& mma_type, const std::string& activation, int64_t num_shared_experts) { + return mega::get_symm_buffer_size_for_mega_moe( + mega::checked_int(num_ranks), mega::checked_int(num_experts), + mega::checked_int(num_max_tokens_per_rank), mega::checked_int(num_topk), + mega::checked_int(hidden), mega::checked_int(intermediate_hidden), + mma_type, activation, mega::checked_int(num_shared_experts)); +} +static mega::SymmBufferSlice _slice_symm_buffer_for_mega_moe( + const torch::Tensor& buffer, + const std::vector& layout_info) { + return mega::slice_symm_buffer_from_layout( + buffer, mega::SymmBufferLayoutInfo::from_int_list(layout_info)); +} +} // namespace deep_gemm::torch_registration + +TORCH_LIBRARY_FRAGMENT(deep_gemm, m) { + m.def("get_token_alignment_for_mega_moe() -> int"); + m.def("get_block_m_for_mega_moe(int num_ranks, int num_experts, int num_max_tokens_per_rank, int num_tokens, int num_topk, str mma_type) -> int"); + + m.def("get_symm_buffer_size_for_mega_moe(int num_ranks, int num_experts, int num_max_tokens_per_rank, int num_topk, int hidden, int intermediate_hidden, str mma_type, str activation, int num_shared_experts) -> (int, int[])", + TORCH_FN(deep_gemm::torch_registration::get_symm_buffer_size_for_mega_moe)); + m.def("_slice_symm_buffer_for_mega_moe(Tensor(a) buffer, int[] layout_info) -> (Tensor(a), Tensor(a), Tensor(a), Tensor(a), Tensor(a), Tensor(a), Tensor(a), Tensor(a), Tensor(a), Tensor(a), Tensor(a), Tensor(a))"); + + m.def("fp8_fp4_mega_moe(Tensor(a!) y, Tensor l1_weights_tuple, Tensor l1_weights_tuple_sf, Tensor l2_weights_tuple, Tensor l2_weights_tuple_sf, Tensor? shared_l1_weights_tuple_opt, Tensor? shared_l1_weights_tuple_opt_sf, Tensor? shared_l2_weights_tuple_opt, Tensor? shared_l2_weights_tuple_opt_sf, Tensor(b!)? cumulative_local_expert_recv_stats, Tensor(c!) sym_buffer, int[] sym_buffer_ptrs, int rank_idx, int num_max_tokens_per_rank, int num_experts, int num_topk, int[3] recipe, str activation, float? activation_clamp_opt, bool fast_math, float activation_alpha, float activation_beta) -> ()"); + m.def("bf16_mega_moe(Tensor(a!) y, Tensor l1_weights, Tensor l2_weights, Tensor? shared_l1_weights_opt, Tensor? shared_l2_weights_opt, Tensor(b!)? cumulative_local_expert_recv_stats, Tensor(c!) sym_buffer, int[] sym_buffer_ptrs, int rank_idx, int num_max_tokens_per_rank, int num_experts, int num_topk, str activation, float? activation_clamp_opt, bool fast_math, float activation_alpha, float activation_beta) -> ()"); +} + +TORCH_LIBRARY_IMPL(deep_gemm, CatchAll, m) { + m.impl("get_token_alignment_for_mega_moe", TORCH_FN(deep_gemm::torch_registration::get_token_alignment_for_mega_moe)); + m.impl("get_block_m_for_mega_moe", TORCH_FN(deep_gemm::torch_registration::get_block_m_for_mega_moe)); +} + +TORCH_LIBRARY_IMPL(deep_gemm, CUDA, m) { + m.impl("_slice_symm_buffer_for_mega_moe", TORCH_FN(deep_gemm::torch_registration::_slice_symm_buffer_for_mega_moe)); + + m.impl("fp8_fp4_mega_moe", TORCH_FN(deep_gemm::torch_registration::fp8_fp4_mega_moe)); + m.impl("bf16_mega_moe", TORCH_FN(deep_gemm::torch_registration::bf16_mega_moe)); +} diff --git a/csrc/apis/nvfp4_mega_moe.hpp b/csrc/apis/nvfp4_mega_moe.hpp index 5a1c4986d1..1f4c4e5a3b 100644 --- a/csrc/apis/nvfp4_mega_moe.hpp +++ b/csrc/apis/nvfp4_mega_moe.hpp @@ -1,9 +1,12 @@ #pragma once +#include +#include "../torch_library_utils.hpp" + #include -#include #include -#include +#include +#include #include #include @@ -14,6 +17,13 @@ namespace deep_gemm::nvfp4_mega { +static int checked_int(const int64_t value) { + // Preserve pybind behavior by rejecting values that do not fit exactly in a C++ int. + DG_HOST_ASSERT(value >= std::numeric_limits::min()); + DG_HOST_ASSERT(value <= std::numeric_limits::max()); + return static_cast(value); +} + static int get_token_alignment_for_nvfp4_mega_moe() { return layout::kLCMCandidateBlockM; } @@ -27,10 +37,72 @@ static int get_block_m_for_nvfp4_mega_moe( return block_m; } -static std::tuple(const torch::Tensor&)>> -get_symm_buffer_size_for_nvfp4_mega_moe( +struct SymmBufferLayoutInfo { + int64_t num_bytes = 0; + int64_t input_token_base = 0; + int64_t input_sf_base = 0; + int64_t input_topk_idx_base = 0; + int64_t input_topk_weights_base = 0; + int64_t shared_l1_token_base = 0; + int64_t shared_l1_sf_base = 0; + int64_t shared_l2_token_base = 0; + int64_t shared_l2_sf_base = 0; + int64_t l1_token_base = 0; + int64_t l1_sf_base = 0; + int64_t l2_token_base = 0; + int64_t l2_sf_base = 0; + bool shared_with_sf = false; + int num_max_tokens_per_rank = 0; + int num_topk = 0; + int hidden = 0; + int intermediate_hidden = 0; + int num_shared_experts = 0; + int shared_intermediate_hidden = 0; + int num_ring_tokens = 0; + int num_sf_ring_tokens = 0; + + std::vector to_int_list() const { + return { + num_bytes, input_token_base, input_sf_base, input_topk_idx_base, + input_topk_weights_base, shared_l1_token_base, shared_l1_sf_base, + shared_l2_token_base, shared_l2_sf_base, l1_token_base, l1_sf_base, + l2_token_base, l2_sf_base, static_cast(shared_with_sf), + num_max_tokens_per_rank, num_topk, hidden, intermediate_hidden, + num_shared_experts, shared_intermediate_hidden, num_ring_tokens, + num_sf_ring_tokens, + }; + } + + static SymmBufferLayoutInfo from_int_list(const std::vector& values) { + DG_HOST_ASSERT(static_cast(values.size()) == 22); + SymmBufferLayoutInfo info; + info.num_bytes = values[0]; + info.input_token_base = values[1]; + info.input_sf_base = values[2]; + info.input_topk_idx_base = values[3]; + info.input_topk_weights_base = values[4]; + info.shared_l1_token_base = values[5]; + info.shared_l1_sf_base = values[6]; + info.shared_l2_token_base = values[7]; + info.shared_l2_sf_base = values[8]; + info.l1_token_base = values[9]; + info.l1_sf_base = values[10]; + info.l2_token_base = values[11]; + info.l2_sf_base = values[12]; + info.shared_with_sf = values[13] != 0; + info.num_max_tokens_per_rank = checked_int(values[14]); + info.num_topk = checked_int(values[15]); + info.hidden = checked_int(values[16]); + info.intermediate_hidden = checked_int(values[17]); + info.num_shared_experts = checked_int(values[18]); + info.shared_intermediate_hidden = checked_int(values[19]); + info.num_ring_tokens = checked_int(values[20]); + info.num_sf_ring_tokens = checked_int(values[21]); + return info; + } +}; + +static SymmBufferLayoutInfo build_symm_buffer_layout( const int& num_ranks, const int& num_experts, const int& num_max_tokens_per_rank, const int& num_topk, const int& hidden, const int& intermediate_hidden, @@ -59,8 +131,6 @@ get_symm_buffer_size_for_nvfp4_mega_moe( num_ring_tokens = math::align(num_ring_tokens, layout::kLCMCandidateBlockM); const bool shared_with_sf = not shared_bf16; - constexpr int packed_sf_k = 64; // Four per-16 E4M3 scale bytes per int32. - // Compute num_sf_ring_tokens (max across all candidate block sizes) int num_sf_ring_tokens = 0; for (auto block_m: layout::kCandidateBlockM) { @@ -82,68 +152,107 @@ get_symm_buffer_size_for_nvfp4_mega_moe( DG_HOST_ASSERT(shared_intermediate_hidden % 128 == 0); DG_HOST_ASSERT(num_sf_ring_tokens % 4 == 0); - // Slice function: creates tensor views from the raw buffer. - // NOTES: `x_sf` is K-major, while `l1_acts_sf` and `l2_acts_sf` are M-major - auto slice_input_buffers = [=](const torch::Tensor& buffer) { + SymmBufferLayoutInfo layout_info; + layout_info.num_bytes = mega_buffer.get_num_bytes(); + layout_info.input_token_base = reinterpret_cast(mega_buffer.input_token_buffer.base); + layout_info.input_sf_base = reinterpret_cast(mega_buffer.input_sf_buffer.base); + layout_info.input_topk_idx_base = reinterpret_cast(mega_buffer.input_topk_idx_buffer.base); + layout_info.input_topk_weights_base = reinterpret_cast(mega_buffer.input_topk_weights_buffer.base); + layout_info.shared_l1_token_base = reinterpret_cast(mega_buffer.shared_l1_token_buffer.base); + layout_info.shared_l1_sf_base = reinterpret_cast(mega_buffer.shared_l1_sf_buffer.base); + layout_info.shared_l2_token_base = reinterpret_cast(mega_buffer.shared_l2_token_buffer.base); + layout_info.shared_l2_sf_base = reinterpret_cast(mega_buffer.shared_l2_sf_buffer.base); + layout_info.l1_token_base = reinterpret_cast(mega_buffer.l1_token_buffer.base); + layout_info.l1_sf_base = reinterpret_cast(mega_buffer.l1_sf_buffer.base); + layout_info.l2_token_base = reinterpret_cast(mega_buffer.l2_token_buffer.base); + layout_info.l2_sf_base = reinterpret_cast(mega_buffer.l2_sf_buffer.base); + layout_info.shared_with_sf = shared_with_sf; + layout_info.num_max_tokens_per_rank = num_max_tokens_per_rank; + layout_info.num_topk = num_topk; + layout_info.hidden = hidden; + layout_info.intermediate_hidden = intermediate_hidden; + layout_info.num_shared_experts = num_shared_experts; + layout_info.shared_intermediate_hidden = shared_intermediate_hidden; + layout_info.num_ring_tokens = num_ring_tokens; + layout_info.num_sf_ring_tokens = num_sf_ring_tokens; + return layout_info; +} + +static std::tuple> get_symm_buffer_size_for_nvfp4_mega_moe( + const int& num_ranks, const int& num_experts, + const int& num_max_tokens_per_rank, const int& num_topk, + const int& hidden, const int& intermediate_hidden, + const int& num_shared_experts = 0, const bool shared_bf16 = false) { + const auto layout_info = build_symm_buffer_layout( + num_ranks, num_experts, num_max_tokens_per_rank, num_topk, + hidden, intermediate_hidden, num_shared_experts, shared_bf16); + return std::make_tuple(layout_info.num_bytes, layout_info.to_int_list()); +} + +using SymmBufferSlice = std::tuple; + +static SymmBufferSlice slice_symm_buffer_from_layout( + const torch::Tensor& buffer, const SymmBufferLayoutInfo& layout_info) { + // `x_sf` is K-major, while `l1_acts_sf` and `l2_acts_sf` are M-major. auto x = torch::from_blob( - math::advance_ptr(buffer.data_ptr(), reinterpret_cast(mega_buffer.input_token_buffer.base)), - {num_max_tokens_per_rank, hidden / 2}, + math::advance_ptr(buffer.data_ptr(), layout_info.input_token_base), + {layout_info.num_max_tokens_per_rank, layout_info.hidden / 2}, torch::TensorOptions().dtype(kPackedFP4).device(buffer.device())); auto x_sf = torch::from_blob( - math::advance_ptr(buffer.data_ptr(), reinterpret_cast(mega_buffer.input_sf_buffer.base)), - {num_max_tokens_per_rank, hidden / packed_sf_k}, + math::advance_ptr(buffer.data_ptr(), layout_info.input_sf_base), + {layout_info.num_max_tokens_per_rank, layout_info.hidden / 64}, torch::TensorOptions().dtype(torch::kInt).device(buffer.device())); auto topk_idx = torch::from_blob( - math::advance_ptr(buffer.data_ptr(), reinterpret_cast(mega_buffer.input_topk_idx_buffer.base)), - {num_max_tokens_per_rank, num_topk}, + math::advance_ptr(buffer.data_ptr(), layout_info.input_topk_idx_base), + {layout_info.num_max_tokens_per_rank, layout_info.num_topk}, torch::TensorOptions().dtype(torch::kInt64).device(buffer.device())); auto topk_weights = torch::from_blob( - math::advance_ptr(buffer.data_ptr(), reinterpret_cast(mega_buffer.input_topk_weights_buffer.base)), - {num_max_tokens_per_rank, num_topk}, + math::advance_ptr(buffer.data_ptr(), layout_info.input_topk_weights_base), + {layout_info.num_max_tokens_per_rank, layout_info.num_topk}, torch::TensorOptions().dtype(torch::kFloat32).device(buffer.device())); - auto shared_l1_acts = num_shared_experts > 0 ? torch::from_blob( - math::advance_ptr(buffer.data_ptr(), reinterpret_cast(mega_buffer.shared_l1_token_buffer.base)), - {num_max_tokens_per_rank, hidden}, - torch::TensorOptions().dtype(shared_with_sf ? torch::kFloat8_e4m3fn : torch::kBFloat16).device(buffer.device())) : x; - auto shared_l1_acts_sf = (shared_with_sf and num_shared_experts > 0) ? torch::from_blob( - math::advance_ptr(buffer.data_ptr(), reinterpret_cast(mega_buffer.shared_l1_sf_buffer.base)), - {layout::get_num_max_shared_sf_tokens(num_max_tokens_per_rank), hidden / 128}, - {1, layout::get_num_max_shared_sf_tokens(num_max_tokens_per_rank)}, + auto shared_l1_acts = layout_info.num_shared_experts > 0 ? torch::from_blob( + math::advance_ptr(buffer.data_ptr(), layout_info.shared_l1_token_base), + {layout_info.num_max_tokens_per_rank, layout_info.hidden}, + torch::TensorOptions().dtype(layout_info.shared_with_sf ? torch::kFloat8_e4m3fn : torch::kBFloat16).device(buffer.device())) : x; + auto shared_l1_acts_sf = (layout_info.shared_with_sf and layout_info.num_shared_experts > 0) ? torch::from_blob( + math::advance_ptr(buffer.data_ptr(), layout_info.shared_l1_sf_base), + {layout::get_num_max_shared_sf_tokens(layout_info.num_max_tokens_per_rank), layout_info.hidden / 128}, + {1, layout::get_num_max_shared_sf_tokens(layout_info.num_max_tokens_per_rank)}, torch::TensorOptions().dtype(torch::kInt).device(buffer.device())) : torch::Tensor(); - auto shared_l2_acts = num_shared_experts > 0 ? torch::from_blob( - math::advance_ptr(buffer.data_ptr(), reinterpret_cast(mega_buffer.shared_l2_token_buffer.base)), - {num_max_tokens_per_rank, shared_intermediate_hidden}, - torch::TensorOptions().dtype(shared_with_sf ? torch::kFloat8_e4m3fn : torch::kBFloat16).device(buffer.device())) : torch::Tensor(); - auto shared_l2_acts_sf = (shared_with_sf and num_shared_experts > 0) ? torch::from_blob( - math::advance_ptr(buffer.data_ptr(), reinterpret_cast(mega_buffer.shared_l2_sf_buffer.base)), - {layout::get_num_max_shared_sf_tokens(num_max_tokens_per_rank), shared_intermediate_hidden / 128}, - {1, layout::get_num_max_shared_sf_tokens(num_max_tokens_per_rank)}, + auto shared_l2_acts = layout_info.num_shared_experts > 0 ? torch::from_blob( + math::advance_ptr(buffer.data_ptr(), layout_info.shared_l2_token_base), + {layout_info.num_max_tokens_per_rank, layout_info.shared_intermediate_hidden}, + torch::TensorOptions().dtype(layout_info.shared_with_sf ? torch::kFloat8_e4m3fn : torch::kBFloat16).device(buffer.device())) : torch::Tensor(); + auto shared_l2_acts_sf = (layout_info.shared_with_sf and layout_info.num_shared_experts > 0) ? torch::from_blob( + math::advance_ptr(buffer.data_ptr(), layout_info.shared_l2_sf_base), + {layout::get_num_max_shared_sf_tokens(layout_info.num_max_tokens_per_rank), layout_info.shared_intermediate_hidden / 128}, + {1, layout::get_num_max_shared_sf_tokens(layout_info.num_max_tokens_per_rank)}, torch::TensorOptions().dtype(torch::kInt).device(buffer.device())) : torch::Tensor(); auto l1_acts = torch::from_blob( - math::advance_ptr(buffer.data_ptr(), reinterpret_cast(mega_buffer.l1_token_buffer.base)), - {num_ring_tokens, hidden / 2}, + math::advance_ptr(buffer.data_ptr(), layout_info.l1_token_base), + {layout_info.num_ring_tokens, layout_info.hidden / 2}, torch::TensorOptions().dtype(kPackedFP4).device(buffer.device())); auto l1_acts_sf = torch::from_blob( - math::advance_ptr(buffer.data_ptr(), reinterpret_cast(mega_buffer.l1_sf_buffer.base)), - {num_sf_ring_tokens, hidden / packed_sf_k}, - {1, num_sf_ring_tokens}, + math::advance_ptr(buffer.data_ptr(), layout_info.l1_sf_base), + {layout_info.num_sf_ring_tokens, layout_info.hidden / 64}, + {1, layout_info.num_sf_ring_tokens}, torch::TensorOptions().dtype(torch::kInt).device(buffer.device())); auto l2_acts = torch::from_blob( - math::advance_ptr(buffer.data_ptr(), reinterpret_cast(mega_buffer.l2_token_buffer.base)), - {num_ring_tokens, intermediate_hidden / 2}, + math::advance_ptr(buffer.data_ptr(), layout_info.l2_token_base), + {layout_info.num_ring_tokens, layout_info.intermediate_hidden / 2}, torch::TensorOptions().dtype(kPackedFP4).device(buffer.device())); auto l2_acts_sf = torch::from_blob( - math::advance_ptr(buffer.data_ptr(), reinterpret_cast(mega_buffer.l2_sf_buffer.base)), - {num_sf_ring_tokens, intermediate_hidden / packed_sf_k}, - {1, num_sf_ring_tokens}, + math::advance_ptr(buffer.data_ptr(), layout_info.l2_sf_base), + {layout_info.num_sf_ring_tokens, layout_info.intermediate_hidden / 64}, + {1, layout_info.num_sf_ring_tokens}, torch::TensorOptions().dtype(torch::kInt).device(buffer.device())); return std::make_tuple(x, x_sf, topk_idx, topk_weights, shared_l1_acts, shared_l1_acts_sf, shared_l2_acts, shared_l2_acts_sf, l1_acts, l1_acts_sf, l2_acts, l2_acts_sf); - }; - return {mega_buffer.get_num_bytes(), slice_input_buffers}; } static void nvfp4_mega_moe( @@ -261,19 +370,20 @@ static void nvfp4_mega_moe( // Check buffer bytes const auto num_ranks = static_cast(sym_buffer_ptrs.size()); const auto num_experts_ = num_experts_per_rank * num_ranks; - const auto [num_required_bytes, slice] = get_symm_buffer_size_for_nvfp4_mega_moe( + const auto layout_info = build_symm_buffer_layout( num_ranks, num_experts, num_max_tokens_per_rank, num_topk, hidden, intermediate_hidden, num_shared_experts, shared_bf16 ); - DG_HOST_ASSERT(sym_buffer.nbytes() >= static_cast(num_required_bytes)); + DG_HOST_ASSERT(sym_buffer.nbytes() >= static_cast(layout_info.num_bytes)); DG_HOST_ASSERT(num_experts == num_experts_); // Already registered tensors const auto [x, x_sf, topk_idx, topk_weights, shared_l1_acts, shared_l1_acts_sf, shared_l2_acts, shared_l2_acts_sf, - l1_acts, l1_acts_sf, l2_acts, l2_acts_sf] = slice(sym_buffer); + l1_acts, l1_acts_sf, l2_acts, l2_acts_sf] = + slice_symm_buffer_from_layout(sym_buffer, layout_info); // Dispatch into different architectures if (arch_major == 10) { @@ -305,24 +415,97 @@ static void nvfp4_mega_moe( sym_buffer.zero_(); } -static void register_apis(pybind11::module_& m) { - m.def("get_token_alignment_for_nvfp4_mega_moe", &get_token_alignment_for_nvfp4_mega_moe); - m.def("get_block_m_for_nvfp4_mega_moe", &get_block_m_for_nvfp4_mega_moe); - m.def("get_symm_buffer_size_for_nvfp4_mega_moe", &get_symm_buffer_size_for_nvfp4_mega_moe, - pybind11::arg("num_ranks"), pybind11::arg("num_experts"), - pybind11::arg("num_max_tokens_per_rank"), pybind11::arg("num_topk"), - pybind11::arg("hidden"), pybind11::arg("intermediate_hidden"), - pybind11::arg("num_shared_experts") = 0, pybind11::arg("shared_bf16") = false); - m.def("nvfp4_mega_moe", &nvfp4_mega_moe, - pybind11::arg("y"), pybind11::arg("l1_weights"), pybind11::arg("l2_weights"), - pybind11::arg("shared_l1_weights"), pybind11::arg("shared_l2_weights"), - pybind11::arg("cumulative_local_expert_recv_stats"), pybind11::arg("sym_buffer"), - pybind11::arg("sym_buffer_ptrs"), pybind11::arg("rank_idx"), - pybind11::arg("num_max_tokens_per_rank"), pybind11::arg("num_experts"), pybind11::arg("num_topk"), - pybind11::arg("activation_clamp"), - pybind11::arg("fast_math"), pybind11::arg("activation_alpha"), pybind11::arg("activation_beta"), - pybind11::arg("l1_alpha") = pybind11::none(), pybind11::arg("l2_alpha") = pybind11::none(), - pybind11::arg("l2_activation_scale") = 1.0f); +} // namespace deep_gemm::nvfp4_mega + +namespace deep_gemm::torch_registration { + +static int64_t get_token_alignment_for_nvfp4_mega_moe() { + return nvfp4_mega::get_token_alignment_for_nvfp4_mega_moe(); } -} // namespace deep_gemm::nvfp4_mega +static int64_t get_block_m_for_nvfp4_mega_moe( + int64_t num_ranks, int64_t num_experts, int64_t num_max_tokens_per_rank, + int64_t num_tokens, int64_t num_topk) { + return nvfp4_mega::get_block_m_for_nvfp4_mega_moe( + nvfp4_mega::checked_int(num_ranks), nvfp4_mega::checked_int(num_experts), + nvfp4_mega::checked_int(num_max_tokens_per_rank), + nvfp4_mega::checked_int(num_tokens), nvfp4_mega::checked_int(num_topk)); +} + +static std::tuple> get_symm_buffer_size_for_nvfp4_mega_moe( + int64_t num_ranks, int64_t num_experts, int64_t num_max_tokens_per_rank, + int64_t num_topk, int64_t hidden, int64_t intermediate_hidden, + int64_t num_shared_experts, bool shared_bf16) { + return nvfp4_mega::get_symm_buffer_size_for_nvfp4_mega_moe( + nvfp4_mega::checked_int(num_ranks), nvfp4_mega::checked_int(num_experts), + nvfp4_mega::checked_int(num_max_tokens_per_rank), + nvfp4_mega::checked_int(num_topk), nvfp4_mega::checked_int(hidden), + nvfp4_mega::checked_int(intermediate_hidden), + nvfp4_mega::checked_int(num_shared_experts), shared_bf16); +} + +static nvfp4_mega::SymmBufferSlice _slice_symm_buffer_for_nvfp4_mega_moe( + const torch::Tensor& buffer, const std::vector& layout_info) { + return nvfp4_mega::slice_symm_buffer_from_layout( + buffer, nvfp4_mega::SymmBufferLayoutInfo::from_int_list(layout_info)); +} + +static void nvfp4_mega_moe( + const torch::Tensor& y, + const torch::Tensor& l1_weights, const torch::Tensor& l1_weights_sf, + const torch::Tensor& l2_weights, const torch::Tensor& l2_weights_sf, + const std::optional& shared_l1_weights, + const std::optional& shared_l1_weights_sf, + const std::optional& shared_l2_weights, + const std::optional& shared_l2_weights_sf, + const std::optional& cumulative_local_expert_recv_stats, + const torch::Tensor& sym_buffer, const std::vector& sym_buffer_ptrs, + int64_t rank_idx, int64_t num_max_tokens_per_rank, + int64_t num_experts, int64_t num_topk, + const std::optional& activation_clamp, bool fast_math, + double activation_alpha, double activation_beta, + const std::optional& l1_alpha, + const std::optional& l2_alpha, + double l2_activation_scale) { + DG_HOST_ASSERT(shared_l1_weights.has_value() == shared_l2_weights.has_value()); + std::optional>> shared_l1; + std::optional>> shared_l2; + if (shared_l1_weights.has_value()) { + shared_l1 = std::make_tuple(*shared_l1_weights, shared_l1_weights_sf); + shared_l2 = std::make_tuple(*shared_l2_weights, shared_l2_weights_sf); + } + nvfp4_mega::nvfp4_mega_moe( + y, std::make_tuple(l1_weights, l1_weights_sf), + std::make_tuple(l2_weights, l2_weights_sf), shared_l1, shared_l2, + cumulative_local_expert_recv_stats, sym_buffer, sym_buffer_ptrs, + nvfp4_mega::checked_int(rank_idx), + nvfp4_mega::checked_int(num_max_tokens_per_rank), + nvfp4_mega::checked_int(num_experts), nvfp4_mega::checked_int(num_topk), + activation_clamp.has_value() + ? std::optional(static_cast(*activation_clamp)) + : std::nullopt, + fast_math, static_cast(activation_alpha), + static_cast(activation_beta), l1_alpha, l2_alpha, + static_cast(l2_activation_scale)); +} + +} // namespace deep_gemm::torch_registration + +TORCH_LIBRARY_FRAGMENT(deep_gemm, m) { + m.def("get_token_alignment_for_nvfp4_mega_moe() -> int"); + m.def("get_block_m_for_nvfp4_mega_moe(int num_ranks, int num_experts, int num_max_tokens_per_rank, int num_tokens, int num_topk) -> int"); + m.def("get_symm_buffer_size_for_nvfp4_mega_moe(int num_ranks, int num_experts, int num_max_tokens_per_rank, int num_topk, int hidden, int intermediate_hidden, int num_shared_experts=0, bool shared_bf16=False) -> (int, int[])", + TORCH_FN(deep_gemm::torch_registration::get_symm_buffer_size_for_nvfp4_mega_moe)); + m.def("_slice_symm_buffer_for_nvfp4_mega_moe(Tensor(a) buffer, int[] layout_info) -> (Tensor(a), Tensor(a), Tensor(a), Tensor(a), Tensor(a), Tensor(a), Tensor(a), Tensor(a), Tensor(a), Tensor(a), Tensor(a), Tensor(a))"); + m.def("nvfp4_mega_moe(Tensor(a!) y, Tensor l1_weights_tuple, Tensor l1_weights_tuple_sf, Tensor l2_weights_tuple, Tensor l2_weights_tuple_sf, Tensor? shared_l1_weights_tuple_opt, Tensor? shared_l1_weights_tuple_opt_sf, Tensor? shared_l2_weights_tuple_opt, Tensor? shared_l2_weights_tuple_opt_sf, Tensor(b!)? cumulative_local_expert_recv_stats, Tensor(c!) sym_buffer, int[] sym_buffer_ptrs, int rank_idx, int num_max_tokens_per_rank, int num_experts, int num_topk, float? activation_clamp_opt, bool fast_math, float activation_alpha, float activation_beta, Tensor? l1_alpha=None, Tensor? l2_alpha=None, float l2_activation_scale=1.0) -> ()"); +} + +TORCH_LIBRARY_IMPL(deep_gemm, CatchAll, m) { + m.impl("get_token_alignment_for_nvfp4_mega_moe", TORCH_FN(deep_gemm::torch_registration::get_token_alignment_for_nvfp4_mega_moe)); + m.impl("get_block_m_for_nvfp4_mega_moe", TORCH_FN(deep_gemm::torch_registration::get_block_m_for_nvfp4_mega_moe)); +} + +TORCH_LIBRARY_IMPL(deep_gemm, CUDA, m) { + m.impl("_slice_symm_buffer_for_nvfp4_mega_moe", TORCH_FN(deep_gemm::torch_registration::_slice_symm_buffer_for_nvfp4_mega_moe)); + m.impl("nvfp4_mega_moe", TORCH_FN(deep_gemm::torch_registration::nvfp4_mega_moe)); +} diff --git a/csrc/apis/sm120_dispatch.hpp b/csrc/apis/sm120_dispatch.hpp index 89ac82c1e4..4b2be00799 100644 --- a/csrc/apis/sm120_dispatch.hpp +++ b/csrc/apis/sm120_dispatch.hpp @@ -17,7 +17,7 @@ // // UNVALIDATED: no sm120 hardware was available while porting; none of this has been run. -#include +#include #include "../jit_kernels/impls/sm120_fp8_fp4_gemm_1d1d.hpp" #include "../jit_kernels/impls/sm120_bf16_gemm.hpp" diff --git a/csrc/jit_kernels/impls/epilogue.hpp b/csrc/jit_kernels/impls/epilogue.hpp index c132a28942..aa5515be4f 100644 --- a/csrc/jit_kernels/impls/epilogue.hpp +++ b/csrc/jit_kernels/impls/epilogue.hpp @@ -3,7 +3,7 @@ #include #include -#include +#include #include diff --git a/csrc/jit_kernels/impls/runtime_utils.hpp b/csrc/jit_kernels/impls/runtime_utils.hpp index 15cd7a14f7..5d73246f96 100644 --- a/csrc/jit_kernels/impls/runtime_utils.hpp +++ b/csrc/jit_kernels/impls/runtime_utils.hpp @@ -8,7 +8,7 @@ #include #include -#include +#include #include #include diff --git a/csrc/jit_kernels/impls/sm100_bf16_gemm.hpp b/csrc/jit_kernels/impls/sm100_bf16_gemm.hpp index 030e77c503..fa80318da0 100644 --- a/csrc/jit_kernels/impls/sm100_bf16_gemm.hpp +++ b/csrc/jit_kernels/impls/sm100_bf16_gemm.hpp @@ -1,7 +1,7 @@ #pragma once #include -#include +#include #include "../../runtime/runtime.hpp" #include "../../utils/exception.hpp" diff --git a/csrc/jit_kernels/impls/sm100_bf16_mega_gate.hpp b/csrc/jit_kernels/impls/sm100_bf16_mega_gate.hpp index e3c263ea2d..24b587bdb6 100644 --- a/csrc/jit_kernels/impls/sm100_bf16_mega_gate.hpp +++ b/csrc/jit_kernels/impls/sm100_bf16_mega_gate.hpp @@ -2,7 +2,7 @@ #include #include -#include +#include #include diff --git a/csrc/jit_kernels/impls/sm100_bf16_mega_moe.hpp b/csrc/jit_kernels/impls/sm100_bf16_mega_moe.hpp index a085d41f2e..9dd2c54780 100644 --- a/csrc/jit_kernels/impls/sm100_bf16_mega_moe.hpp +++ b/csrc/jit_kernels/impls/sm100_bf16_mega_moe.hpp @@ -1,7 +1,7 @@ #pragma once #include -#include +#include #include #include diff --git a/csrc/jit_kernels/impls/sm100_bmk_bnk_mn.hpp b/csrc/jit_kernels/impls/sm100_bmk_bnk_mn.hpp index e9a14c20f1..51f688065b 100644 --- a/csrc/jit_kernels/impls/sm100_bmk_bnk_mn.hpp +++ b/csrc/jit_kernels/impls/sm100_bmk_bnk_mn.hpp @@ -2,7 +2,7 @@ #include #include -#include +#include #include "../../runtime/runtime.hpp" #include "../../utils/exception.hpp" diff --git a/csrc/jit_kernels/impls/sm100_fp8_fp4_gemm_1d1d.hpp b/csrc/jit_kernels/impls/sm100_fp8_fp4_gemm_1d1d.hpp index 0e637da1ea..9d56d2a034 100644 --- a/csrc/jit_kernels/impls/sm100_fp8_fp4_gemm_1d1d.hpp +++ b/csrc/jit_kernels/impls/sm100_fp8_fp4_gemm_1d1d.hpp @@ -1,7 +1,7 @@ #pragma once #include -#include +#include #include "../../runtime/runtime.hpp" #include "../../utils/exception.hpp" diff --git a/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp b/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp index 77ba9d4cdd..2e27aece4b 100644 --- a/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp +++ b/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp @@ -1,7 +1,7 @@ #pragma once #include -#include +#include #include #include diff --git a/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe_situ.hpp b/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe_situ.hpp index b997e5e950..6a573c0f93 100644 --- a/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe_situ.hpp +++ b/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe_situ.hpp @@ -1,7 +1,7 @@ #pragma once #include -#include +#include #include #include diff --git a/csrc/jit_kernels/impls/sm100_nvfp4_mega_moe.hpp b/csrc/jit_kernels/impls/sm100_nvfp4_mega_moe.hpp index 0f3f26c356..194c601ba5 100644 --- a/csrc/jit_kernels/impls/sm100_nvfp4_mega_moe.hpp +++ b/csrc/jit_kernels/impls/sm100_nvfp4_mega_moe.hpp @@ -1,7 +1,7 @@ #pragma once #include -#include +#include #include #include diff --git a/csrc/jit_kernels/impls/sm100_tf32_hc_prenorm_gemm.hpp b/csrc/jit_kernels/impls/sm100_tf32_hc_prenorm_gemm.hpp index 004f73070b..b9532d6e55 100644 --- a/csrc/jit_kernels/impls/sm100_tf32_hc_prenorm_gemm.hpp +++ b/csrc/jit_kernels/impls/sm100_tf32_hc_prenorm_gemm.hpp @@ -2,7 +2,7 @@ #include #include -#include +#include #include "../../runtime/jit.hpp" #include "../../utils/exception.hpp" diff --git a/csrc/jit_kernels/impls/sm120_bf16_gemm.hpp b/csrc/jit_kernels/impls/sm120_bf16_gemm.hpp index 4fa3daf63c..5e90cf0bf0 100644 --- a/csrc/jit_kernels/impls/sm120_bf16_gemm.hpp +++ b/csrc/jit_kernels/impls/sm120_bf16_gemm.hpp @@ -2,7 +2,7 @@ #include #include -#include +#include #include "../../runtime/runtime.hpp" #include "../../utils/exception.hpp" diff --git a/csrc/jit_kernels/impls/sm120_bmk_bnk_mn.hpp b/csrc/jit_kernels/impls/sm120_bmk_bnk_mn.hpp index dc977d90ef..7becfdbea5 100644 --- a/csrc/jit_kernels/impls/sm120_bmk_bnk_mn.hpp +++ b/csrc/jit_kernels/impls/sm120_bmk_bnk_mn.hpp @@ -2,7 +2,7 @@ #include #include -#include +#include #include "../../runtime/runtime.hpp" #include "../../utils/exception.hpp" diff --git a/csrc/jit_kernels/impls/sm120_fp8_fp4_gemm_1d1d.hpp b/csrc/jit_kernels/impls/sm120_fp8_fp4_gemm_1d1d.hpp index 84c1a976f8..3a0476ae3f 100644 --- a/csrc/jit_kernels/impls/sm120_fp8_fp4_gemm_1d1d.hpp +++ b/csrc/jit_kernels/impls/sm120_fp8_fp4_gemm_1d1d.hpp @@ -2,7 +2,7 @@ #include #include -#include +#include #include "../../runtime/runtime.hpp" #include "../../utils/exception.hpp" diff --git a/csrc/jit_kernels/impls/sm120_mqa_logits.hpp b/csrc/jit_kernels/impls/sm120_mqa_logits.hpp index d38708b249..192e8fbf36 100644 --- a/csrc/jit_kernels/impls/sm120_mqa_logits.hpp +++ b/csrc/jit_kernels/impls/sm120_mqa_logits.hpp @@ -3,7 +3,7 @@ #include #include #include -#include +#include #include "../../runtime/runtime.hpp" #include "../../utils/exception.hpp" diff --git a/csrc/jit_kernels/impls/sm120_tf32_hc_prenorm_gemm.hpp b/csrc/jit_kernels/impls/sm120_tf32_hc_prenorm_gemm.hpp index 33c3b0e103..823d029308 100644 --- a/csrc/jit_kernels/impls/sm120_tf32_hc_prenorm_gemm.hpp +++ b/csrc/jit_kernels/impls/sm120_tf32_hc_prenorm_gemm.hpp @@ -1,7 +1,7 @@ #pragma once #include -#include +#include #include "../../runtime/jit.hpp" #include "../../utils/exception.hpp" diff --git a/csrc/jit_kernels/impls/sm90_bf16_gemm.hpp b/csrc/jit_kernels/impls/sm90_bf16_gemm.hpp index 699db4191e..ce52720417 100644 --- a/csrc/jit_kernels/impls/sm90_bf16_gemm.hpp +++ b/csrc/jit_kernels/impls/sm90_bf16_gemm.hpp @@ -1,7 +1,7 @@ #pragma once #include -#include +#include #include "../../runtime/runtime.hpp" #include "../../utils/exception.hpp" diff --git a/csrc/jit_kernels/impls/sm90_bmk_bnk_mn.hpp b/csrc/jit_kernels/impls/sm90_bmk_bnk_mn.hpp index 2bda372666..9aac59329f 100644 --- a/csrc/jit_kernels/impls/sm90_bmk_bnk_mn.hpp +++ b/csrc/jit_kernels/impls/sm90_bmk_bnk_mn.hpp @@ -2,7 +2,7 @@ #include #include -#include +#include #include "../../runtime/runtime.hpp" #include "../../utils/exception.hpp" diff --git a/csrc/jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp b/csrc/jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp index 29f516c6d6..ddbbf5627d 100644 --- a/csrc/jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp +++ b/csrc/jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp @@ -1,7 +1,7 @@ #pragma once #include -#include +#include #include "../../runtime/runtime.hpp" #include "../../utils/exception.hpp" diff --git a/csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp b/csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp index 89d0d2b2fe..5adc816865 100644 --- a/csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp +++ b/csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp @@ -1,7 +1,7 @@ #pragma once #include -#include +#include #include "../../runtime/runtime.hpp" #include "../../utils/exception.hpp" diff --git a/csrc/jit_kernels/impls/sm90_tf32_hc_prenorm_gemm.hpp b/csrc/jit_kernels/impls/sm90_tf32_hc_prenorm_gemm.hpp index 2dcee9a0e3..9d376a24cf 100644 --- a/csrc/jit_kernels/impls/sm90_tf32_hc_prenorm_gemm.hpp +++ b/csrc/jit_kernels/impls/sm90_tf32_hc_prenorm_gemm.hpp @@ -2,7 +2,7 @@ #include #include -#include +#include #include "../../runtime/jit.hpp" #include "../../utils/exception.hpp" diff --git a/csrc/jit_kernels/impls/smxx_layout.hpp b/csrc/jit_kernels/impls/smxx_layout.hpp index a40cee1fbc..e651274eda 100644 --- a/csrc/jit_kernels/impls/smxx_layout.hpp +++ b/csrc/jit_kernels/impls/smxx_layout.hpp @@ -1,7 +1,7 @@ #pragma once #include -#include +#include #include "../../runtime/runtime.hpp" #include "../../utils/exception.hpp" diff --git a/csrc/python_api.cpp b/csrc/python_api.cpp index b4bbfde983..9361bc226c 100644 --- a/csrc/python_api.cpp +++ b/csrc/python_api.cpp @@ -1,9 +1,13 @@ -#include -#include +#include + +#include +#include +#include +#include "utils/registration.h" #include -#include +#include "runtime/runtime.hpp" #include "apis/config.hpp" #include "apis/attention.hpp" #include "apis/einsum.hpp" @@ -15,29 +19,18 @@ #include "apis/mega_mhc.hpp" #include "apis/mega_gate.hpp" -#ifndef TORCH_EXTENSION_NAME -#define TORCH_EXTENSION_NAME _C -#endif - -// ReSharper disable once CppParameterMayBeConstPtrOrRef -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.doc() = "DeepGEMM C++ library"; +namespace deep_gemm { +// Preserve the opaque runtime handle exposed by DeepJIT's former Python binding. +struct JitRuntimeHandle : torch::CustomClassHolder { + std::shared_ptr> value; + JitRuntimeHandle() : value(jit.get()) {} +}; +} +TORCH_LIBRARY(deep_gemm, m) { // Register JIT objects - deep_jit::register_python_api(m, deep_gemm::jit); - - // Register config APIs - deep_gemm::config::register_apis(m); - - // Register kernels - // TODO: make SM80 incompatible issues raise errors - deep_gemm::attention::register_apis(m); - deep_gemm::einsum::register_apis(m); - deep_gemm::hyperconnection::register_apis(m); - deep_gemm::gemm::register_apis(m); - deep_gemm::layout::register_apis(m); - deep_gemm::mega::register_apis(m); - deep_gemm::nvfp4_mega::register_apis(m); - deep_gemm::mega_mhc::register_apis(m); - deep_gemm::mega_gate::register_apis(m); + m.class_("Runtime"); + m.def("get_jit() -> __torch__.torch.classes.deep_gemm.Runtime", []() { return c10::make_intrusive(); }); } + +REGISTER_EXTENSION(_C_extension) diff --git a/csrc/torch_library_utils.hpp b/csrc/torch_library_utils.hpp new file mode 100644 index 0000000000..84797cd3a6 --- /dev/null +++ b/csrc/torch_library_utils.hpp @@ -0,0 +1,66 @@ +#pragma once + +#include +#include +#include +#include +#include + +#include + +#include "utils/exception.hpp" + +namespace deep_gemm::torch_utils { + +inline std::optional> list_to_recipe3( + const c10::optional>& recipe) { + if (not recipe.has_value()) { + return std::nullopt; + } + DG_HOST_ASSERT(recipe->size() == 3); + return std::make_tuple(static_cast((*recipe)[0]), + static_cast((*recipe)[1]), + static_cast((*recipe)[2])); +} + +inline std::optional> list_to_recipe2( + const c10::optional>& recipe) { + if (not recipe.has_value()) { + return std::nullopt; + } + DG_HOST_ASSERT(recipe->size() == 2); + return std::make_tuple(static_cast((*recipe)[0]), static_cast((*recipe)[1])); +} + +inline std::variant, std::tuple> list_to_recipe_variant( + const std::vector& recipe) { + DG_HOST_ASSERT(recipe.size() == 2 or recipe.size() == 3); + if (recipe.size() == 2) { + return std::make_tuple(static_cast(recipe[0]), static_cast(recipe[1])); + } + return std::make_tuple(static_cast(recipe[0]), + static_cast(recipe[1]), + static_cast(recipe[2])); +} + +inline std::tuple list_to_tuple3(const std::vector& values) { + DG_HOST_ASSERT(values.size() == 3); + return std::make_tuple(static_cast(values[0]), + static_cast(values[1]), + static_cast(values[2])); +} + +inline std::optional> list_to_optional_vector_int( + const c10::optional>& values) { + if (not values.has_value()) { + return std::nullopt; + } + std::vector out; + out.reserve(values->size()); + for (const auto value : *values) { + out.push_back(static_cast(value)); + } + return out; +} + +} // namespace deep_gemm::torch_utils diff --git a/csrc/utils/layout.hpp b/csrc/utils/layout.hpp index 136135d6fb..dce87941e2 100644 --- a/csrc/utils/layout.hpp +++ b/csrc/utils/layout.hpp @@ -1,7 +1,7 @@ #pragma once #include -#include +#include #include "math.hpp" #include "exception.hpp" diff --git a/csrc/utils/math.hpp b/csrc/utils/math.hpp index 19a86c3890..902ebab1b8 100644 --- a/csrc/utils/math.hpp +++ b/csrc/utils/math.hpp @@ -1,7 +1,7 @@ // TODO: merge this file with `math.cuh` (the device part) #pragma once -#include +#include #include "exception.hpp" diff --git a/csrc/utils/registration.h b/csrc/utils/registration.h new file mode 100644 index 0000000000..f625822732 --- /dev/null +++ b/csrc/utils/registration.h @@ -0,0 +1,17 @@ +#pragma once + +#include + +#define _CONCAT(A, B) A##B +#define CONCAT(A, B) _CONCAT(A, B) + +#define _STRINGIFY(A) #A +#define STRINGIFY(A) _STRINGIFY(A) + +// Empty PyInit so the .so is importable; ops still register via TORCH_LIBRARY. +#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); \ + } diff --git a/deep_gemm/_C.py b/deep_gemm/_C.py new file mode 100644 index 0000000000..13a862b99b --- /dev/null +++ b/deep_gemm/_C.py @@ -0,0 +1,720 @@ +"""Compatibility facade for the former ``deep_gemm._C`` extension. + +Keeping this module name preserves the existing Python API after moving +operator registration from pybind to ``TORCH_LIBRARY``. It loads the compiled +extension and forwards calls to ``torch.ops.deep_gemm``. +""" + +import torch +from pathlib import Path + + +def _load_extension(): + so_files = list(Path(__file__).parent.glob('_C_extension*.so')) + assert len(so_files) == 1, ( + f'Expected one _C_extension*.so file, found {len(so_files)}: {so_files}' + ) + torch.ops.load_library(str(so_files[0])) + + +_load_extension() +_torch_ops = torch.ops.deep_gemm + + +def _bind_guarded_ops(*names): + """Bind ops when all are registered (matches one C++ #if guard group).""" + present = [name for name in names if hasattr(_torch_ops, name)] + if not present: + return + assert len(present) == len(names), ( + f'Guard group mismatch: {sorted(set(names) - set(present))} missing while ' + f'{present} are registered — the C++ #if guards for these ops have diverged.' + ) + globals().update({name: getattr(_torch_ops, name) for name in names}) + + +init = _torch_ops.init +set_num_sms = _torch_ops.set_num_sms +get_num_sms = _torch_ops.get_num_sms +set_tc_util = _torch_ops.set_tc_util +get_tc_util = _torch_ops.get_tc_util +set_pdl = _torch_ops.set_pdl +get_pdl = _torch_ops.get_pdl +set_ignore_compile_dims = _torch_ops.set_ignore_compile_dims +get_mk_alignment_for_contiguous_layout = _torch_ops.get_mk_alignment_for_contiguous_layout +get_theoretical_mk_alignment_for_contiguous_layout = _torch_ops.get_theoretical_mk_alignment_for_contiguous_layout +cublaslt_gemm_nt = _torch_ops.cublaslt_gemm_nt +cublaslt_gemm_nn = _torch_ops.cublaslt_gemm_nn +cublaslt_gemm_tn = _torch_ops.cublaslt_gemm_tn +cublaslt_gemm_tt = _torch_ops.cublaslt_gemm_tt + + +def set_block_size_multiple_of(value): + if isinstance(value, int): + return _torch_ops.set_block_size_multiple_of([value]) + return _torch_ops.set_block_size_multiple_of(list(value)) + + +def set_mk_alignment_for_contiguous_layout(value): + return _torch_ops.set_mk_alignment_for_contiguous_layout(value) + + +def _unpack_ab_pair(a, b): + return a[0], a[1], b[0], b[1] + + +def _unpack_kv(kv): + return kv[0], kv[1] + + +def _as_int_list(value): + """Reject a bare scalar instead of letting int[N] silently broadcast it into a list.""" + return None if value is None else list(value) + + +def _register_deep_gemm_kernels(): + """Export compatibility wrappers for the registered DeepGEMM kernels.""" + def fp8_fp4_gemm_nt(a, b, d, c=None, recipe=None, recipe_a=None, recipe_b=None, + compiled_dims='nk', disable_ue8m0_cast=False, alpha=None): + a_tensor, sfa, b_tensor, sfb = _unpack_ab_pair(a, b) + return _torch_ops.fp8_fp4_gemm_nt( + a_tensor, sfa, b_tensor, sfb, d, c, _as_int_list(recipe), _as_int_list(recipe_a), _as_int_list(recipe_b), + compiled_dims, disable_ue8m0_cast, alpha, + ) + + def fp8_fp4_gemm_nn(a, b, d, c=None, recipe=None, recipe_a=None, recipe_b=None, + compiled_dims='nk', disable_ue8m0_cast=False, alpha=None): + a_tensor, sfa, b_tensor, sfb = _unpack_ab_pair(a, b) + return _torch_ops.fp8_fp4_gemm_nn( + a_tensor, sfa, b_tensor, sfb, d, c, _as_int_list(recipe), _as_int_list(recipe_a), _as_int_list(recipe_b), + compiled_dims, disable_ue8m0_cast, alpha, + ) + + def fp8_fp4_gemm_tn(a, b, d, c=None, recipe=None, recipe_a=None, recipe_b=None, + compiled_dims='mn', disable_ue8m0_cast=False, alpha=None): + a_tensor, sfa, b_tensor, sfb = _unpack_ab_pair(a, b) + return _torch_ops.fp8_fp4_gemm_tn( + a_tensor, sfa, b_tensor, sfb, d, c, _as_int_list(recipe), _as_int_list(recipe_a), _as_int_list(recipe_b), + compiled_dims, disable_ue8m0_cast, alpha, + ) + + def fp8_fp4_gemm_tt(a, b, d, c=None, recipe=None, recipe_a=None, recipe_b=None, + compiled_dims='mn', disable_ue8m0_cast=False, alpha=None): + a_tensor, sfa, b_tensor, sfb = _unpack_ab_pair(a, b) + return _torch_ops.fp8_fp4_gemm_tt( + a_tensor, sfa, b_tensor, sfb, d, c, _as_int_list(recipe), _as_int_list(recipe_a), _as_int_list(recipe_b), + compiled_dims, disable_ue8m0_cast, alpha, + ) + + def m_grouped_fp8_fp4_gemm_nt_contiguous(a, b, d, grouped_layout, recipe=None, recipe_a=None, recipe_b=None, + compiled_dims='nk', disable_ue8m0_cast=False, use_psum_layout=False, + ensure_zero_padding=True, expected_m_for_psum_layout=None): + a_tensor, sfa, b_tensor, sfb = _unpack_ab_pair(a, b) + return _torch_ops.m_grouped_fp8_fp4_gemm_nt_contiguous( + a_tensor, sfa, b_tensor, sfb, d, grouped_layout, _as_int_list(recipe), _as_int_list(recipe_a), _as_int_list(recipe_b), + compiled_dims, disable_ue8m0_cast, use_psum_layout, ensure_zero_padding, + expected_m_for_psum_layout, + ) + + def m_grouped_fp8_fp4_gemm_nn_contiguous(a, b, d, grouped_layout, recipe=None, recipe_a=None, recipe_b=None, + compiled_dims='nk', disable_ue8m0_cast=False, use_psum_layout=False, + ensure_zero_padding=True): + a_tensor, sfa, b_tensor, sfb = _unpack_ab_pair(a, b) + return _torch_ops.m_grouped_fp8_fp4_gemm_nn_contiguous( + a_tensor, sfa, b_tensor, sfb, d, grouped_layout, _as_int_list(recipe), _as_int_list(recipe_a), _as_int_list(recipe_b), + compiled_dims, disable_ue8m0_cast, use_psum_layout, ensure_zero_padding, + ) + + def m_grouped_fp8_fp4_gemm_nt_masked(a, b, d, masked_m, expected_m, recipe=None, recipe_a=None, recipe_b=None, + compiled_dims='nk', disable_ue8m0_cast=False): + a_tensor, sfa, b_tensor, sfb = _unpack_ab_pair(a, b) + return _torch_ops.m_grouped_fp8_fp4_gemm_nt_masked( + a_tensor, sfa, b_tensor, sfb, d, masked_m, expected_m, _as_int_list(recipe), _as_int_list(recipe_a), _as_int_list(recipe_b), + compiled_dims, disable_ue8m0_cast, + ) + + def k_grouped_fp8_gemm_tn_contiguous(a, b, d, ks_cpu, grouped_layout, c=None, recipe=(1, 1, 128), + compiled_dims='mn', use_psum_layout=False): + a_tensor, sfa, b_tensor, sfb = _unpack_ab_pair(a, b) + return _torch_ops.k_grouped_fp8_gemm_tn_contiguous( + a_tensor, sfa, b_tensor, sfb, d, ks_cpu, grouped_layout, c, list(recipe), + compiled_dims, use_psum_layout, + ) + + def k_grouped_fp8_gemm_nt_contiguous(a, b, d, ks_cpu, grouped_layout, c=None, recipe=(1, 1, 128), + compiled_dims='mn', use_psum_layout=False): + a_tensor, sfa, b_tensor, sfb = _unpack_ab_pair(a, b) + return _torch_ops.k_grouped_fp8_gemm_nt_contiguous( + a_tensor, sfa, b_tensor, sfb, d, ks_cpu, grouped_layout, c, list(recipe), + compiled_dims, use_psum_layout, + ) + + def fp8_gemm_nt_skip_head_mid(a, b, d, head_splits, recipe=None, compiled_dims='nk', disable_ue8m0_cast=False): + a_tensor, sfa, b_tensor, sfb = _unpack_ab_pair(a, b) + return _torch_ops.fp8_gemm_nt_skip_head_mid( + a_tensor, sfa, b_tensor, sfb, d, list(head_splits), _as_int_list(recipe), compiled_dims, disable_ue8m0_cast, + ) + + def fp8_einsum(expr, a, b, d, c=None, recipe=(1, 128, 128)): + d_tensor, sfd = d if isinstance(d, (tuple, list)) else (d, None) + return _torch_ops.fp8_einsum(expr, a[0], a[1], b[0], b[1], d_tensor, c, list(recipe), sfd) + + def fp8_fp4_mqa_logits(q, kv, weights, cu_seq_len_k_start, cu_seq_len_k_end, clean_logits=True, + max_seqlen_k=0, logits_dtype=torch.float32, schedule_meta=None): + q_fp, q_sf = q[0], q[1] + kv_fp, kv_sf = _unpack_kv(kv) + return _torch_ops.fp8_fp4_mqa_logits( + q_fp, q_sf, kv_fp, kv_sf, weights, cu_seq_len_k_start, cu_seq_len_k_end, + clean_logits, max_seqlen_k, logits_dtype, schedule_meta, + ) + + def fp8_fp4_paged_mqa_logits(q, kv_cache, weights, context_lens, block_table, schedule_meta, max_context_len, + clean_logits=False, logits_dtype=torch.float32, indices=None): + q_fp, q_sf = q[0], q[1] + return _torch_ops.fp8_fp4_paged_mqa_logits( + q_fp, q_sf, kv_cache, weights, context_lens, block_table, schedule_meta, max_context_len, + clean_logits, logits_dtype, indices, + ) + + def fp8_mqa_logits(q, kv, weights, cu_seq_len_k_start, cu_seq_len_k_end, clean_logits=True, max_seqlen_k=0): + kv_fp, kv_sf = _unpack_kv(kv) + return _torch_ops.fp8_mqa_logits(q, kv_fp, kv_sf, weights, cu_seq_len_k_start, cu_seq_len_k_end, clean_logits, max_seqlen_k) + + def fp8_paged_mqa_logits(q, kv_cache, weights, context_lens, block_table, schedule_meta, max_context_len, + clean_logits=False, indices=None): + return _torch_ops.fp8_paged_mqa_logits( + q, kv_cache, weights, context_lens, block_table, schedule_meta, max_context_len, clean_logits, indices, + ) + + globals().update({ + 'fp8_fp4_gemm_nt': fp8_fp4_gemm_nt, + 'fp8_fp4_gemm_nn': fp8_fp4_gemm_nn, + 'fp8_fp4_gemm_tn': fp8_fp4_gemm_tn, + 'fp8_fp4_gemm_tt': fp8_fp4_gemm_tt, + 'fp8_gemm_nt': fp8_fp4_gemm_nt, + 'fp8_gemm_nn': fp8_fp4_gemm_nn, + 'fp8_gemm_tn': fp8_fp4_gemm_tn, + 'fp8_gemm_tt': fp8_fp4_gemm_tt, + 'm_grouped_fp8_fp4_gemm_nt_contiguous': m_grouped_fp8_fp4_gemm_nt_contiguous, + 'm_grouped_fp8_fp4_gemm_nn_contiguous': m_grouped_fp8_fp4_gemm_nn_contiguous, + 'm_grouped_fp8_fp4_gemm_nt_masked': m_grouped_fp8_fp4_gemm_nt_masked, + 'm_grouped_fp8_gemm_nt_contiguous': m_grouped_fp8_fp4_gemm_nt_contiguous, + 'm_grouped_fp8_gemm_nn_contiguous': m_grouped_fp8_fp4_gemm_nn_contiguous, + 'm_grouped_fp8_gemm_nt_masked': m_grouped_fp8_fp4_gemm_nt_masked, + 'k_grouped_fp8_gemm_tn_contiguous': k_grouped_fp8_gemm_tn_contiguous, + 'k_grouped_fp8_gemm_nt_contiguous': k_grouped_fp8_gemm_nt_contiguous, + 'fp8_gemm_nt_skip_head_mid': fp8_gemm_nt_skip_head_mid, + 'fp8_einsum': fp8_einsum, + 'fp8_fp4_mqa_logits': fp8_fp4_mqa_logits, + 'fp8_fp4_paged_mqa_logits': fp8_fp4_paged_mqa_logits, + 'fp8_mqa_logits': fp8_mqa_logits, + 'fp8_paged_mqa_logits': fp8_paged_mqa_logits, + }) + + # BF16 GEMMs + _bind_guarded_ops( + 'bf16_gemm_nt', + 'bf16_gemm_nn', + 'bf16_gemm_tn', + 'bf16_gemm_tt', + 'm_grouped_bf16_gemm_nt_contiguous', + 'm_grouped_bf16_gemm_nn_contiguous', + 'm_grouped_bf16_gemm_nt_masked', + 'k_grouped_bf16_gemm_tn_contiguous', + ) + + # Einsum, hyperconnection, and attention metadata + _bind_guarded_ops( + 'einsum', # einsum.hpp + 'tf32_hc_prenorm_gemm', # hyperconnection.hpp + 'get_paged_mqa_logits_metadata', # attention.hpp + ) + + # Layout helpers + _bind_guarded_ops( + 'transform_sf_into_required_layout', + 'get_tma_aligned_size', + 'get_mn_major_tma_aligned_tensor', + 'get_mn_major_tma_aligned_packed_ue8m0_tensor', + 'get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor', + ) + + +_register_deep_gemm_kernels() + + +def get_mqa_logits_metadata(cu_seq_len_k_start, cu_seq_len_k_end, num_kv_tokens, num_heads): + return _torch_ops.get_mqa_logits_metadata(cu_seq_len_k_start, cu_seq_len_k_end, num_kv_tokens, num_heads) + + +def 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=False, +): + return _torch_ops.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, + ) + + +def get_paged_sparse_mqa_logits_metadata( + context_lens, block_table, indices, page_kv, sparse_kv_block_indices, qk_dtype, sparse_block_kv +): + return _torch_ops.get_paged_sparse_mqa_logits_metadata( + context_lens, block_table, indices, page_kv, sparse_kv_block_indices, qk_dtype, sparse_block_kv + ) + + +def fp8_fp4_sparse_mqa_logits(q, kv, weights, metadata, num_max_sparse_blocks, sparse_block_kv, use_unaligned_ks=False): + return _torch_ops.fp8_fp4_sparse_mqa_logits( + q[0], q[1], kv[0], kv[1], weights, metadata, num_max_sparse_blocks, sparse_block_kv, use_unaligned_ks + ) + + +def fp8_fp4_paged_sparse_mqa_logits(q, kv_cache, weights, metadata, num_max_sparse_blocks, sparse_block_kv): + return _torch_ops.fp8_fp4_paged_sparse_mqa_logits( + q[0], q[1], kv_cache, weights, metadata, num_max_sparse_blocks, sparse_block_kv + ) + + +def k_grouped_fp4_gemm_nt_contiguous( + a, b, d, ks_cpu, grouped_layout, c=None, recipe=(1, 1, 32), compiled_dims='mn', use_psum_layout=False +): + return _torch_ops.k_grouped_fp4_gemm_nt_contiguous( + a[0], a[1], b[0], b[1], d, ks_cpu, grouped_layout, c, list(recipe), compiled_dims, use_psum_layout + ) + + +def cublaslt_nvfp4_gemm_nt(a, b, d, c=None): + return _torch_ops.cublaslt_nvfp4_gemm_nt(a[0], a[1], b[0], b[1], d, c) + + +def batched_syrk(a, d): + return _torch_ops.batched_syrk(a, d) + + +def batched_symm(a, b, d): + return _torch_ops.batched_symm(a, b, d) + + +def mega_mhc( + x, + residual, + shifted_prev_mix, + post_mix, + comb_res_mix, + fn, + mix_scales, + mix_bases, + hc_mult, + hc_norm_eps, + hc_pre_eps, + hc_post_scale, + sinkhorn_eps, + num_sinkhorn_iters, + rmsnorm_weight, + rmsnorm_eps, + rmsnorm_scale, + new_residual, + new_prev_mix, + new_post_mix, + new_comb_res_mix, + y_bf16=None, + y_fp8=None, + y_gemm_sf=None, + y_routed_sf=None, + y_shared_sf=None, + shared_sf_block_m=0, +): + return _torch_ops.mega_mhc( + x, + residual, + shifted_prev_mix, + post_mix, + comb_res_mix, + fn, + mix_scales, + mix_bases, + hc_mult, + hc_norm_eps, + hc_pre_eps, + hc_post_scale, + sinkhorn_eps, + num_sinkhorn_iters, + rmsnorm_weight, + rmsnorm_eps, + rmsnorm_scale, + new_residual, + new_prev_mix, + new_post_mix, + new_comb_res_mix, + y_bf16, + y_fp8, + y_gemm_sf, + y_routed_sf, + y_shared_sf, + shared_sf_block_m, + ) + + +def get_token_alignment_for_mega_moe(): + return _torch_ops.get_token_alignment_for_mega_moe() + + +def get_block_m_for_mega_moe(num_ranks, num_experts, num_max_tokens_per_rank, num_tokens, num_topk, mma_type): + return _torch_ops.get_block_m_for_mega_moe( + num_ranks, num_experts, num_max_tokens_per_rank, num_tokens, num_topk, mma_type + ) + + +def get_token_alignment_for_nvfp4_mega_moe(): + return _torch_ops.get_token_alignment_for_nvfp4_mega_moe() + + +def get_block_m_for_nvfp4_mega_moe( + num_ranks, num_experts, num_max_tokens_per_rank, num_tokens, num_topk, +): + return _torch_ops.get_block_m_for_nvfp4_mega_moe( + num_ranks, num_experts, num_max_tokens_per_rank, num_tokens, num_topk + ) + + +def nvfp4_mega_moe( + y, + l1_weights_tuple, + l2_weights_tuple, + shared_l1_weights_tuple_opt, + shared_l2_weights_tuple_opt, + cumulative_local_expert_recv_stats, + sym_buffer, + sym_buffer_ptrs, + rank_idx, + num_max_tokens_per_rank, + num_experts, + num_topk, + activation_clamp_opt, + fast_math, + activation_alpha, + activation_beta, + l1_alpha=None, + l2_alpha=None, + l2_activation_scale=1.0, +): + return _torch_ops.nvfp4_mega_moe( + y, + l1_weights_tuple[0], + l1_weights_tuple[1], + l2_weights_tuple[0], + l2_weights_tuple[1], + shared_l1_weights_tuple_opt[0] if shared_l1_weights_tuple_opt is not None else None, + shared_l1_weights_tuple_opt[1] if shared_l1_weights_tuple_opt is not None else None, + shared_l2_weights_tuple_opt[0] if shared_l2_weights_tuple_opt is not None else None, + shared_l2_weights_tuple_opt[1] if shared_l2_weights_tuple_opt is not None else None, + cumulative_local_expert_recv_stats, + sym_buffer, + sym_buffer_ptrs, + rank_idx, + num_max_tokens_per_rank, + num_experts, + num_topk, + activation_clamp_opt, + fast_math, + activation_alpha, + activation_beta, + l1_alpha, + l2_alpha, + l2_activation_scale, + ) + + +def fp8_fp4_mega_moe( + y, + l1_weights_tuple, + l2_weights_tuple, + shared_l1_weights_tuple_opt, + shared_l2_weights_tuple_opt, + cumulative_local_expert_recv_stats, + sym_buffer, + sym_buffer_ptrs, + rank_idx, + num_max_tokens_per_rank, + num_experts, + num_topk, + recipe, + activation, + activation_clamp_opt, + fast_math, + activation_alpha, + activation_beta, +): + return _torch_ops.fp8_fp4_mega_moe( + y, + l1_weights_tuple[0], + l1_weights_tuple[1], + l2_weights_tuple[0], + l2_weights_tuple[1], + shared_l1_weights_tuple_opt[0] if shared_l1_weights_tuple_opt is not None else None, + shared_l1_weights_tuple_opt[1] if shared_l1_weights_tuple_opt is not None else None, + shared_l2_weights_tuple_opt[0] if shared_l2_weights_tuple_opt is not None else None, + shared_l2_weights_tuple_opt[1] if shared_l2_weights_tuple_opt is not None else None, + cumulative_local_expert_recv_stats, + sym_buffer, + sym_buffer_ptrs, + rank_idx, + num_max_tokens_per_rank, + num_experts, + num_topk, + list(recipe), + activation, + activation_clamp_opt, + fast_math, + activation_alpha, + activation_beta, + ) + + +def bf16_mega_moe( + y, + l1_weights, + l2_weights, + shared_l1_weights_opt, + shared_l2_weights_opt, + cumulative_local_expert_recv_stats, + sym_buffer, + sym_buffer_ptrs, + rank_idx, + num_max_tokens_per_rank, + num_experts, + num_topk, + activation, + activation_clamp_opt, + fast_math, + activation_alpha, + activation_beta, +): + return _torch_ops.bf16_mega_moe( + y, + l1_weights, + l2_weights, + shared_l1_weights_opt, + shared_l2_weights_opt, + cumulative_local_expert_recv_stats, + sym_buffer, + sym_buffer_ptrs, + rank_idx, + num_max_tokens_per_rank, + num_experts, + num_topk, + activation, + activation_clamp_opt, + fast_math, + activation_alpha, + activation_beta, + ) + + +use_deterministic_algorithms = _torch_ops.use_deterministic_algorithms +fp4_gemm_nt = fp8_fp4_gemm_nt +m_grouped_fp4_gemm_nt_contiguous = m_grouped_fp8_fp4_gemm_nt_contiguous +m_grouped_fp4_gemm_nt_masked = m_grouped_fp8_fp4_gemm_nt_masked + + +def get_bf16_mega_gate_config(num_tokens, hidden, num_routed_experts, num_topk): + return _torch_ops.get_bf16_mega_gate_config(num_tokens, hidden, num_routed_experts, num_topk) + + +def bf16_mega_gate( + x, + weight, + num_topk, + use_shared_as_routed, + num_shared_experts, + routed_scaling_factor, + ep_rank, + scoring_func='identity', + mask=None, + bias=None, + image_bias=None, + image_token_mask=None, + fix_routing_mask=None, + to_physical_map=None, + logical_count=None, + unmapped_topk_idx=None, + force_random=None, + out=None, +): + if out is None: + shape = (x.shape[0], num_topk + (num_shared_experts if use_shared_as_routed else 0)) + out = ( + torch.empty(shape, dtype=torch.int64, device=x.device), + torch.empty(shape, dtype=torch.float32, device=x.device), + ) + _torch_ops.bf16_mega_gate( + x, + weight, + num_topk, + use_shared_as_routed, + num_shared_experts, + routed_scaling_factor, + ep_rank, + scoring_func, + mask, + bias, + image_bias, + image_token_mask, + fix_routing_mask, + to_physical_map, + logical_count, + unmapped_topk_idx, + force_random, + *out, + ) + return tuple(out) + + +def get_symm_buffer_size_for_mega_moe( + num_ranks, + num_experts, + num_max_tokens_per_rank, + num_topk, + hidden, + intermediate_hidden, + mma_type, + activation, + num_shared_experts, +): + num_bytes, layout_info = _torch_ops.get_symm_buffer_size_for_mega_moe( + num_ranks, + num_experts, + num_max_tokens_per_rank, + num_topk, + hidden, + intermediate_hidden, + mma_type, + activation, + num_shared_experts, + ) + + def slice_input_buffers(buffer): + return _torch_ops._slice_symm_buffer_for_mega_moe(buffer, layout_info) + + return num_bytes, slice_input_buffers + + +def get_symm_buffer_size_for_nvfp4_mega_moe( + num_ranks, + num_experts, + num_max_tokens_per_rank, + num_topk, + hidden, + intermediate_hidden, + num_shared_experts=0, + shared_bf16=False, +): + num_bytes, layout_info = _torch_ops.get_symm_buffer_size_for_nvfp4_mega_moe( + num_ranks, + num_experts, + num_max_tokens_per_rank, + num_topk, + hidden, + intermediate_hidden, + num_shared_experts, + shared_bf16, + ) + + def slice_input_buffers(buffer): + return _torch_ops._slice_symm_buffer_for_nvfp4_mega_moe(buffer, layout_info) + + return num_bytes, slice_input_buffers + + +Runtime = torch.classes.deep_gemm.Runtime +get_jit = _torch_ops.get_jit + + +_PUBLIC_API = [ + # Runtime and configuration + 'Runtime', 'get_jit', 'init', + 'set_num_sms', 'get_num_sms', + 'set_tc_util', 'get_tc_util', + 'set_pdl', 'get_pdl', + 'use_deterministic_algorithms', + 'set_ignore_compile_dims', + 'set_block_size_multiple_of', + 'set_mk_alignment_for_contiguous_layout', + 'get_mk_alignment_for_contiguous_layout', + 'get_theoretical_mk_alignment_for_contiguous_layout', + # cuBLASLt kernels + 'cublaslt_gemm_nt', 'cublaslt_gemm_nn', + 'cublaslt_gemm_tn', 'cublaslt_gemm_tt', + 'cublaslt_nvfp4_gemm_nt', + 'batched_syrk', 'batched_symm', + # Mega kernels + 'mega_mhc', + 'get_token_alignment_for_mega_moe', + 'get_block_m_for_mega_moe', + 'get_symm_buffer_size_for_mega_moe', + 'fp8_fp4_mega_moe', 'bf16_mega_moe', + 'get_token_alignment_for_nvfp4_mega_moe', + 'get_block_m_for_nvfp4_mega_moe', + 'get_symm_buffer_size_for_nvfp4_mega_moe', + 'nvfp4_mega_moe', + 'get_bf16_mega_gate_config', 'bf16_mega_gate', + # FP8/FP4 GEMMs + 'fp8_fp4_gemm_nt', 'fp8_fp4_gemm_nn', + 'fp8_fp4_gemm_tn', 'fp8_fp4_gemm_tt', + 'fp8_gemm_nt', 'fp8_gemm_nn', + 'fp8_gemm_tn', 'fp8_gemm_tt', + 'fp4_gemm_nt', + 'm_grouped_fp8_fp4_gemm_nt_contiguous', + 'm_grouped_fp8_fp4_gemm_nn_contiguous', + 'm_grouped_fp8_fp4_gemm_nt_masked', + 'm_grouped_fp8_gemm_nt_contiguous', + 'm_grouped_fp8_gemm_nn_contiguous', + 'm_grouped_fp8_gemm_nt_masked', + 'm_grouped_fp4_gemm_nt_contiguous', + 'm_grouped_fp4_gemm_nt_masked', + 'k_grouped_fp8_gemm_tn_contiguous', + 'k_grouped_fp8_gemm_nt_contiguous', + 'k_grouped_fp4_gemm_nt_contiguous', + 'fp8_gemm_nt_skip_head_mid', + # BF16 GEMMs + 'bf16_gemm_nt', 'bf16_gemm_nn', + 'bf16_gemm_tn', 'bf16_gemm_tt', + 'm_grouped_bf16_gemm_nt_contiguous', + 'm_grouped_bf16_gemm_nn_contiguous', + 'm_grouped_bf16_gemm_nt_masked', + 'k_grouped_bf16_gemm_tn_contiguous', + # Einsum + 'einsum', 'fp8_einsum', + # Attention + 'fp8_fp4_mqa_logits', + 'get_mqa_logits_metadata', + 'get_paged_mqa_logits_metadata', + 'get_sparse_mqa_logits_metadata', + 'get_paged_sparse_mqa_logits_metadata', + 'fp8_fp4_sparse_mqa_logits', + 'fp8_fp4_paged_sparse_mqa_logits', + 'fp8_fp4_paged_mqa_logits', + 'fp8_mqa_logits', 'fp8_paged_mqa_logits', + # Hyperconnection + 'tf32_hc_prenorm_gemm', + # Layout + 'transform_sf_into_required_layout', + 'get_tma_aligned_size', + 'get_mn_major_tma_aligned_tensor', + 'get_mn_major_tma_aligned_packed_ue8m0_tensor', + 'get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor', +] + +__all__ = _PUBLIC_API diff --git a/deepjit_main.txt b/deepjit_main.txt new file mode 100644 index 0000000000..fff9238646 --- /dev/null +++ b/deepjit_main.txt @@ -0,0 +1,227 @@ +Using /home/dev/.venv/bin/python + +{"description": "bf16_gemm_nt m=4096 n=7168 k=2048", "first_call_s": 0.06649942696094513, "warm_s": 0.000301715576171875, "tflops": 398.58427534246135, "gbps": 347.5379074902879} +{"description": "bf16_gemm_nt m=4096 n=7168 k=2048", "first_call_s": 0.06561460718512535, "warm_s": 0.00030224401855468753, "tflops": 397.8873919923101, "gbps": 346.9302734307949} +bf16_gemm_nt m=4096 n=7168 k=2048 (trial 1/5) + cold first call: 66.499 ms + disk-cache first call: 65.615 ms + warm call: 302.2 us + throughput: 397.89 TFLOP/s + bandwidth: 346.93 GB/s + +{"description": "bf16_gemm_nt m=4096 n=7168 k=2048", "first_call_s": 0.06516649201512337, "warm_s": 0.00030251123046875, "tflops": 397.53593313430065, "gbps": 346.6238256263085} +{"description": "bf16_gemm_nt m=4096 n=7168 k=2048", "first_call_s": 0.06440368853509426, "warm_s": 0.00030269241333007814, "tflops": 397.29797970476596, "gbps": 346.41634670128167} +bf16_gemm_nt m=4096 n=7168 k=2048 (trial 2/5) + cold first call: 65.166 ms + disk-cache first call: 64.404 ms + warm call: 302.7 us + throughput: 397.30 TFLOP/s + bandwidth: 346.42 GB/s + +{"description": "bf16_gemm_nt m=4096 n=7168 k=2048", "first_call_s": 0.06394599564373493, "warm_s": 0.00030183917236328124, "tflops": 398.4210642588865, "gbps": 347.39559871903475} +{"description": "bf16_gemm_nt m=4096 n=7168 k=2048", "first_call_s": 0.06866894848644733, "warm_s": 0.0003019266662597656, "tflops": 398.30560770850855, "gbps": 347.29492859628607} +bf16_gemm_nt m=4096 n=7168 k=2048 (trial 3/5) + cold first call: 63.946 ms + disk-cache first call: 68.669 ms + warm call: 301.9 us + throughput: 398.31 TFLOP/s + bandwidth: 347.29 GB/s + +{"description": "bf16_gemm_nt m=4096 n=7168 k=2048", "first_call_s": 0.06693107075989246, "warm_s": 0.00030319757080078127, "tflops": 396.63604154341107, "gbps": 345.839182428337} +{"description": "bf16_gemm_nt m=4096 n=7168 k=2048", "first_call_s": 0.06670908443629742, "warm_s": 0.0003029931640625, "tflops": 396.9036221001789, "gbps": 346.07249415821957} +bf16_gemm_nt m=4096 n=7168 k=2048 (trial 4/5) + cold first call: 66.931 ms + disk-cache first call: 66.709 ms + warm call: 303.0 us + throughput: 396.90 TFLOP/s + bandwidth: 346.07 GB/s + +{"description": "bf16_gemm_nt m=4096 n=7168 k=2048", "first_call_s": 0.06635340861976147, "warm_s": 0.0003020225219726563, "tflops": 398.1791937320744, "gbps": 347.1847043562311} +{"description": "bf16_gemm_nt m=4096 n=7168 k=2048", "first_call_s": 0.06536657176911831, "warm_s": 0.0003014137573242188, "tflops": 398.98339530216623, "gbps": 347.8859124774748} +bf16_gemm_nt m=4096 n=7168 k=2048 (trial 5/5) + cold first call: 66.353 ms + disk-cache first call: 65.367 ms + warm call: 301.4 us + throughput: 398.98 TFLOP/s + bandwidth: 347.89 GB/s + +{"description": "fp8_gemm_nt m=4096 n=7168 k=2048", "first_call_s": 4.181794486939907, "warm_s": 0.00017804121398925783, "tflops": 675.4564383910339, "gbps": 460.8745029392509} +{"description": "fp8_gemm_nt m=4096 n=7168 k=2048", "first_call_s": 0.010474476963281631, "warm_s": 0.00017821011352539063, "tflops": 674.8162711364077, "gbps": 460.437706798886} +fp8_gemm_nt m=4096 n=7168 k=2048 (trial 1/5) + cold first call: 4181.794 ms + disk-cache first call: 10.474 ms + warm call: 178.2 us + throughput: 674.82 TFLOP/s + bandwidth: 460.44 GB/s + +{"description": "fp8_gemm_nt m=4096 n=7168 k=2048", "first_call_s": 4.260368146002293, "warm_s": 0.00017857107543945314, "tflops": 673.4522037908396, "gbps": 459.50698229300696} +{"description": "fp8_gemm_nt m=4096 n=7168 k=2048", "first_call_s": 0.009816871955990791, "warm_s": 0.00017852511596679686, "tflops": 673.6255772009496, "gbps": 459.62527768507937} +fp8_gemm_nt m=4096 n=7168 k=2048 (trial 2/5) + cold first call: 4260.368 ms + disk-cache first call: 9.817 ms + warm call: 178.5 us + throughput: 673.63 TFLOP/s + bandwidth: 459.63 GB/s + +{"description": "fp8_gemm_nt m=4096 n=7168 k=2048", "first_call_s": 4.1381640918552876, "warm_s": 0.0001778419189453125, "tflops": 676.2133753458902, "gbps": 461.390972874243} +{"description": "fp8_gemm_nt m=4096 n=7168 k=2048", "first_call_s": 0.008485311642289162, "warm_s": 0.00017830213928222657, "tflops": 674.4679832340497, "gbps": 460.20006451027103} +fp8_gemm_nt m=4096 n=7168 k=2048 (trial 3/5) + cold first call: 4138.164 ms + disk-cache first call: 8.485 ms + warm call: 178.3 us + throughput: 674.47 TFLOP/s + bandwidth: 460.20 GB/s + +{"description": "fp8_gemm_nt m=4096 n=7168 k=2048", "first_call_s": 4.143187824636698, "warm_s": 0.00017789657592773436, "tflops": 676.0056154023558, "gbps": 461.24921501205546} +{"description": "fp8_gemm_nt m=4096 n=7168 k=2048", "first_call_s": 0.009114265441894531, "warm_s": 0.00017864364624023437, "tflops": 673.1786258229378, "gbps": 459.3203157623388} +fp8_gemm_nt m=4096 n=7168 k=2048 (trial 4/5) + cold first call: 4143.188 ms + disk-cache first call: 9.114 ms + warm call: 178.6 us + throughput: 673.18 TFLOP/s + bandwidth: 459.32 GB/s + +{"description": "fp8_gemm_nt m=4096 n=7168 k=2048", "first_call_s": 4.193111157044768, "warm_s": 0.0001781525115966797, "tflops": 675.0344590159643, "gbps": 460.58657980508247} +{"description": "fp8_gemm_nt m=4096 n=7168 k=2048", "first_call_s": 0.009256750345230103, "warm_s": 0.00017856927490234376, "tflops": 673.4589942965691, "gbps": 459.5116155613791} +fp8_gemm_nt m=4096 n=7168 k=2048 (trial 5/5) + cold first call: 4193.111 ms + disk-cache first call: 9.257 ms + warm call: 178.6 us + throughput: 673.46 TFLOP/s + bandwidth: 459.51 GB/s + +{"description": "grouped_fp8_nt groups=4 m=4480 n=4096 k=2048", "first_call_s": 5.832242663949728, "warm_s": 0.0002526413421630859, "tflops": 297.50446635721727, "gbps": 315.63505528134965} +{"description": "grouped_fp8_nt groups=4 m=4480 n=4096 k=2048", "first_call_s": 0.002241833135485649, "warm_s": 0.0002523389434814453, "tflops": 297.8609906303532, "gbps": 316.0133069427055} +grouped_fp8_nt groups=4 m=4480 n=4096 k=2048 (trial 1/5) + cold first call: 5832.243 ms + disk-cache first call: 2.242 ms + warm call: 252.3 us + throughput: 297.86 TFLOP/s + bandwidth: 316.01 GB/s + +{"description": "grouped_fp8_nt groups=4 m=4480 n=4096 k=2048", "first_call_s": 5.70818018913269, "warm_s": 0.00025179936218261716, "tflops": 298.49927747429683, "gbps": 316.6904924173989} +{"description": "grouped_fp8_nt groups=4 m=4480 n=4096 k=2048", "first_call_s": 0.00203133188188076, "warm_s": 0.00025269235229492186, "tflops": 297.44441015878925, "gbps": 315.5713391235961} +grouped_fp8_nt groups=4 m=4480 n=4096 k=2048 (trial 2/5) + cold first call: 5708.180 ms + disk-cache first call: 2.031 ms + warm call: 252.7 us + throughput: 297.44 TFLOP/s + bandwidth: 315.57 GB/s + +{"description": "grouped_fp8_nt groups=4 m=4480 n=4096 k=2048", "first_call_s": 5.686355398967862, "warm_s": 0.0002534339141845703, "tflops": 296.57407108214113, "gbps": 314.6479596330795} +{"description": "grouped_fp8_nt groups=4 m=4480 n=4096 k=2048", "first_call_s": 0.0020052585750818253, "warm_s": 0.00025268537902832034, "tflops": 297.4526186241114, "gbps": 315.58004783119117} +grouped_fp8_nt groups=4 m=4480 n=4096 k=2048 (trial 3/5) + cold first call: 5686.355 ms + disk-cache first call: 2.005 ms + warm call: 252.7 us + throughput: 297.45 TFLOP/s + bandwidth: 315.58 GB/s + +{"description": "grouped_fp8_nt groups=4 m=4480 n=4096 k=2048", "first_call_s": 5.885648449882865, "warm_s": 0.00025210426330566407, "tflops": 298.13826507515205, "gbps": 316.3074791135768} +{"description": "grouped_fp8_nt groups=4 m=4480 n=4096 k=2048", "first_call_s": 0.00233672559261322, "warm_s": 0.00025262342834472655, "tflops": 297.5255627416909, "gbps": 315.6574373267728} +grouped_fp8_nt groups=4 m=4480 n=4096 k=2048 (trial 4/5) + cold first call: 5885.648 ms + disk-cache first call: 2.337 ms + warm call: 252.6 us + throughput: 297.53 TFLOP/s + bandwidth: 315.66 GB/s + +{"description": "grouped_fp8_nt groups=4 m=4480 n=4096 k=2048", "first_call_s": 5.876402880996466, "warm_s": 0.00025255162048339844, "tflops": 297.61015801892586, "gbps": 315.7471880297909} +{"description": "grouped_fp8_nt groups=4 m=4480 n=4096 k=2048", "first_call_s": 0.002096550539135933, "warm_s": 0.0002526886444091797, "tflops": 297.4487747786956, "gbps": 315.57596973322126} +grouped_fp8_nt groups=4 m=4480 n=4096 k=2048 (trial 5/5) + cold first call: 5876.403 ms + disk-cache first call: 2.097 ms + warm call: 252.7 us + throughput: 297.45 TFLOP/s + bandwidth: 315.58 GB/s + +{"description": "tf32_hc_prenorm_gemm m=4096 n=24 k=7168 splits=16", "first_call_s": 1.8277694415301085, "warm_s": 6.259897613525391e-05, "tflops": 22.5129264247875, "gbps": 1053.7230490396494} +{"description": "tf32_hc_prenorm_gemm m=4096 n=24 k=7168 splits=16", "first_call_s": 0.007721079513430595, "warm_s": 6.246540832519532e-05, "tflops": 22.561065104437443, "gbps": 1055.9761917604299} +tf32_hc_prenorm_gemm m=4096 n=24 k=7168 splits=16 (trial 1/5) + cold first call: 1827.769 ms + disk-cache first call: 7.721 ms + warm call: 62.5 us + throughput: 22.56 TFLOP/s + bandwidth: 1055.98 GB/s + +{"description": "tf32_hc_prenorm_gemm m=4096 n=24 k=7168 splits=16", "first_call_s": 1.8242545556277037, "warm_s": 6.245753479003906e-05, "tflops": 22.563909202269024, "gbps": 1056.1093104577646} +{"description": "tf32_hc_prenorm_gemm m=4096 n=24 k=7168 splits=16", "first_call_s": 0.009904064238071442, "warm_s": 6.237184143066407e-05, "tflops": 22.59491000544916, "gbps": 1057.5603106624153} +tf32_hc_prenorm_gemm m=4096 n=24 k=7168 splits=16 (trial 2/5) + cold first call: 1824.255 ms + disk-cache first call: 9.904 ms + warm call: 62.4 us + throughput: 22.59 TFLOP/s + bandwidth: 1057.56 GB/s + +{"description": "tf32_hc_prenorm_gemm m=4096 n=24 k=7168 splits=16", "first_call_s": 1.7698302939534187, "warm_s": 6.237913513183594e-05, "tflops": 22.59226808806385, "gbps": 1057.436655070511} +{"description": "tf32_hc_prenorm_gemm m=4096 n=24 k=7168 splits=16", "first_call_s": 0.008888935670256615, "warm_s": 6.251801681518556e-05, "tflops": 22.54208012013724, "gbps": 1055.0875949087674} +tf32_hc_prenorm_gemm m=4096 n=24 k=7168 splits=16 (trial 3/5) + cold first call: 1769.830 ms + disk-cache first call: 8.889 ms + warm call: 62.5 us + throughput: 22.54 TFLOP/s + bandwidth: 1055.09 GB/s + +{"description": "tf32_hc_prenorm_gemm m=4096 n=24 k=7168 splits=16", "first_call_s": 1.7741438802331686, "warm_s": 6.251232147216796e-05, "tflops": 22.544133873310866, "gbps": 1055.1837213303286} +{"description": "tf32_hc_prenorm_gemm m=4096 n=24 k=7168 splits=16", "first_call_s": 0.009329654276371002, "warm_s": 6.241676712036133e-05, "tflops": 22.578646876766367, "gbps": 1056.7991109312384} +tf32_hc_prenorm_gemm m=4096 n=24 k=7168 splits=16 (trial 4/5) + cold first call: 1774.144 ms + disk-cache first call: 9.330 ms + warm call: 62.4 us + throughput: 22.58 TFLOP/s + bandwidth: 1056.80 GB/s + +{"description": "tf32_hc_prenorm_gemm m=4096 n=24 k=7168 splits=16", "first_call_s": 1.8082384653389454, "warm_s": 6.237747192382812e-05, "tflops": 22.59287047928043, "gbps": 1057.4648501393115} +{"description": "tf32_hc_prenorm_gemm m=4096 n=24 k=7168 splits=16", "first_call_s": 0.007910523563623428, "warm_s": 6.238182449340821e-05, "tflops": 22.591294106008668, "gbps": 1057.3910676012706} +tf32_hc_prenorm_gemm m=4096 n=24 k=7168 splits=16 (trial 5/5) + cold first call: 1808.238 ms + disk-cache first call: 7.911 ms + warm call: 62.4 us + throughput: 22.59 TFLOP/s + bandwidth: 1057.39 GB/s + +{"description": "fp8_mqa_logits seq=512 kv=4096 h=32 d=128", "first_call_s": 1.756322830915451, "warm_s": 5.014611053466797e-05, "tflops": 321.14215910856444, "gbps": 53.99134591162817} +{"description": "fp8_mqa_logits seq=512 kv=4096 h=32 d=128", "first_call_s": 0.013768279924988747, "warm_s": 5.049683380126953e-05, "tflops": 318.91168209431646, "gbps": 53.616351683656106} +fp8_mqa_logits seq=512 kv=4096 h=32 d=128 (trial 1/5) + cold first call: 1756.323 ms + disk-cache first call: 13.768 ms + warm call: 50.5 us + throughput: 318.91 TFLOP/s + bandwidth: 53.62 GB/s + +{"description": "fp8_mqa_logits seq=512 kv=4096 h=32 d=128", "first_call_s": 1.7796995379030704, "warm_s": 5.023462295532226e-05, "tflops": 320.5763129211226, "gbps": 53.896214218786135} +{"description": "fp8_mqa_logits seq=512 kv=4096 h=32 d=128", "first_call_s": 0.012739954516291618, "warm_s": 5.0546241760253906e-05, "tflops": 318.59995218602194, "gbps": 53.56394275249476} +fp8_mqa_logits seq=512 kv=4096 h=32 d=128 (trial 2/5) + cold first call: 1779.700 ms + disk-cache first call: 12.740 ms + warm call: 50.5 us + throughput: 318.60 TFLOP/s + bandwidth: 53.56 GB/s + +{"description": "fp8_mqa_logits seq=512 kv=4096 h=32 d=128", "first_call_s": 1.8266457580029964, "warm_s": 4.98721923828125e-05, "tflops": 322.906001091501, "gbps": 54.287888112435844} +{"description": "fp8_mqa_logits seq=512 kv=4096 h=32 d=128", "first_call_s": 0.0115666463971138, "warm_s": 4.9816574096679686e-05, "tflops": 323.26651320395285, "gbps": 54.348498448440154} +fp8_mqa_logits seq=512 kv=4096 h=32 d=128 (trial 3/5) + cold first call: 1826.646 ms + disk-cache first call: 11.567 ms + warm call: 49.8 us + throughput: 323.27 TFLOP/s + bandwidth: 54.35 GB/s + +{"description": "fp8_mqa_logits seq=512 kv=4096 h=32 d=128", "first_call_s": 1.5486265439540148, "warm_s": 5.0309185028076176e-05, "tflops": 320.10119422552333, "gbps": 53.816335893516126} +{"description": "fp8_mqa_logits seq=512 kv=4096 h=32 d=128", "first_call_s": 0.01149073801934719, "warm_s": 5.05354232788086e-05, "tflops": 318.66815716873646, "gbps": 53.57540957088091} +fp8_mqa_logits seq=512 kv=4096 h=32 d=128 (trial 4/5) + cold first call: 1548.627 ms + disk-cache first call: 11.491 ms + warm call: 50.5 us + throughput: 318.67 TFLOP/s + bandwidth: 53.58 GB/s + +{"description": "fp8_mqa_logits seq=512 kv=4096 h=32 d=128", "first_call_s": 1.574046803638339, "warm_s": 5.008639907836914e-05, "tflops": 321.52501486086794, "gbps": 54.055712724799804} +{"description": "fp8_mqa_logits seq=512 kv=4096 h=32 d=128", "first_call_s": 0.011414902284741402, "warm_s": 5.030559921264648e-05, "tflops": 320.1240112442903, "gbps": 53.820171956511835} +fp8_mqa_logits seq=512 kv=4096 h=32 d=128 (trial 5/5) + cold first call: 1574.047 ms + disk-cache first call: 11.415 ms + warm call: 50.3 us + throughput: 320.12 TFLOP/s + bandwidth: 53.82 GB/s + diff --git a/deepjit_torch_library.txt b/deepjit_torch_library.txt new file mode 100644 index 0000000000..13a7705f64 --- /dev/null +++ b/deepjit_torch_library.txt @@ -0,0 +1,227 @@ +Using /home/dev/.venv/bin/python + +{"description": "bf16_gemm_nt m=4096 n=7168 k=2048", "first_call_s": 0.06517756544053555, "warm_s": 0.00030330630493164063, "tflops": 396.4938490649051, "gbps": 345.7152004262914} +{"description": "bf16_gemm_nt m=4096 n=7168 k=2048", "first_call_s": 0.0673441793769598, "warm_s": 0.000304598388671875, "tflops": 394.8119516073595, "gbps": 344.24870222460896} +bf16_gemm_nt m=4096 n=7168 k=2048 (trial 1/5) + cold first call: 65.178 ms + disk-cache first call: 67.344 ms + warm call: 304.6 us + throughput: 394.81 TFLOP/s + bandwidth: 344.25 GB/s + +{"description": "bf16_gemm_nt m=4096 n=7168 k=2048", "first_call_s": 0.07014842890202999, "warm_s": 0.0003013708801269531, "tflops": 399.0401602083805, "gbps": 347.9354075477648} +{"description": "bf16_gemm_nt m=4096 n=7168 k=2048", "first_call_s": 0.06667943298816681, "warm_s": 0.00030103283691406253, "tflops": 399.48826022036593, "gbps": 348.32611975129566} +bf16_gemm_nt m=4096 n=7168 k=2048 (trial 2/5) + cold first call: 70.148 ms + disk-cache first call: 66.679 ms + warm call: 301.0 us + throughput: 399.49 TFLOP/s + bandwidth: 348.33 GB/s + +{"description": "bf16_gemm_nt m=4096 n=7168 k=2048", "first_call_s": 0.06572775915265083, "warm_s": 0.0003022747497558594, "tflops": 397.8469401930879, "gbps": 346.89500226099324} +{"description": "bf16_gemm_nt m=4096 n=7168 k=2048", "first_call_s": 0.06833427399396896, "warm_s": 0.0003026773681640625, "tflops": 397.31772817191626, "gbps": 346.43356599811335} +bf16_gemm_nt m=4096 n=7168 k=2048 (trial 3/5) + cold first call: 65.728 ms + disk-cache first call: 68.334 ms + warm call: 302.7 us + throughput: 397.32 TFLOP/s + bandwidth: 346.43 GB/s + +{"description": "bf16_gemm_nt m=4096 n=7168 k=2048", "first_call_s": 0.06797941401600838, "warm_s": 0.0003017897033691406, "tflops": 398.4863729459401, "gbps": 347.4525433750175} +{"description": "bf16_gemm_nt m=4096 n=7168 k=2048", "first_call_s": 0.06298637948930264, "warm_s": 0.00030273855590820314, "tflops": 397.23742463931535, "gbps": 346.363546874403} +bf16_gemm_nt m=4096 n=7168 k=2048 (trial 4/5) + cold first call: 67.979 ms + disk-cache first call: 62.986 ms + warm call: 302.7 us + throughput: 397.24 TFLOP/s + bandwidth: 346.36 GB/s + +{"description": "bf16_gemm_nt m=4096 n=7168 k=2048", "first_call_s": 0.06577018648386002, "warm_s": 0.0003020613708496094, "tflops": 398.12798289879544, "gbps": 347.1400520532187} +{"description": "bf16_gemm_nt m=4096 n=7168 k=2048", "first_call_s": 0.06585092470049858, "warm_s": 0.00030162841796875, "tflops": 398.69944979938646, "gbps": 347.6383316470655} +bf16_gemm_nt m=4096 n=7168 k=2048 (trial 5/5) + cold first call: 65.770 ms + disk-cache first call: 65.851 ms + warm call: 301.6 us + throughput: 398.70 TFLOP/s + bandwidth: 347.64 GB/s + +{"description": "fp8_gemm_nt m=4096 n=7168 k=2048", "first_call_s": 4.044099438935518, "warm_s": 0.00017790074157714844, "tflops": 675.9897863373911, "gbps": 461.23841459320823} +{"description": "fp8_gemm_nt m=4096 n=7168 k=2048", "first_call_s": 0.009746948257088661, "warm_s": 0.00017849063110351562, "tflops": 673.7557234489007, "gbps": 459.71407850764115} +fp8_gemm_nt m=4096 n=7168 k=2048 (trial 1/5) + cold first call: 4044.099 ms + disk-cache first call: 9.747 ms + warm call: 178.5 us + throughput: 673.76 TFLOP/s + bandwidth: 459.71 GB/s + +{"description": "fp8_gemm_nt m=4096 n=7168 k=2048", "first_call_s": 4.082718379795551, "warm_s": 0.00017787826538085937, "tflops": 676.0752025016121, "gbps": 461.29669537934177} +{"description": "fp8_gemm_nt m=4096 n=7168 k=2048", "first_call_s": 0.010123221203684807, "warm_s": 0.00017828665161132812, "tflops": 674.5265739252848, "gbps": 460.2400418562034} +fp8_gemm_nt m=4096 n=7168 k=2048 (trial 2/5) + cold first call: 4082.718 ms + disk-cache first call: 10.123 ms + warm call: 178.3 us + throughput: 674.53 TFLOP/s + bandwidth: 460.24 GB/s + +{"description": "fp8_gemm_nt m=4096 n=7168 k=2048", "first_call_s": 4.202683877199888, "warm_s": 0.00017806866455078124, "tflops": 675.3523119375378, "gbps": 460.80345582981465} +{"description": "fp8_gemm_nt m=4096 n=7168 k=2048", "first_call_s": 0.00910363532602787, "warm_s": 0.00017734188842773437, "tflops": 678.120016394236, "gbps": 462.69190391212464} +fp8_gemm_nt m=4096 n=7168 k=2048 (trial 3/5) + cold first call: 4202.684 ms + disk-cache first call: 9.104 ms + warm call: 177.3 us + throughput: 678.12 TFLOP/s + bandwidth: 462.69 GB/s + +{"description": "fp8_gemm_nt m=4096 n=7168 k=2048", "first_call_s": 4.061393616721034, "warm_s": 0.00017800653076171873, "tflops": 675.5880459744478, "gbps": 460.964300853878} +{"description": "fp8_gemm_nt m=4096 n=7168 k=2048", "first_call_s": 0.008923590183258057, "warm_s": 0.0001780452423095703, "tflops": 675.441156011928, "gbps": 460.86407553272426} +fp8_gemm_nt m=4096 n=7168 k=2048 (trial 4/5) + cold first call: 4061.394 ms + disk-cache first call: 8.924 ms + warm call: 178.0 us + throughput: 675.44 TFLOP/s + bandwidth: 460.86 GB/s + +{"description": "fp8_gemm_nt m=4096 n=7168 k=2048", "first_call_s": 4.094407599419355, "warm_s": 0.00017833222961425782, "tflops": 674.3541789845102, "gbps": 460.12241408907755} +{"description": "fp8_gemm_nt m=4096 n=7168 k=2048", "first_call_s": 0.009091893211007118, "warm_s": 0.0001779968566894531, "tflops": 675.6247639687996, "gbps": 460.9893541162854} +fp8_gemm_nt m=4096 n=7168 k=2048 (trial 5/5) + cold first call: 4094.408 ms + disk-cache first call: 9.092 ms + warm call: 178.0 us + throughput: 675.62 TFLOP/s + bandwidth: 460.99 GB/s + +{"description": "grouped_fp8_nt groups=4 m=4480 n=4096 k=2048", "first_call_s": 5.702284090220928, "warm_s": 0.0002529977569580078, "tflops": 297.0853519957304, "gbps": 315.1903991513867} +{"description": "grouped_fp8_nt groups=4 m=4480 n=4096 k=2048", "first_call_s": 0.00243414007127285, "warm_s": 0.0002533164215087891, "tflops": 296.71162742756565, "gbps": 314.7938989704749} +grouped_fp8_nt groups=4 m=4480 n=4096 k=2048 (trial 1/5) + cold first call: 5702.284 ms + disk-cache first call: 2.434 ms + warm call: 253.3 us + throughput: 296.71 TFLOP/s + bandwidth: 314.79 GB/s + +{"description": "grouped_fp8_nt groups=4 m=4480 n=4096 k=2048", "first_call_s": 6.077576534822583, "warm_s": 0.00025338925170898435, "tflops": 296.62634532866025, "gbps": 314.70341958933454} +{"description": "grouped_fp8_nt groups=4 m=4480 n=4096 k=2048", "first_call_s": 0.0023938026279211044, "warm_s": 0.00025203514099121096, "tflops": 298.2200314781544, "gbps": 316.3942285444267} +grouped_fp8_nt groups=4 m=4480 n=4096 k=2048 (trial 2/5) + cold first call: 6077.577 ms + disk-cache first call: 2.394 ms + warm call: 252.0 us + throughput: 298.22 TFLOP/s + bandwidth: 316.39 GB/s + +{"description": "grouped_fp8_nt groups=4 m=4480 n=4096 k=2048", "first_call_s": 6.167798668146133, "warm_s": 0.00025283853149414066, "tflops": 297.27244196457383, "gbps": 315.3888908022232} +{"description": "grouped_fp8_nt groups=4 m=4480 n=4096 k=2048", "first_call_s": 0.002321135252714157, "warm_s": 0.00025162509155273436, "tflops": 298.70601225096, "gbps": 316.90982607467015} +grouped_fp8_nt groups=4 m=4480 n=4096 k=2048 (trial 3/5) + cold first call: 6167.799 ms + disk-cache first call: 2.321 ms + warm call: 251.6 us + throughput: 298.71 TFLOP/s + bandwidth: 316.91 GB/s + +{"description": "grouped_fp8_nt groups=4 m=4480 n=4096 k=2048", "first_call_s": 6.098514148965478, "warm_s": 0.0002517086029052734, "tflops": 298.6069081964831, "gbps": 316.80468239700906} +{"description": "grouped_fp8_nt groups=4 m=4480 n=4096 k=2048", "first_call_s": 0.002363240346312523, "warm_s": 0.00025327122497558597, "tflops": 296.76457594914393, "gbps": 314.8500742936224} +grouped_fp8_nt groups=4 m=4480 n=4096 k=2048 (trial 4/5) + cold first call: 6098.514 ms + disk-cache first call: 2.363 ms + warm call: 253.3 us + throughput: 296.76 TFLOP/s + bandwidth: 314.85 GB/s + +{"description": "grouped_fp8_nt groups=4 m=4480 n=4096 k=2048", "first_call_s": 5.893278189003468, "warm_s": 0.00025232019042968753, "tflops": 297.88312838541907, "gbps": 316.0367938221786} +{"description": "grouped_fp8_nt groups=4 m=4480 n=4096 k=2048", "first_call_s": 0.0023332349956035614, "warm_s": 0.0002533895111083984, "tflops": 296.6260416669189, "gbps": 314.7030974217662} +grouped_fp8_nt groups=4 m=4480 n=4096 k=2048 (trial 5/5) + cold first call: 5893.278 ms + disk-cache first call: 2.333 ms + warm call: 253.4 us + throughput: 296.63 TFLOP/s + bandwidth: 314.70 GB/s + +{"description": "tf32_hc_prenorm_gemm m=4096 n=24 k=7168 splits=16", "first_call_s": 1.7681599371135235, "warm_s": 6.237580871582032e-05, "tflops": 22.593472902621688, "gbps": 1057.4930467117154} +{"description": "tf32_hc_prenorm_gemm m=4096 n=24 k=7168 splits=16", "first_call_s": 0.007659167051315308, "warm_s": 6.233484649658203e-05, "tflops": 22.608319795530008, "gbps": 1058.1879591797317} +tf32_hc_prenorm_gemm m=4096 n=24 k=7168 splits=16 (trial 1/5) + cold first call: 1768.160 ms + disk-cache first call: 7.659 ms + warm call: 62.3 us + throughput: 22.61 TFLOP/s + bandwidth: 1058.19 GB/s + +{"description": "tf32_hc_prenorm_gemm m=4096 n=24 k=7168 splits=16", "first_call_s": 1.8053273316472769, "warm_s": 6.24769287109375e-05, "tflops": 22.55690497400017, "gbps": 1055.7814758338527} +{"description": "tf32_hc_prenorm_gemm m=4096 n=24 k=7168 splits=16", "first_call_s": 0.008736524730920792, "warm_s": 6.23813133239746e-05, "tflops": 22.591479225211792, "gbps": 1057.3997321510265} +tf32_hc_prenorm_gemm m=4096 n=24 k=7168 splits=16 (trial 2/5) + cold first call: 1805.327 ms + disk-cache first call: 8.737 ms + warm call: 62.4 us + throughput: 22.59 TFLOP/s + bandwidth: 1057.40 GB/s + +{"description": "tf32_hc_prenorm_gemm m=4096 n=24 k=7168 splits=16", "first_call_s": 1.7553360406309366, "warm_s": 6.2342529296875e-05, "tflops": 22.605533652460302, "gbps": 1058.0575530692565} +{"description": "tf32_hc_prenorm_gemm m=4096 n=24 k=7168 splits=16", "first_call_s": 0.008079558610916138, "warm_s": 6.244057464599609e-05, "tflops": 22.57003802399132, "gbps": 1056.3961714633213} +tf32_hc_prenorm_gemm m=4096 n=24 k=7168 splits=16 (trial 3/5) + cold first call: 1755.336 ms + disk-cache first call: 8.080 ms + warm call: 62.4 us + throughput: 22.57 TFLOP/s + bandwidth: 1056.40 GB/s + +{"description": "tf32_hc_prenorm_gemm m=4096 n=24 k=7168 splits=16", "first_call_s": 1.9040807969868183, "warm_s": 6.248499298095703e-05, "tflops": 22.553993795430127, "gbps": 1055.6452174060837} +{"description": "tf32_hc_prenorm_gemm m=4096 n=24 k=7168 splits=16", "first_call_s": 0.008027838543057442, "warm_s": 6.247455978393555e-05, "tflops": 22.557760292732432, "gbps": 1055.8215092371277} +tf32_hc_prenorm_gemm m=4096 n=24 k=7168 splits=16 (trial 4/5) + cold first call: 1904.081 ms + disk-cache first call: 8.028 ms + warm call: 62.5 us + throughput: 22.56 TFLOP/s + bandwidth: 1055.82 GB/s + +{"description": "tf32_hc_prenorm_gemm m=4096 n=24 k=7168 splits=16", "first_call_s": 1.833843868225813, "warm_s": 6.23891830444336e-05, "tflops": 22.58862955452239, "gbps": 1057.2663526147128} +{"description": "tf32_hc_prenorm_gemm m=4096 n=24 k=7168 splits=16", "first_call_s": 0.008773460984230042, "warm_s": 6.236211013793946e-05, "tflops": 22.59843582718391, "gbps": 1057.7253376144256} +tf32_hc_prenorm_gemm m=4096 n=24 k=7168 splits=16 (trial 5/5) + cold first call: 1833.844 ms + disk-cache first call: 8.773 ms + warm call: 62.4 us + throughput: 22.60 TFLOP/s + bandwidth: 1057.73 GB/s + +{"description": "fp8_mqa_logits seq=512 kv=4096 h=32 d=128", "first_call_s": 1.7449549660086632, "warm_s": 5.0295745849609374e-05, "tflops": 320.18672625221785, "gbps": 53.83071578450462} +{"description": "fp8_mqa_logits seq=512 kv=4096 h=32 d=128", "first_call_s": 0.012757135555148125, "warm_s": 5.0264190673828125e-05, "tflops": 320.38773512740846, "gbps": 53.8645099762789} +fp8_mqa_logits seq=512 kv=4096 h=32 d=128 (trial 1/5) + cold first call: 1744.955 ms + disk-cache first call: 12.757 ms + warm call: 50.3 us + throughput: 320.39 TFLOP/s + bandwidth: 53.86 GB/s + +{"description": "fp8_mqa_logits seq=512 kv=4096 h=32 d=128", "first_call_s": 1.8783042784780264, "warm_s": 5.031340789794922e-05, "tflops": 320.074327715265, "gbps": 53.81181901833281} +{"description": "fp8_mqa_logits seq=512 kv=4096 h=32 d=128", "first_call_s": 0.013591498136520386, "warm_s": 5.0268672943115234e-05, "tflops": 320.35916735306614, "gbps": 53.85970707967161} +fp8_mqa_logits seq=512 kv=4096 h=32 d=128 (trial 2/5) + cold first call: 1878.304 ms + disk-cache first call: 13.591 ms + warm call: 50.3 us + throughput: 320.36 TFLOP/s + bandwidth: 53.86 GB/s + +{"description": "fp8_mqa_logits seq=512 kv=4096 h=32 d=128", "first_call_s": 1.7959023490548134, "warm_s": 5.0507137298583984e-05, "tflops": 318.8466238503581, "gbps": 53.60541390406433} +{"description": "fp8_mqa_logits seq=512 kv=4096 h=32 d=128", "first_call_s": 0.012327706441283226, "warm_s": 5.0550209045410155e-05, "tflops": 318.5749478015702, "gbps": 53.5597389432721} +fp8_mqa_logits seq=512 kv=4096 h=32 d=128 (trial 3/5) + cold first call: 1795.902 ms + disk-cache first call: 12.328 ms + warm call: 50.6 us + throughput: 318.57 TFLOP/s + bandwidth: 53.56 GB/s + +{"description": "fp8_mqa_logits seq=512 kv=4096 h=32 d=128", "first_call_s": 1.694915683940053, "warm_s": 5.023040008544922e-05, "tflops": 320.60326377262976, "gbps": 53.90074527366343} +{"description": "fp8_mqa_logits seq=512 kv=4096 h=32 d=128", "first_call_s": 0.011795859783887863, "warm_s": 5.0212608337402343e-05, "tflops": 320.7168625813935, "gbps": 53.919843832993465} +fp8_mqa_logits seq=512 kv=4096 h=32 d=128 (trial 4/5) + cold first call: 1694.916 ms + disk-cache first call: 11.796 ms + warm call: 50.2 us + throughput: 320.72 TFLOP/s + bandwidth: 53.92 GB/s + +{"description": "fp8_mqa_logits seq=512 kv=4096 h=32 d=128", "first_call_s": 1.6631217785179615, "warm_s": 5.0175872802734376e-05, "tflops": 320.95167076241466, "gbps": 53.95932046153575} +{"description": "fp8_mqa_logits seq=512 kv=4096 h=32 d=128", "first_call_s": 0.012903977185487747, "warm_s": 5.030771255493164e-05, "tflops": 320.1105633736736, "gbps": 53.81791106172228} +fp8_mqa_logits seq=512 kv=4096 h=32 d=128 (trial 5/5) + cold first call: 1663.122 ms + disk-cache first call: 12.904 ms + warm call: 50.3 us + throughput: 320.11 TFLOP/s + bandwidth: 53.82 GB/s + diff --git a/gil_production_results_torch_deepjit_main.json b/gil_production_results_torch_deepjit_main.json new file mode 100644 index 0000000000..0d22f2ed0d --- /dev/null +++ b/gil_production_results_torch_deepjit_main.json @@ -0,0 +1,154 @@ +{ + "created_at": "2026-09-23T13:01:53.136694+00:00", + "build_label": "torch_deepjit_main", + "deep_gemm_path": "/workspaces/vllm_container/DeepGEMM/deep_gemm/__init__.py", + "device": "NVIDIA H200 MIG 3g.71gb", + "shape": { + "m": 128, + "n": 4096, + "k": 7168 + }, + "warmups": 20, + "calls": 1000, + "trials": 3, + "summaries": { + "no_worker": { + "median_foreground_us_per_call": 33.422213047742844, + "median_worker_progress_ratio": null + }, + "io_bound": { + "median_foreground_us_per_call": 33.45991484820843, + "median_worker_progress_ratio": 0.84375 + }, + "scheduler": { + "median_foreground_us_per_call": 33.470259979367256, + "median_worker_progress_ratio": 0.8214285714285714 + }, + "cpu_bound": { + "median_foreground_us_per_call": 38.60392048954964, + "median_worker_progress_ratio": 0.7272727272727273 + } + }, + "results": [ + { + "profile": "no_worker", + "trial": 1, + "foreground_seconds": 0.033422213047742844, + "foreground_us_per_call": 33.422213047742844, + "foreground_calls_per_second": 29920.220979129168, + "worker_tasks": 0, + "standalone_worker_tasks": 0, + "worker_progress_ratio": null + }, + { + "profile": "no_worker", + "trial": 2, + "foreground_seconds": 0.033437011763453484, + "foreground_us_per_call": 33.437011763453484, + "foreground_calls_per_second": 29906.978741832303, + "worker_tasks": 0, + "standalone_worker_tasks": 0, + "worker_progress_ratio": null + }, + { + "profile": "no_worker", + "trial": 3, + "foreground_seconds": 0.03340873308479786, + "foreground_us_per_call": 33.40873308479786, + "foreground_calls_per_second": 29932.29337556158, + "worker_tasks": 0, + "standalone_worker_tasks": 0, + "worker_progress_ratio": null + }, + { + "profile": "io_bound", + "trial": 1, + "foreground_seconds": 0.03344629891216755, + "foreground_us_per_call": 33.44629891216755, + "foreground_calls_per_second": 29898.67436830825, + "worker_tasks": 27, + "standalone_worker_tasks": 32, + "worker_progress_ratio": 0.84375 + }, + { + "profile": "io_bound", + "trial": 2, + "foreground_seconds": 0.03345991484820843, + "foreground_us_per_call": 33.45991484820843, + "foreground_calls_per_second": 29886.5076177426, + "worker_tasks": 27, + "standalone_worker_tasks": 32, + "worker_progress_ratio": 0.84375 + }, + { + "profile": "io_bound", + "trial": 3, + "foreground_seconds": 0.033468952402472496, + "foreground_us_per_call": 33.468952402472496, + "foreground_calls_per_second": 29878.437423877233, + "worker_tasks": 27, + "standalone_worker_tasks": 32, + "worker_progress_ratio": 0.84375 + }, + { + "profile": "scheduler", + "trial": 1, + "foreground_seconds": 0.033488623797893524, + "foreground_us_per_call": 33.488623797893524, + "foreground_calls_per_second": 29860.886671099972, + "worker_tasks": 23, + "standalone_worker_tasks": 28, + "worker_progress_ratio": 0.8214285714285714 + }, + { + "profile": "scheduler", + "trial": 2, + "foreground_seconds": 0.033470259979367256, + "foreground_us_per_call": 33.470259979367256, + "foreground_calls_per_second": 29877.27016809699, + "worker_tasks": 23, + "standalone_worker_tasks": 28, + "worker_progress_ratio": 0.8214285714285714 + }, + { + "profile": "scheduler", + "trial": 3, + "foreground_seconds": 0.03346007503569126, + "foreground_us_per_call": 33.46007503569126, + "foreground_calls_per_second": 29886.36453843328, + "worker_tasks": 23, + "standalone_worker_tasks": 27, + "worker_progress_ratio": 0.8518518518518519 + }, + { + "profile": "cpu_bound", + "trial": 1, + "foreground_seconds": 0.03860392048954964, + "foreground_us_per_call": 38.60392048954964, + "foreground_calls_per_second": 25904.104746840603, + "worker_tasks": 206, + "standalone_worker_tasks": 277, + "worker_progress_ratio": 0.7436823104693141 + }, + { + "profile": "cpu_bound", + "trial": 2, + "foreground_seconds": 0.03857681713998318, + "foreground_us_per_call": 38.57681713998318, + "foreground_calls_per_second": 25922.304485912184, + "worker_tasks": 162, + "standalone_worker_tasks": 273, + "worker_progress_ratio": 0.5934065934065934 + }, + { + "profile": "cpu_bound", + "trial": 3, + "foreground_seconds": 0.03863685578107834, + "foreground_us_per_call": 38.63685578107834, + "foreground_calls_per_second": 25882.023259504746, + "worker_tasks": 200, + "standalone_worker_tasks": 275, + "worker_progress_ratio": 0.7272727272727273 + } + ] +} diff --git a/gil_production_results_torch_deepjit_torch_gil.json b/gil_production_results_torch_deepjit_torch_gil.json new file mode 100644 index 0000000000..d89d97bfc0 --- /dev/null +++ b/gil_production_results_torch_deepjit_torch_gil.json @@ -0,0 +1,154 @@ +{ + "created_at": "2026-09-23T12:29:54.576527+00:00", + "build_label": "torch_deepjit_gil_enabled", + "deep_gemm_path": "/workspaces/vllm_container/DeepGEMM/deep_gemm/__init__.py", + "device": "NVIDIA H200 MIG 3g.71gb", + "shape": { + "m": 128, + "n": 4096, + "k": 7168 + }, + "warmups": 20, + "calls": 1000, + "trials": 3, + "summaries": { + "no_worker": { + "median_foreground_us_per_call": 33.43009576201439, + "median_worker_progress_ratio": null + }, + "io_bound": { + "median_foreground_us_per_call": 33.46824273467064, + "median_worker_progress_ratio": 1.0 + }, + "scheduler": { + "median_foreground_us_per_call": 33.4872305393219, + "median_worker_progress_ratio": 0.9642857142857143 + }, + "cpu_bound": { + "median_foreground_us_per_call": 6706.36167190969, + "median_worker_progress_ratio": 1.1055547192533495 + } + }, + "results": [ + { + "profile": "no_worker", + "trial": 1, + "foreground_seconds": 0.03343009576201439, + "foreground_us_per_call": 33.43009576201439, + "foreground_calls_per_second": 29913.165882589838, + "worker_tasks": 0, + "standalone_worker_tasks": 0, + "worker_progress_ratio": null + }, + { + "profile": "no_worker", + "trial": 2, + "foreground_seconds": 0.03343972936272621, + "foreground_us_per_call": 33.43972936272621, + "foreground_calls_per_second": 29904.54824418094, + "worker_tasks": 0, + "standalone_worker_tasks": 0, + "worker_progress_ratio": null + }, + { + "profile": "no_worker", + "trial": 3, + "foreground_seconds": 0.03341834805905819, + "foreground_us_per_call": 33.41834805905819, + "foreground_calls_per_second": 29923.681393010855, + "worker_tasks": 0, + "standalone_worker_tasks": 0, + "worker_progress_ratio": null + }, + { + "profile": "io_bound", + "trial": 1, + "foreground_seconds": 0.03344297595322132, + "foreground_us_per_call": 33.44297595322132, + "foreground_calls_per_second": 29901.645158575586, + "worker_tasks": 32, + "standalone_worker_tasks": 32, + "worker_progress_ratio": 1.0 + }, + { + "profile": "io_bound", + "trial": 2, + "foreground_seconds": 0.03347584418952465, + "foreground_us_per_call": 33.47584418952465, + "foreground_calls_per_second": 29872.28624731509, + "worker_tasks": 32, + "standalone_worker_tasks": 32, + "worker_progress_ratio": 1.0 + }, + { + "profile": "io_bound", + "trial": 3, + "foreground_seconds": 0.03346824273467064, + "foreground_us_per_call": 33.46824273467064, + "foreground_calls_per_second": 29879.07097267684, + "worker_tasks": 32, + "standalone_worker_tasks": 32, + "worker_progress_ratio": 1.0 + }, + { + "profile": "scheduler", + "trial": 1, + "foreground_seconds": 0.0334872305393219, + "foreground_us_per_call": 33.4872305393219, + "foreground_calls_per_second": 29862.129053215205, + "worker_tasks": 27, + "standalone_worker_tasks": 28, + "worker_progress_ratio": 0.9642857142857143 + }, + { + "profile": "scheduler", + "trial": 2, + "foreground_seconds": 0.033480679616332054, + "foreground_us_per_call": 33.480679616332054, + "foreground_calls_per_second": 29867.97196052719, + "worker_tasks": 27, + "standalone_worker_tasks": 28, + "worker_progress_ratio": 0.9642857142857143 + }, + { + "profile": "scheduler", + "trial": 3, + "foreground_seconds": 0.03350985236465931, + "foreground_us_per_call": 33.50985236465931, + "foreground_calls_per_second": 29841.969732299858, + "worker_tasks": 27, + "standalone_worker_tasks": 27, + "worker_progress_ratio": 1.0 + }, + { + "profile": "cpu_bound", + "trial": 1, + "foreground_seconds": 6.762459525838494, + "foreground_us_per_call": 6762.459525838494, + "foreground_calls_per_second": 147.87519188530854, + "worker_tasks": 37172, + "standalone_worker_tasks": 35825, + "worker_progress_ratio": 1.0375994417306351 + }, + { + "profile": "cpu_bound", + "trial": 2, + "foreground_seconds": 6.70636167190969, + "foreground_us_per_call": 6706.36167190969, + "foreground_calls_per_second": 149.11214887031915, + "worker_tasks": 36754, + "standalone_worker_tasks": 32567, + "worker_progress_ratio": 1.1285657260416986 + }, + { + "profile": "cpu_bound", + "trial": 3, + "foreground_seconds": 6.638097869232297, + "foreground_us_per_call": 6638.097869232297, + "foreground_calls_per_second": 150.6455643920253, + "worker_tasks": 36721, + "standalone_worker_tasks": 33215, + "worker_progress_ratio": 1.1055547192533495 + } + ] +} diff --git a/kernel_benchmark_comparison.txt b/kernel_benchmark_comparison.txt new file mode 100644 index 0000000000..de700d6314 --- /dev/null +++ b/kernel_benchmark_comparison.txt @@ -0,0 +1,37 @@ +Baseline: /workspaces/vllm_container/DeepGEMM/kernel_benchmark_results_main.json (NVIDIA H200 MIG 3g.71gb) +Candidate: /workspaces/vllm_container/DeepGEMM/kernel_benchmark_results_torch_library.json (NVIDIA H200 MIG 3g.71gb) + +kernel and shape baseline wall ms/call candidate wall ms/call delta ms relative +------------------------------------------------------------------------------------------------------------------------------------------ +batched_symm (576x5120x576 x16) 0.169921 0.169532 -0.000389 -0.23% +batched_syrk (576x576x5120 x16) 0.151197 0.151113 -0.000083 -0.06% +bf16_gemm_nn (128x4096x7168) 0.033421 0.033426 +0.000005 +0.01% +bf16_gemm_nt (128x4096x7168) 0.033100 0.033080 -0.000020 -0.06% +bf16_gemm_tn (128x4096x7168) 0.032508 0.032626 +0.000118 +0.36% +bf16_gemm_tt (128x4096x7168) 0.032946 0.032949 +0.000003 +0.01% +cublaslt_gemm_nn (128x4096x7168) 0.033420 0.033424 +0.000003 +0.01% +cublaslt_gemm_nt (128x4096x7168) 0.033428 0.033407 -0.000021 -0.06% +cublaslt_gemm_tn (128x4096x7168) 0.032470 0.032670 +0.000200 +0.62% +cublaslt_gemm_tt (128x4096x7168) 0.032938 0.032924 -0.000013 -0.04% +fp8_fp4_gemm_nn (128x4096x7168) 0.027151 0.026539 -0.000612 -2.25% +fp8_fp4_gemm_nt (128x4096x7168) 0.026879 0.026662 -0.000217 -0.81% +fp8_fp4_gemm_tn (128x4096x7168) 0.028118 0.027942 -0.000177 -0.63% +fp8_fp4_gemm_tt (128x4096x7168) 0.028052 0.027755 -0.000298 -1.06% +k_grouped_bf16_gemm_tn_contiguous (768x2048x256 x8) 0.084497 0.084451 -0.000047 -0.06% +k_grouped_fp8_gemm_nt_contiguous (768x2048x256 x8) 0.089690 0.089682 -0.000008 -0.01% +m_grouped_bf16_gemm_nn_contiguous (192x6144x7168 x32) 1.821838 1.822937 +0.001098 +0.06% +m_grouped_bf16_gemm_nt_contiguous (192x6144x7168 x32) 1.839028 1.841066 +0.002038 +0.11% +m_grouped_bf16_gemm_nt_masked (192x6144x7168 x32) 1.794968 1.794944 -0.000024 -0.00% +m_grouped_fp8_fp4_gemm_nn_contiguous (192x6144x7168 x32) 2.169581 2.165532 -0.004049 -0.19% +m_grouped_fp8_fp4_gemm_nt_contiguous (192x6144x7168 x32) 2.157507 2.166435 +0.008928 +0.41% +m_grouped_fp8_fp4_gemm_nt_masked (192x6144x7168 x32) 2.131780 2.148537 +0.016757 +0.79% +------------------------------------------------------------------------------------------------------------------------------------------ +TOTAL (one call per compared case) 12.784439 12.807630 +0.023191 +0.18% + +Compared 22 matching passed cases. +Negative delta and relative values mean the candidate is faster. + +Not compared: + - cublaslt_nvfp4_gemm_nt (128x4096x7168): status baseline=skipped candidate=skipped + - k_grouped_fp4_gemm_nt_contiguous (768x2048x256 x8): status baseline=skipped candidate=skipped + - k_grouped_fp8_gemm_tn_contiguous (768x2048x256 x8): status baseline=skipped candidate=skipped diff --git a/kernel_benchmark_results_main.json b/kernel_benchmark_results_main.json new file mode 100644 index 0000000000..8271071fed --- /dev/null +++ b/kernel_benchmark_results_main.json @@ -0,0 +1,364 @@ +{ + "created_at": "2026-09-24T18:55:17.843764+00:00", + "device": "NVIDIA H200 MIG 3g.71gb", + "binding": { + "package_module": "/workspaces/vllm_container/DeepGEMM/deep_gemm/__init__.py", + "binding_module": "/workspaces/vllm_container/DeepGEMM/deep_gemm/_C.cpython-312-x86_64-linux-gnu.so", + "native_modules": { + "deep_gemm._C": "/workspaces/vllm_container/DeepGEMM/deep_gemm/_C.cpython-312-x86_64-linux-gnu.so" + } + }, + "warmups": 10, + "repetitions": 1000, + "results": [ + { + "kernel": "bf16_gemm_nt", + "m": 128, + "n": 4096, + "k": 7168, + "groups_or_batches": 1, + "min_sm": null, + "supported_sms": null, + "status": "passed", + "runtime binding path": "native extension; dispatcher event not observed", + "total wall milliseconds": 33.1000629812479, + "wall milliseconds per call": 0.0331000629812479, + "effective tflops": 227.0748781432268 + }, + { + "kernel": "bf16_gemm_nn", + "m": 128, + "n": 4096, + "k": 7168, + "groups_or_batches": 1, + "min_sm": null, + "supported_sms": null, + "status": "passed", + "runtime binding path": "native extension; dispatcher event not observed", + "total wall milliseconds": 33.42066332697868, + "wall milliseconds per call": 0.033420663326978683, + "effective tflops": 224.8965765419918 + }, + { + "kernel": "bf16_gemm_tn", + "m": 128, + "n": 4096, + "k": 7168, + "groups_or_batches": 1, + "min_sm": null, + "supported_sms": null, + "status": "passed", + "runtime binding path": "native extension; dispatcher event not observed", + "total wall milliseconds": 32.507847994565964, + "wall milliseconds per call": 0.032507847994565964, + "effective tflops": 231.21163754845944 + }, + { + "kernel": "bf16_gemm_tt", + "m": 128, + "n": 4096, + "k": 7168, + "groups_or_batches": 1, + "min_sm": null, + "supported_sms": null, + "status": "passed", + "runtime binding path": "native extension; dispatcher event not observed", + "total wall milliseconds": 32.94606879353523, + "wall milliseconds per call": 0.03294606879353523, + "effective tflops": 228.13625549992318 + }, + { + "kernel": "cublaslt_gemm_nt", + "m": 128, + "n": 4096, + "k": 7168, + "groups_or_batches": 1, + "min_sm": null, + "supported_sms": null, + "status": "passed", + "runtime binding path": "native extension; dispatcher event not observed", + "total wall milliseconds": 33.42802822589874, + "wall milliseconds per call": 0.03342802822589874, + "effective tflops": 224.84702708778812 + }, + { + "kernel": "cublaslt_gemm_nn", + "m": 128, + "n": 4096, + "k": 7168, + "groups_or_batches": 1, + "min_sm": null, + "supported_sms": null, + "status": "passed", + "runtime binding path": "native extension; dispatcher event not observed", + "total wall milliseconds": 33.42042677104473, + "wall milliseconds per call": 0.03342042677104473, + "effective tflops": 224.89816840136785 + }, + { + "kernel": "cublaslt_gemm_tn", + "m": 128, + "n": 4096, + "k": 7168, + "groups_or_batches": 1, + "min_sm": null, + "supported_sms": null, + "status": "passed", + "runtime binding path": "native extension; dispatcher event not observed", + "total wall milliseconds": 32.4699804186821, + "wall milliseconds per call": 0.0324699804186821, + "effective tflops": 231.48128428421975 + }, + { + "kernel": "cublaslt_gemm_tt", + "m": 128, + "n": 4096, + "k": 7168, + "groups_or_batches": 1, + "min_sm": null, + "supported_sms": null, + "status": "passed", + "runtime binding path": "native extension; dispatcher event not observed", + "total wall milliseconds": 32.93773718178272, + "wall milliseconds per call": 0.03293773718178272, + "effective tflops": 228.19396264285794 + }, + { + "kernel": "fp8_fp4_gemm_nt", + "m": 128, + "n": 4096, + "k": 7168, + "groups_or_batches": 1, + "min_sm": null, + "supported_sms": null, + "status": "passed", + "runtime binding path": "native extension; dispatcher event not observed", + "total wall milliseconds": 26.87935158610344, + "wall milliseconds per call": 0.02687935158610344, + "effective tflops": 279.6270119806705 + }, + { + "kernel": "fp8_fp4_gemm_nn", + "m": 128, + "n": 4096, + "k": 7168, + "groups_or_batches": 1, + "min_sm": null, + "supported_sms": null, + "status": "passed", + "runtime binding path": "native extension; dispatcher event not observed", + "total wall milliseconds": 27.150673791766167, + "wall milliseconds per call": 0.027150673791766167, + "effective tflops": 276.8326423736634 + }, + { + "kernel": "fp8_fp4_gemm_tn", + "m": 128, + "n": 4096, + "k": 7168, + "groups_or_batches": 1, + "min_sm": null, + "supported_sms": null, + "status": "passed", + "runtime binding path": "native extension; dispatcher event not observed", + "total wall milliseconds": 28.118368238210678, + "wall milliseconds per call": 0.028118368238210678, + "effective tflops": 267.3054390754467 + }, + { + "kernel": "fp8_fp4_gemm_tt", + "m": 128, + "n": 4096, + "k": 7168, + "groups_or_batches": 1, + "min_sm": null, + "supported_sms": null, + "status": "passed", + "runtime binding path": "native extension; dispatcher event not observed", + "total wall milliseconds": 28.052490204572678, + "wall milliseconds per call": 0.028052490204572678, + "effective tflops": 267.9331750296745 + }, + { + "kernel": "m_grouped_fp8_fp4_gemm_nt_contiguous", + "m": 192, + "n": 6144, + "k": 7168, + "groups_or_batches": 32, + "min_sm": null, + "supported_sms": null, + "status": "passed", + "runtime binding path": "native extension; dispatcher event not observed", + "total wall milliseconds": 2157.507175579667, + "wall milliseconds per call": 2.157507175579667, + "effective tflops": 250.82923728891078 + }, + { + "kernel": "m_grouped_fp8_fp4_gemm_nn_contiguous", + "m": 192, + "n": 6144, + "k": 7168, + "groups_or_batches": 32, + "min_sm": null, + "supported_sms": null, + "status": "passed", + "runtime binding path": "native extension; dispatcher event not observed", + "total wall milliseconds": 2169.5808079093695, + "wall milliseconds per call": 2.1695808079093695, + "effective tflops": 249.43338239494892 + }, + { + "kernel": "m_grouped_fp8_fp4_gemm_nt_masked", + "m": 192, + "n": 6144, + "k": 7168, + "groups_or_batches": 32, + "min_sm": null, + "supported_sms": null, + "status": "passed", + "runtime binding path": "native extension; dispatcher event not observed", + "total wall milliseconds": 2131.780432537198, + "wall milliseconds per call": 2.131780432537198, + "effective tflops": 253.85629356392784 + }, + { + "kernel": "k_grouped_fp8_gemm_nt_contiguous", + "m": 768, + "n": 2048, + "k": 256, + "groups_or_batches": 8, + "min_sm": null, + "supported_sms": [ + 9, + 12 + ], + "status": "passed", + "runtime binding path": "native extension; dispatcher event not observed", + "total wall milliseconds": 89.68985080718994, + "wall milliseconds per call": 0.08968985080718994, + "effective tflops": 71.83032289628409 + }, + { + "kernel": "k_grouped_fp8_gemm_tn_contiguous", + "m": 768, + "n": 2048, + "k": 256, + "groups_or_batches": 8, + "min_sm": null, + "supported_sms": [ + 10, + 12 + ], + "status": "skipped", + "reason": "requires one of: SM100, SM120" + }, + { + "kernel": "k_grouped_fp4_gemm_nt_contiguous", + "m": 768, + "n": 2048, + "k": 256, + "groups_or_batches": 8, + "min_sm": null, + "supported_sms": [ + 10 + ], + "status": "skipped", + "reason": "requires one of: SM100" + }, + { + "kernel": "m_grouped_bf16_gemm_nt_contiguous", + "m": 192, + "n": 6144, + "k": 7168, + "groups_or_batches": 32, + "min_sm": null, + "supported_sms": null, + "status": "passed", + "runtime binding path": "native extension; dispatcher event not observed", + "total wall milliseconds": 1839.027689769864, + "wall milliseconds per call": 1.839027689769864, + "effective tflops": 294.26739048378414 + }, + { + "kernel": "m_grouped_bf16_gemm_nn_contiguous", + "m": 192, + "n": 6144, + "k": 7168, + "groups_or_batches": 32, + "min_sm": null, + "supported_sms": null, + "status": "passed", + "runtime binding path": "native extension; dispatcher event not observed", + "total wall milliseconds": 1821.8383397907019, + "wall milliseconds per call": 1.8218383397907019, + "effective tflops": 297.0438526165668 + }, + { + "kernel": "m_grouped_bf16_gemm_nt_masked", + "m": 192, + "n": 6144, + "k": 7168, + "groups_or_batches": 32, + "min_sm": null, + "supported_sms": null, + "status": "passed", + "runtime binding path": "native extension; dispatcher event not observed", + "total wall milliseconds": 1794.968007132411, + "wall milliseconds per call": 1.794968007132411, + "effective tflops": 301.49054308803585 + }, + { + "kernel": "k_grouped_bf16_gemm_tn_contiguous", + "m": 768, + "n": 2048, + "k": 256, + "groups_or_batches": 8, + "min_sm": null, + "supported_sms": null, + "status": "passed", + "runtime binding path": "native extension; dispatcher event not observed", + "total wall milliseconds": 84.49743129312992, + "wall milliseconds per call": 0.08449743129312992, + "effective tflops": 76.24434074984484 + }, + { + "kernel": "cublaslt_nvfp4_gemm_nt", + "m": 128, + "n": 4096, + "k": 7168, + "groups_or_batches": 1, + "min_sm": 10, + "supported_sms": null, + "status": "skipped", + "reason": "requires SM100+" + }, + { + "kernel": "batched_syrk", + "m": 576, + "n": 576, + "k": 5120, + "groups_or_batches": 16, + "min_sm": null, + "supported_sms": null, + "status": "passed", + "runtime binding path": "native extension; dispatcher event not observed", + "total wall milliseconds": 151.19662694633007, + "wall milliseconds per call": 0.15119662694633007, + "effective tflops": 359.5197917959863 + }, + { + "kernel": "batched_symm", + "m": 576, + "n": 5120, + "k": 576, + "groups_or_batches": 16, + "min_sm": null, + "supported_sms": null, + "status": "passed", + "runtime binding path": "native extension; dispatcher event not observed", + "total wall milliseconds": 169.92095112800598, + "wall milliseconds per call": 0.16992095112800598, + "effective tflops": 319.90275171571125 + } + ] +} diff --git a/kernel_benchmark_results_torch_library.json b/kernel_benchmark_results_torch_library.json new file mode 100644 index 0000000000..406aba2253 --- /dev/null +++ b/kernel_benchmark_results_torch_library.json @@ -0,0 +1,362 @@ +{ + "created_at": "2026-09-24T18:50:54.125014+00:00", + "device": "NVIDIA H200 MIG 3g.71gb", + "binding": { + "package_module": "/workspaces/vllm_container/DeepGEMM/deep_gemm/__init__.py", + "binding_module": "/workspaces/vllm_container/DeepGEMM/deep_gemm/_C.py", + "native_modules": {} + }, + "warmups": 10, + "repetitions": 1000, + "results": [ + { + "kernel": "bf16_gemm_nt", + "m": 128, + "n": 4096, + "k": 7168, + "groups_or_batches": 1, + "min_sm": null, + "supported_sms": null, + "status": "passed", + "runtime binding path": "torch dispatcher confirmed", + "total wall milliseconds": 33.07993337512016, + "wall milliseconds per call": 0.03307993337512016, + "effective tflops": 227.21305640999338 + }, + { + "kernel": "bf16_gemm_nn", + "m": 128, + "n": 4096, + "k": 7168, + "groups_or_batches": 1, + "min_sm": null, + "supported_sms": null, + "status": "passed", + "runtime binding path": "torch dispatcher confirmed", + "total wall milliseconds": 33.42551365494728, + "wall milliseconds per call": 0.03342551365494728, + "effective tflops": 224.8639421248665 + }, + { + "kernel": "bf16_gemm_tn", + "m": 128, + "n": 4096, + "k": 7168, + "groups_or_batches": 1, + "min_sm": null, + "supported_sms": null, + "status": "passed", + "runtime binding path": "torch dispatcher confirmed", + "total wall milliseconds": 32.625505700707436, + "wall milliseconds per call": 0.032625505700707436, + "effective tflops": 230.37781657548445 + }, + { + "kernel": "bf16_gemm_tt", + "m": 128, + "n": 4096, + "k": 7168, + "groups_or_batches": 1, + "min_sm": null, + "supported_sms": null, + "status": "passed", + "runtime binding path": "torch dispatcher confirmed", + "total wall milliseconds": 32.949307933449745, + "wall milliseconds per call": 0.032949307933449745, + "effective tflops": 228.1138281623709 + }, + { + "kernel": "cublaslt_gemm_nt", + "m": 128, + "n": 4096, + "k": 7168, + "groups_or_batches": 1, + "min_sm": null, + "supported_sms": null, + "status": "passed", + "runtime binding path": "torch dispatcher confirmed", + "total wall milliseconds": 33.406682312488556, + "wall milliseconds per call": 0.033406682312488556, + "effective tflops": 224.9906978996891 + }, + { + "kernel": "cublaslt_gemm_nn", + "m": 128, + "n": 4096, + "k": 7168, + "groups_or_batches": 1, + "min_sm": null, + "supported_sms": null, + "status": "passed", + "runtime binding path": "torch dispatcher confirmed", + "total wall milliseconds": 33.4236491471529, + "wall milliseconds per call": 0.0334236491471529, + "effective tflops": 224.87648595486306 + }, + { + "kernel": "cublaslt_gemm_tn", + "m": 128, + "n": 4096, + "k": 7168, + "groups_or_batches": 1, + "min_sm": null, + "supported_sms": null, + "status": "passed", + "runtime binding path": "torch dispatcher confirmed", + "total wall milliseconds": 32.67024829983711, + "wall milliseconds per call": 0.03267024829983711, + "effective tflops": 230.06230926128205 + }, + { + "kernel": "cublaslt_gemm_tt", + "m": 128, + "n": 4096, + "k": 7168, + "groups_or_batches": 1, + "min_sm": null, + "supported_sms": null, + "status": "passed", + "runtime binding path": "torch dispatcher confirmed", + "total wall milliseconds": 32.924335449934006, + "wall milliseconds per call": 0.032924335449934006, + "effective tflops": 228.28684817129894 + }, + { + "kernel": "fp8_fp4_gemm_nt", + "m": 128, + "n": 4096, + "k": 7168, + "groups_or_batches": 1, + "min_sm": null, + "supported_sms": null, + "status": "passed", + "runtime binding path": "torch dispatcher confirmed", + "total wall milliseconds": 26.662005111575127, + "wall milliseconds per call": 0.026662005111575127, + "effective tflops": 281.90650840198424 + }, + { + "kernel": "fp8_fp4_gemm_nn", + "m": 128, + "n": 4096, + "k": 7168, + "groups_or_batches": 1, + "min_sm": null, + "supported_sms": null, + "status": "passed", + "runtime binding path": "torch dispatcher confirmed", + "total wall milliseconds": 26.538513600826263, + "wall milliseconds per call": 0.026538513600826263, + "effective tflops": 283.2183023153937 + }, + { + "kernel": "fp8_fp4_gemm_tn", + "m": 128, + "n": 4096, + "k": 7168, + "groups_or_batches": 1, + "min_sm": null, + "supported_sms": null, + "status": "passed", + "runtime binding path": "torch dispatcher confirmed", + "total wall milliseconds": 27.941623702645302, + "wall milliseconds per call": 0.027941623702645302, + "effective tflops": 268.9962776675868 + }, + { + "kernel": "fp8_fp4_gemm_tt", + "m": 128, + "n": 4096, + "k": 7168, + "groups_or_batches": 1, + "min_sm": null, + "supported_sms": null, + "status": "passed", + "runtime binding path": "torch dispatcher confirmed", + "total wall milliseconds": 27.75469981133938, + "wall milliseconds per call": 0.02775469981133938, + "effective tflops": 270.8079287144445 + }, + { + "kernel": "m_grouped_fp8_fp4_gemm_nt_contiguous", + "m": 192, + "n": 6144, + "k": 7168, + "groups_or_batches": 32, + "min_sm": null, + "supported_sms": null, + "status": "passed", + "runtime binding path": "torch dispatcher confirmed", + "total wall milliseconds": 2166.4346866309643, + "wall milliseconds per call": 2.1664346866309643, + "effective tflops": 249.7956124112702 + }, + { + "kernel": "m_grouped_fp8_fp4_gemm_nn_contiguous", + "m": 192, + "n": 6144, + "k": 7168, + "groups_or_batches": 32, + "min_sm": null, + "supported_sms": null, + "status": "passed", + "runtime binding path": "torch dispatcher confirmed", + "total wall milliseconds": 2165.5323058366776, + "wall milliseconds per call": 2.1655323058366776, + "effective tflops": 249.89970264466433 + }, + { + "kernel": "m_grouped_fp8_fp4_gemm_nt_masked", + "m": 192, + "n": 6144, + "k": 7168, + "groups_or_batches": 32, + "min_sm": null, + "supported_sms": null, + "status": "passed", + "runtime binding path": "torch dispatcher confirmed", + "total wall milliseconds": 2148.5373228788376, + "wall milliseconds per call": 2.1485373228788376, + "effective tflops": 251.8764154261415 + }, + { + "kernel": "k_grouped_fp8_gemm_nt_contiguous", + "m": 768, + "n": 2048, + "k": 256, + "groups_or_batches": 8, + "min_sm": null, + "supported_sms": [ + 9, + 12 + ], + "status": "passed", + "runtime binding path": "torch dispatcher confirmed", + "total wall milliseconds": 89.68152664601803, + "wall milliseconds per call": 0.08968152664601803, + "effective tflops": 71.83699012427608 + }, + { + "kernel": "k_grouped_fp8_gemm_tn_contiguous", + "m": 768, + "n": 2048, + "k": 256, + "groups_or_batches": 8, + "min_sm": null, + "supported_sms": [ + 10, + 12 + ], + "status": "skipped", + "reason": "requires one of: SM100, SM120" + }, + { + "kernel": "k_grouped_fp4_gemm_nt_contiguous", + "m": 768, + "n": 2048, + "k": 256, + "groups_or_batches": 8, + "min_sm": null, + "supported_sms": [ + 10 + ], + "status": "skipped", + "reason": "requires one of: SM100" + }, + { + "kernel": "m_grouped_bf16_gemm_nt_contiguous", + "m": 192, + "n": 6144, + "k": 7168, + "groups_or_batches": 32, + "min_sm": null, + "supported_sms": null, + "status": "passed", + "runtime binding path": "torch dispatcher confirmed", + "total wall milliseconds": 1841.0661090165377, + "wall milliseconds per call": 1.8410661090165377, + "effective tflops": 293.9415790914105 + }, + { + "kernel": "m_grouped_bf16_gemm_nn_contiguous", + "m": 192, + "n": 6144, + "k": 7168, + "groups_or_batches": 32, + "min_sm": null, + "supported_sms": null, + "status": "passed", + "runtime binding path": "torch dispatcher confirmed", + "total wall milliseconds": 1822.9365292936563, + "wall milliseconds per call": 1.8229365292936563, + "effective tflops": 296.8649048388364 + }, + { + "kernel": "m_grouped_bf16_gemm_nt_masked", + "m": 192, + "n": 6144, + "k": 7168, + "groups_or_batches": 32, + "min_sm": null, + "supported_sms": null, + "status": "passed", + "runtime binding path": "torch dispatcher confirmed", + "total wall milliseconds": 1794.9435133486986, + "wall milliseconds per call": 1.7949435133486986, + "effective tflops": 301.49465722538827 + }, + { + "kernel": "k_grouped_bf16_gemm_tn_contiguous", + "m": 768, + "n": 2048, + "k": 256, + "groups_or_batches": 8, + "min_sm": null, + "supported_sms": null, + "status": "passed", + "runtime binding path": "torch dispatcher confirmed", + "total wall milliseconds": 84.45078693330288, + "wall milliseconds per call": 0.08445078693330288, + "effective tflops": 76.28645247661323 + }, + { + "kernel": "cublaslt_nvfp4_gemm_nt", + "m": 128, + "n": 4096, + "k": 7168, + "groups_or_batches": 1, + "min_sm": 10, + "supported_sms": null, + "status": "skipped", + "reason": "requires SM100+" + }, + { + "kernel": "batched_syrk", + "m": 576, + "n": 576, + "k": 5120, + "groups_or_batches": 16, + "min_sm": null, + "supported_sms": null, + "status": "passed", + "runtime binding path": "torch dispatcher confirmed", + "total wall milliseconds": 151.1132325977087, + "wall milliseconds per call": 0.1511132325977087, + "effective tflops": 359.71819876761884 + }, + { + "kernel": "batched_symm", + "m": 576, + "n": 5120, + "k": 576, + "groups_or_batches": 16, + "min_sm": null, + "supported_sms": null, + "status": "passed", + "runtime binding path": "torch dispatcher confirmed", + "total wall milliseconds": 169.5318277925253, + "wall milliseconds per call": 0.1695318277925253, + "effective tflops": 320.63701871087045 + } + ] +} diff --git a/scripts/benchmark_gil_production.py b/scripts/benchmark_gil_production.py new file mode 100644 index 0000000000..e5bbedecfa --- /dev/null +++ b/scripts/benchmark_gil_production.py @@ -0,0 +1,182 @@ +#!/usr/bin/env python3 +"""Benchmark DeepGEMM under representative Python-thread contention. + +Run this on the same TORCH_LIBRARY branch twice: once with DeepJIT's normal +GIL support and once with ``DJ_DISABLE_GIL=1``. Change ``BUILD_LABEL`` between +runs. Comparing those two reports isolates the DeepJIT macro. +""" + +from __future__ import annotations + +import json +import statistics +import threading +import time +from datetime import datetime, timezone +from pathlib import Path + +import torch + +import deep_gemm + + +REPO_ROOT = Path(__file__).resolve().parents[1] + +# Edit only the label between the two builds being compared. +BUILD_LABEL = "torch_deepjit_torch_temp" +M, N, K = 128, 4096, 7168 +WARMUPS = 20 +CALLS = 1_000 +TRIALS = 3 +PYTHON_BURST_ITERS = 2_000 +SCHEDULER_WAIT_SECONDS = 0.001 +IO_WAIT_SECONDS = 0.001 +OUTPUT_FILE = REPO_ROOT / f"gil_production_results_{BUILD_LABEL}.json" + + +def make_kernel_call(): + a = torch.randn((M, K), device="cuda", dtype=torch.bfloat16) + b = torch.randn((N, K), device="cuda", dtype=torch.bfloat16) + d = torch.empty((M, N), device="cuda", dtype=torch.bfloat16) + return lambda: deep_gemm.bf16_gemm_nt(a, b, d) + + +def run_python_burst(state: int) -> int: + for value in range(PYTHON_BURST_ITERS): + state = (state * 1_664_525 + value + 1_013_904_223) & 0xFFFFFFFF + return state + + +def worker_loop(profile: str, stop: threading.Event, ready: threading.Event, + result: list[int]) -> None: + tasks = 0 + state = 1 + ready.set() + while not stop.is_set(): + if profile == "io_bound": + tasks += 1 + stop.wait(IO_WAIT_SECONDS) + elif profile == "scheduler": + state = run_python_burst(state) + tasks += 1 + stop.wait(SCHEDULER_WAIT_SECONDS) + elif profile == "cpu_bound": + state = run_python_burst(state) + tasks += 1 + else: + raise ValueError(f"unknown worker profile: {profile}") + result[:] = [tasks, state] + + +def start_worker(profile: str): + stop = threading.Event() + ready = threading.Event() + result = [0, 0] + worker = threading.Thread( + target=worker_loop, args=(profile, stop, ready, result), daemon=True + ) + worker.start() + ready.wait() + return worker, stop, result + + +def run_foreground(call, profile: str | None) -> tuple[float, int]: + torch.cuda.synchronize() + worker = stop = result = None + if profile is not None: + worker, stop, result = start_worker(profile) + + start = time.perf_counter() + for _ in range(CALLS): + call() + torch.cuda.synchronize() + elapsed = time.perf_counter() - start + + tasks = 0 + if worker is not None and stop is not None and result is not None: + stop.set() + worker.join() + tasks = result[0] + return elapsed, tasks + + +def run_worker_alone(profile: str, duration: float) -> int: + worker, stop, result = start_worker(profile) + time.sleep(duration) + stop.set() + worker.join() + return result[0] + + +def main() -> None: + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required") + + call = make_kernel_call() + for _ in range(WARMUPS): + call() + torch.cuda.synchronize() + + results = [] + for profile in (None, "io_bound", "scheduler", "cpu_bound"): + profile_name = profile or "no_worker" + for trial in range(1, TRIALS + 1): + elapsed, worker_tasks = run_foreground(call, profile) + standalone_tasks = ( + run_worker_alone(profile, elapsed) if profile is not None else 0 + ) + worker_progress = ( + worker_tasks / standalone_tasks if standalone_tasks else None + ) + row = { + "profile": profile_name, + "trial": trial, + "foreground_seconds": elapsed, + "foreground_us_per_call": elapsed / CALLS * 1e6, + "foreground_calls_per_second": CALLS / elapsed, + "worker_tasks": worker_tasks, + "standalone_worker_tasks": standalone_tasks, + "worker_progress_ratio": worker_progress, + } + results.append(row) + progress = "n/a" if worker_progress is None else f"{worker_progress:.1%}" + print( + f"{profile_name:10} trial {trial}: " + f"{row['foreground_us_per_call']:9.2f} us/call, " + f"worker progress={progress}" + ) + + summaries = {} + for profile in ("no_worker", "io_bound", "scheduler", "cpu_bound"): + rows = [row for row in results if row["profile"] == profile] + progress_values = [ + row["worker_progress_ratio"] for row in rows + if row["worker_progress_ratio"] is not None + ] + summaries[profile] = { + "median_foreground_us_per_call": statistics.median( + row["foreground_us_per_call"] for row in rows + ), + "median_worker_progress_ratio": ( + statistics.median(progress_values) if progress_values else None + ), + } + + report = { + "created_at": datetime.now(timezone.utc).isoformat(), + "build_label": BUILD_LABEL, + "deep_gemm_path": str(Path(deep_gemm.__file__).resolve()), + "device": torch.cuda.get_device_name(torch.cuda.current_device()), + "shape": {"m": M, "n": N, "k": K}, + "warmups": WARMUPS, + "calls": CALLS, + "trials": TRIALS, + "summaries": summaries, + "results": results, + } + OUTPUT_FILE.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8") + print(f"\nSaved results to {OUTPUT_FILE}") + + +if __name__ == "__main__": + main() diff --git a/scripts/benchmark_kernels.py b/scripts/benchmark_kernels.py new file mode 100644 index 0000000000..0bd9acc1b1 --- /dev/null +++ b/scripts/benchmark_kernels.py @@ -0,0 +1,431 @@ +#!/usr/bin/env python3 +"""Benchmark an explicit, editable list of DeepGEMM kernels. + +Edit the configuration values in ``main`` and run: + + python scripts/benchmark_kernels.py + +Inputs are created once and only the kernel call is timed. Results are printed +and written to a JSON file. +""" + +from __future__ import annotations + +import json +import random +import sys +import time +from dataclasses import asdict, dataclass +from datetime import datetime, timezone +from pathlib import Path +from typing import Callable + +import torch + + +REPO_ROOT = Path(__file__).resolve().parents[1] +TESTS_DIR = REPO_ROOT / "tests" + + +@dataclass(frozen=True) +class BenchmarkCase: + kernel: str + m: int + n: int + k: int + groups_or_batches: int = 1 + min_sm: int | None = None + supported_sms: tuple[int, ...] | None = None + + +def unsupported_reason(case: BenchmarkCase, sm: int) -> str | None: + if case.supported_sms is not None and sm not in case.supported_sms: + supported = ", ".join(f"SM{value}0" for value in case.supported_sms) + return f"requires one of: {supported}" + if case.min_sm is not None and sm < case.min_sm: + return f"requires SM{case.min_sm}0+" + return None + + +def prioritize_repo_on_path() -> None: + repo_root = str(REPO_ROOT) + tests_dir = str(TESTS_DIR) + if repo_root in sys.path: + sys.path.remove(repo_root) + if tests_dir in sys.path: + sys.path.remove(tests_dir) + sys.path[:0] = [repo_root, tests_dir] + + +def report_loaded_binding() -> tuple[bool, dict[str, object]]: + """Print imported module paths and return binding metadata.""" + prioritize_repo_on_path() + import deep_gemm + from deep_gemm import _C + + binding_path = str(getattr(_C, "__file__", "")) + is_python_facade = binding_path.endswith(".py") + package_path = str(getattr(deep_gemm, "__file__", "")) + print(f"Loaded DeepGEMM package: {deep_gemm.__file__}") + print(f"Loaded binding module: {binding_path}") + native_modules = sorted( + (name, str(module.__file__)) + for name, module in sys.modules.items() + if name.startswith("deep_gemm") + and isinstance(getattr(module, "__file__", None), str) + and module.__file__.endswith((".so", ".pyd", ".dylib")) + ) + if native_modules: + for name, path in native_modules: + print(f"Loaded native module: {name} ({path})") + else: + print("Loaded native module: none found") + return is_python_facade, { + "package_module": package_path, + "binding_module": binding_path, + "native_modules": dict(native_modules), + } + + +def make_call(case: BenchmarkCase) -> Callable[[], None]: + """Allocate inputs for one case and return only its kernel invocation.""" + prioritize_repo_on_path() + + import deep_gemm + from deep_gemm import _C + + def get_kernel(name: str): + kernel = getattr(deep_gemm, name, None) + source = "deep_gemm" if kernel is not None else "deep_gemm._C" + if kernel is None: + kernel = getattr(_C, name) + callable_module = getattr(kernel, "__module__", type(kernel).__module__) + callable_name = getattr(kernel, "__name__", type(kernel).__name__) + print( + f"Selected binding: {name} via {source} " + f"({callable_module}.{callable_name}; {type(kernel).__name__})" + ) + return kernel + + m, n, k = case.m, case.n, case.k + count = case.groups_or_batches + if case.kernel.startswith(("bf16_gemm_", "cublaslt_gemm_")): + layout = case.kernel.rsplit("_", 1)[-1] + a = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) + b = torch.randn((n, k), device="cuda", dtype=torch.bfloat16) + d = torch.empty((m, n), device="cuda", dtype=torch.bfloat16) + inputs = { + "nt": (a, b), + "nn": (a, b.T.contiguous()), + "tn": (a.T.contiguous(), b.T.contiguous()), + "tt": (a.T.contiguous(), b), + }[layout] + kernel = get_kernel(case.kernel) + return lambda: kernel(*inputs, d) + + if case.kernel.startswith("fp8_fp4_gemm_"): + from deep_gemm.testing import get_arch_major + from generators import KernelType, MajorTypeAB, generate_normal, get_ue8m0_usage + + layout = case.kernel.rsplit("_", 1)[-1] + kernel_type = ( + KernelType.Kernel1D2D if get_arch_major() == 9 else KernelType.Kernel1D1D + ) + use_ue8m0 = get_ue8m0_usage(kernel_type) + a, b, _, d, _ = generate_normal( + m, n, k, + MajorTypeAB.KMajor, MajorTypeAB.KMajor, + False, torch.bfloat16, kernel_type, + use_ue8m0=use_ue8m0, + ) + if layout == "nn": + b = (b[0].T, b[1].T) + elif layout == "tn": + a, b = (a[0].T, a[1].T), (b[0].T, b[1].T) + elif layout == "tt": + a = (a[0].T, a[1].T) + kernel = get_kernel(case.kernel) + return lambda: kernel(a, b, d, disable_ue8m0_cast=not use_ue8m0) + + from deep_gemm.testing import get_arch_major + from generators import ( + KernelType, + MajorTypeAB, + QuantConfig, + generate_k_grouped_contiguous, + generate_m_grouped_contiguous, + generate_m_grouped_masked, + generate_normal, + get_ue8m0_usage, + ) + + def set_mk_alignment(expected_m=None): + alignment = get_kernel( + "get_theoretical_mk_alignment_for_contiguous_layout" + )(expected_m) + get_kernel("set_mk_alignment_for_contiguous_layout")(alignment) + + if case.kernel.startswith("m_grouped_bf16_gemm_") and case.kernel.endswith("_contiguous"): + layout = case.kernel.removeprefix("m_grouped_bf16_gemm_").removesuffix("_contiguous") + set_mk_alignment() + major_b = MajorTypeAB.KMajor if layout == "nt" else MajorTypeAB.MNMajor + _, a, b, grouped_layout, d, _, _ = generate_m_grouped_contiguous( + count, m, n, k, MajorTypeAB.KMajor, major_b, use_bf16=True + ) + if not major_b.is_k_major(): + b = b.mT.contiguous() + kernel = get_kernel(case.kernel) + return lambda: kernel(a, b, d, grouped_layout) + + if case.kernel.startswith("m_grouped_fp8_fp4_gemm_") and case.kernel.endswith("_contiguous"): + layout = case.kernel.removeprefix("m_grouped_fp8_fp4_gemm_").removesuffix("_contiguous") + kernel_type = KernelType.Kernel1D2D if get_arch_major() == 9 else KernelType.Kernel1D1D + use_ue8m0 = get_ue8m0_usage(kernel_type) + set_mk_alignment() + _, a, b, grouped_layout, d, _, _ = generate_m_grouped_contiguous( + count, m, n, k, MajorTypeAB.KMajor, MajorTypeAB.KMajor, + use_ue8m0=use_ue8m0, + ) + if layout == "nn": + b = b[0].mT, b[1].mT + kernel = get_kernel(case.kernel) + return lambda: kernel( + a, b, d, grouped_layout, disable_ue8m0_cast=not use_ue8m0 + ) + + if case.kernel == "m_grouped_bf16_gemm_nt_masked": + set_mk_alignment(int(m * 1.2)) + a, b, masked_m, d, _, _ = generate_m_grouped_masked( + count, 2 * m, m, n, k, use_bf16=True + ) + kernel = get_kernel(case.kernel) + return lambda: kernel(a, b, d, masked_m, m) + + if case.kernel == "m_grouped_fp8_fp4_gemm_nt_masked": + kernel_type = KernelType.Kernel1D2D if get_arch_major() == 9 else KernelType.Kernel1D1D + use_ue8m0 = get_ue8m0_usage(kernel_type) + set_mk_alignment(int(m * 1.2)) + a, b, masked_m, d, _, _ = generate_m_grouped_masked( + count, 2 * m, m, n, k, use_ue8m0=use_ue8m0 + ) + kernel = get_kernel(case.kernel) + return lambda: kernel( + a, b, d, masked_m, m, disable_ue8m0_cast=not use_ue8m0 + ) + + if case.kernel == "k_grouped_bf16_gemm_tn_contiguous": + _, a, b, c, d, _, grouped_layout, ks = generate_k_grouped_contiguous( + count, m, n, MajorTypeAB.MNMajor, MajorTypeAB.MNMajor, [k] * count, + use_bf16=True, + ) + kernel = get_kernel(case.kernel) + return lambda: kernel(a, b, d, ks, grouped_layout, c=c) + + if case.kernel in { + "k_grouped_fp8_gemm_nt_contiguous", + "k_grouped_fp8_gemm_tn_contiguous", + }: + layout = "nt" if "_nt_" in case.kernel else "tn" + if layout == "tn" and get_arch_major() < 10: + raise RuntimeError(f"{case.kernel} requires SM100+") + use_ue8m0 = get_arch_major() >= 10 + _, a, b, c, d, _, grouped_layout, ks = generate_k_grouped_contiguous( + count, m, n, MajorTypeAB.MNMajor, MajorTypeAB.MNMajor, [k] * count, + use_ue8m0=use_ue8m0, + ) + if layout == "nt": + a = a[0].T.contiguous(), a[1].T.contiguous() + b = b[0].T.contiguous(), b[1].T.contiguous() + kernel = get_kernel(case.kernel) + return lambda: kernel(a, b, d, ks, grouped_layout, c=c) + + if case.kernel == "k_grouped_fp4_gemm_nt_contiguous": + if get_arch_major() < 10: + raise RuntimeError(f"{case.kernel} requires SM100+") + _, a, b, _, d, _, grouped_layout, ks = generate_k_grouped_contiguous( + count, m, n, MajorTypeAB.KMajor, MajorTypeAB.KMajor, [k] * count, + use_ue8m0=True, gran_k=32, + quant_config=QuantConfig((32, 32, True, True)), + k_alignment=256, accumulate=False, out_dtype=torch.bfloat16, + ) + kernel = get_kernel(case.kernel) + return lambda: kernel(a, b, d, ks, grouped_layout, recipe=(1, 1, 32)) + + if case.kernel == "cublaslt_nvfp4_gemm_nt": + if get_arch_major() < 10: + raise RuntimeError(f"{case.kernel} requires SM100+") + from utils import to_cublaslt_vec16_sf_layout + a, b, _, d, _ = generate_normal( + m, n, k, MajorTypeAB.KMajor, MajorTypeAB.KMajor, + False, torch.bfloat16, KernelType.Kernel1D1D, use_ue8m0=True, + quant_config=QuantConfig((32, 32, True, True)), + ) + a = a[0], to_cublaslt_vec16_sf_layout(a[1]) + b = b[0], to_cublaslt_vec16_sf_layout(b[1]) + kernel = get_kernel(case.kernel) + return lambda: kernel(a, b, d) + + if case.kernel == "batched_syrk": + a = torch.randn(count, m, k, device="cuda", dtype=torch.bfloat16) + d = torch.empty(count, m, m, device="cuda", dtype=a.dtype) + kernel = get_kernel(case.kernel) + return lambda: kernel(a, d) + + if case.kernel == "batched_symm": + a = torch.randn(count, m, m, device="cuda", dtype=torch.bfloat16) + b = torch.randn(count, m, n, device="cuda", dtype=torch.bfloat16) + d = torch.empty(count, m, n, device="cuda", dtype=a.dtype) + kernel = get_kernel(case.kernel) + return lambda: kernel(a, b, d) + + raise ValueError(f"no input builder for kernel: {case.kernel}") + + +def benchmark(call: Callable[[], None], warmups: int, repetitions: int) -> float: + for _ in range(warmups): + call() + torch.cuda.synchronize() + + start = time.perf_counter() + for _ in range(repetitions): + call() + torch.cuda.synchronize() + return (time.perf_counter() - start) * 1e3 + + +def verify_runtime_dispatch(call: Callable[[], None], case: BenchmarkCase, + require_torch_dispatch: bool) -> str: + """Profile one untimed call to see whether this API reached torch.ops.""" + op_name = f"deep_gemm::{case.kernel}" + with torch.profiler.profile( + activities=[torch.profiler.ProfilerActivity.CPU] + ) as profile: + call() + torch.cuda.synchronize() + recorded = {event.key for event in profile.key_averages()} + if op_name in recorded: + print(f"Runtime path check: confirmed dispatcher call {op_name}") + return "torch dispatcher confirmed" + + if require_torch_dispatch: + raise RuntimeError( + f"{case.kernel} is exported through the Python _C facade, but the " + f"profiler did not observe the expected dispatcher call {op_name}. " + "Check the Python wrapper and registered operator name." + ) + print( + f"Runtime path check: no {op_name} dispatcher event observed; " + "this is consistent with a direct native-extension binding." + ) + return "native extension; dispatcher event not observed" + + +def main() -> None: + random.seed(0) + + # Edit these values to change the benchmark. To add another case, copy one + # line and change the kernel name and/or dimensions. + dense_m, dense_n, dense_k = 128, 4096, 7168 + m_grouped_groups = 32 + m_grouped_m, m_grouped_n, m_grouped_k = 192, 6144, 7168 + k_grouped_groups = 8 + k_grouped_m, k_grouped_n, k_grouped_k = 768, 2048, 256 + nvfp4_m, nvfp4_n, nvfp4_k = 128, 4096, 7168 + syrk_batch, syrk_m, syrk_k = 16, 576, 5120 + symm_batch, symm_m, symm_n = 16, 576, 5120 + cases = [ + BenchmarkCase("bf16_gemm_nt", dense_m, dense_n, dense_k), + BenchmarkCase("bf16_gemm_nn", dense_m, dense_n, dense_k), + BenchmarkCase("bf16_gemm_tn", dense_m, dense_n, dense_k), + BenchmarkCase("bf16_gemm_tt", dense_m, dense_n, dense_k), + BenchmarkCase("cublaslt_gemm_nt", dense_m, dense_n, dense_k), + BenchmarkCase("cublaslt_gemm_nn", dense_m, dense_n, dense_k), + BenchmarkCase("cublaslt_gemm_tn", dense_m, dense_n, dense_k), + BenchmarkCase("cublaslt_gemm_tt", dense_m, dense_n, dense_k), + BenchmarkCase("fp8_fp4_gemm_nt", dense_m, dense_n, dense_k), + BenchmarkCase("fp8_fp4_gemm_nn", dense_m, dense_n, dense_k), + BenchmarkCase("fp8_fp4_gemm_tn", dense_m, dense_n, dense_k), + BenchmarkCase("fp8_fp4_gemm_tt", dense_m, dense_n, dense_k), + BenchmarkCase("m_grouped_fp8_fp4_gemm_nt_contiguous", m_grouped_m, m_grouped_n, m_grouped_k, m_grouped_groups), + BenchmarkCase("m_grouped_fp8_fp4_gemm_nn_contiguous", m_grouped_m, m_grouped_n, m_grouped_k, m_grouped_groups), + BenchmarkCase("m_grouped_fp8_fp4_gemm_nt_masked", m_grouped_m, m_grouped_n, m_grouped_k, m_grouped_groups), + BenchmarkCase( + "k_grouped_fp8_gemm_nt_contiguous", k_grouped_m, k_grouped_n, k_grouped_k, k_grouped_groups, + supported_sms=(9, 12), + ), + BenchmarkCase( + "k_grouped_fp8_gemm_tn_contiguous", k_grouped_m, k_grouped_n, k_grouped_k, k_grouped_groups, + supported_sms=(10, 12), + ), + BenchmarkCase( + "k_grouped_fp4_gemm_nt_contiguous", k_grouped_m, k_grouped_n, k_grouped_k, k_grouped_groups, + supported_sms=(10,), + ), + BenchmarkCase("m_grouped_bf16_gemm_nt_contiguous", m_grouped_m, m_grouped_n, m_grouped_k, m_grouped_groups), + BenchmarkCase("m_grouped_bf16_gemm_nn_contiguous", m_grouped_m, m_grouped_n, m_grouped_k, m_grouped_groups), + BenchmarkCase("m_grouped_bf16_gemm_nt_masked", m_grouped_m, m_grouped_n, m_grouped_k, m_grouped_groups), + BenchmarkCase("k_grouped_bf16_gemm_tn_contiguous", k_grouped_m, k_grouped_n, k_grouped_k, k_grouped_groups), + BenchmarkCase("cublaslt_nvfp4_gemm_nt", nvfp4_m, nvfp4_n, nvfp4_k, min_sm=10), + BenchmarkCase("batched_syrk", syrk_m, syrk_m, syrk_k, syrk_batch), + BenchmarkCase("batched_symm", symm_m, symm_n, symm_m, symm_batch), + ] + warmups = 10 + repetitions = 1000 + FILE_INDEX = 1 + output_file = REPO_ROOT / f"kernel_benchmark_results_temp.json" # _no_gil{FILE_INDEX}.json" + + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required to benchmark DeepGEMM") + + is_python_facade, binding_info = report_loaded_binding() + sm = torch.cuda.get_device_capability()[0] + kernel_results = [] + for case in cases: + reason = unsupported_reason(case, sm) + if reason is not None: + result = { + **asdict(case), + "status": "skipped", + "reason": reason, + } + kernel_results.append(result) + print(f"skip {case.kernel:44} {reason} (device is SM{sm}0)") + continue + + call = make_call(case) + runtime_binding_path = verify_runtime_dispatch(call, case, is_python_facade) + total_wall_milliseconds = benchmark(call, warmups, repetitions) + wall_milliseconds_per_call = total_wall_milliseconds / repetitions + tflops = ( + 2 * case.m * case.n * case.k * case.groups_or_batches + ) / (wall_milliseconds_per_call * 1e9) + result = { + **asdict(case), + "status": "passed", + "runtime binding path": runtime_binding_path, + "total wall milliseconds": total_wall_milliseconds, + "wall milliseconds per call": wall_milliseconds_per_call, + "effective tflops": tflops, + } + kernel_results.append(result) + print( + f"{case.kernel:24} m={case.m:5} n={case.n:5} k={case.k:5} " + f"{wall_milliseconds_per_call:8.3f} ms/call " + f"{total_wall_milliseconds:8.3f} ms total " + f"{tflops:8.2f} effective TFLOP/s" + ) + + report = { + "created_at": datetime.now(timezone.utc).isoformat(), + "device": torch.cuda.get_device_name(torch.cuda.current_device()), + "binding": binding_info, + "warmups": warmups, + "repetitions": repetitions, + "results": kernel_results, + } + output_file.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8") + print(f"\nSaved results to {output_file}") + + +if __name__ == "__main__": + main() diff --git a/scripts/benchmark_sm90_deepjit.py b/scripts/benchmark_sm90_deepjit.py new file mode 100644 index 0000000000..695e9eace8 --- /dev/null +++ b/scripts/benchmark_sm90_deepjit.py @@ -0,0 +1,273 @@ +#!/usr/bin/env python3 +"""Minimal SM90 benchmark for DeepGEMM kernels that use DeepJIT. + +Run: + python scripts/benchmark_sm90_deepjit.py +""" + +from __future__ import annotations + +import json +import os +from pathlib import Path +import random +import subprocess +import sys +import tempfile +import time + + +REPO_ROOT = Path(__file__).resolve().parents[1] +TESTS_DIR = REPO_ROOT / "tests" +PYTHON = os.environ.get("PYTHON", sys.executable) + +CASES = ( + "bf16_gemm_nt", + "fp8_gemm_nt", + "grouped_fp8_gemm_nt", + "tf32_hc_prenorm_gemm", + "mqa_logits", +) + +WARMUP_ITERS = 10 +MEASURE_ITERS = 500 +TRIALS = 5 +SEED = 0 + + +def build_case(name: str): + if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + if str(TESTS_DIR) not in sys.path: + sys.path.insert(0, str(TESTS_DIR)) + + import torch + import deep_gemm + from deep_gemm.testing.numeric import count_bytes + from deep_gemm.utils import per_custom_dims_cast_to_fp8 + from generators import ( + KernelType, + MajorTypeAB, + generate_m_grouped_contiguous, + generate_normal, + ) + + if name == "bf16_gemm_nt": + m, n, k = 4096, 7168, 2048 + a, b, _, d, _ = generate_normal( + m=m, + n=n, + k=k, + major_a=MajorTypeAB.KMajor, + major_b=MajorTypeAB.KMajor, + accumulate=False, + out_dtype=torch.bfloat16, + kernel_type=KernelType.KernelNoSF, + use_bf16=True, + ) + return { + "description": f"bf16_gemm_nt m={m} n={n} k={k}", + "run": lambda: deep_gemm.bf16_gemm_nt(a, b, d), + "flops": float(2 * m * n * k), + "bytes": float(count_bytes(a, b, d)), + } + + if name == "fp8_gemm_nt": + m, n, k = 4096, 7168, 2048 + a, b, _, d, _ = generate_normal( + m=m, + n=n, + k=k, + major_a=MajorTypeAB.KMajor, + major_b=MajorTypeAB.KMajor, + accumulate=False, + out_dtype=torch.bfloat16, + kernel_type=KernelType.Kernel1D2D, + use_ue8m0=False, + ) + return { + "description": f"fp8_gemm_nt m={m} n={n} k={k}", + "run": lambda: deep_gemm.fp8_gemm_nt(a, b, d, disable_ue8m0_cast=True), + "flops": float(2 * m * n * k), + "bytes": float(count_bytes(a, b, d)), + } + + if name == "grouped_fp8_gemm_nt": + groups, expected_m_per_group, n, k = 4, 1024, 4096, 2048 + deep_gemm.set_mk_alignment_for_contiguous_layout( + deep_gemm.get_theoretical_mk_alignment_for_contiguous_layout() + ) + m, a, b, grouped_layout, d, _, _ = generate_m_grouped_contiguous( + num_groups=groups, + expected_m_per_group=expected_m_per_group, + n=n, + k=k, + major_a=MajorTypeAB.KMajor, + major_b=MajorTypeAB.KMajor, + use_ue8m0=False, + use_psum_layout=False, + ) + return { + "description": f"grouped_fp8_nt groups={groups} m={m} n={n} k={k}", + "run": lambda: deep_gemm.m_grouped_fp8_gemm_nt_contiguous( + a, + b, + d, + grouped_layout, + disable_ue8m0_cast=True, + ensure_zero_padding=False, + ), + "flops": float(2 * m * n * k), + "bytes": float(count_bytes(a, b, d, grouped_layout)), + } + + if name == "tf32_hc_prenorm_gemm": + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cudnn.allow_tf32 = True + m, n, k, splits = 4096, 24, 7168, 16 + a = torch.randn((m, k), dtype=torch.bfloat16, device="cuda") + b = torch.randn((n, k), dtype=torch.float32, device="cuda") + d = torch.empty((splits, m, n), dtype=torch.float32, device="cuda") + sqr_sum = torch.empty((splits, m), dtype=torch.float32, device="cuda") + return { + "description": f"tf32_hc_prenorm_gemm m={m} n={n} k={k} splits={splits}", + "run": lambda: deep_gemm.tf32_hc_prenorm_gemm( + a, b, d, sqr_sum, num_splits=splits + ), + "flops": float(2 * m * n * k), + "bytes": float(count_bytes(a, b, d, sqr_sum)), + } + + if name == "mqa_logits": + seq_len, seq_len_kv, heads, dim = 512, 4096, 32, 128 + q = torch.randn((seq_len, heads, dim), dtype=torch.bfloat16, device="cuda") + kv = torch.randn((seq_len_kv, dim), dtype=torch.bfloat16, device="cuda") + weights = torch.randn((seq_len, heads), dtype=torch.float32, device="cuda") + q_in = (q.to(torch.float8_e4m3fn), None) + kv_in = per_custom_dims_cast_to_fp8(kv, (0,), False) + ks = torch.zeros(seq_len, dtype=torch.int32, device="cuda") + ke = torch.arange(seq_len, dtype=torch.int32, device="cuda") + (seq_len_kv - seq_len) + + element_size = torch.empty((), dtype=torch.float32).element_size() + stride = ((weights.size(1) * element_size + 15) // 16) * 16 // element_size + packed_weights = torch.empty( + (weights.size(0), stride), dtype=torch.float32, device="cuda" + ) + packed_weights[:, : weights.size(1)].copy_(weights) + active_k = int((ke - ks).sum().item()) + + return { + "description": f"fp8_mqa_logits seq={seq_len} kv={seq_len_kv} h={heads} d={dim}", + "run": lambda: deep_gemm.fp8_fp4_mqa_logits( + q=q_in, + kv=kv_in, + weights=packed_weights, + cu_seq_len_k_start=ks, + cu_seq_len_k_end=ke, + clean_logits=False, + ), + "flops": float(2 * active_k * heads * dim), + "bytes": float(count_bytes(q_in, kv_in, packed_weights, ks, ke)), + } + + raise ValueError(f"unknown case: {name}") + + +def run_child() -> int: + import torch + from deep_gemm.testing import get_arch_major + + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required") + if get_arch_major() != 9: + raise RuntimeError(f"expected SM90, got arch major {get_arch_major()}") + + random.seed(SEED) + torch.manual_seed(SEED) + torch.cuda.manual_seed(SEED) + case = build_case(os.environ["DG_BENCH_CASE"]) + run = case["run"] + + torch.cuda.synchronize() + start = time.perf_counter() + run() + torch.cuda.synchronize() + first_call_s = time.perf_counter() - start + + for _ in range(WARMUP_ITERS): + run() + + start_event = torch.cuda.Event(enable_timing=True) + end_event = torch.cuda.Event(enable_timing=True) + start_event.record() + for _ in range(MEASURE_ITERS): + run() + end_event.record() + torch.cuda.synchronize() + warm_s = start_event.elapsed_time(end_event) / MEASURE_ITERS / 1e3 + + print(json.dumps({ + "description": case["description"], + "first_call_s": first_call_s, + "warm_s": warm_s, + "tflops": case["flops"] / warm_s / 1e12 if warm_s > 0 else 0.0, + "gbps": case["bytes"] / warm_s / 1e9 if warm_s > 0 else 0.0, + })) + return 0 + + +def run_case(name: str, cache_dir: Path) -> dict: + env = os.environ.copy() + env["DG_BENCH_CASE"] = name + env["DG_JIT_CACHE_DIR"] = str(cache_dir) + result = subprocess.run( + [PYTHON, str(Path(__file__).resolve())], + cwd=str(REPO_ROOT), + env=env, + capture_output=True, + text=True, + check=True, + ) + if result.stdout: + print(result.stdout, end="") + lines = [line.strip() for line in result.stdout.splitlines() if line.strip()] + return json.loads(lines[-1]) + + +def fmt_time(seconds: float) -> str: + if seconds < 1e-3: + return f"{seconds * 1e6:8.1f} us" + return f"{seconds * 1e3:8.3f} ms" + + +def main() -> int: + if "DG_BENCH_CASE" in os.environ: + return run_child() + + with tempfile.TemporaryDirectory(prefix="deepjit_sm90_", dir="/tmp") as root: + print(f"Using {PYTHON}") + print() + for name in CASES: + for trial in range(1, TRIALS + 1): + cache_dir = Path(root) / name / f"trial_{trial}" + cold = run_case(name, cache_dir) + cached = run_case(name, cache_dir) + + if cold["description"] != cached["description"]: + raise RuntimeError( + "cold and disk-cache processes generated different cases: " + f"{cold['description']!r} != {cached['description']!r}" + ) + + print(f"{cold['description']} (trial {trial}/{TRIALS})") + print(f" cold first call: {fmt_time(cold['first_call_s'])}") + print(f" disk-cache first call: {fmt_time(cached['first_call_s'])}") + print(f" warm call: {fmt_time(cached['warm_s'])}") + print(f" throughput: {cached['tflops']:8.2f} TFLOP/s") + print(f" bandwidth: {cached['gbps']:8.2f} GB/s") + print() + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/generate_pyi.py b/scripts/generate_pyi.py index df7490d410..f4547d1155 100644 --- a/scripts/generate_pyi.py +++ b/scripts/generate_pyi.py @@ -1,108 +1,11 @@ +"""Generate deep_gemm/_C.pyi from TORCH_LIBRARY schemas and deep_gemm/_C.py wrappers.""" +import ast import re from pathlib import Path - -def build_cpp_function_index(root_path): - func_index = {} - extensions = {'.cpp', '.cc', '.cxx', '.c', '.hpp', '.h'} - - pattern = re.compile( - r'([\w:\s*<&>,\[\]\(\)]+?)' - r'\s+' - r'([a-zA-Z_][a-zA-Z0-9_:]*)' - r'\s*\(', - ) - - for file_path in Path(root_path).rglob('*'): - if file_path.suffix.lower() not in extensions: - continue - if not file_path.is_file(): - continue - - try: - with open(file_path, 'r', encoding='utf-8', errors='ignore') as f: - content = f.read() - except Exception as e: - print(f'Failed to read file {file_path}: {e}') - continue - - # Remove the compile directives and comments - lines = content.split('\n') - clean_lines = [line for line in lines if not line.strip().startswith(('#', '//'))] - content = '\n'.join(clean_lines) - - for match in pattern.finditer(content): - return_type_part = match.group(1).strip() - full_func_name = match.group(2).strip() - - if not return_type_part or not re.match(r'^[a-zA-Z_]', return_type_part): - continue - - first_token = return_type_part.split()[0] - if first_token in {'return', 'if', 'else', 'for', 'while', 'switch', 'case', 'throw', 'catch', 'auto'}: - continue - - # Extract base name - if '::' in full_func_name: - base_name = full_func_name.split('::')[-1] - else: - base_name = full_func_name - - if not re.match(r'^[a-zA-Z_][a-zA-Z0-9_]*$', base_name): - continue - - # Find matching ')' - paren_start = match.end() - 1 - paren_count = 0 - pos = paren_start - while pos < len(content): - ch = content[pos] - if ch == '(': - paren_count += 1 - elif ch == ')': - paren_count -= 1 - if paren_count == 0: - break - elif paren_count < 0: - pos = -1 - break - pos += 1 - else: - continue - - if pos == -1: - continue - - # Check context before match: should be at statement boundary - match_start = match.start() - context_before = content[max(0, match_start - 50):match_start] - if context_before and re.search(r'[a-zA-Z0-9_]$', context_before.rstrip()): - continue - - # Check for definition or header declaration - is_header = file_path.suffix.lower() in {'.h', '.hpp', '.cuh'} - after_paren = content[pos+1:pos+500] - has_brace = '{' in after_paren - has_semicolon = ';' in after_paren.split('{')[0] - - if has_brace or (is_header and has_semicolon): - sig_start = match.start(1) - full_signature = content[sig_start:pos+1].strip() - if base_name not in func_index: - func_index[base_name] = full_signature - - return func_index - - class BracketTracker: - """ - Tracks nesting levels of various brackets in C++ code: - - () → paren - - [] → bracket - - {} → brace - - <> → angle (treated as template brackets only at top level) - Provides is_top_level() to check if currently outside all brackets. - """ + """Track () [] {} <> nesting for top-level comma/default splitting.""" + def __init__(self): self.paren = 0 # () self.bracket = 0 # [] @@ -110,9 +13,6 @@ def __init__(self): self.angle = 0 # <> def update(self, char: str): - """ - Update internal counters based on the given character. - """ if char == '(': self.paren += 1 elif char == ')': @@ -133,32 +33,382 @@ def update(self, char: str): self.angle -= 1 def _in_top_level_of_other_brackets(self): - """ - Check if not inside parentheses, square brackets, or braces (for correct template bracket recognition). - """ return self.paren == 0 and self.bracket == 0 and self.brace == 0 def is_top_level(self): - """ - Check if completely at top level (all bracket counters are zero). - """ - return (self.paren == 0 and - self.bracket == 0 and - self.brace == 0 and - self.angle == 0) - - -def extract_m_def_statements(root_path): - """ - Scan all c files under root_path and extract all m.def(...) statements. - """ - results = [] - extensions = {'.hpp', '.cpp', '.h', '.cc'} + return self.paren == 0 and self.bracket == 0 and self.brace == 0 and self.angle == 0 + + +def split_top_level_commas(value: str) -> list[str]: + """Split on commas not nested inside brackets.""" + parts = [] + current = [] + tracker = BracketTracker() + for ch in value: + if ch in '()[]{}<>': + tracker.update(ch) + if ch == ',' and tracker.is_top_level(): + parts.append(''.join(current).strip()) + current = [] + else: + current.append(ch) + if current: + parts.append(''.join(current).strip()) + return parts + + +def find_top_level_equals(value: str) -> int: + """Return index of top-level '=', or -1.""" + tracker = BracketTracker() + for i, ch in enumerate(value): + if ch in '()[]{}<>': + tracker.update(ch) + elif ch == '=' and tracker.is_top_level(): + return i + return -1 + + +def schema_type_to_python(type_str: str) -> str: + """Map a TORCH schema type to a Python annotation.""" + type_str = type_str.strip() + optional = type_str.endswith('?') + if optional: + type_str = type_str[:-1].strip() + + if type_str.startswith('Tensor'): + py_type = 'torch.Tensor' + elif type_str == 'int': + py_type = 'int' + elif type_str == 'bool': + py_type = 'bool' + elif type_str == 'float': + py_type = 'float' + elif type_str == 'str': + py_type = 'str' + elif type_str == 'int[]': + py_type = 'list[int]' + elif re.match(r'^int\[\d+\]$', type_str): + # Fixed-size int[N] maps directly to a same-arity tuple. + n = int(re.match(r'^int\[(\d+)\]$', type_str).group(1)) + py_type = f"tuple[{', '.join(['int'] * n)}]" + elif type_str == 'ScalarType': + py_type = 'torch.dtype' + else: + print(f'Warning: unrecognized schema type {type_str!r}, using Any') + py_type = 'Any' + + if optional: + return f'Optional[{py_type}]' + return py_type + + +def schema_return_to_python(return_str: str) -> str: + """Map a TORCH schema return type to a Python annotation.""" + return_str = re.sub(r"Tensor\([^)]*\)", "Tensor", return_str.strip()) + if return_str == 'int[]': + return 'list[int]' + if return_str == 'Dict(str, int)': + return 'dict[str, int]' + if return_str.startswith('__torch__.torch.classes.'): + return 'Any' + if return_str == '()': + return 'None' + if return_str in {'int', 'bool', 'float', 'str', 'Tensor'}: + return { + 'int': 'int', + 'bool': 'bool', + 'float': 'float', + 'str': 'str', + 'Tensor': 'torch.Tensor', + }[return_str] + if return_str.startswith('(') and return_str.endswith(')'): + inner = return_str[1:-1].strip() + if not inner: + return 'tuple[()]' + parts = split_top_level_commas(inner) + py_parts = [schema_return_to_python(part) for part in parts] + return f'tuple[{", ".join(py_parts)}]' + print(f'Warning: unrecognized schema return type {return_str!r}, using Any') + return 'Any' + - # Regex: match m.def( ... ), supports multi-line - pattern = re.compile(r'm\.def\s*\(') +_SCALAR_TYPE_DEFAULTS = { + 'float': 'torch.float32', + 'float32': 'torch.float32', + 'double': 'torch.float64', + 'float64': 'torch.float64', + 'half': 'torch.float16', + 'float16': 'torch.float16', + 'bfloat16': 'torch.bfloat16', + 'byte': 'torch.uint8', + 'char': 'torch.int8', + 'short': 'torch.int16', + 'int': 'torch.int32', + 'long': 'torch.int64', +} + + +def schema_default_to_python(default_str: str) -> str: + """Convert a TORCH schema default literal to a Python expression string.""" + default_str = default_str.strip() + if default_str in {'None', 'True', 'False'}: + return default_str + if (default_str.startswith("'") and default_str.endswith("'")) or ( + default_str.startswith('"') and default_str.endswith('"')): + return default_str + if default_str in _SCALAR_TYPE_DEFAULTS: + return _SCALAR_TYPE_DEFAULTS[default_str] + if re.match(r'^[+-]?\d+$', default_str): + return default_str + if re.match(r'^[+-]?\d*\.\d+([eE][+-]?\d+)?$', default_str): + return default_str + print(f'Warning: unrecognized schema default {default_str!r}, using None') + return 'None' + + +def parse_schema_arg(arg_str: str) -> dict: + """Parse one TORCH schema argument such as 'Tensor? c=None'.""" + arg_str = arg_str.strip() + if not arg_str: + raise ValueError('empty schema argument') + + default = None + eq_pos = find_top_level_equals(arg_str) + if eq_pos != -1: + default = schema_default_to_python(arg_str[eq_pos + 1:].strip()) + arg_str = arg_str[:eq_pos].strip() + + match = re.match(r'^(.+?)\s+([a-zA-Z_][a-zA-Z0-9_]*)$', arg_str) + if not match: + raise ValueError(f'could not parse schema argument: {arg_str!r}') + return { + 'name': match.group(2), + 'py_type': schema_type_to_python(match.group(1)), + 'default': default, + } + + +def parse_torch_schema(schema: str) -> dict: + """Parse a TORCH schema into name, parameters, and return type.""" + arrow = schema.rfind(' -> ') + if arrow == -1: + raise ValueError(f'schema missing return type: {schema!r}') + + signature = schema[:arrow].strip() + return_type = schema_return_to_python(schema[arrow + 4:].strip()) + + open_paren = signature.find('(') + if open_paren == -1: + raise ValueError(f'schema missing argument list: {schema!r}') + + name = signature[:open_paren].strip() + paren_depth = 0 + close_paren = -1 + for i in range(open_paren, len(signature)): + if signature[i] == '(': + paren_depth += 1 + elif signature[i] == ')': + paren_depth -= 1 + if paren_depth == 0: + close_paren = i + break + if close_paren == -1: + raise ValueError(f'unclosed argument list in schema: {schema!r}') + + args_blob = signature[open_paren + 1:close_paren].strip() + parameters = [] + if args_blob: + for arg in split_top_level_commas(args_blob): + parameters.append(parse_schema_arg(arg)) + + return { + 'python_function_name': name, + 'parameters': parameters, + 'return_type': return_type, + } + + +def _merge_named_pairs(parameters: list[dict], pairs: tuple[tuple[str, str], ...]) -> list[dict]: + """Replace (tensor, scale_factor) arg pairs with one tuple-typed parameter.""" + by_name = {param['name']: param for param in parameters} + sf_of = dict(pairs) + drop = set(sf_of.values()) + + out = [] + for param in parameters: + if param['name'] in drop: + continue + sf_name = sf_of.get(param['name']) + if sf_name is None: + out.append(dict(param)) + continue + base_optional = param['py_type'] == 'Optional[torch.Tensor]' + sf_optional = by_name[sf_name]['py_type'] == 'Optional[torch.Tensor]' + if base_optional and sf_optional: + py_type = 'Optional[tuple[torch.Tensor, torch.Tensor]]' + elif sf_optional: + py_type = 'tuple[torch.Tensor, Optional[torch.Tensor]]' + else: + py_type = 'tuple[torch.Tensor, torch.Tensor]' + out.append({ + 'name': param['name'], + 'py_type': py_type, + 'default': None, + }) + return out + + +def _is_tensor_schema_param(param: dict) -> bool: + py_type = param['py_type'] + return py_type in {'torch.Tensor', 'Optional[torch.Tensor]'} + + +def _is_tensor_scale_factor_pair(base_name: str, sf_name: str) -> bool: + if base_name == 'a' and sf_name == 'sfa': + return True + if base_name == 'b' and sf_name == 'sfb': + return True + return sf_name == f'{base_name}_sf' + + +def detect_tensor_sf_pairs(parameters: list[dict]) -> list[tuple[str, str]]: + """Detect consecutive (tensor, scale_factor) arg pairs.""" + pairs = [] + i = 0 + while i < len(parameters) - 1: + left, right = parameters[i], parameters[i + 1] + if ( + _is_tensor_schema_param(left) + and _is_tensor_schema_param(right) + and _is_tensor_scale_factor_pair(left['name'], right['name']) + ): + pairs.append((left['name'], right['name'])) + i += 2 + else: + i += 1 + return pairs + + +def _maybe_widen_int_list_value_param(parameters: list[dict]) -> None: + """Single int[] value param in a Python wrapper usually accepts int | list[int].""" + if len(parameters) == 1 and parameters[0]['name'] == 'value': + if parameters[0]['py_type'] == 'list[int]': + parameters[0]['py_type'] = 'int | list[int]' + + +def _promote_transform_sf_recipe_type(op_name: str, parameters: list[dict]) -> None: + """transform_sf_into_required_layout's recipe is a std::variant, which int[N] can't express, so promote it here.""" + if op_name != 'transform_sf_into_required_layout': + return + recipe = next((p for p in parameters if p['name'] == 'recipe'), None) + if recipe is not None and recipe['py_type'] == 'list[int]': + recipe['py_type'] = 'tuple[int, int] | tuple[int, int, int]' + + +def adjust_for_c_py_wrapper( + name: str, + parameters: list[dict], + wrapper_defaults: dict[str, dict[str, str]] | None = None, +) -> list[dict]: + """Adjust flat schema params to match deep_gemm._C Python wrappers.""" + pairs = detect_tensor_sf_pairs(parameters) + if pairs: + parameters = _merge_named_pairs(parameters, tuple(pairs)) + + _maybe_widen_int_list_value_param(parameters) + _promote_transform_sf_recipe_type(name, parameters) + + return parameters + + +def format_ast_default(node: ast.AST) -> str: + """Convert an AST default value node to a Python expression string.""" + if isinstance(node, ast.Constant): + if node.value is None: + return 'None' + if isinstance(node.value, bool): + return 'True' if node.value else 'False' + if isinstance(node.value, str): + return f'"{node.value}"' + if isinstance(node.value, (int, float)): + return repr(node.value) + if isinstance(node, ast.Tuple): + elts = ', '.join(format_ast_default(element) for element in node.elts) + return f'({elts})' + if isinstance(node, ast.List): + elts = ', '.join(format_ast_default(element) for element in node.elts) + return f'[{elts}]' + if isinstance(node, ast.Name): + return node.id + if isinstance(node, ast.Attribute): + return ast.unparse(node) + if isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.USub): + return f'-{format_ast_default(node.operand)}' + return ast.unparse(node) + + +def extract_function_defaults(func_def: ast.FunctionDef) -> dict[str, str]: + """Extract {param_name: default_expr} from a Python function definition.""" + defaults: dict[str, str] = {} + args = func_def.args + pos_args = args.args + if args.defaults: + first_default_idx = len(pos_args) - len(args.defaults) + for idx, default_node in enumerate(args.defaults): + defaults[pos_args[first_default_idx + idx].arg] = format_ast_default(default_node) + for arg, default_node in zip(args.kwonlyargs, args.kw_defaults): + if default_node is not None: + defaults[arg.arg] = format_ast_default(default_node) + return defaults + + +def _is_globals_update_call(node: ast.Call) -> bool: + if not isinstance(node.func, ast.Attribute) or node.func.attr != 'update': + return False + base = node.func.value + if isinstance(base, ast.Name): + return base.id == 'globals' + if isinstance(base, ast.Call) and isinstance(base.func, ast.Name): + return base.func.id == 'globals' + return False + + +def parse_c_py_metadata(c_py_path: Path) -> dict[str, dict[str, str]]: + """Parse wrapper defaults from deep_gemm/_C.py without importing it.""" + source = c_py_path.read_text(encoding='utf-8') + module = ast.parse(source, filename=str(c_py_path)) + + func_defaults: dict[str, dict[str, str]] = {} + + for node in ast.walk(module): + if isinstance(node, ast.FunctionDef): + func_defaults[node.name] = extract_function_defaults(node) + + for node in ast.walk(module): + if not isinstance(node, ast.Call): + continue + if not _is_globals_update_call(node): + continue + if not node.args or not isinstance(node.args[0], ast.Dict): + continue + alias_dict = node.args[0] + for key_node, value_node in zip(alias_dict.keys, alias_dict.values): + if not isinstance(key_node, ast.Constant) or not isinstance(key_node.value, str): + continue + alias_name = key_node.value + if isinstance(value_node, ast.Name): + if value_node.id in func_defaults: + func_defaults[alias_name] = func_defaults[value_node.id] + + return func_defaults - for file_path in Path(root_path).rglob('*'): + +def extract_m_def_statements(root_path) -> list[str]: + """Scan C++ sources under root_path for m.def(...) registrations.""" + statements = [] + extensions = {'.hpp', '.cpp', '.h', '.cc'} + + for file_path in sorted(Path(root_path).rglob('*')): if file_path.suffix.lower() not in extensions: continue if not file_path.is_file(): @@ -177,15 +427,11 @@ def extract_m_def_statements(root_path): while i < len(lines): line = lines[i] if 'm.def(' in line: - # Found a potential starting line - start_i = i - # Check if it's a comment stripped = line.lstrip() if stripped.startswith('//') or stripped.startswith('/*'): i += 1 continue - # Try to match the complete m.def(...) call paren_count = 0 j = i found_start = False @@ -201,7 +447,6 @@ def extract_m_def_statements(root_path): if found_start: paren_count -= 1 if paren_count == 0: - # Found complete statement full_stmt = ''.join(lines[i:j+1]).rstrip() m_def_list.append(full_stmt) i = j @@ -209,28 +454,16 @@ def extract_m_def_statements(root_path): if paren_count <= 0 and found_start: break j += 1 - else: - pass i += 1 if m_def_list: - results.append({ - 'file': str(file_path), - 'm_def_statements': m_def_list - }) + statements.extend(m_def_list) - return results + return statements def parse_m_def_statement(m_def_str): - result = { - 'python_function_name': None, - 'num_args': 0, - 'default_args': {}, - 'is_lambda': False, - } - - # Extract top-level arguments + """Parse a TORCH_LIBRARY m.def(...) statement.""" start = m_def_str.find('m.def(') if start == -1: raise ValueError(f'[{m_def_str}] Could not find m.def start position') @@ -252,634 +485,155 @@ def parse_m_def_statement(m_def_str): raise ValueError(f'[{m_def_str}] m.def parentheses not closed') args_content = m_def_str[content_start:content_end] + args_list = split_top_level_commas(args_content) + if not args_list: + raise ValueError(f'[{m_def_str}] m.def has no arguments') - # Split arguments using BracketTracker - args_list = [] - current = [] - tracker = BracketTracker() - - for ch in args_content: - if ch in '()[]{}<>': - tracker.update(ch) - if ch == ',' and tracker.is_top_level(): - args_list.append(''.join(current).strip()) - current = [] - else: - current.append(ch) - - if current: - args_list.append(''.join(current).strip()) - - if len(args_list) < 2: - raise ValueError(f'[{m_def_str}] m.def has insufficient arguments') - - # Extract Python function name first = args_list[0].strip() str_match = re.match(r'^"([^"\\]*(?:\\.[^"\\]*)*)"', first) - if str_match: - result['python_function_name'] = str_match.group(1) - else: + if not str_match: raise ValueError(f'[{m_def_str}] m.def first argument should be a string literal') - cpp_func_part = args_list[1].strip() - if cpp_func_part.startswith('&'): - cpp_func_part = cpp_func_part[1:].strip() + return parse_torch_schema(str_match.group(1)) - if cpp_func_part.startswith('['): - result['is_lambda'] = True - result['cpp_function_name'] = None - else: - if '::' in cpp_func_part: - cpp_func_name = cpp_func_part.split('::')[-1] - else: - cpp_func_name = cpp_func_part - match = re.match(r'^([a-zA-Z_][a-zA-Z0-9_]*)', cpp_func_name) - if match: - result['cpp_function_name'] = match.group(1) - else: - result['cpp_function_name'] = cpp_func_name +def apply_wrapper_defaults(name: str, parameters: list[dict], wrapper_defaults: dict[str, dict[str, str]]) -> list[dict]: + """Overlay public API defaults from deep_gemm/_C.py onto schema-derived parameters.""" + by_name = wrapper_defaults.get(name, {}) + if not by_name: + return parameters - # Parse py::arg arguments - py_args = args_list[2:] - result['num_args'] = len(py_args) + out = [] + for param in parameters: + param = dict(param) + if param['name'] in by_name: + param['default'] = by_name[param['name']] + out.append(param) + return out - for idx, arg_expr in enumerate(py_args): - expr = arg_expr.strip() - # Find top-level '=' - eq_pos = -1 - p_depth = b_depth = br_depth = angle_depth = 0 - i = 0 - while i < len(expr): - ch = expr[i] - if ch == '(': - p_depth += 1 - elif ch == ')': - p_depth -= 1 - elif ch == '[': - b_depth += 1 - elif ch == ']': - b_depth -= 1 - elif ch == '{': - br_depth += 1 - elif ch == '}': - br_depth -= 1 - elif ch == '<' and p_depth == 0 and b_depth == 0 and br_depth == 0: - angle_depth += 1 - elif ch == '>' and angle_depth > 0 and p_depth == 0 and b_depth == 0 and br_depth == 0: - angle_depth -= 1 - elif ch == '=' and all(d == 0 for d in [p_depth, b_depth, br_depth, angle_depth]): - eq_pos = i - break - i += 1 - if eq_pos != -1: - default_val = expr[eq_pos + 1:].strip() - if not default_val: - raise ValueError(f'[{expr}] Default value is empty (arg {idx})') - result['default_args'][idx] = default_val - - return result - - -def extract_cpp_signature_from_content(cpp_func_name, content): - """ - Search for the C++ function signature of cpp_func_name in the given file content. - """ - if not cpp_func_name: - return None - - # Build regex: match function starting with cpp_func_name (after word boundary) - # Note: function name may be preceded by return type (with templates, namespaces, etc.), followed by '(' - pattern = re.compile( - r'^\s*' # leading whitespace - r'([\w:\s*<&>,\[\]\(\)]+?)' # return type (non-greedy, allows templates, pointers, etc.) - r'\s+' # at least one space - r'\b' + re.escape(cpp_func_name) + r'\b' # function name (word boundary) - r'\s*\(', # optional whitespace + start of param list - re.MULTILINE +def generate_pyi_function(parsed, wrapper_defaults=None): + """Generate a typed .pyi stub for one registered op.""" + py_name = parsed['python_function_name'] + parameters = adjust_for_c_py_wrapper( + py_name, + parsed['parameters'], + wrapper_defaults=wrapper_defaults, ) + if wrapper_defaults: + parameters = apply_wrapper_defaults(py_name, parameters, wrapper_defaults) + return_type = parsed['return_type'] - for match in pattern.finditer(content): - # Find '(' position after function name - paren_start = match.end() - 1 - if content[paren_start] != '(': - paren_start = content.find('(', match.end(0) - 1) - if paren_start == -1: - continue - - # From '(', match to corresponding ')' - paren_count = 0 - pos = paren_start - while pos < len(content): - ch = content[pos] - if ch == '(': - paren_count += 1 - elif ch == ')': - paren_count -= 1 - if paren_count == 0: - start_sig = match.start(1) - full_signature = content[start_sig:pos+1].strip() - return full_signature - pos += 1 - - return None - - -def parse_mdef_and_attach_cpp_signatures(item, func_index): - """ - Enhance item by parsing m.def and extracting C++ function signature from global index - """ - statements_with_parsed_signatures = [] - - for stmt in item['m_def_statements']: - parsed = parse_m_def_statement(stmt,) - cpp_func_name = parsed.get('cpp_function_name') - - cpp_sig = None - if cpp_func_name and cpp_func_name in func_index: - cpp_sig = func_index[cpp_func_name] - else: - if not parsed['is_lambda']: - print(f'Warning: C++ function "{cpp_func_name}" not found in any .cpp file') - - parsed['cpp_signature'] = cpp_sig - statements_with_parsed_signatures.append({ - 'raw': stmt, - 'parsed': parsed - }) - - return { - 'm_def_statements': statements_with_parsed_signatures - } - - -def parse_cpp_signature(cpp_sig): - """ - Parse a C++ function signature and extract return type, parameter types, and names. - """ - if not cpp_sig or not cpp_sig.strip(): - return None - - # Find function name: last identifier before '(' - paren_pos = cpp_sig.find('(') - if paren_pos == -1: - return None - - before_paren = cpp_sig[:paren_pos].strip() - if not before_paren: - return None - - # Function name is the last word in before_paren (may include templates like func) - tokens = before_paren.split() - if len(tokens) < 2: - return None - - # Heuristic: function name is usually the last token (may include <>) - func_name_part = tokens[-1] - return_type = ' '.join(tokens[:-1]).strip() - - # Now extract parameter list content - param_list_str = cpp_sig[paren_pos+1:cpp_sig.rfind(')')].strip() - parameters = [] - - if param_list_str and param_list_str != 'void': # 'void' means no parameters - # Split parameters (handle commas not inside templates/brackets) - param_decls = split_cpp_parameters(param_list_str) - for decl in param_decls: - decl = decl.strip() - if not decl: - continue - # Try to split type and name from right to left - param_info = parse_parameter_declaration(decl) - if param_info: - parameters.append(param_info) - - return { - 'return_type': return_type, - 'parameters': parameters, - 'num_parameters': len(parameters) - } - - -def split_cpp_parameters(param_str: str): - """ - Split a C++ parameter list string by top-level commas, - e.g., 'int a, std::vector b' → ['int a', 'std::vector b'] - """ - if not param_str.strip() or param_str == 'void': - return [] - params = [] - current = [] - tracker = BracketTracker() - - for ch in param_str: - if ch in '()[]{}<>': - tracker.update(ch) - if ch == ',' and tracker.is_top_level(): - param = ''.join(current).strip() - if param: # Only add non-empty parameters - params.append(param) - current = [] - else: - current.append(ch) - - if current: - final_param = ''.join(current).strip() - if final_param: # Only add non-empty parameters - params.append(final_param) - return params - - -def parse_parameter_declaration(decl: str): - """ - Parse a single parameter declaration, e.g., 'const std::string& name' → {'type': 'const std::string&', 'name': 'name'} - Improved version that better handles template types. - """ - decl = decl.strip() - if not decl: - return None - - # Remove possible default value (starting from top-level '=') - tracker = BracketTracker() - eq_pos = -1 - for i, ch in enumerate(decl): - if ch in '()[]{}<>': - tracker.update(ch) - elif ch == '=' and tracker.is_top_level(): - eq_pos = i - break - - if eq_pos != -1: - decl = decl[:eq_pos].strip() - - # Now decl is 'type name' or just 'type' - # Instead of simple splitting, we'll use a more robust approach - # to find the parameter name - - # First, let's handle the case where there's no explicit parameter name - # (this sometimes happens in function declarations) - if not re.search(r'[a-zA-Z_][a-zA-Z0-9_]*$', decl): - # No parameter name found, just return the type - return { - 'type': decl, - 'name': None - } - - # Use bracket tracking to find where the type ends and name begins - tracker = BracketTracker() - name_start = -1 - - # Scan from the end to find the start of the parameter name - # We look for the first identifier that's outside all brackets - i = len(decl) - 1 - while i >= 0: - ch = decl[i] - - if ch in '()[]{}<>': - tracker.update(ch) - - # If we're at top level and find an identifier character - if tracker.is_top_level() and re.match(r'[a-zA-Z0-9_]', ch): - # Track back to find the start of this identifier - name_start = i - while name_start > 0 and re.match(r'[a-zA-Z0-9_]', decl[name_start - 1]): - name_start -= 1 - - # Check if this might be part of a type keyword (like 'int', 'bool', etc.) - potential_name = decl[name_start:i+1] - type_keywords = {'int', 'long', 'short', 'char', 'bool', 'float', 'double', - 'void', 'auto', 'const', 'static', 'volatile', 'mutable', - 'unsigned', 'signed'} - - # If it's not a type keyword and looks like a parameter name, use it - if (potential_name not in type_keywords and - re.match(r'^[a-zA-Z_][a-zA-Z0-9_]*$', potential_name)): - break - - i -= 1 - - if name_start != -1 and i >= 0: - param_name = decl[name_start:i+1] - param_type = decl[:name_start].strip() - - # Clean up the type - remove trailing &, * and whitespace - param_type = param_type.rstrip('&* \t') - - return { - 'type': param_type, - 'name': param_name - } - - # Fallback: if we can't find a clear parameter name, just return the type - return { - 'type': decl, - 'name': None - } - - -def extract_cpp_signature_details(item): - """ - For each m.def entry in item, parse cpp_signature to extract return type and parameter details. - """ - statements_with_parsed_signatures = [] - for stmt_info in item['m_def_statements']: - parsed = stmt_info['parsed'] - cpp_sig = parsed.get('cpp_signature') - - cpp_params_info = None - if cpp_sig: - try: - cpp_params_info = parse_cpp_signature(cpp_sig) - except Exception as e: - print(f'Failed to parse C++ signature: {e}') - - parsed['cpp_parsed_signature'] = cpp_params_info - statements_with_parsed_signatures.append({ - 'raw': stmt_info['raw'], - 'parsed': parsed - }) - - return { - 'm_def_statements': statements_with_parsed_signatures - } - - -def cpp_type_to_python_type(cpp_type: str) -> str: - if not cpp_type: - return 'Any' - - original = cpp_type.strip() - if not original: - return 'Any' - - # Remove C++ specifiers that don't affect Python type - cleaned = re.sub(r'\b(static|inline|constexpr|thread_local|extern|mutable|const|volatile|endif)\b', '', original) - cleaned = cleaned.replace('&', '').replace('*', '').strip() - cleaned = re.sub(r'\s+', ' ', cleaned).strip() - - # Handle void - if cleaned == 'void': - return 'None' - - # Handle template types — ORDER MATTERS! Must come before internal type checks. - - # std::pair - if cleaned.startswith('std::pair<'): - inner = cleaned[10:-1].strip() # len('std::pair<') == 10 - args = split_template_args(inner) - if len(args) == 2: - t1 = cpp_type_to_python_type(args[0]) - t2 = cpp_type_to_python_type(args[1]) - return f'tuple[{t1}, {t2}]' - else: - print(f'Warning: std::pair with unexpected number of args: {cleaned}') - return 'Any' - - # std::tuple - if cleaned.startswith('std::tuple<'): - inner = cleaned[11:-1].strip() # len('std::tuple<') == 11 - args = split_template_args(inner) - py_types = [cpp_type_to_python_type(arg) for arg in args] - return f"tuple[{', '.join(py_types)}]" - - # std::vector - if cleaned.startswith('std::vector<'): - inner = cleaned[12:-1].strip() # len('std::vector<') == 12 - args = split_template_args(inner) - if len(args) == 1: - inner_py = cpp_type_to_python_type(args[0]) - return f'list[{inner_py}]' - else: - print(f'Warning: std::vector with unexpected args: {cleaned}') - return 'Any' - - # std::optional - if cleaned.startswith('std::optional<'): - inner = cleaned[14:-1].strip() # len('std::optional<') == 14 - args = split_template_args(inner) - if len(args) == 1: - inner_py = cpp_type_to_python_type(args[0]) - return f'Optional[{inner_py}]' - else: - print(f'Warning: std::optional with unexpected args: {cleaned}') - return 'Any' - - # std::string - if re.search(r'\bstd::string\b', original): - return 'str' - - # C-style strings: char*, const char*, char[], etc. - if re.search(r'\b(?:const\s+)?char\s*[\*\[]', original): - return 'str' - - # Boolean - if re.search(r'\bbool\b', cleaned): - return 'bool' - - # Integer types (including fixed-width and common aliases) - if re.search(r'\b(int|long|short|size_t|ssize_t|ptrdiff_t|' - r'int8_t|int16_t|int32_t|int64_t|' - r'uint8_t|uint16_t|uint32_t|uint64_t)\b', cleaned): - return 'int' - - # Floating-point - if re.search(r'\b(float|double|long\s+double)\b', cleaned): - return 'float' - - # torch::Tensor - if re.search(r'\btorch::Tensor\b', original): - return 'torch.Tensor' - - # Unrecognized type - print(f'Warning: Unrecognized C++ type: {original}') - return 'Any' - - -def split_template_args(template_args: str): - """ - Split template arguments, e.g., 'int, std::vector' → ['int', 'std::vector'] - """ - if not template_args.strip(): - return [] - args = [] - current = [] - tracker = BracketTracker() - - for ch in template_args: - if ch in '()[]{}<>': - tracker.update(ch) - if ch == ',' and tracker.is_top_level(): - args.append(''.join(current).strip()) - current = [] + param_lines = [] + for param in parameters: + name = param['name'] + if param['default'] is not None: + param_lines.append(f' {name}: {param["py_type"]} = {param["default"]}') else: - current.append(ch) - - if current: - args.append(''.join(current).strip()) - return args - - -def cpp_default_to_python_default(cpp_default: str): - """ - Convert C++ default value string to valid Python expression string. - """ - if not cpp_default: - return 'None' - - s = cpp_default.strip() - - # Handle string literals: 'bf16' → 'bf16' - # Match: starts and ends with unescaped double quotes - string_match = re.match(r'^"([^"\\]*(?:\\.[^"\\]*)*)"$', s) - if string_match: - return s - - # Handle boolean literals - if s == 'false': - return 'False' - if s == 'true': - return 'True' - - # Handle null-like values: nullptr, nullopt, NULL, etc. - if s in ('nullptr', 'NULL') or 'nullopt' in s: - return 'None' - - # Handle std::tuple({128, 128}) → (128, 128) - tuple_match = re.match(r'std::tuple\s*<[^>]*>\s*\(\s*({.*?})\s*\)', s) - if tuple_match: - inner = tuple_match.group(1) # {128, 128} - inner_py = inner.replace('{', '(').replace('}', ')') - return inner_py - - # Handle std::make_tuple(1, 2, 3) → (1, 2, 3) - make_tuple_match = re.match(r'std::make_tuple\s*\(\s*(.*?)\s*\)', s) - if make_tuple_match: - inner = make_tuple_match.group(1) - # Ensure it's a valid tuple even with one element: add comma if needed? - # But in C++ default args, it's usually multi-element, so we assume valid. - return f'({inner})' - - # Handle std::vector({1,2,3}) → [1, 2, 3] - vector_match = re.match(r'std::vector\s*<[^>]*>\s*\(\s*({.*?})\s*\)', s) - if vector_match: - inner = vector_match.group(1) - inner_py = inner.replace('{', '[').replace('}', ']') - return inner_py - - # Handle numeric literals: integers and floats - if re.match(r'^[+-]?\d+$', s): # integer - return s - if re.match(r'^[+-]?\d*\.\d+([eE][+-]?\d+)?$', s): # float - return s - - # Fallback: unrecognized → warn and return None - print(f'Warning: Unrecognized default value: {s}') - return 'None' - + param_lines.append(f' {name}: {param["py_type"]}') -def generate_pyi_function(item_entry): - parsed = item_entry['parsed'] - py_name = parsed['python_function_name'] - - if parsed.get('is_lambda'): - return f'def {py_name}(*args, **kwargs) -> Any: ...' - - sig_info = parsed.get('cpp_parsed_signature') - default_args = parsed.get('default_args', {}) + if param_lines: + params_block = ',\n'.join(param_lines) + return f'def {py_name}(\n{params_block}\n) -> {return_type}: ...' + return f'def {py_name}() -> {return_type}: ...' - if not sig_info: - return f'def {py_name}(*args, **kwargs) -> Any: ...' - return_type = cpp_type_to_python_type(sig_info['return_type']) - params = sig_info['parameters'] - num_params = len(params) +def generate_pyi_file_content( + parsed_ops, + module_name: str = 'my_module', + wrapper_defaults=None, +): + """Assemble the full .pyi file from parsed TORCH ops.""" + decls = [] - # Build parameter list - param_lines = [] - for i in range(num_params): - param_info = params[i] if i < len(params) else {'type': 'Any', 'name': f'arg{i}'} - param_type = cpp_type_to_python_type(param_info['type']) - param_name = param_info['name'] or f'arg{i}' - - # Replace invalid Python identifiers (e.g., keywords) - if param_name in {'def', 'class', 'from', 'import', 'None', 'True', 'False'}: - param_name = f'{param_name}_' - - # Check for default value - if i in default_args: - cpp_default = default_args[i] - py_default = cpp_default_to_python_default(cpp_default) - param_str = f' {param_name}: {param_type} = {py_default}' - else: - param_str = f' {param_name}: {param_type}' + for parsed in parsed_ops: + name = parsed['python_function_name'] + try: + decl = generate_pyi_function(parsed, wrapper_defaults=wrapper_defaults) + except Exception as e: + decl = f'# ERROR: failed to generate stub for {name}: {e}' + decls.append(decl) - param_lines.append(param_str) + lines = [ + f'# Stubs for module: {module_name}', + '', + 'from typing import Any, Callable, Optional', + 'import torch', + '', + '', + ] - if param_lines: - params_block = ',\n'.join(param_lines) - func_def = f'def {py_name}(\n{params_block}\n) -> {return_type}: ...' - else: - func_def = f'def {py_name}() -> {return_type}: ...' - - return func_def - - -def generate_pyi_file_content(enhanced_results, module_name: str = 'my_module'): - function_decls = [] - has_optional = False - has_torch = False - has_numpy = False - - for item in enhanced_results: - for stmt in item['m_def_statements']: - try: - decl = generate_pyi_function(stmt) - function_decls.append(decl) - - if 'Optional[' in decl: - has_optional = True - if 'torch.Tensor' in decl: - has_torch = True - if 'numpy.ndarray' in decl or 'py::array' in str(stmt): - has_numpy = True - except Exception as e: - func_name = stmt['parsed'].get('python_function_name', 'unknown') - function_decls.append(f'# ERROR: failed to generate stub for {func_name}: {e}') - - imports = ['from typing import Any'] - if has_optional: - imports[0] += ', Optional' - - if has_torch: - imports.append('import torch') - if has_numpy: - imports.append('import numpy') - - lines = [f'# Stubs for module: {module_name}', ''] - lines.extend(imports) - lines.append('') - lines.append('') - - for decl in function_decls: - lines.append(decl) - lines.append('') - lines.append('') + for decl in decls: + lines.extend([decl, '', '']) return '\n'.join(lines) -def generate_pyi_file(name, root, output_dir='.'): - func_index = build_cpp_function_index(root) - results = extract_m_def_statements(root) - - cpp_results = [] - for item in results: - enhanced_item = parse_mdef_and_attach_cpp_signatures(item, func_index) - cpp_item = extract_cpp_signature_details(enhanced_item) - cpp_results.append(cpp_item) +def apply_python_wrappers(parsed_ops, c_py_path): + """Use public wrapper signatures when packing differs from dispatcher schemas.""" + module = ast.parse(Path(c_py_path).read_text()) + by_name = {op['python_function_name']: op for op in parsed_ops} + for node in ast.walk(module): + if not isinstance(node, ast.FunctionDef) or node.name not in by_name: + continue + op = by_name[node.name] + params = adjust_for_c_py_wrapper(node.name, op['parameters']) + types = {p['name']: p['py_type'] for p in params} + defaults = extract_function_defaults(node) + op['parameters'] = [dict(name=arg.arg, py_type=types.get(arg.arg, 'Any'), + default=defaults.get(arg.arg)) for arg in node.args.args] + if node.name == 'fp8_einsum': + next(p for p in op['parameters'] if p['name'] == 'd')['py_type'] = 'torch.Tensor | tuple[torch.Tensor, torch.Tensor]' + elif node.name == 'bf16_mega_gate': + op['return_type'] = 'tuple[torch.Tensor, torch.Tensor]' + next(p for p in op['parameters'] if p['name'] == 'out')['py_type'] = 'Optional[tuple[torch.Tensor, torch.Tensor]]' + elif node.name in { + 'get_symm_buffer_size_for_mega_moe', + 'get_symm_buffer_size_for_nvfp4_mega_moe', + }: + op['return_type'] = 'tuple[int, Callable[[torch.Tensor], tuple[Optional[torch.Tensor], ...]]]' + if node.name == 'nvfp4_mega_moe': + for name in ('shared_l1_weights_tuple_opt', 'shared_l2_weights_tuple_opt'): + next(p for p in op['parameters'] if p['name'] == name)['py_type'] = ( + 'Optional[tuple[torch.Tensor, Optional[torch.Tensor]]]' + ) + elif node.name == 'set_block_size_multiple_of': + op['parameters'][0]['py_type'] = 'int | tuple[int, int] | list[int]' + aliases = {} + for node in ast.walk(module): + if isinstance(node, ast.Assign) and isinstance(node.value, ast.Name): + for target in node.targets: + if isinstance(target, ast.Name): + aliases[target.id] = node.value.id + elif isinstance(node, ast.Call) and _is_globals_update_call(node) and node.args and isinstance(node.args[0], ast.Dict): + for key, value in zip(node.args[0].keys, node.args[0].values): + if isinstance(key, ast.Constant) and isinstance(value, ast.Name): + aliases[key.value] = value.id + for alias, target in aliases.items(): + if alias not in by_name and target in by_name: + op = dict(by_name[target], python_function_name=alias) + parsed_ops.append(op) + + +def generate_pyi_file(name, root, output_dir='.', c_py_path=None): + """Generate stubs/.pyi from csrc/ schemas and optional _C.py wrapper defaults.""" + m_def_statements = extract_m_def_statements(root) + parsed_ops = [parse_m_def_statement(stmt) for stmt in m_def_statements] + + wrapper_defaults = {} + if c_py_path is not None: + c_py_path = Path(c_py_path) + if c_py_path.is_file(): + wrapper_defaults = parse_c_py_metadata(c_py_path) + apply_python_wrappers(parsed_ops, c_py_path) + else: + print(f'Warning: wrapper file not found: {c_py_path}') - pyi_content = generate_pyi_file_content(cpp_results, module_name=name) + pyi_content = generate_pyi_file_content( + parsed_ops, + module_name=name, + wrapper_defaults=wrapper_defaults, + ) output_path = Path(output_dir) / f'{name}.pyi' output_path.parent.mkdir(parents=True, exist_ok=True) diff --git a/setup.py b/setup.py index 7b6deb1c51..5f7994f79d 100644 --- a/setup.py +++ b/setup.py @@ -25,6 +25,8 @@ # Compiler flags cxx_flags = ['-std=c++20', '-O3', '-fPIC', '-Wno-psabi', '-Wno-deprecated-declarations', + '-DPy_LIMITED_API=0x030a0000', + '-DDJ_DISABLE_GIL=1', f'-D_GLIBCXX_USE_CXX11_ABI={int(torch.compiled_with_cxx11_abi())}'] # Sources @@ -80,7 +82,7 @@ def get_platform(): def get_wheel_url(): torch_version = parse(torch.__version__) torch_version = f'{torch_version.major}.{torch_version.minor}' - python_version = f'cp{sys.version_info.major}{sys.version_info.minor}' + python_version = 'cp310-abi3' platform_name = get_platform() deep_gemm_version = get_package_version() cxx11_abi = int(torch._C._GLIBCXX_USE_CXX11_ABI) @@ -100,12 +102,13 @@ def get_ext_modules(): if DG_SKIP_CUDA_BUILD: return [] - return [CUDAExtension(name='deep_gemm._C', + return [CUDAExtension(name='deep_gemm._C_extension', sources=sources, include_dirs=build_include_dirs, libraries=build_libraries, library_dirs=build_library_dirs, - extra_compile_args=cxx_flags)] + extra_compile_args=cxx_flags, + py_limited_api=True)] class CustomBuildPy(build_py): @@ -134,7 +137,7 @@ def prepare_agent_files(self): shutil.copytree(os.path.join(current_dir, name), dst) def generate_pyi_file(self): - generate_pyi_file(name='_C', root='./csrc', output_dir='./stubs') + generate_pyi_file(name='_C', root='./csrc', output_dir='./stubs', c_py_path='./deep_gemm/_C.py') pyi_source = os.path.join(current_dir, 'stubs', '_C.pyi') pyi_target = os.path.join(self.build_lib, 'deep_gemm', '_C.pyi') @@ -215,6 +218,7 @@ def run(self): }, ext_modules=get_ext_modules(), zip_safe=False, + options={'bdist_wheel': {'py_limited_api': 'cp310'}}, cmdclass={ 'build_py': CustomBuildPy, 'bdist_wheel': CachedWheelsCommand,