From a3e7b0ec4118b57ccd98f4825afc50db6b544ead Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Mon, 6 Jul 2026 13:33:25 +0000 Subject: [PATCH 01/28] migrated off of pybind and onto TORCH_LIBRARY Signed-off-by: Chris Leonard --- csrc/apis/attention.hpp | 154 +++++-- csrc/apis/einsum.hpp | 48 +- csrc/apis/gemm.hpp | 413 ++++++++++++----- csrc/apis/hyperconnection.hpp | 33 +- csrc/apis/layout.hpp | 121 +++-- csrc/apis/mega.hpp | 424 ++++++++++++++---- csrc/apis/runtime.hpp | 101 +++-- csrc/jit/device_runtime.hpp | 2 + csrc/jit_kernels/impls/runtime_utils.hpp | 2 +- csrc/jit_kernels/impls/sm100_bf16_gemm.hpp | 2 +- .../jit_kernels/impls/sm100_bf16_mega_moe.hpp | 2 +- csrc/jit_kernels/impls/sm100_bmk_bnk_mn.hpp | 2 +- .../impls/sm100_fp8_fp4_gemm_1d1d.hpp | 2 +- .../impls/sm100_fp8_fp4_mega_moe.hpp | 2 +- .../impls/sm100_tf32_hc_prenorm_gemm.hpp | 2 +- csrc/jit_kernels/impls/sm90_bf16_gemm.hpp | 2 +- csrc/jit_kernels/impls/sm90_bmk_bnk_mn.hpp | 2 +- csrc/jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp | 2 +- csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp | 2 +- .../impls/sm90_tf32_hc_prenorm_gemm.hpp | 2 +- csrc/jit_kernels/impls/smxx_layout.hpp | 2 +- csrc/python_api.cpp | 22 +- csrc/torch_library_macros.hpp | 17 + csrc/torch_library_utils.hpp | 64 +++ csrc/utils/layout.hpp | 2 +- csrc/utils/math.hpp | 2 +- csrc/utils/torch_compat.hpp | 20 + deep_gemm/_C.py | 364 +++++++++++++++ deep_gemm/__init__.py | 73 +-- .../include/deep_gemm/layout/mqa_logits.cuh | 1 + deep_gemm/mega/__init__.py | 12 +- scripts/generate_pyi.py | 34 +- setup.py | 9 +- 33 files changed, 1495 insertions(+), 447 deletions(-) create mode 100644 csrc/torch_library_macros.hpp create mode 100644 csrc/torch_library_utils.hpp create mode 100644 csrc/utils/torch_compat.hpp create mode 100644 deep_gemm/_C.py diff --git a/csrc/apis/attention.hpp b/csrc/apis/attention.hpp index 1abfd5c9b0..d48d776022 100644 --- a/csrc/apis/attention.hpp +++ b/csrc/apis/attention.hpp @@ -15,6 +15,7 @@ #endif #include "layout.hpp" +#include "../torch_library_macros.hpp" namespace deep_gemm::attention { @@ -463,41 +464,126 @@ static torch::Tensor fp8_paged_mqa_logits(const torch::Tensor& q, } #endif -static void register_apis(pybind11::module_& m) { +} // namespace deep_gemm::attention + +namespace deep_gemm::torch_registration { + #if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE - m.def("fp8_gemm_nt_skip_head_mid", &fp8_gemm_nt_skip_head_mid, - py::arg("a"), py::arg("b"), py::arg("d"), py::arg("head_splits"), - py::arg("recipe") = std::nullopt, - py::arg("compiled_dims") = "nk", - py::arg("disable_ue8m0_cast") = false); - m.def("fp8_fp4_mqa_logits", &fp8_fp4_mqa_logits, - py::arg("q"), py::arg("kv"), py::arg("weights"), - py::arg("cu_seq_len_k_start"), py::arg("cu_seq_len_k_end"), - py::arg("clean_logits") = true, - py::arg("max_seqlen_k") = 0, - py::arg("logits_dtype") = torch::kFloat32); - m.def("get_paged_mqa_logits_metadata", &get_paged_mqa_logits_metadata, - py::arg("context_lens"), py::arg("block_kv"), py::arg("num_sms"), - py::arg("indices") = std::nullopt); - m.def("fp8_fp4_paged_mqa_logits", &fp8_fp4_paged_mqa_logits, - py::arg("q"), py::arg("kv_cache"), py::arg("weights"), - py::arg("context_lens"), py::arg("block_table"), py::arg("schedule_meta"), - py::arg("max_context_len"), - py::arg("clean_logits") = false, - py::arg("logits_dtype") = torch::kFloat32, - py::arg("indices") = std::nullopt); - // Legacy API - m.def("fp8_mqa_logits", &fp8_mqa_logits, - py::arg("q"), py::arg("kv"), py::arg("weights"), - py::arg("cu_seq_len_k_start"), py::arg("cu_seq_len_k_end"), - py::arg("clean_logits") = true, - py::arg("max_seqlen_k") = 0); - m.def("fp8_paged_mqa_logits", &fp8_paged_mqa_logits, - py::arg("q"), py::arg("kv_cache"), py::arg("weights"), - py::arg("context_lens"), py::arg("block_table"), py::arg("schedule_meta"), - py::arg("max_context_len"), py::arg("clean_logits") = false, - py::arg("indices") = std::nullopt); +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 c10::List& 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, + const int64_t& logits_dtype) { + return attention::fp8_fp4_mqa_logits( + std::make_tuple(q, q_sf), + std::make_tuple(kv, kv_sf), + weights, cu_seq_len_k_start, cu_seq_len_k_end, + clean_logits, static_cast(max_seqlen_k), + static_cast(logits_dtype)); +} + +static torch::Tensor get_paged_mqa_logits_metadata( + const torch::Tensor& context_lens, const int64_t& block_kv, + const int64_t& num_sms, const c10::optional& 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, + const int64_t& 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, + static_cast(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& fused_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, fused_kv_cache, weights, + context_lens, block_table, schedule_meta, + static_cast(max_context_len), clean_logits, indices); +} #endif + +} // 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[] head_splits, int[]? 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, int logits_dtype=6) -> 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, int logits_dtype=6, Tensor? indices=None) -> Tensor"); + m.def( + "fp8_mqa_logits(Tensor q, Tensor kv, Tensor kv_sf, Tensor weights, Tensor cu_seq_len_k_start, Tensor cu_seq_len_k_end, bool clean_logits=True, int max_seqlen_k=0) -> Tensor"); + m.def( + "fp8_paged_mqa_logits(Tensor q, Tensor kv_cache, Tensor weights, Tensor context_lens, Tensor block_table, Tensor schedule_meta, int max_context_len, bool clean_logits=False, Tensor? indices=None) -> Tensor"); } -} // namespace deep_gemm::attention +TORCH_LIBRARY_IMPL(deep_gemm, CUDA, m) { + using namespace deep_gemm::torch_registration; + +#if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE + m.impl("fp8_gemm_nt_skip_head_mid", TORCH_FN(fp8_gemm_nt_skip_head_mid)); + m.impl("fp8_fp4_mqa_logits", TORCH_FN(fp8_fp4_mqa_logits)); + m.impl("get_paged_mqa_logits_metadata", TORCH_FN(get_paged_mqa_logits_metadata)); + m.impl("fp8_fp4_paged_mqa_logits", TORCH_FN(fp8_fp4_paged_mqa_logits)); + m.impl("fp8_mqa_logits", TORCH_FN(fp8_mqa_logits)); + m.impl("fp8_paged_mqa_logits", TORCH_FN(fp8_paged_mqa_logits)); +#endif +} diff --git a/csrc/apis/einsum.hpp b/csrc/apis/einsum.hpp index ff3ac590c0..bc08c7a14b 100644 --- a/csrc/apis/einsum.hpp +++ b/csrc/apis/einsum.hpp @@ -1,7 +1,6 @@ #pragma once -#include -#include +#include "../utils/torch_compat.hpp" #include "../utils/exception.hpp" #include "../utils/format.hpp" @@ -19,6 +18,7 @@ #include "../jit_kernels/impls/sm120_fp8_fp4_gemm_1d1d.hpp" #include "../jit_kernels/impls/smxx_cublaslt.hpp" #endif +#include "../torch_library_macros.hpp" namespace deep_gemm::einsum { @@ -268,17 +268,41 @@ static void fp8_einsum(const std::string& expr, } #endif -static void register_apis(pybind11::module_& m) { +} // namespace deep_gemm::einsum + +namespace deep_gemm::torch_registration { + #if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE - m.def("einsum", &einsum, - py::arg("expr"), py::arg("a"), py::arg("b"), - py::arg("d"), py::arg("c") = std::nullopt, - py::arg("use_cublaslt") = false); - m.def("fp8_einsum", &fp8_einsum, - py::arg("expr"), py::arg("a"), py::arg("b"), - py::arg("d"), py::arg("c") = std::nullopt, - py::arg("recipe") = std::make_tuple(1, 128, 128)); +static void einsum(const std::string& expr, + const torch::Tensor& a, const torch::Tensor& b, + const torch::Tensor& d, const c10::optional& c, + const bool& use_cublaslt) { + einsum::einsum(expr, a, b, d, c, use_cublaslt); +} + +static void fp8_einsum(const std::string& expr, + const torch::Tensor& a, const torch::Tensor& sfa, + const torch::Tensor& b, const torch::Tensor& sfb, + const torch::Tensor& d, const c10::optional& c, + const c10::optional>& recipe) { + einsum::fp8_einsum(expr, {a, sfa}, {b, sfb}, d, c, list_to_tuple3(recipe.value())); +} #endif + +} // 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, bool use_cublaslt=False) -> ()"); + m.def( + "fp8_einsum(str expr, Tensor a, Tensor sfa, Tensor b, Tensor sfb, Tensor(d!) d, Tensor? c=None, int[]? recipe=None) -> ()"); } -} // namespace deep_gemm::einsum +TORCH_LIBRARY_IMPL(deep_gemm, CUDA, m) { + using namespace deep_gemm::torch_registration; + +#if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE + m.impl("einsum", TORCH_FN(einsum)); + m.impl("fp8_einsum", TORCH_FN(fp8_einsum)); +#endif +} diff --git a/csrc/apis/gemm.hpp b/csrc/apis/gemm.hpp index 902ae699a9..ee0fc50659 100644 --- a/csrc/apis/gemm.hpp +++ b/csrc/apis/gemm.hpp @@ -15,6 +15,7 @@ #include "../jit_kernels/impls/smxx_cublaslt.hpp" #include "layout.hpp" +#include "../torch_library_macros.hpp" namespace deep_gemm::gemm { @@ -782,130 +783,306 @@ static void cublaslt_gemm_tt(const torch::Tensor& a, const torch::Tensor& b, cublaslt_gemm_nt(a.transpose(0, 1), b, d, c); } -static void register_apis(pybind11::module_& m) { +} // namespace deep_gemm::gemm + +namespace deep_gemm::torch_registration { #if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE - // 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); - 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); - 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); - 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); - 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); - - // FP8 GEMM alias names - 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_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_fp8_gemm_nt_masked") = m.attr("m_grouped_fp8_fp4_gemm_nt_masked"); +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) { + 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); +} + +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) { + 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); +} + +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) { + 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); +} + +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) { + 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); +} + +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 c10::optional>& 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.value()), 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 c10::optional>& 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.value()), compiled_dims, use_psum_layout); +} #endif #if DG_TENSORMAP_COMPATIBLE - // BF16 GEMMs - m.def("bf16_gemm_nt", &bf16_gemm_nt, - py::arg("a"), py::arg("b"), py::arg("d"), - py::arg("c") = std::nullopt, - py::arg("compiled_dims") = "nk"); - m.def("bf16_gemm_nn", &bf16_gemm_nn, - py::arg("a"), py::arg("b"), py::arg("d"), - py::arg("c") = std::nullopt, - py::arg("compiled_dims") = "nk"); - m.def("bf16_gemm_tn", &bf16_gemm_tn, - py::arg("a"), py::arg("b"), py::arg("d"), - py::arg("c") = std::nullopt, - py::arg("compiled_dims") = "mn"); - m.def("bf16_gemm_tt", &bf16_gemm_tt, - py::arg("a"), py::arg("b"), py::arg("d"), - py::arg("c") = std::nullopt, - py::arg("compiled_dims") = "mn"); - m.def("m_grouped_bf16_gemm_nt_contiguous", &m_grouped_bf16_gemm_nt_contiguous, - py::arg("a"), py::arg("b"), py::arg("d"), py::arg("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); +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) { + gemm::bf16_gemm_nt(a, b, d, c, compiled_dims); +} + +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) { + gemm::bf16_gemm_nn(a, b, d, c, compiled_dims); +} + +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) { + gemm::bf16_gemm_tn(a, b, d, c, compiled_dims); +} + +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) { + gemm::bf16_gemm_tt(a, b, d, c, compiled_dims); +} + +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_nn_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) { + gemm::m_grouped_bf16_gemm_nn_contiguous( + a, b, d, grouped_layout, compiled_dims, use_psum_layout, ensure_zero_padding); +} + +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); +} #endif - // 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); +static void cublaslt_gemm_nt( + const torch::Tensor& a, const torch::Tensor& b, + const torch::Tensor& d, const c10::optional& c) { + gemm::cublaslt_gemm_nt(a, b, d, c); +} + +static void cublaslt_gemm_nn( + const torch::Tensor& a, const torch::Tensor& b, + const torch::Tensor& d, const c10::optional& c) { + gemm::cublaslt_gemm_nn(a, b, d, c); +} + +static void cublaslt_gemm_tn( + const torch::Tensor& a, const torch::Tensor& b, + const torch::Tensor& d, const c10::optional& c) { + gemm::cublaslt_gemm_tn(a, b, d, c); } -} // namespace deep_gemm::gemm +static void cublaslt_gemm_tt( + const torch::Tensor& a, const torch::Tensor& b, + const torch::Tensor& d, const c10::optional& c) { + gemm::cublaslt_gemm_tt(a, b, 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[]? recipe=None, int[]? recipe_a=None, int[]? recipe_b=None, str compiled_dims='nk', bool disable_ue8m0_cast=False) -> ()"); + m.def( + "fp8_fp4_gemm_nn(Tensor a, Tensor sfa, Tensor b, Tensor sfb, Tensor(d!) d, Tensor? c=None, int[]? recipe=None, int[]? recipe_a=None, int[]? recipe_b=None, str compiled_dims='nk', bool disable_ue8m0_cast=False) -> ()"); + m.def( + "fp8_fp4_gemm_tn(Tensor a, Tensor sfa, Tensor b, Tensor sfb, Tensor(d!) d, Tensor? c=None, int[]? recipe=None, int[]? recipe_a=None, int[]? recipe_b=None, str compiled_dims='mn', bool disable_ue8m0_cast=False) -> ()"); + m.def( + "fp8_fp4_gemm_tt(Tensor a, Tensor sfa, Tensor b, Tensor sfb, Tensor(d!) d, Tensor? c=None, int[]? recipe=None, int[]? recipe_a=None, int[]? recipe_b=None, str compiled_dims='mn', bool disable_ue8m0_cast=False) -> ()"); + m.def( + "m_grouped_fp8_fp4_gemm_nt_contiguous(Tensor a, Tensor sfa, Tensor b, Tensor sfb, Tensor(d!) d, Tensor grouped_layout, int[]? recipe=None, int[]? recipe_a=None, int[]? 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[]? recipe=None, int[]? recipe_a=None, int[]? 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[]? recipe=None, int[]? recipe_a=None, int[]? 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[]? recipe=None, 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[]? recipe=None, 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') -> ()"); + m.def( + "bf16_gemm_nn(Tensor a, Tensor b, Tensor(d!) d, Tensor? c=None, str compiled_dims='nk') -> ()"); + m.def( + "bf16_gemm_tn(Tensor a, Tensor b, Tensor(d!) d, Tensor? c=None, str compiled_dims='mn') -> ()"); + m.def( + "bf16_gemm_tt(Tensor a, Tensor b, Tensor(d!) d, Tensor? c=None, str compiled_dims='mn') -> ()"); + 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) -> ()"); +} + +TORCH_LIBRARY_IMPL(deep_gemm, CUDA, m) { + using namespace deep_gemm::torch_registration; + +#if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE + 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)); +#endif + +#if DG_TENSORMAP_COMPATIBLE + 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(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)); +#endif + + m.impl("cublaslt_gemm_nt", TORCH_FN(cublaslt_gemm_nt)); + m.impl("cublaslt_gemm_nn", TORCH_FN(cublaslt_gemm_nn)); + m.impl("cublaslt_gemm_tn", TORCH_FN(cublaslt_gemm_tn)); + m.impl("cublaslt_gemm_tt", TORCH_FN(cublaslt_gemm_tt)); +} diff --git a/csrc/apis/hyperconnection.hpp b/csrc/apis/hyperconnection.hpp index a695f5938e..897a214950 100644 --- a/csrc/apis/hyperconnection.hpp +++ b/csrc/apis/hyperconnection.hpp @@ -7,6 +7,7 @@ #include "../jit_kernels/impls/sm100_tf32_hc_prenorm_gemm.hpp" #include "../jit_kernels/impls/sm120_tf32_hc_prenorm_gemm.hpp" #endif +#include "../torch_library_macros.hpp" namespace deep_gemm::hyperconnection { @@ -59,15 +60,35 @@ static void tf32_hc_prenorm_gemm(const torch::Tensor& a, DG_HOST_UNREACHABLE("Unsupported architecture"); } } - #endif -static void register_apis(pybind11::module_& m) { +} // namespace deep_gemm::hyperconnection + +namespace deep_gemm::torch_registration { + #if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE - 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); +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); +} #endif + +} // 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; + +#if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE + m.impl("tf32_hc_prenorm_gemm", TORCH_FN(tf32_hc_prenorm_gemm)); +#endif +} diff --git a/csrc/apis/layout.hpp b/csrc/apis/layout.hpp index 5f5c870a62..512e67876c 100644 --- a/csrc/apis/layout.hpp +++ b/csrc/apis/layout.hpp @@ -7,6 +7,7 @@ #if DG_TENSORMAP_COMPATIBLE #include "../jit_kernels/impls/smxx_layout.hpp" #endif +#include "../torch_library_macros.hpp" namespace deep_gemm::layout { @@ -137,34 +138,102 @@ static torch::Tensor transform_k_grouped_sf_into_required_layout(const torch::Te #endif -static void register_apis(pybind11::module_& m) { +} // namespace deep_gemm::layout + +namespace deep_gemm::torch_registration { + #if DG_TENSORMAP_COMPATIBLE - 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); +static torch::Tensor transform_sf_into_required_layout( + const torch::Tensor& sf, const int64_t& mn, const int64_t& k, + const c10::List& 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_mn_major_tma_aligned_tensor(const torch::Tensor& sf) { + return ::deep_gemm::get_mn_major_tma_aligned_tensor(sf); +} + +static torch::Tensor get_mn_major_tma_aligned_packed_ue8m0_tensor( + const torch::Tensor& sf, const c10::optional& psum_layout) { + return ::deep_gemm::get_mn_major_tma_aligned_packed_ue8m0_tensor(sf, 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); +} #endif - 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, - const std::optional& num_groups) { - return heuristics_runtime->get_theoretical_mk_alignment_for_contiguous_layout(expected_m, num_groups); - }, py::arg("expected_m") = std::nullopt, py::arg("num_groups") = std::nullopt); +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)); +} + +static int64_t get_mk_alignment_for_contiguous_layout() { + return heuristics_runtime->get_mk_alignment_for_contiguous_layout(); } -} // namespace deep_gemm::layout +static int64_t get_theoretical_mk_alignment_for_contiguous_layout( + const c10::optional& expected_m, + const c10::optional& num_groups) { + return HeuristicsRuntime::get_theoretical_mk_alignment_for_contiguous_layout( + expected_m.has_value() ? std::make_optional(static_cast(expected_m.value())) : std::nullopt, + num_groups.has_value() ? std::make_optional(static_cast(num_groups.value())) : std::nullopt); +} + +} // namespace deep_gemm::torch_registration + +TORCH_LIBRARY_FRAGMENT(deep_gemm, m) { +#if DG_TENSORMAP_COMPATIBLE + m.def( + "transform_sf_into_required_layout(Tensor 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"); + m.def("get_tma_aligned_size(int x, int element_size) -> int", DEEP_GEMM_IMPL(get_tma_aligned_size)); + m.def("get_mn_major_tma_aligned_tensor(Tensor sf) -> Tensor"); + 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"); +#endif + + m.def("set_mk_alignment_for_contiguous_layout(int new_value) -> ()", + DEEP_GEMM_IMPL(set_mk_alignment_for_contiguous_layout)); + m.def("get_mk_alignment_for_contiguous_layout() -> int", + DEEP_GEMM_IMPL(get_mk_alignment_for_contiguous_layout)); + m.def("get_theoretical_mk_alignment_for_contiguous_layout(int? expected_m=None, int? num_groups=None) -> int", + DEEP_GEMM_IMPL(get_theoretical_mk_alignment_for_contiguous_layout)); +} + +TORCH_LIBRARY_IMPL(deep_gemm, CUDA, m) { + using namespace deep_gemm::torch_registration; + +#if DG_TENSORMAP_COMPATIBLE + m.impl("transform_sf_into_required_layout", + TORCH_FN(transform_sf_into_required_layout)); + m.impl("get_mn_major_tma_aligned_tensor", + TORCH_FN(get_mn_major_tma_aligned_tensor)); + m.impl("get_mn_major_tma_aligned_packed_ue8m0_tensor", + TORCH_FN(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)); +#endif +} diff --git a/csrc/apis/mega.hpp b/csrc/apis/mega.hpp index 93a9138ce4..3712f0574c 100644 --- a/csrc/apis/mega.hpp +++ b/csrc/apis/mega.hpp @@ -1,12 +1,13 @@ #pragma once #include -#include #include -#include +#include +#include "../utils/torch_compat.hpp" #include #include +#include "../utils/math.hpp" #if DG_TENSORMAP_COMPATIBLE #include "../jit/compiler.hpp" @@ -15,6 +16,7 @@ #include "../jit_kernels/impls/sm100_bf16_mega_moe.hpp" #include "../jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp" #include "../jit_kernels/impls/sm100_fp4_fp4_mega_moe.hpp" +#include "../torch_library_macros.hpp" namespace deep_gemm::mega { @@ -33,10 +35,33 @@ 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; + MmaKind mma_kind = MmaKind::BF16; + bool with_sf = false; + int sf_gran_k = 0; + 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; +}; + +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, @@ -104,76 +129,135 @@ 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. + 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.mma_kind = mma_kind; + layout_info.with_sf = with_sf; + layout_info.sf_gran_k = sf_gran_k; + 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 int64_t 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) { + return build_symm_buffer_layout( + num_ranks, num_experts, num_max_tokens_per_rank, num_topk, + hidden, intermediate_hidden, mma_type, activation, num_shared_experts).num_bytes; +} + +using SymmBufferSlice = std::tuple; + +static SymmBufferSlice slice_symm_buffer_from_layout( + const torch::Tensor& buffer, const SymmBufferLayoutInfo& layout_info) { // NOTES: `x_sf` is K-major, while `l1_acts_sf` and `l2_acts_sf` are M-major - // NOTES: for NVFP4, token views are packed E2M1 bytes (2 elements each) and SF - // views pack 4 E4M3 bytes per `int` - const auto is_fp4 = mma_kind == MmaKind::NVFP4; - const auto token_dtype = is_fp4 ? torch::kUInt8 : (with_sf ? torch::kFloat8_e4m3fn : torch::kBFloat16); - const auto hidden_cols = is_fp4 ? hidden / 2 : hidden; - const auto intermediate_cols = is_fp4 ? intermediate_hidden / 2 : intermediate_hidden; - const auto hidden_sf_cols = hidden / (sf_gran_k * 4); - const auto intermediate_sf_cols = intermediate_hidden / (sf_gran_k * 4); - auto slice_input_buffers = [=](const torch::Tensor& buffer) { - 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_cols}, - torch::TensorOptions().dtype(token_dtype).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_sf_cols}, - 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}, - 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}, - torch::TensorOptions().dtype(torch::kFloat32).device(buffer.device())); - - // NOTES: NVFP4 shared experts are BF16 and SF-free; their L1 activations come from - // a caller-provided BF16 tensor, so `shared_l1_acts`/SF views stay undefined - const bool shared_with_sf = with_sf and not is_fp4; - auto shared_l1_acts = is_fp4 ? torch::Tensor() : 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)}, - 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)}, - 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_cols}, - torch::TensorOptions().dtype(token_dtype).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_sf_cols}, - {1, 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_cols}, - torch::TensorOptions().dtype(token_dtype).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_sf_cols}, - {1, 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}; + // NVFP4 token views pack two E2M1 elements per byte, while SF views pack + // four E4M3 values per int32. + const bool is_fp4 = layout_info.mma_kind == MmaKind::NVFP4; + const auto token_dtype = is_fp4 + ? torch::kUInt8 + : (layout_info.with_sf ? torch::kFloat8_e4m3fn : torch::kBFloat16); + const auto hidden_cols = is_fp4 ? layout_info.hidden / 2 : layout_info.hidden; + const auto intermediate_cols = is_fp4 + ? layout_info.intermediate_hidden / 2 + : layout_info.intermediate_hidden; + const auto hidden_sf_cols = layout_info.with_sf + ? layout_info.hidden / (layout_info.sf_gran_k * 4) : 0; + const auto intermediate_sf_cols = layout_info.with_sf + ? layout_info.intermediate_hidden / (layout_info.sf_gran_k * 4) : 0; + auto x = torch::from_blob( + math::advance_ptr(buffer.data_ptr(), layout_info.input_token_base), + {layout_info.num_max_tokens_per_rank, hidden_cols}, + torch::TensorOptions().dtype(token_dtype).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, hidden_sf_cols}, + torch::TensorOptions().dtype(torch::kInt).device(buffer.device())) : torch::Tensor(); + auto topk_idx = torch::from_blob( + 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(), 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())); + + // NVFP4 shared experts consume caller-provided BF16 inputs and do not use SFs. + const bool shared_with_sf = layout_info.with_sf and not is_fp4; + auto shared_l1_acts = is_fp4 ? torch::Tensor() : x; + auto shared_l1_acts_sf = (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 = 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(shared_with_sf ? torch::kFloat8_e4m3fn : torch::kBFloat16).device(buffer.device())) : torch::Tensor(); + auto shared_l2_acts_sf = (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(), layout_info.l1_token_base), + {layout_info.num_ring_tokens, hidden_cols}, + torch::TensorOptions().dtype(token_dtype).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, hidden_sf_cols}, + {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(), layout_info.l2_token_base), + {layout_info.num_ring_tokens, intermediate_cols}, + torch::TensorOptions().dtype(token_dtype).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, intermediate_sf_cols}, + {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); +} + +static SymmBufferSlice slice_symm_buffer_for_mega_moe( + const torch::Tensor& buffer, + 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) { + return slice_symm_buffer_from_layout( + buffer, + build_symm_buffer_layout( + num_ranks, num_experts, num_max_tokens_per_rank, num_topk, + hidden, intermediate_hidden, mma_type, activation, num_shared_experts)); } static void fp8_fp4_mega_moe( @@ -272,22 +356,22 @@ static void fp8_fp4_mega_moe( DG_HOST_ASSERT(cumulative_local_expert_recv_stats->is_contiguous()); } - // Check buffer bytes + // Check buffer bytes and slice views from one shared layout plan. 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, "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) { @@ -434,15 +518,15 @@ static void fp4_fp4_mega_moe( DG_HOST_ASSERT(is_local_cuda_tensor(x_bf16)); } - // Check buffer bytes + // Check buffer bytes and slice views from one shared layout plan. 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, "fp4xfp4", 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_); // Check the optional per-local-expert scales (e.g. modelopt's `weight_scale_2`) @@ -478,7 +562,8 @@ static void fp4_fp4_mega_moe( // 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) { @@ -580,22 +665,22 @@ static void bf16_mega_moe( DG_HOST_ASSERT(cumulative_local_expert_recv_stats->is_contiguous()); } - // Check buffer bytes + // Check buffer bytes and slice views from one shared layout plan. 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) { @@ -622,15 +707,156 @@ static void bf16_mega_moe( sym_buffer.zero_(); } -static void register_apis(pybind11::module_& m) { -#if DG_TENSORMAP_COMPATIBLE - 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("fp4_fp4_mega_moe", &fp4_fp4_mega_moe); - m.def("bf16_mega_moe", &bf16_mega_moe); -#endif +} // namespace deep_gemm::mega + +namespace deep_gemm::torch_registration { + +static int64_t get_token_alignment_for_mega_moe() { + return static_cast(mega::get_token_alignment_for_mega_moe()); } -} // namespace deep_gemm::mega +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 static_cast(mega::get_block_m_for_mega_moe( + static_cast(num_ranks), static_cast(num_experts), + static_cast(num_max_tokens_per_rank), static_cast(num_tokens), + static_cast(num_topk), mma_type)); +} + +static int64_t get_symm_buffer_size_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_topk, + const int64_t& hidden, const int64_t& intermediate_hidden, + const std::string& mma_type, const std::string& activation, + const int64_t& num_shared_experts) { + return mega::get_symm_buffer_size_for_mega_moe( + static_cast(num_ranks), static_cast(num_experts), + static_cast(num_max_tokens_per_rank), static_cast(num_topk), + static_cast(hidden), static_cast(intermediate_hidden), + mma_type, activation, static_cast(num_shared_experts)); +} + +static mega::SymmBufferSlice slice_symm_buffer_for_mega_moe( + const torch::Tensor& buffer, + const int64_t& num_ranks, const int64_t& num_experts, + const int64_t& num_max_tokens_per_rank, const int64_t& num_topk, + const int64_t& hidden, const int64_t& intermediate_hidden, + const std::string& mma_type, const std::string& activation, + const int64_t& num_shared_experts) { + return mega::slice_symm_buffer_for_mega_moe( + buffer, + static_cast(num_ranks), static_cast(num_experts), + static_cast(num_max_tokens_per_rank), static_cast(num_topk), + static_cast(hidden), static_cast(intermediate_hidden), + mma_type, activation, static_cast(num_shared_experts)); +} + +static void fp8_fp4_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 c10::optional& shared_l1_weights, + const c10::optional& shared_l1_weights_sf, + const c10::optional& shared_l2_weights, + const c10::optional& shared_l2_weights_sf, + const c10::optional& cumulative_local_expert_recv_stats, + const torch::Tensor& sym_buffer, + const c10::List& 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 c10::List& recipe, + const std::string& activation, + const c10::optional& activation_clamp, + const bool& fast_math) { + std::optional> shared_l1_opt = std::nullopt; + std::optional> shared_l2_opt = std::nullopt; + if (shared_l1_weights.has_value()) { + DG_HOST_ASSERT(shared_l1_weights_sf.has_value() and shared_l2_weights.has_value() and shared_l2_weights_sf.has_value()); + shared_l1_opt = std::make_tuple(shared_l1_weights.value(), shared_l1_weights_sf.value()); + shared_l2_opt = std::make_tuple(shared_l2_weights.value(), shared_l2_weights_sf.value()); + } else { + DG_HOST_ASSERT(not shared_l1_weights_sf.has_value() and not shared_l2_weights.has_value() and not shared_l2_weights_sf.has_value()); + } + + mega::fp8_fp4_mega_moe( + y, + std::make_tuple(l1_weights, l1_weights_sf), + std::make_tuple(l2_weights, l2_weights_sf), + shared_l1_opt, + shared_l2_opt, + cumulative_local_expert_recv_stats, + sym_buffer, + std::vector(sym_buffer_ptrs.begin(), sym_buffer_ptrs.end()), + static_cast(rank_idx), + static_cast(num_max_tokens_per_rank), + static_cast(num_experts), static_cast(num_topk), + list_to_tuple3(recipe), + activation, + activation_clamp.has_value() + ? std::make_optional(static_cast(activation_clamp.value())) + : std::nullopt, + fast_math); +} + +static void bf16_mega_moe( + const torch::Tensor& y, + const torch::Tensor& l1_weights, + const torch::Tensor& l2_weights, + const c10::optional& shared_l1_weights, + const c10::optional& shared_l2_weights, + const c10::optional& cumulative_local_expert_recv_stats, + const torch::Tensor& sym_buffer, + const c10::List& 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 c10::optional& activation_clamp, + const bool& fast_math) { + mega::bf16_mega_moe( + y, l1_weights, l2_weights, + shared_l1_weights, + shared_l2_weights, + cumulative_local_expert_recv_stats, + sym_buffer, + std::vector(sym_buffer_ptrs.begin(), sym_buffer_ptrs.end()), + static_cast(rank_idx), + static_cast(num_max_tokens_per_rank), + static_cast(num_experts), static_cast(num_topk), + activation, + activation_clamp.has_value() + ? std::make_optional(static_cast(activation_clamp.value())) + : std::nullopt, + fast_math); +} + +} // namespace deep_gemm::torch_registration + +TORCH_LIBRARY_FRAGMENT(deep_gemm, m) { + m.def( + "get_token_alignment_for_mega_moe() -> int", + DEEP_GEMM_IMPL(get_token_alignment_for_mega_moe)); + 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", + DEEP_GEMM_IMPL(get_block_m_for_mega_moe)); + 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=0) -> int", + DEEP_GEMM_IMPL(get_symm_buffer_size_for_mega_moe)); + m.def( + "slice_symm_buffer_for_mega_moe(Tensor buffer, 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=0) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor)"); + m.def( + "fp8_fp4_mega_moe(Tensor(y!) y, Tensor l1_weights, Tensor l1_weights_sf, Tensor l2_weights, Tensor l2_weights_sf, Tensor? shared_l1_weights, Tensor? shared_l1_weights_sf, Tensor? shared_l2_weights, Tensor? shared_l2_weights_sf, Tensor? cumulative_local_expert_recv_stats, Tensor(sym_buffer!) sym_buffer, int[] sym_buffer_ptrs, int rank_idx, int num_max_tokens_per_rank, int num_experts, int num_topk, int[] recipe, str activation, float? activation_clamp, bool fast_math) -> ()"); + m.def( + "bf16_mega_moe(Tensor(y!) y, Tensor l1_weights, Tensor l2_weights, Tensor? shared_l1_weights, Tensor? shared_l2_weights, Tensor? cumulative_local_expert_recv_stats, Tensor(sym_buffer!) 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, bool fast_math) -> ()"); +} + +TORCH_LIBRARY_IMPL(deep_gemm, CUDA, m) { + using namespace deep_gemm::torch_registration; + + m.impl("slice_symm_buffer_for_mega_moe", TORCH_FN(slice_symm_buffer_for_mega_moe)); + m.impl("fp8_fp4_mega_moe", TORCH_FN(fp8_fp4_mega_moe)); + m.impl("bf16_mega_moe", TORCH_FN(bf16_mega_moe)); +} diff --git a/csrc/apis/runtime.hpp b/csrc/apis/runtime.hpp index 58fef941b7..61a094b283 100644 --- a/csrc/apis/runtime.hpp +++ b/csrc/apis/runtime.hpp @@ -2,50 +2,73 @@ #if DG_TENSORMAP_COMPATIBLE #include "../jit/compiler.hpp" +#include "../jit/kernel_runtime.hpp" #endif #include "../jit/device_runtime.hpp" #include "../jit_kernels/heuristics/runtime.hpp" -namespace deep_gemm::runtime { - -static void register_apis(pybind11::module_& m) { - m.def("set_num_sms", [&](const int& new_num_sms) { - device_runtime->set_num_sms(new_num_sms); - }); - m.def("get_num_sms", [&]() { - return device_runtime->get_num_sms(); - }); - m.def("set_tc_util", [&](const int& new_tc_util) { - device_runtime->set_tc_util(new_tc_util); - }); - m.def("get_tc_util", [&]() { - return device_runtime->get_tc_util(); - }); - m.def("set_pdl", [&](const bool& new_enable_pdl) { - device_runtime->set_pdl(new_enable_pdl); - }); - m.def("get_pdl", [&]() { - return device_runtime->get_pdl(); - }); - m.def("set_ignore_compile_dims", [&](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("init", [&](const std::string& library_root_path, const std::string& cuda_home_path_by_python) { +#include "../torch_library_macros.hpp" + +namespace deep_gemm::torch_registration { + +static void set_num_sms(const int64_t& new_num_sms) { + device_runtime->set_num_sms(static_cast(new_num_sms)); +} + +static int64_t get_num_sms() { + return device_runtime->get_num_sms(); +} + +static void set_tc_util(const int64_t& new_tc_util) { + device_runtime->set_tc_util(static_cast(new_tc_util)); +} + +static int64_t get_tc_util() { + return device_runtime->get_tc_util(); +} + +static void set_pdl(const bool& new_enable_pdl) { + device_runtime->set_pdl(new_enable_pdl); +} + +static bool get_pdl() { + return device_runtime->get_pdl(); +} + +static void set_ignore_compile_dims(const bool& new_value) { + heuristics_runtime->set_ignore_compile_dims(new_value); +} + +static void set_block_size_multiple_of(const c10::List& value) { + if (value.size() == 1) { + const int v = static_cast(value[0]); + heuristics_runtime->set_block_size_multiple_of(v, v); + } else { + DG_HOST_ASSERT(value.size() == 2); + heuristics_runtime->set_block_size_multiple_of( + static_cast(value[0]), static_cast(value[1])); + } +} + +static void init(const std::string& library_root_path, + const std::string& cuda_home_path_by_python) { #if DG_TENSORMAP_COMPATIBLE - Compiler::prepare_init(library_root_path, cuda_home_path_by_python); - KernelRuntime::prepare_init(cuda_home_path_by_python); - IncludeParser::prepare_init(library_root_path); + Compiler::prepare_init(library_root_path, cuda_home_path_by_python); + KernelRuntime::prepare_init(cuda_home_path_by_python); + IncludeParser::prepare_init(library_root_path); #endif - }); } -} // namespace deep_gemm::runtime +} // namespace deep_gemm::torch_registration + +TORCH_LIBRARY_FRAGMENT(deep_gemm, m) { + m.def("set_num_sms(int new_num_sms) -> ()", DEEP_GEMM_IMPL(set_num_sms)); + m.def("get_num_sms() -> int", DEEP_GEMM_IMPL(get_num_sms)); + m.def("set_tc_util(int new_tc_util) -> ()", DEEP_GEMM_IMPL(set_tc_util)); + m.def("get_tc_util() -> int", DEEP_GEMM_IMPL(get_tc_util)); + m.def("set_pdl(bool new_enable_pdl) -> ()", DEEP_GEMM_IMPL(set_pdl)); + m.def("get_pdl() -> bool", DEEP_GEMM_IMPL(get_pdl)); + m.def("set_ignore_compile_dims(bool new_value) -> ()", DEEP_GEMM_IMPL(set_ignore_compile_dims)); + m.def("set_block_size_multiple_of(int[] value) -> ()", DEEP_GEMM_IMPL(set_block_size_multiple_of)); + m.def("init(str library_root_path, str cuda_home_path_by_python) -> ()", DEEP_GEMM_IMPL(init)); +} diff --git a/csrc/jit/device_runtime.hpp b/csrc/jit/device_runtime.hpp index d433558bfc..14a3f79442 100644 --- a/csrc/jit/device_runtime.hpp +++ b/csrc/jit/device_runtime.hpp @@ -4,6 +4,8 @@ #include #include +#include "../utils/torch_compat.hpp" + #include "../utils/exception.hpp" #include "../utils/lazy_init.hpp" diff --git a/csrc/jit_kernels/impls/runtime_utils.hpp b/csrc/jit_kernels/impls/runtime_utils.hpp index b2959494a5..739543d6f9 100644 --- a/csrc/jit_kernels/impls/runtime_utils.hpp +++ b/csrc/jit_kernels/impls/runtime_utils.hpp @@ -1,7 +1,7 @@ #pragma once #include -#include +#include "../../utils/torch_compat.hpp" #include "../heuristics/sm90.hpp" #include "../../jit/handle.hpp" diff --git a/csrc/jit_kernels/impls/sm100_bf16_gemm.hpp b/csrc/jit_kernels/impls/sm100_bf16_gemm.hpp index f9b2f361cb..5c2d8b8295 100644 --- a/csrc/jit_kernels/impls/sm100_bf16_gemm.hpp +++ b/csrc/jit_kernels/impls/sm100_bf16_gemm.hpp @@ -1,6 +1,6 @@ #pragma once -#include +#include "../../utils/torch_compat.hpp" #include "../../jit/compiler.hpp" #include "../../jit/device_runtime.hpp" diff --git a/csrc/jit_kernels/impls/sm100_bf16_mega_moe.hpp b/csrc/jit_kernels/impls/sm100_bf16_mega_moe.hpp index 273874cfd9..3be0ef94f3 100644 --- a/csrc/jit_kernels/impls/sm100_bf16_mega_moe.hpp +++ b/csrc/jit_kernels/impls/sm100_bf16_mega_moe.hpp @@ -1,6 +1,6 @@ #pragma once -#include +#include "../../utils/torch_compat.hpp" #include "../../jit/compiler.hpp" #include "../../jit/kernel_runtime.hpp" diff --git a/csrc/jit_kernels/impls/sm100_bmk_bnk_mn.hpp b/csrc/jit_kernels/impls/sm100_bmk_bnk_mn.hpp index 65c9d501c2..2ec5d12a7f 100644 --- a/csrc/jit_kernels/impls/sm100_bmk_bnk_mn.hpp +++ b/csrc/jit_kernels/impls/sm100_bmk_bnk_mn.hpp @@ -1,6 +1,6 @@ #pragma once -#include +#include "../../utils/torch_compat.hpp" #include "../../jit/compiler.hpp" #include "../../jit/device_runtime.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 9e4ba58b60..d2fb50e951 100644 --- a/csrc/jit_kernels/impls/sm100_fp8_fp4_gemm_1d1d.hpp +++ b/csrc/jit_kernels/impls/sm100_fp8_fp4_gemm_1d1d.hpp @@ -1,6 +1,6 @@ #pragma once -#include +#include "../../utils/torch_compat.hpp" #include "../../jit/compiler.hpp" #include "../../jit/device_runtime.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 b323baa3c1..c11d2dd6c7 100644 --- a/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp +++ b/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp @@ -1,6 +1,6 @@ #pragma once -#include +#include "../../utils/torch_compat.hpp" #include "../../jit/compiler.hpp" #include "../../jit/kernel_runtime.hpp" 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 0071e2c57f..e91a5d41fe 100644 --- a/csrc/jit_kernels/impls/sm100_tf32_hc_prenorm_gemm.hpp +++ b/csrc/jit_kernels/impls/sm100_tf32_hc_prenorm_gemm.hpp @@ -1,6 +1,6 @@ #pragma once -#include +#include "../../utils/torch_compat.hpp" #include "../../jit/compiler.hpp" #include "../../jit/device_runtime.hpp" diff --git a/csrc/jit_kernels/impls/sm90_bf16_gemm.hpp b/csrc/jit_kernels/impls/sm90_bf16_gemm.hpp index 24edd46562..e901351b66 100644 --- a/csrc/jit_kernels/impls/sm90_bf16_gemm.hpp +++ b/csrc/jit_kernels/impls/sm90_bf16_gemm.hpp @@ -1,6 +1,6 @@ #pragma once -#include +#include "../../utils/torch_compat.hpp" #include "../../jit/compiler.hpp" #include "../../jit/kernel_runtime.hpp" diff --git a/csrc/jit_kernels/impls/sm90_bmk_bnk_mn.hpp b/csrc/jit_kernels/impls/sm90_bmk_bnk_mn.hpp index 473677b70c..19a1556e6c 100644 --- a/csrc/jit_kernels/impls/sm90_bmk_bnk_mn.hpp +++ b/csrc/jit_kernels/impls/sm90_bmk_bnk_mn.hpp @@ -1,6 +1,6 @@ #pragma once -#include +#include "../../utils/torch_compat.hpp" #include "../../jit/compiler.hpp" #include "../../jit/device_runtime.hpp" diff --git a/csrc/jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp b/csrc/jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp index 120b91faf8..960578047a 100644 --- a/csrc/jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp +++ b/csrc/jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp @@ -1,6 +1,6 @@ #pragma once -#include +#include "../../utils/torch_compat.hpp" #include "../../jit/compiler.hpp" #include "../../jit/device_runtime.hpp" diff --git a/csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp b/csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp index 892edaee7d..c296524b60 100644 --- a/csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp +++ b/csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp @@ -1,6 +1,6 @@ #pragma once -#include +#include "../../utils/torch_compat.hpp" #include "../../jit/compiler.hpp" #include "../../jit/device_runtime.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 c17d1b554e..4a10d69775 100644 --- a/csrc/jit_kernels/impls/sm90_tf32_hc_prenorm_gemm.hpp +++ b/csrc/jit_kernels/impls/sm90_tf32_hc_prenorm_gemm.hpp @@ -1,6 +1,6 @@ #pragma once -#include +#include "../../utils/torch_compat.hpp" #include "../../jit/compiler.hpp" #include "../../jit/device_runtime.hpp" diff --git a/csrc/jit_kernels/impls/smxx_layout.hpp b/csrc/jit_kernels/impls/smxx_layout.hpp index 82de55e198..ef5b6d4080 100644 --- a/csrc/jit_kernels/impls/smxx_layout.hpp +++ b/csrc/jit_kernels/impls/smxx_layout.hpp @@ -1,6 +1,6 @@ #pragma once -#include +#include "../../utils/torch_compat.hpp" #include "../../jit/kernel_runtime.hpp" #include "../../jit/compiler.hpp" diff --git a/csrc/python_api.cpp b/csrc/python_api.cpp index 55c0fa2b33..efd5a622d5 100644 --- a/csrc/python_api.cpp +++ b/csrc/python_api.cpp @@ -1,5 +1,4 @@ -#include -#include +#include "utils/torch_compat.hpp" #include "apis/attention.hpp" #include "apis/einsum.hpp" @@ -9,22 +8,3 @@ #include "apis/mega.hpp" #include "apis/sm90_mega.hpp" #include "apis/runtime.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"; - - // 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::mega::register_sm90_apis(m); - deep_gemm::runtime::register_apis(m); -} diff --git a/csrc/torch_library_macros.hpp b/csrc/torch_library_macros.hpp new file mode 100644 index 0000000000..5ab5688185 --- /dev/null +++ b/csrc/torch_library_macros.hpp @@ -0,0 +1,17 @@ +#pragma once + +#include + +#include "torch_library_utils.hpp" + +namespace deep_gemm::torch_registration { + +using deep_gemm::torch_library_utils::list_to_optional_vector_int; +using deep_gemm::torch_library_utils::list_to_recipe2; +using deep_gemm::torch_library_utils::list_to_recipe3; +using deep_gemm::torch_library_utils::list_to_recipe_variant; +using deep_gemm::torch_library_utils::list_to_tuple3; + +} // namespace deep_gemm::torch_registration + +#define DEEP_GEMM_IMPL(fn) TORCH_FN(deep_gemm::torch_registration::fn) diff --git a/csrc/torch_library_utils.hpp b/csrc/torch_library_utils.hpp new file mode 100644 index 0000000000..77cfa0abcc --- /dev/null +++ b/csrc/torch_library_utils.hpp @@ -0,0 +1,64 @@ +#pragma once + +#include +#include +#include +#include +#include + +#include "utils/exception.hpp" + +namespace deep_gemm::torch_library_utils { + +inline std::optional> list_to_recipe3( + const c10::optional>& recipe) { + if (not recipe.has_value() or recipe->empty()) { + 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() or recipe->empty()) { + 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 c10::List& 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 c10::List& 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_library_utils diff --git a/csrc/utils/layout.hpp b/csrc/utils/layout.hpp index 07a81c4e37..09a9126d68 100644 --- a/csrc/utils/layout.hpp +++ b/csrc/utils/layout.hpp @@ -1,7 +1,7 @@ #pragma once #include -#include +#include "torch_compat.hpp" #include "math.hpp" #include "exception.hpp" diff --git a/csrc/utils/math.hpp b/csrc/utils/math.hpp index 0aa28eb400..f77049584e 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 "torch_compat.hpp" #include "exception.hpp" diff --git a/csrc/utils/torch_compat.hpp b/csrc/utils/torch_compat.hpp new file mode 100644 index 0000000000..9bc017ac14 --- /dev/null +++ b/csrc/utils/torch_compat.hpp @@ -0,0 +1,20 @@ +#pragma once + +#include + +// torch/library.h declares `namespace torch` for op registration but does not +// re-export ATen types. DeepGEMM csrc uses torch::Tensor throughout; under +// Py_LIMITED_API we cannot include torch/python.h or torch/types.h (autograd +// pulls the full Python C-API). Re-export at:: into torch:: instead. +namespace torch { +using namespace at; + +constexpr auto kUInt8 = at::kByte; +constexpr auto kInt8 = at::kChar; +constexpr auto kInt16 = at::kShort; +constexpr auto kInt32 = at::kInt; +constexpr auto kInt64 = at::kLong; +constexpr auto kFloat16 = at::kHalf; +constexpr auto kFloat32 = at::kFloat; +constexpr auto kFloat64 = at::kDouble; +} // namespace torch diff --git a/deep_gemm/_C.py b/deep_gemm/_C.py new file mode 100644 index 0000000000..6975cc7e44 --- /dev/null +++ b/deep_gemm/_C.py @@ -0,0 +1,364 @@ +import torch +from pathlib import Path + +_SCALAR_TYPE = { + torch.float32: 6, + torch.bfloat16: 15, +} + + +def _as_scalar_type(dtype): + if isinstance(dtype, int): + return dtype + return _SCALAR_TYPE.get(dtype, 6) + + +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).""" + bound = {} + for name in names: + op = getattr(_torch_ops, name, None) + if op is None: + return + bound[name] = op + globals().update(bound) + + +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_q(q): + if isinstance(q, tuple): + q_fp = q[0] + q_sf = q[1] if len(q) > 1 else None + else: + q_fp, q_sf = q, None + return q_fp, q_sf + + +def _unpack_kv(kv): + return kv[0], kv[1] + + +def _fp8_fp4_gemm(name, a, b, d, c=None, 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 getattr(_torch_ops, name)( + a_tensor, sfa, b_tensor, sfb, d, c, recipe, recipe_a, recipe_b, + compiled_dims, disable_ue8m0_cast, + ) + + +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): + return _fp8_fp4_gemm('fp8_fp4_gemm_nt', a, b, d, c, recipe, recipe_a, recipe_b, + compiled_dims, disable_ue8m0_cast) + + +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): + return _fp8_fp4_gemm('fp8_fp4_gemm_nn', a, b, d, c, recipe, recipe_a, recipe_b, + compiled_dims, disable_ue8m0_cast) + + +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): + return _fp8_fp4_gemm('fp8_fp4_gemm_tn', a, b, d, c, recipe, recipe_a, recipe_b, + compiled_dims, disable_ue8m0_cast) + + +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): + return _fp8_fp4_gemm('fp8_fp4_gemm_tt', a, b, d, c, recipe, recipe_a, recipe_b, + compiled_dims, disable_ue8m0_cast) + + +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 + + +def _m_grouped_fp8_fp4_gemm(name, 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 getattr(_torch_ops, name)( + a_tensor, sfa, b_tensor, sfb, d, grouped_layout, recipe, recipe_a, recipe_b, + compiled_dims, disable_ue8m0_cast, use_psum_layout, ensure_zero_padding, + expected_m_for_psum_layout, + ) + + +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): + return _m_grouped_fp8_fp4_gemm( + 'm_grouped_fp8_fp4_gemm_nt_contiguous', a, b, d, grouped_layout, recipe, recipe_a, 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): + return _m_grouped_fp8_fp4_gemm( + 'm_grouped_fp8_fp4_gemm_nn_contiguous', a, b, d, grouped_layout, recipe, recipe_a, recipe_b, + compiled_dims, disable_ue8m0_cast, use_psum_layout, ensure_zero_padding, None, + ) + + +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, recipe, recipe_a, recipe_b, + compiled_dims, disable_ue8m0_cast, + ) + + +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 + + +def _k_grouped_fp8_gemm(name, a, b, d, ks_cpu, grouped_layout, c=None, recipe=None, + compiled_dims='mn', use_psum_layout=False): + a_tensor, sfa, b_tensor, sfb = _unpack_ab_pair(a, b) + return getattr(_torch_ops, name)( + a_tensor, sfa, b_tensor, sfb, d, ks_cpu, grouped_layout, c, recipe, compiled_dims, use_psum_layout, + ) + + +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): + return _k_grouped_fp8_gemm( + 'k_grouped_fp8_gemm_tn_contiguous', a, b, 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): + return _k_grouped_fp8_gemm( + 'k_grouped_fp8_gemm_nt_contiguous', a, b, 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), recipe, compiled_dims, disable_ue8m0_cast, + ) + + +def fp8_einsum(expr, a, b, d, c=None, recipe=(1, 128, 128)): + return _torch_ops.fp8_einsum(expr, a[0], a[1], b[0], b[1], d, c, list(recipe) if recipe is not None else None) + + +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): + q_fp, q_sf = _unpack_q(q) + 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, _as_scalar_type(logits_dtype), + ) + + +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 = _unpack_q(q) + 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, _as_scalar_type(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, + ) + + +def get_symm_buffer_size_for_mega_moe(*args, **kwargs): + return _torch_ops.get_symm_buffer_size_for_mega_moe(*args, **kwargs) + + +def slice_symm_buffer_for_mega_moe(buffer, *args, **kwargs): + return _torch_ops.slice_symm_buffer_for_mega_moe(buffer, *args, **kwargs) + + +# DG_TENSORMAP_COMPATIBLE — layout.hpp (schema and impl conditional) +_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', +) + +# DG_TENSORMAP_COMPATIBLE — gemm.hpp (BF16 impl conditional) +_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', +) + +# DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE — hyperconnection.hpp, einsum.hpp, attention.hpp +_bind_guarded_ops( + 'tf32_hc_prenorm_gemm', + 'einsum', + 'get_paged_mqa_logits_metadata', +) + +# DG_TENSORMAP_COMPATIBLE — mega.hpp (C++ impl conditional; matches legacy pybind export guard) +_bind_guarded_ops( + 'get_token_alignment_for_mega_moe', + 'get_block_m_for_mega_moe', +) + + +def fp8_fp4_mega_moe(y, l1_weights, l2_weights, shared_l1_weights, shared_l2_weights, + 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, fast_math): + shared_l1_w = shared_l1_sf = shared_l2_w = shared_l2_sf = None + if shared_l1_weights is not None: + shared_l1_w, shared_l1_sf = shared_l1_weights + shared_l2_w, shared_l2_sf = shared_l2_weights + return _torch_ops.fp8_fp4_mega_moe( + y, l1_weights[0], l1_weights[1], l2_weights[0], l2_weights[1], + shared_l1_w, shared_l1_sf, shared_l2_w, shared_l2_sf, + cumulative_local_expert_recv_stats, sym_buffer, list(sym_buffer_ptrs), rank_idx, + num_max_tokens_per_rank, num_experts, num_topk, list(recipe), activation, + activation_clamp, fast_math, + ) + + +def bf16_mega_moe(y, l1_weights, l2_weights, shared_l1_weights, shared_l2_weights, + cumulative_local_expert_recv_stats, sym_buffer, + sym_buffer_ptrs, rank_idx, num_max_tokens_per_rank, num_experts, num_topk, + activation, activation_clamp, fast_math): + return _torch_ops.bf16_mega_moe( + y, l1_weights, l2_weights, shared_l1_weights, shared_l2_weights, + cumulative_local_expert_recv_stats, sym_buffer, + list(sym_buffer_ptrs), rank_idx, num_max_tokens_per_rank, num_experts, num_topk, + activation, activation_clamp, fast_math, + ) + + +_PUBLIC_API = ( + # Runtime + 'init', + 'set_num_sms', 'get_num_sms', + 'set_tc_util', 'get_tc_util', + 'set_pdl', 'get_pdl', + '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 GEMMs + 'cublaslt_gemm_nt', 'cublaslt_gemm_nn', + 'cublaslt_gemm_tn', 'cublaslt_gemm_tt', + # 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', + 'fp8_gemm_nt_skip_head_mid', + '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', + 'k_grouped_fp8_gemm_tn_contiguous', + 'k_grouped_fp8_gemm_nt_contiguous', + # BF16 GEMMs (guarded) + '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_paged_mqa_logits_metadata', + 'fp8_fp4_paged_mqa_logits', + 'fp8_mqa_logits', + 'fp8_paged_mqa_logits', + # Hyperconnection (guarded) + 'tf32_hc_prenorm_gemm', + # Layout (guarded) + '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', + # Mega MoE + 'get_token_alignment_for_mega_moe', + 'get_block_m_for_mega_moe', + 'get_symm_buffer_size_for_mega_moe', + 'slice_symm_buffer_for_mega_moe', + 'fp8_fp4_mega_moe', + 'bf16_mega_moe', +) + +# Only export names actually bound (guarded ops may be absent on older CUDA builds). +__all__ = [name for name in _PUBLIC_API if name in globals()] diff --git a/deep_gemm/__init__.py b/deep_gemm/__init__.py index 48819ca4ee..cc4943e649 100644 --- a/deep_gemm/__init__.py +++ b/deep_gemm/__init__.py @@ -1,6 +1,9 @@ import os import subprocess -import torch + +# Load the C++ extension and re-export ops through deep_gemm._C. +from ._C import * # noqa: F403 +from . import _C # Set some default environment provided at setup try: @@ -12,74 +15,12 @@ except ImportError: pass -# Configs -from . import _C -from ._C import ( - set_num_sms, - get_num_sms, - set_tc_util, - get_tc_util, - set_ignore_compile_dims, - set_block_size_multiple_of, - set_pdl, - get_pdl, -) - -# cuBLASLt Kernels -from ._C import ( - cublaslt_gemm_nt, cublaslt_gemm_nn, - cublaslt_gemm_tn, cublaslt_gemm_tt, -) - +# Some alias for legacy supports +# TODO: remove these later try: - # DeepGEMM Kernels - from ._C import ( - # FP8 FP4 GEMMs - fp8_fp4_gemm_nt, fp8_fp4_gemm_nn, - fp8_fp4_gemm_tn, fp8_fp4_gemm_tt, - m_grouped_fp8_fp4_gemm_nt_contiguous, - m_grouped_fp8_fp4_gemm_nn_contiguous, - m_grouped_fp8_fp4_gemm_nt_masked, - # FP8 GEMMs - fp8_gemm_nt, fp8_gemm_nn, - fp8_gemm_tn, fp8_gemm_tt, - fp8_gemm_nt_skip_head_mid, - m_grouped_fp8_gemm_nt_contiguous, - m_grouped_fp8_gemm_nn_contiguous, - m_grouped_fp8_gemm_nt_masked, - k_grouped_fp8_gemm_nt_contiguous, - k_grouped_fp8_gemm_tn_contiguous, - # BF16 GEMMs - bf16_gemm_nt, bf16_gemm_nn, - bf16_gemm_tn, bf16_gemm_tt, - m_grouped_bf16_gemm_nt_contiguous, - m_grouped_bf16_gemm_nn_contiguous, - m_grouped_bf16_gemm_nt_masked, - k_grouped_bf16_gemm_tn_contiguous, - # Einsum kernels - einsum, - fp8_einsum, - # Attention kernels - fp8_fp4_mqa_logits, - get_paged_mqa_logits_metadata, - fp8_fp4_paged_mqa_logits, - # Attention kernels (legacy) - fp8_mqa_logits, - fp8_paged_mqa_logits, - # Hyperconnection kernels - tf32_hc_prenorm_gemm, - # Layout kernels - transform_sf_into_required_layout, - # MegaMoE - get_block_m_for_mega_moe, - ) - - # Some alias for legacy supports - # TODO: remove these later fp8_m_grouped_gemm_nt_masked = m_grouped_fp8_gemm_nt_masked bf16_m_grouped_gemm_nt_masked = m_grouped_bf16_gemm_nt_masked -except ImportError: - # Expected behavior for CUDA runtime version before 12.1 +except NameError: pass # Mega kernels diff --git a/deep_gemm/include/deep_gemm/layout/mqa_logits.cuh b/deep_gemm/include/deep_gemm/layout/mqa_logits.cuh index 14485b54c1..3358eda73e 100644 --- a/deep_gemm/include/deep_gemm/layout/mqa_logits.cuh +++ b/deep_gemm/include/deep_gemm/layout/mqa_logits.cuh @@ -1,5 +1,6 @@ #pragma once +#include #include #include diff --git a/deep_gemm/mega/__init__.py b/deep_gemm/mega/__init__.py index 3da1297cf4..e43a5ab409 100644 --- a/deep_gemm/mega/__init__.py +++ b/deep_gemm/mega/__init__.py @@ -36,12 +36,20 @@ def __init__(self, group: dist.ProcessGroup, self.activation = activation # Allocate a symmetric buffer - num_bytes, slice_input_buffers = _C.get_symm_buffer_size_for_mega_moe( + num_bytes = _C.get_symm_buffer_size_for_mega_moe( group.size(), num_experts, num_max_tokens_per_rank, num_topk, hidden, intermediate_hidden, mma_type, activation, - num_shared_experts + num_shared_experts, + ) + slice_input_buffers = lambda buffer: _C.slice_symm_buffer_for_mega_moe( + buffer, + group.size(), num_experts, + num_max_tokens_per_rank, num_topk, + hidden, intermediate_hidden, + mma_type, activation, + num_shared_experts, ) allocator = torch if group.size() == 1 else symm_mem self.buffer = allocator.empty(num_bytes, dtype=torch.int8, device='cuda') diff --git a/scripts/generate_pyi.py b/scripts/generate_pyi.py index df7490d410..f54fce87b4 100644 --- a/scripts/generate_pyi.py +++ b/scripts/generate_pyi.py @@ -148,6 +148,19 @@ def is_top_level(self): self.angle == 0) +def is_torch_schema_string(schema_or_name: str) -> bool: + """Return True if the m.def string literal is a TORCH_LIBRARY schema.""" + return ' -> ' in schema_or_name or '(' in schema_or_name + + +def extract_torch_op_name(schema: str) -> str: + """Extract the operator name from a TORCH schema string.""" + paren_pos = schema.find('(') + if paren_pos == -1: + return schema.strip() + return schema[:paren_pos].strip() + + def extract_m_def_statements(root_path): """ Scan all c files under root_path and extract all m.def(...) statements. @@ -270,17 +283,26 @@ def parse_m_def_statement(m_def_str): if current: args_list.append(''.join(current).strip()) - if len(args_list) < 2: - raise ValueError(f'[{m_def_str}] m.def has insufficient arguments') + if len(args_list) < 1: + raise ValueError(f'[{m_def_str}] m.def has no 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') + schema_or_name = str_match.group(1) + if is_torch_schema_string(schema_or_name): + result['is_torch_schema'] = True + result['python_function_name'] = extract_torch_op_name(schema_or_name) + return result + + if len(args_list) < 2: + raise ValueError(f'[{m_def_str}] m.def has insufficient arguments') + + result['python_function_name'] = schema_or_name + cpp_func_part = args_list[1].strip() if cpp_func_part.startswith('&'): cpp_func_part = cpp_func_part[1:].strip() @@ -400,7 +422,7 @@ def parse_mdef_and_attach_cpp_signatures(item, func_index): if cpp_func_name and cpp_func_name in func_index: cpp_sig = func_index[cpp_func_name] else: - if not parsed['is_lambda']: + if not parsed['is_lambda'] and not parsed.get('is_torch_schema'): print(f'Warning: C++ function "{cpp_func_name}" not found in any .cpp file') parsed['cpp_signature'] = cpp_sig diff --git a/setup.py b/setup.py index c4d74ae929..48482efc9c 100644 --- a/setup.py +++ b/setup.py @@ -26,7 +26,8 @@ # Compiler flags cxx_flags = ['-std=c++17', '-O3', '-fPIC', '-Wno-psabi', '-Wno-deprecated-declarations', - f'-D_GLIBCXX_USE_CXX11_ABI={int(torch.compiled_with_cxx11_abi())}'] + f'-D_GLIBCXX_USE_CXX11_ABI={int(torch.compiled_with_cxx11_abi())}', + '-DPy_LIMITED_API=0x03090000'] if DG_JIT_USE_RUNTIME_API: cxx_flags.append('-DDG_JIT_USE_RUNTIME_API') @@ -103,12 +104,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): @@ -207,6 +209,7 @@ def run(self): }, ext_modules=get_ext_modules(), zip_safe=False, + options={'bdist_wheel': {'py_limited_api': 'cp39'}}, cmdclass={ 'build_py': CustomBuildPy, 'bdist_wheel': CachedWheelsCommand, From c91c2098a3bba47ae4becf8d02eddd85330d424a Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Mon, 6 Jul 2026 17:01:19 +0000 Subject: [PATCH 02/28] added #if DG..._COMPATABLE guards to the function schema registers to match what was in the legacy code Signed-off-by: Chris Leonard --- csrc/apis/attention.hpp | 2 ++ csrc/apis/einsum.hpp | 2 ++ csrc/apis/gemm.hpp | 6 ++++++ csrc/apis/hyperconnection.hpp | 2 ++ csrc/apis/mega.hpp | 4 ++++ 5 files changed, 16 insertions(+) diff --git a/csrc/apis/attention.hpp b/csrc/apis/attention.hpp index d48d776022..001235fac0 100644 --- a/csrc/apis/attention.hpp +++ b/csrc/apis/attention.hpp @@ -561,6 +561,7 @@ static torch::Tensor fp8_paged_mqa_logits( } // namespace deep_gemm::torch_registration TORCH_LIBRARY_FRAGMENT(deep_gemm, m) { +#if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE m.def( "fp8_gemm_nt_skip_head_mid(Tensor a, Tensor sfa, Tensor b, Tensor sfb, Tensor(d!) d, int[] head_splits, int[]? recipe=None, str compiled_dims='nk', bool disable_ue8m0_cast=False) -> ()"); m.def( @@ -573,6 +574,7 @@ TORCH_LIBRARY_FRAGMENT(deep_gemm, m) { "fp8_mqa_logits(Tensor q, Tensor kv, Tensor kv_sf, Tensor weights, Tensor cu_seq_len_k_start, Tensor cu_seq_len_k_end, bool clean_logits=True, int max_seqlen_k=0) -> Tensor"); m.def( "fp8_paged_mqa_logits(Tensor q, Tensor kv_cache, Tensor weights, Tensor context_lens, Tensor block_table, Tensor schedule_meta, int max_context_len, bool clean_logits=False, Tensor? indices=None) -> Tensor"); +#endif } TORCH_LIBRARY_IMPL(deep_gemm, CUDA, m) { diff --git a/csrc/apis/einsum.hpp b/csrc/apis/einsum.hpp index bc08c7a14b..a2bd1f84cd 100644 --- a/csrc/apis/einsum.hpp +++ b/csrc/apis/einsum.hpp @@ -292,10 +292,12 @@ static void fp8_einsum(const std::string& expr, } // namespace deep_gemm::torch_registration TORCH_LIBRARY_FRAGMENT(deep_gemm, m) { +#if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE m.def( "einsum(str expr, Tensor a, Tensor b, Tensor(d!) d, Tensor? c=None, bool use_cublaslt=False) -> ()"); m.def( "fp8_einsum(str expr, Tensor a, Tensor sfa, Tensor b, Tensor sfb, Tensor(d!) d, Tensor? c=None, int[]? recipe=None) -> ()"); +#endif } TORCH_LIBRARY_IMPL(deep_gemm, CUDA, m) { diff --git a/csrc/apis/gemm.hpp b/csrc/apis/gemm.hpp index ee0fc50659..da158a90d9 100644 --- a/csrc/apis/gemm.hpp +++ b/csrc/apis/gemm.hpp @@ -1012,6 +1012,7 @@ static void cublaslt_gemm_tt( } // namespace deep_gemm::torch_registration TORCH_LIBRARY_FRAGMENT(deep_gemm, m) { +#if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE // GEMM — FP8/FP4 m.def( "fp8_fp4_gemm_nt(Tensor a, Tensor sfa, Tensor b, Tensor sfb, Tensor(d!) d, Tensor? c=None, int[]? recipe=None, int[]? recipe_a=None, int[]? recipe_b=None, str compiled_dims='nk', bool disable_ue8m0_cast=False) -> ()"); @@ -1031,6 +1032,9 @@ TORCH_LIBRARY_FRAGMENT(deep_gemm, m) { "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[]? recipe=None, 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[]? recipe=None, str compiled_dims='mn', bool use_psum_layout=False) -> ()"); +#endif + +#if DG_TENSORMAP_COMPATIBLE // GEMM — BF16 m.def( "bf16_gemm_nt(Tensor a, Tensor b, Tensor(d!) d, Tensor? c=None, str compiled_dims='nk') -> ()"); @@ -1048,6 +1052,8 @@ TORCH_LIBRARY_FRAGMENT(deep_gemm, m) { "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) -> ()"); +#endif + // 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) -> ()"); diff --git a/csrc/apis/hyperconnection.hpp b/csrc/apis/hyperconnection.hpp index 897a214950..ffe04c0672 100644 --- a/csrc/apis/hyperconnection.hpp +++ b/csrc/apis/hyperconnection.hpp @@ -81,8 +81,10 @@ static void tf32_hc_prenorm_gemm(const torch::Tensor& a, const torch::Tensor& b, } // namespace deep_gemm::torch_registration TORCH_LIBRARY_FRAGMENT(deep_gemm, m) { +#if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE m.def( "tf32_hc_prenorm_gemm(Tensor a, Tensor b, Tensor(d!) d, Tensor(sqr_sum!) sqr_sum, int? num_splits=None) -> ()"); +#endif } TORCH_LIBRARY_IMPL(deep_gemm, CUDA, m) { diff --git a/csrc/apis/mega.hpp b/csrc/apis/mega.hpp index 3712f0574c..adc98750c0 100644 --- a/csrc/apis/mega.hpp +++ b/csrc/apis/mega.hpp @@ -836,6 +836,7 @@ static void bf16_mega_moe( } // namespace deep_gemm::torch_registration TORCH_LIBRARY_FRAGMENT(deep_gemm, m) { +#if DG_TENSORMAP_COMPATIBLE m.def( "get_token_alignment_for_mega_moe() -> int", DEEP_GEMM_IMPL(get_token_alignment_for_mega_moe)); @@ -851,12 +852,15 @@ TORCH_LIBRARY_FRAGMENT(deep_gemm, m) { "fp8_fp4_mega_moe(Tensor(y!) y, Tensor l1_weights, Tensor l1_weights_sf, Tensor l2_weights, Tensor l2_weights_sf, Tensor? shared_l1_weights, Tensor? shared_l1_weights_sf, Tensor? shared_l2_weights, Tensor? shared_l2_weights_sf, Tensor? cumulative_local_expert_recv_stats, Tensor(sym_buffer!) sym_buffer, int[] sym_buffer_ptrs, int rank_idx, int num_max_tokens_per_rank, int num_experts, int num_topk, int[] recipe, str activation, float? activation_clamp, bool fast_math) -> ()"); m.def( "bf16_mega_moe(Tensor(y!) y, Tensor l1_weights, Tensor l2_weights, Tensor? shared_l1_weights, Tensor? shared_l2_weights, Tensor? cumulative_local_expert_recv_stats, Tensor(sym_buffer!) 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, bool fast_math) -> ()"); +#endif } TORCH_LIBRARY_IMPL(deep_gemm, CUDA, m) { using namespace deep_gemm::torch_registration; +#if DG_TENSORMAP_COMPATIBLE m.impl("slice_symm_buffer_for_mega_moe", TORCH_FN(slice_symm_buffer_for_mega_moe)); m.impl("fp8_fp4_mega_moe", TORCH_FN(fp8_fp4_mega_moe)); m.impl("bf16_mega_moe", TORCH_FN(bf16_mega_moe)); +#endif } From 926d9939e09b311e23bd4551bb008c06f74b9226 Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Wed, 8 Jul 2026 16:38:03 +0000 Subject: [PATCH 03/28] updated __init__.py import to align with legacy code. Removed pybind helper code from generate_pyi.py Signed-off-by: Chris Leonard --- csrc/apis/runtime.hpp | 6 +- deep_gemm/_C.py | 347 ++++++++++--------- deep_gemm/__init__.py | 72 +++- scripts/generate_pyi.py | 721 ++-------------------------------------- 4 files changed, 275 insertions(+), 871 deletions(-) diff --git a/csrc/apis/runtime.hpp b/csrc/apis/runtime.hpp index 61a094b283..c0173c5b6d 100644 --- a/csrc/apis/runtime.hpp +++ b/csrc/apis/runtime.hpp @@ -53,9 +53,9 @@ static void set_block_size_multiple_of(const c10::List& value) { static void init(const std::string& library_root_path, const std::string& cuda_home_path_by_python) { #if DG_TENSORMAP_COMPATIBLE - Compiler::prepare_init(library_root_path, cuda_home_path_by_python); - KernelRuntime::prepare_init(cuda_home_path_by_python); - IncludeParser::prepare_init(library_root_path); + Compiler::prepare_init(library_root_path, cuda_home_path_by_python); + KernelRuntime::prepare_init(cuda_home_path_by_python); + IncludeParser::prepare_init(library_root_path); #endif } diff --git a/deep_gemm/_C.py b/deep_gemm/_C.py index 6975cc7e44..5fbdbde1b4 100644 --- a/deep_gemm/_C.py +++ b/deep_gemm/_C.py @@ -79,152 +79,174 @@ def _unpack_kv(kv): return kv[0], kv[1] -def _fp8_fp4_gemm(name, a, b, d, c=None, 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 getattr(_torch_ops, name)( - a_tensor, sfa, b_tensor, sfb, d, c, recipe, recipe_a, recipe_b, - compiled_dims, disable_ue8m0_cast, +def _register_deep_gemm_kernels(): + """Export DeepGEMM kernels only when C++ ops are registered.""" + 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): + 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, recipe, recipe_a, recipe_b, + compiled_dims, disable_ue8m0_cast, + ) + + 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): + 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, recipe, recipe_a, recipe_b, + compiled_dims, disable_ue8m0_cast, + ) + + 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): + 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, recipe, recipe_a, recipe_b, + compiled_dims, disable_ue8m0_cast, + ) + + 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): + 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, recipe, recipe_a, recipe_b, + compiled_dims, disable_ue8m0_cast, + ) + + 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, recipe, recipe_a, 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, recipe, recipe_a, 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, recipe, recipe_a, 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), recipe, compiled_dims, disable_ue8m0_cast, + ) + + def fp8_einsum(expr, a, b, d, c=None, recipe=(1, 128, 128)): + return _torch_ops.fp8_einsum(expr, a[0], a[1], b[0], b[1], d, c, list(recipe) if recipe is not None else None) + + 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): + q_fp, q_sf = _unpack_q(q) + 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, _as_scalar_type(logits_dtype), + ) + + 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 = _unpack_q(q) + 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, _as_scalar_type(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, + }) + + # DG_TENSORMAP_COMPATIBLE — gemm.hpp (BF16 impl conditional) + _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', ) - -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): - return _fp8_fp4_gemm('fp8_fp4_gemm_nt', a, b, d, c, recipe, recipe_a, recipe_b, - compiled_dims, disable_ue8m0_cast) - - -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): - return _fp8_fp4_gemm('fp8_fp4_gemm_nn', a, b, d, c, recipe, recipe_a, recipe_b, - compiled_dims, disable_ue8m0_cast) - - -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): - return _fp8_fp4_gemm('fp8_fp4_gemm_tn', a, b, d, c, recipe, recipe_a, recipe_b, - compiled_dims, disable_ue8m0_cast) - - -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): - return _fp8_fp4_gemm('fp8_fp4_gemm_tt', a, b, d, c, recipe, recipe_a, recipe_b, - compiled_dims, disable_ue8m0_cast) - - -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 - - -def _m_grouped_fp8_fp4_gemm(name, 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 getattr(_torch_ops, name)( - a_tensor, sfa, b_tensor, sfb, d, grouped_layout, recipe, recipe_a, recipe_b, - compiled_dims, disable_ue8m0_cast, use_psum_layout, ensure_zero_padding, - expected_m_for_psum_layout, - ) - - -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): - return _m_grouped_fp8_fp4_gemm( - 'm_grouped_fp8_fp4_gemm_nt_contiguous', a, b, d, grouped_layout, recipe, recipe_a, 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): - return _m_grouped_fp8_fp4_gemm( - 'm_grouped_fp8_fp4_gemm_nn_contiguous', a, b, d, grouped_layout, recipe, recipe_a, recipe_b, - compiled_dims, disable_ue8m0_cast, use_psum_layout, ensure_zero_padding, None, - ) - - -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, recipe, recipe_a, recipe_b, - compiled_dims, disable_ue8m0_cast, - ) - - -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 - - -def _k_grouped_fp8_gemm(name, a, b, d, ks_cpu, grouped_layout, c=None, recipe=None, - compiled_dims='mn', use_psum_layout=False): - a_tensor, sfa, b_tensor, sfb = _unpack_ab_pair(a, b) - return getattr(_torch_ops, name)( - a_tensor, sfa, b_tensor, sfb, d, ks_cpu, grouped_layout, c, recipe, compiled_dims, use_psum_layout, - ) - - -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): - return _k_grouped_fp8_gemm( - 'k_grouped_fp8_gemm_tn_contiguous', a, b, 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): - return _k_grouped_fp8_gemm( - 'k_grouped_fp8_gemm_nt_contiguous', a, b, 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), recipe, compiled_dims, disable_ue8m0_cast, - ) - - -def fp8_einsum(expr, a, b, d, c=None, recipe=(1, 128, 128)): - return _torch_ops.fp8_einsum(expr, a[0], a[1], b[0], b[1], d, c, list(recipe) if recipe is not None else None) - - -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): - q_fp, q_sf = _unpack_q(q) - 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, _as_scalar_type(logits_dtype), + # DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE — einsum.hpp, attention.hpp, hyperconnection.hpp + _bind_guarded_ops( + 'einsum', + 'tf32_hc_prenorm_gemm', + 'get_paged_mqa_logits_metadata', ) - -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 = _unpack_q(q) - 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, _as_scalar_type(logits_dtype), indices, + # DG_TENSORMAP_COMPATIBLE — layout.hpp (schema and impl conditional) + _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', ) -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, - ) +_register_deep_gemm_kernels() def get_symm_buffer_size_for_mega_moe(*args, **kwargs): @@ -235,34 +257,6 @@ def slice_symm_buffer_for_mega_moe(buffer, *args, **kwargs): return _torch_ops.slice_symm_buffer_for_mega_moe(buffer, *args, **kwargs) -# DG_TENSORMAP_COMPATIBLE — layout.hpp (schema and impl conditional) -_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', -) - -# DG_TENSORMAP_COMPATIBLE — gemm.hpp (BF16 impl conditional) -_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', -) - -# DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE — hyperconnection.hpp, einsum.hpp, attention.hpp -_bind_guarded_ops( - 'tf32_hc_prenorm_gemm', - 'einsum', - 'get_paged_mqa_logits_metadata', -) - # DG_TENSORMAP_COMPATIBLE — mega.hpp (C++ impl conditional; matches legacy pybind export guard) _bind_guarded_ops( 'get_token_alignment_for_mega_moe', @@ -299,7 +293,7 @@ def bf16_mega_moe(y, l1_weights, l2_weights, shared_l1_weights, shared_l2_weight ) -_PUBLIC_API = ( +_UNCONDITIONAL_API = ( # Runtime 'init', 'set_num_sms', 'get_num_sms', @@ -313,6 +307,14 @@ def bf16_mega_moe(y, l1_weights, l2_weights, shared_l1_weights, shared_l2_weight # cuBLASLt GEMMs 'cublaslt_gemm_nt', 'cublaslt_gemm_nn', 'cublaslt_gemm_tn', 'cublaslt_gemm_tt', + # Mega MoE (imported via deep_gemm.mega; always defined, fails at call if unregistered) + 'get_symm_buffer_size_for_mega_moe', + 'slice_symm_buffer_for_mega_moe', + 'fp8_fp4_mega_moe', + 'bf16_mega_moe', +) + +_DEEP_GEMM_API = ( # FP8/FP4 GEMMs 'fp8_fp4_gemm_nt', 'fp8_fp4_gemm_nn', 'fp8_fp4_gemm_tn', 'fp8_fp4_gemm_tt', @@ -351,14 +353,9 @@ def bf16_mega_moe(y, l1_weights, l2_weights, shared_l1_weights, shared_l2_weight 'get_mn_major_tma_aligned_tensor', 'get_mn_major_tma_aligned_packed_ue8m0_tensor', 'get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor', - # Mega MoE + # Mega helpers (guarded) 'get_token_alignment_for_mega_moe', 'get_block_m_for_mega_moe', - 'get_symm_buffer_size_for_mega_moe', - 'slice_symm_buffer_for_mega_moe', - 'fp8_fp4_mega_moe', - 'bf16_mega_moe', ) -# Only export names actually bound (guarded ops may be absent on older CUDA builds). -__all__ = [name for name in _PUBLIC_API if name in globals()] +__all__ = list(_UNCONDITIONAL_API) + [name for name in _DEEP_GEMM_API if name in globals()] diff --git a/deep_gemm/__init__.py b/deep_gemm/__init__.py index cc4943e649..d8fee3ec5a 100644 --- a/deep_gemm/__init__.py +++ b/deep_gemm/__init__.py @@ -1,10 +1,6 @@ import os import subprocess -# Load the C++ extension and re-export ops through deep_gemm._C. -from ._C import * # noqa: F403 -from . import _C - # Set some default environment provided at setup try: # noinspection PyUnresolvedReferences @@ -15,12 +11,74 @@ except ImportError: pass -# Some alias for legacy supports -# TODO: remove these later +# Configs +from . import _C +from ._C import ( + set_num_sms, + get_num_sms, + set_tc_util, + get_tc_util, + set_ignore_compile_dims, + set_block_size_multiple_of, + set_pdl, + get_pdl, +) + +# cuBLASLt Kernels +from ._C import ( + cublaslt_gemm_nt, cublaslt_gemm_nn, + cublaslt_gemm_tn, cublaslt_gemm_tt, +) + try: + # DeepGEMM Kernels + from ._C import ( + # FP8 FP4 GEMMs + fp8_fp4_gemm_nt, fp8_fp4_gemm_nn, + fp8_fp4_gemm_tn, fp8_fp4_gemm_tt, + m_grouped_fp8_fp4_gemm_nt_contiguous, + m_grouped_fp8_fp4_gemm_nn_contiguous, + m_grouped_fp8_fp4_gemm_nt_masked, + # FP8 GEMMs + fp8_gemm_nt, fp8_gemm_nn, + fp8_gemm_tn, fp8_gemm_tt, + fp8_gemm_nt_skip_head_mid, + m_grouped_fp8_gemm_nt_contiguous, + m_grouped_fp8_gemm_nn_contiguous, + m_grouped_fp8_gemm_nt_masked, + k_grouped_fp8_gemm_nt_contiguous, + k_grouped_fp8_gemm_tn_contiguous, + # BF16 GEMMs + bf16_gemm_nt, bf16_gemm_nn, + bf16_gemm_tn, bf16_gemm_tt, + m_grouped_bf16_gemm_nt_contiguous, + m_grouped_bf16_gemm_nn_contiguous, + m_grouped_bf16_gemm_nt_masked, + k_grouped_bf16_gemm_tn_contiguous, + # Einsum kernels + einsum, + fp8_einsum, + # Attention kernels + fp8_fp4_mqa_logits, + get_paged_mqa_logits_metadata, + fp8_fp4_paged_mqa_logits, + # Attention kernels (legacy) + fp8_mqa_logits, + fp8_paged_mqa_logits, + # Hyperconnection kernels + tf32_hc_prenorm_gemm, + # Layout kernels + transform_sf_into_required_layout, + # MegaMoE + get_block_m_for_mega_moe, + ) + + # Some alias for legacy supports + # TODO: remove these later fp8_m_grouped_gemm_nt_masked = m_grouped_fp8_gemm_nt_masked bf16_m_grouped_gemm_nt_masked = m_grouped_bf16_gemm_nt_masked -except NameError: +except ImportError: + # Expected behavior for CUDA runtime version before 12.1 pass # Mega kernels diff --git a/scripts/generate_pyi.py b/scripts/generate_pyi.py index f54fce87b4..72bb2c0dce 100644 --- a/scripts/generate_pyi.py +++ b/scripts/generate_pyi.py @@ -2,98 +2,6 @@ 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: @@ -148,11 +56,6 @@ def is_top_level(self): self.angle == 0) -def is_torch_schema_string(schema_or_name: str) -> bool: - """Return True if the m.def string literal is a TORCH_LIBRARY schema.""" - return ' -> ' in schema_or_name or '(' in schema_or_name - - def extract_torch_op_name(schema: str) -> str: """Extract the operator name from a TORCH schema string.""" paren_pos = schema.find('(') @@ -163,14 +66,12 @@ def extract_torch_op_name(schema: str) -> str: def extract_m_def_statements(root_path): """ - Scan all c files under root_path and extract all m.def(...) statements. + Scan all C++ files under root_path and extract all m.def(...) statements. + Supports multi-line m.def(...) calls. """ results = [] extensions = {'.hpp', '.cpp', '.h', '.cc'} - # Regex: match m.def( ... ), supports multi-line - pattern = re.compile(r'm\.def\s*\(') - for file_path in Path(root_path).rglob('*'): if file_path.suffix.lower() not in extensions: continue @@ -191,7 +92,6 @@ def extract_m_def_statements(root_path): 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('/*'): @@ -222,8 +122,6 @@ 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: @@ -236,13 +134,12 @@ def extract_m_def_statements(root_path): def parse_m_def_statement(m_def_str): - result = { - 'python_function_name': None, - 'num_args': 0, - 'default_args': {}, - 'is_lambda': False, - } + """ + Parse a TORCH_LIBRARY m.def(...) statement. + DeepGEMM registers ops via TORCH_LIBRARY_FRAGMENT, so the first argument is + always a schema string such as "fp8_fp4_gemm_nt(Tensor a, ...) -> ()". + """ # Extract top-level arguments start = m_def_str.find('m.def(') if start == -1: @@ -283,625 +180,77 @@ def parse_m_def_statement(m_def_str): if current: args_list.append(''.join(current).strip()) - if len(args_list) < 1: + if not args_list: raise ValueError(f'[{m_def_str}] m.def has no arguments') - # Extract Python function name + # Extract operator name from the TORCH schema string first = args_list[0].strip() str_match = re.match(r'^"([^"\\]*(?:\\.[^"\\]*)*)"', first) if not str_match: raise ValueError(f'[{m_def_str}] m.def first argument should be a string literal') - schema_or_name = str_match.group(1) - if is_torch_schema_string(schema_or_name): - result['is_torch_schema'] = True - result['python_function_name'] = extract_torch_op_name(schema_or_name) - return result - - if len(args_list) < 2: - raise ValueError(f'[{m_def_str}] m.def has insufficient arguments') - - result['python_function_name'] = schema_or_name - - cpp_func_part = args_list[1].strip() - if cpp_func_part.startswith('&'): - cpp_func_part = cpp_func_part[1:].strip() - - 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 - - # Parse py::arg arguments - py_args = args_list[2:] - result['num_args'] = len(py_args) - - 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 - ) - - 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'] and not parsed.get('is_torch_schema'): - 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 + schema = str_match.group(1) return { - 'type': decl, - 'name': None + 'python_function_name': extract_torch_op_name(schema), + 'schema': schema, } -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'] +def generate_pyi_function(item_entry): """ - if not template_args.strip(): - return [] - args = [] - current = [] - tracker = BracketTracker() + Generate a .pyi stub for one registered op. - for ch in template_args: - if ch in '()[]{}<>': - tracker.update(ch) - if ch == ',' and tracker.is_top_level(): - args.append(''.join(current).strip()) - current = [] - 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. + Typed stubs require parsing the TORCH schema in item_entry['parsed']['schema']. + Until then, emit a generic signature that matches deep_gemm._C wrappers. """ - 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' - - -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 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) - - # 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}' - - param_lines.append(param_str) - - 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 + py_name = item_entry['parsed']['python_function_name'] + return f'def {py_name}(*args, **kwargs) -> Any: ...' 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' - + lines = [ + f'# Stubs for module: {module_name}', + '', + 'from typing import Any', + ] 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('') + lines.append('import torch') + lines.extend(['', '']) for decl in function_decls: - lines.append(decl) - lines.append('') - lines.append('') + 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 = [] + enhanced_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) + statements = [] + for stmt in item['m_def_statements']: + statements.append({ + 'raw': stmt, + 'parsed': parse_m_def_statement(stmt), + }) + enhanced_results.append({'m_def_statements': statements}) - pyi_content = generate_pyi_file_content(cpp_results, module_name=name) + pyi_content = generate_pyi_file_content(enhanced_results, module_name=name) output_path = Path(output_dir) / f'{name}.pyi' output_path.parent.mkdir(parents=True, exist_ok=True) From c38efbb776cf229ac09fbe92344c2eb7b0d0822e Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Wed, 8 Jul 2026 17:45:23 +0000 Subject: [PATCH 04/28] generate_pyi.py stopped creating the function type hints in the past updates so I updated it to inlcude the type hints Signed-off-by: Chris Leonard --- scripts/generate_pyi.py | 419 ++++++++++++++++++++++++++++++++++++---- 1 file changed, 378 insertions(+), 41 deletions(-) diff --git a/scripts/generate_pyi.py b/scripts/generate_pyi.py index 72bb2c0dce..6977b00e61 100644 --- a/scripts/generate_pyi.py +++ b/scripts/generate_pyi.py @@ -1,6 +1,28 @@ import re from pathlib import Path +_TENSOR_PAIR = 'tuple[torch.Tensor, torch.Tensor]' +_Q_TYPE = f'torch.Tensor | tuple[torch.Tensor, Optional[torch.Tensor]]' + +_MERGE_AB_PAIR_OPS = frozenset({ + 'fp8_fp4_gemm_nt', 'fp8_fp4_gemm_nn', 'fp8_fp4_gemm_tn', 'fp8_fp4_gemm_tt', + 'm_grouped_fp8_fp4_gemm_nt_contiguous', 'm_grouped_fp8_fp4_gemm_nn_contiguous', + 'm_grouped_fp8_fp4_gemm_nt_masked', + 'k_grouped_fp8_gemm_tn_contiguous', 'k_grouped_fp8_gemm_nt_contiguous', + 'fp8_gemm_nt_skip_head_mid', +}) +_MERGE_EINSUM_AB_OPS = frozenset({'fp8_einsum'}) +_MERGE_MEGA_WEIGHT_OPS = frozenset({'fp8_fp4_mega_moe'}) +_PYI_ALIASES = { + '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_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', +} + class BracketTracker: """ @@ -56,6 +78,35 @@ def is_top_level(self): self.angle == 0) +def split_top_level_commas(value: str) -> list[str]: + """Split a string on top-level commas.""" + 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 '=' in a schema argument, 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 extract_torch_op_name(schema: str) -> str: """Extract the operator name from a TORCH schema string.""" paren_pos = schema.find('(') @@ -64,6 +115,232 @@ def extract_torch_op_name(schema: str) -> str: return schema[:paren_pos].strip() +def schema_type_to_python(type_str: str) -> str: + """Map a TORCH_LIBRARY schema type to a Python type annotation string.""" + 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]' + 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_LIBRARY return type to a Python annotation.""" + return_str = return_str.strip() + 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' + + +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 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_LIBRARY 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, + 'schema': schema, + } + + +def _merge_named_pairs(parameters: list[dict], pairs: tuple[tuple[str, str], ...]) -> list[dict]: + """Replace (left, right) arg pairs with a single tuple-typed parameter.""" + drop = {right for left, right in pairs} + merged_left = {left for left, _ in pairs} + out = [] + for param in parameters: + if param['name'] in drop: + continue + if param['name'] in merged_left: + out.append({ + 'name': param['name'], + 'py_type': _TENSOR_PAIR, + 'default': None, + }) + continue + out.append(dict(param)) + return out + + +def adjust_for_c_py_wrapper(name: str, parameters: list[dict]) -> list[dict]: + """ + Adjust parsed schema parameters to match deep_gemm._C Python wrappers. + + TORCH_LIBRARY registers flat tensor/scales args; _C.py preserves the legacy + pybind API by accepting (tensor, scale_factor) tuples for many kernels. + """ + if name in _MERGE_AB_PAIR_OPS: + parameters = _merge_named_pairs(parameters, (('a', 'sfa'), ('b', 'sfb'))) + + elif name in _MERGE_EINSUM_AB_OPS: + parameters = _merge_named_pairs(parameters, (('a', 'sfa'), ('b', 'sfb'))) + for param in parameters: + if param['name'] == 'recipe': + param['py_type'] = 'tuple[int, int, int]' + param['default'] = '(1, 128, 128)' + + elif name in _MERGE_MEGA_WEIGHT_OPS: + parameters = _merge_named_pairs( + parameters, + (('l1_weights', 'l1_weights_sf'), ('l2_weights', 'l2_weights_sf')), + ) + for param in parameters: + if param['name'] == 'recipe': + param['py_type'] = 'tuple[int, int, int]' + + elif name == 'fp8_fp4_mqa_logits': + parameters = _merge_named_pairs(parameters, (('kv', 'kv_sf'),)) + out = [] + for param in parameters: + if param['name'] == 'q_sf': + continue + if param['name'] == 'q': + param['py_type'] = _Q_TYPE + if param['name'] == 'logits_dtype': + param['py_type'] = 'torch.dtype' + param['default'] = 'torch.float32' + out.append(param) + return out + + elif name == 'fp8_fp4_paged_mqa_logits': + out = [] + for param in parameters: + if param['name'] == 'q_sf': + continue + if param['name'] == 'q': + param['py_type'] = _Q_TYPE + if param['name'] == 'logits_dtype': + param['py_type'] = 'torch.dtype' + param['default'] = 'torch.float32' + out.append(param) + return out + + elif name == 'fp8_mqa_logits': + parameters = _merge_named_pairs(parameters, (('kv', 'kv_sf'),)) + + elif name == 'set_block_size_multiple_of': + for param in parameters: + if param['name'] == 'value': + param['py_type'] = 'int | list[int]' + + if name in {'k_grouped_fp8_gemm_tn_contiguous', 'k_grouped_fp8_gemm_nt_contiguous'}: + for param in parameters: + if param['name'] == 'recipe': + param['py_type'] = 'tuple[int, int, int]' + param['default'] = '(1, 1, 128)' + + return parameters + + +def sanitize_param_name(name: str) -> str: + if name in {'def', 'class', 'from', 'import', 'None', 'True', 'False'}: + return f'{name}_' + return name + + def extract_m_def_statements(root_path): """ Scan all C++ files under root_path and extract all m.def(...) statements. @@ -164,75 +441,89 @@ def parse_m_def_statement(m_def_str): args_content = m_def_str[content_start:content_end] # 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()) + args_list = split_top_level_commas(args_content) if not args_list: raise ValueError(f'[{m_def_str}] m.def has no arguments') - # Extract operator name from the TORCH schema string + # Extract operator schema from the first string literal first = args_list[0].strip() str_match = re.match(r'^"([^"\\]*(?:\\.[^"\\]*)*)"', first) if not str_match: raise ValueError(f'[{m_def_str}] m.def first argument should be a string literal') - schema = str_match.group(1) - return { - 'python_function_name': extract_torch_op_name(schema), - 'schema': schema, - } + return parse_torch_schema(str_match.group(1)) def generate_pyi_function(item_entry): - """ - Generate a .pyi stub for one registered op. + """Generate a typed .pyi stub for one registered op.""" + parsed = item_entry['parsed'] + py_name = parsed['python_function_name'] + parameters = adjust_for_c_py_wrapper(py_name, parsed['parameters']) + return_type = parsed['return_type'] + + param_lines = [] + for param in parameters: + name = sanitize_param_name(param['name']) + if param['default'] is not None: + param_lines.append(f' {name}: {param["py_type"]} = {param["default"]}') + else: + param_lines.append(f' {name}: {param["py_type"]}') - Typed stubs require parsing the TORCH schema in item_entry['parsed']['schema']. - Until then, emit a generic signature that matches deep_gemm._C wrappers. - """ - py_name = item_entry['parsed']['python_function_name'] - return f'def {py_name}(*args, **kwargs) -> Any: ...' + 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}: ...' -def generate_pyi_file_content(enhanced_results, module_name: str = 'my_module'): - function_decls = [] - has_torch = False +def _alias_pyi_decl(decl: str, alias_name: str, source_name: str) -> str: + return decl.replace(f'def {source_name}(', f'def {alias_name}(', 1) + +def generate_pyi_file_content(enhanced_results, module_name: str = 'my_module'): + by_name = {} for item in enhanced_results: for stmt in item['m_def_statements']: - try: - decl = generate_pyi_function(stmt) - function_decls.append(decl) - if 'torch.Tensor' in decl: - has_torch = 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}') + name = stmt['parsed']['python_function_name'] + by_name[name] = stmt + + decl_by_name = {} + has_optional = False + has_torch = False + + for name in sorted(by_name): + stmt = by_name[name] + try: + decl = generate_pyi_function(stmt) + decl_by_name[name] = decl + if 'Optional[' in decl: + has_optional = True + if 'torch.' in decl: + has_torch = True + except Exception as e: + decl_by_name[name] = f'# ERROR: failed to generate stub for {name}: {e}' + + for alias_name, source_name in _PYI_ALIASES.items(): + if source_name in decl_by_name and alias_name not in decl_by_name: + decl_by_name[alias_name] = _alias_pyi_decl(decl_by_name[source_name], alias_name, source_name) + if 'Optional[' in decl_by_name[alias_name]: + has_optional = True + if 'torch.' in decl_by_name[alias_name]: + has_torch = True lines = [ f'# Stubs for module: {module_name}', '', 'from typing import Any', ] + if has_optional: + lines[2] += ', Optional' if has_torch: lines.append('import torch') lines.extend(['', '']) - for decl in function_decls: - lines.extend([decl, '', '']) + for name in sorted(decl_by_name): + lines.extend([decl_by_name[name], '', '']) return '\n'.join(lines) @@ -259,3 +550,49 @@ def generate_pyi_file(name, root, output_dir='.'): f.write(pyi_content) print(f'.pyi file generated: {output_path}') + + +def main(argv=None) -> int: + import argparse + import sys + + parser = argparse.ArgumentParser( + description='Generate deep_gemm/_C.pyi stubs from TORCH_LIBRARY schemas.', + ) + parser.add_argument('--name', default='_C', help='Module name for the .pyi file (default: _C)') + parser.add_argument('--root', default='./csrc', help='Root to scan for m.def(...) (default: ./csrc)') + parser.add_argument('--output-dir', default='./stubs', help='Output directory (default: ./stubs)') + parser.add_argument( + '--check', + action='store_true', + help='Verify the output has typed stubs (no generic *args, **kwargs)', + ) + args = parser.parse_args(argv) + + repo_root = Path(__file__).resolve().parent.parent + root = Path(args.root) + output_dir = Path(args.output_dir) + if not root.is_absolute(): + root = repo_root / root + if not output_dir.is_absolute(): + output_dir = repo_root / output_dir + + generate_pyi_file(name=args.name, root=str(root), output_dir=str(output_dir)) + + pyi_path = output_dir / f'{args.name}.pyi' + if args.check: + content = pyi_path.read_text(encoding='utf-8') + if '*args, **kwargs' in content: + print(f'CHECK FAILED: generic stubs found in {pyi_path}', file=sys.stderr) + return 1 + stub_count = content.count('def ') + if stub_count == 0: + print(f'CHECK FAILED: no function stubs in {pyi_path}', file=sys.stderr) + return 1 + print(f'CHECK PASSED: {stub_count} typed stubs in {pyi_path}') + return 0 + + +if __name__ == '__main__': + import sys + raise SystemExit(main()) From 2630e0a5ad2c1cfc236332c1609b7e2244648227 Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Wed, 8 Jul 2026 19:40:24 +0000 Subject: [PATCH 05/28] updated generate_pyi.py to align more with legacy code Signed-off-by: Chris Leonard --- scripts/generate_pyi.py | 1225 +++++++++++++++++++++++++++------------ 1 file changed, 866 insertions(+), 359 deletions(-) diff --git a/scripts/generate_pyi.py b/scripts/generate_pyi.py index 6977b00e61..6408496eac 100644 --- a/scripts/generate_pyi.py +++ b/scripts/generate_pyi.py @@ -1,27 +1,227 @@ import re from pathlib import Path -_TENSOR_PAIR = 'tuple[torch.Tensor, torch.Tensor]' -_Q_TYPE = f'torch.Tensor | tuple[torch.Tensor, Optional[torch.Tensor]]' - -_MERGE_AB_PAIR_OPS = frozenset({ - 'fp8_fp4_gemm_nt', 'fp8_fp4_gemm_nn', 'fp8_fp4_gemm_tn', 'fp8_fp4_gemm_tt', - 'm_grouped_fp8_fp4_gemm_nt_contiguous', 'm_grouped_fp8_fp4_gemm_nn_contiguous', - 'm_grouped_fp8_fp4_gemm_nt_masked', - 'k_grouped_fp8_gemm_tn_contiguous', 'k_grouped_fp8_gemm_nt_contiguous', - 'fp8_gemm_nt_skip_head_mid', -}) -_MERGE_EINSUM_AB_OPS = frozenset({'fp8_einsum'}) -_MERGE_MEGA_WEIGHT_OPS = frozenset({'fp8_fp4_mega_moe'}) -_PYI_ALIASES = { - '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_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', -} + +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 + + +def extract_torch_op_name(schema: str) -> str: + """Extract the operator name from a TORCH schema string or plain name.""" + paren_pos = schema.find('(') + if paren_pos == -1: + return schema.strip() + return schema[:paren_pos].strip() + + +def split_schema_args(args_str: str) -> list[str]: + """Split a TORCH schema argument list by top-level commas.""" + if not args_str.strip(): + return [] + return split_cpp_parameters(args_str) + + +def parse_schema_arg_default(spec: str) -> str | None: + """Extract the default value from a single TORCH schema argument, if present.""" + spec = spec.strip() + if not spec: + return None + + in_quote = None + paren = bracket = angle = 0 + for i, ch in enumerate(spec): + if in_quote: + if ch == in_quote and (i == 0 or spec[i - 1] != '\\'): + in_quote = None + continue + if ch in ('"', "'"): + in_quote = ch + continue + if ch == '(': + paren += 1 + elif ch == ')': + paren -= 1 + elif ch == '[': + bracket += 1 + elif ch == ']': + bracket -= 1 + elif ch == '<': + angle += 1 + elif ch == '>': + angle -= 1 + elif ch == '=' and paren == bracket == angle == 0: + return schema_default_to_python(spec[i + 1:].strip()) + return None + + +def schema_default_to_python(val: str) -> str: + """Convert a TORCH schema default literal to a Python expression string.""" + val = val.strip() + if not val: + return 'None' + if val == 'None': + return 'None' + if val in ('True', 'true'): + return 'True' + if val in ('False', 'false'): + return 'False' + if (val.startswith("'") and val.endswith("'")) or (val.startswith('"') and val.endswith('"')): + return f'"{val[1:-1]}"' + if re.match(r'^[+-]?\d+$', val): + return val + if re.match(r'^[+-]?\d*\.\d+([eE][+-]?\d+)?$', val): + return val + print(f'Warning: Unrecognized schema default value: {val}') + return val + + +def parse_schema_parameter_name(spec: str) -> str | None: + """Extract the parameter name from a TORCH schema argument.""" + spec = spec.strip() + if not spec: + return None + + in_quote = None + paren = bracket = angle = 0 + eq_pos = -1 + for i, ch in enumerate(spec): + if in_quote: + if ch == in_quote and (i == 0 or spec[i - 1] != '\\'): + in_quote = None + continue + if ch in ('"', "'"): + in_quote = ch + continue + if ch == '(': + paren += 1 + elif ch == ')': + paren -= 1 + elif ch == '[': + bracket += 1 + elif ch == ']': + bracket -= 1 + elif ch == '<': + angle += 1 + elif ch == '>': + angle -= 1 + elif ch == '=' and paren == bracket == angle == 0: + eq_pos = i + break + + left = spec[:eq_pos].strip() if eq_pos != -1 else spec + name_match = re.search(r'([a-zA-Z_][a-zA-Z0-9_]*)\s*$', left) + return name_match.group(1) if name_match else None + + +def parse_schema_parameter_defaults(schema: str) -> dict[str, str]: + """Parse parameter defaults from a TORCH schema string, keyed by parameter name.""" + arrow = schema.rfind(' -> ') + if arrow == -1: + return {} + + sig_part = schema[:arrow].strip() + open_paren = sig_part.find('(') + close_paren = sig_part.rfind(')') + if open_paren == -1 or close_paren == -1 or close_paren <= open_paren: + return {} + + defaults = {} + for spec in split_schema_args(sig_part[open_paren + 1:close_paren]): + default_val = parse_schema_arg_default(spec) + if default_val is None: + continue + param_name = parse_schema_parameter_name(spec) + if param_name: + defaults[param_name] = default_val + return defaults class BracketTracker: @@ -78,277 +278,16 @@ def is_top_level(self): self.angle == 0) -def split_top_level_commas(value: str) -> list[str]: - """Split a string on top-level commas.""" - 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 '=' in a schema argument, 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 extract_torch_op_name(schema: str) -> str: - """Extract the operator name from a TORCH schema string.""" - paren_pos = schema.find('(') - if paren_pos == -1: - return schema.strip() - return schema[:paren_pos].strip() - - -def schema_type_to_python(type_str: str) -> str: - """Map a TORCH_LIBRARY schema type to a Python type annotation string.""" - 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]' - 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_LIBRARY return type to a Python annotation.""" - return_str = return_str.strip() - 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' - - -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 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_LIBRARY 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, - 'schema': schema, - } - - -def _merge_named_pairs(parameters: list[dict], pairs: tuple[tuple[str, str], ...]) -> list[dict]: - """Replace (left, right) arg pairs with a single tuple-typed parameter.""" - drop = {right for left, right in pairs} - merged_left = {left for left, _ in pairs} - out = [] - for param in parameters: - if param['name'] in drop: - continue - if param['name'] in merged_left: - out.append({ - 'name': param['name'], - 'py_type': _TENSOR_PAIR, - 'default': None, - }) - continue - out.append(dict(param)) - return out - - -def adjust_for_c_py_wrapper(name: str, parameters: list[dict]) -> list[dict]: - """ - Adjust parsed schema parameters to match deep_gemm._C Python wrappers. - - TORCH_LIBRARY registers flat tensor/scales args; _C.py preserves the legacy - pybind API by accepting (tensor, scale_factor) tuples for many kernels. - """ - if name in _MERGE_AB_PAIR_OPS: - parameters = _merge_named_pairs(parameters, (('a', 'sfa'), ('b', 'sfb'))) - - elif name in _MERGE_EINSUM_AB_OPS: - parameters = _merge_named_pairs(parameters, (('a', 'sfa'), ('b', 'sfb'))) - for param in parameters: - if param['name'] == 'recipe': - param['py_type'] = 'tuple[int, int, int]' - param['default'] = '(1, 128, 128)' - - elif name in _MERGE_MEGA_WEIGHT_OPS: - parameters = _merge_named_pairs( - parameters, - (('l1_weights', 'l1_weights_sf'), ('l2_weights', 'l2_weights_sf')), - ) - for param in parameters: - if param['name'] == 'recipe': - param['py_type'] = 'tuple[int, int, int]' - - elif name == 'fp8_fp4_mqa_logits': - parameters = _merge_named_pairs(parameters, (('kv', 'kv_sf'),)) - out = [] - for param in parameters: - if param['name'] == 'q_sf': - continue - if param['name'] == 'q': - param['py_type'] = _Q_TYPE - if param['name'] == 'logits_dtype': - param['py_type'] = 'torch.dtype' - param['default'] = 'torch.float32' - out.append(param) - return out - - elif name == 'fp8_fp4_paged_mqa_logits': - out = [] - for param in parameters: - if param['name'] == 'q_sf': - continue - if param['name'] == 'q': - param['py_type'] = _Q_TYPE - if param['name'] == 'logits_dtype': - param['py_type'] = 'torch.dtype' - param['default'] = 'torch.float32' - out.append(param) - return out - - elif name == 'fp8_mqa_logits': - parameters = _merge_named_pairs(parameters, (('kv', 'kv_sf'),)) - - elif name == 'set_block_size_multiple_of': - for param in parameters: - if param['name'] == 'value': - param['py_type'] = 'int | list[int]' - - if name in {'k_grouped_fp8_gemm_tn_contiguous', 'k_grouped_fp8_gemm_nt_contiguous'}: - for param in parameters: - if param['name'] == 'recipe': - param['py_type'] = 'tuple[int, int, int]' - param['default'] = '(1, 1, 128)' - - return parameters - - -def sanitize_param_name(name: str) -> str: - if name in {'def', 'class', 'from', 'import', 'None', 'True', 'False'}: - return f'{name}_' - return name - - def extract_m_def_statements(root_path): """ - Scan all C++ files under root_path and extract all m.def(...) statements. - Supports multi-line m.def(...) calls. + Scan all c files under root_path and extract all m.def(...) statements. """ results = [] extensions = {'.hpp', '.cpp', '.h', '.cc'} + # Regex: match m.def( ... ), supports multi-line + pattern = re.compile(r'm\.def\s*\(') + for file_path in Path(root_path).rglob('*'): if file_path.suffix.lower() not in extensions: continue @@ -369,6 +308,7 @@ def extract_m_def_statements(root_path): 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('/*'): @@ -399,6 +339,8 @@ 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: @@ -411,12 +353,13 @@ def extract_m_def_statements(root_path): def parse_m_def_statement(m_def_str): - """ - Parse a TORCH_LIBRARY m.def(...) statement. + result = { + 'python_function_name': None, + 'num_args': 0, + 'default_args': {}, + 'is_lambda': False, + } - DeepGEMM registers ops via TORCH_LIBRARY_FRAGMENT, so the first argument is - always a schema string such as "fp8_fp4_gemm_nt(Tensor a, ...) -> ()". - """ # Extract top-level arguments start = m_def_str.find('m.def(') if start == -1: @@ -441,107 +384,671 @@ def parse_m_def_statement(m_def_str): args_content = m_def_str[content_start:content_end] # Split arguments using BracketTracker - args_list = split_top_level_commas(args_content) + 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 not args_list: raise ValueError(f'[{m_def_str}] m.def has no arguments') - # Extract operator schema from the first string literal + # Extract Python function name from the first string literal (plain name or schema). first = args_list[0].strip() str_match = re.match(r'^"([^"\\]*(?:\\.[^"\\]*)*)"', first) - if not str_match: + if str_match: + schema_or_name = str_match.group(1) + result['python_function_name'] = extract_torch_op_name(schema_or_name) + if '(' in schema_or_name: + result['schema_default_args'] = parse_schema_parameter_defaults(schema_or_name) + else: raise ValueError(f'[{m_def_str}] m.def first argument should be a string literal') - return parse_torch_schema(str_match.group(1)) + if len(args_list) == 1: + result['cpp_function_name'] = result['python_function_name'] + else: + cpp_func_part = args_list[1].strip() + if cpp_func_part.startswith('&'): + cpp_func_part = cpp_func_part[1:].strip() + + if cpp_func_part.startswith('['): + result['is_lambda'] = True + result['cpp_function_name'] = None + elif cpp_func_part.startswith(('DEEP_GEMM_IMPL(', 'TORCH_FN(')): + result['cpp_function_name'] = result['python_function_name'] + 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 + + # Parse py::arg arguments (legacy pybind registrations only). + py_args = args_list[2:] + result['num_args'] = len(py_args) + + 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 + ) + + 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() + if return_type.startswith('static '): + return_type = return_type[len('static '):].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 / c10::optional + if cleaned.startswith('std::optional<') or cleaned.startswith('c10::optional<'): + inner = cleaned[cleaned.index('<') + 1:-1].strip() + 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: optional with unexpected args: {cleaned}') + return 'Any' + + # c10::List + if cleaned.startswith('c10::List<'): + inner = cleaned[10:-1].strip() + args = split_template_args(inner) + if len(args) == 1: + inner_py = cpp_type_to_python_type(args[0]) + return f'list[{inner_py}]' + print(f'Warning: c10::List 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' + + # at::ScalarType + if re.search(r'\bat::ScalarType\b', original): + return 'torch.dtype' + + # mega.hpp type alias + if re.search(r'\bSymmBufferSlice\b', original): + tensor = 'torch.Tensor' + return f'tuple[{", ".join([tensor] * 8)}]' + + # 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 = [] + 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 + + if s == 'torch::kFloat32': + return 'torch.float32' + + # Fallback: unrecognized → warn and return None + print(f'Warning: Unrecognized default value: {s}') + return 'None' def generate_pyi_function(item_entry): - """Generate a typed .pyi stub for one registered op.""" parsed = item_entry['parsed'] py_name = parsed['python_function_name'] - parameters = adjust_for_c_py_wrapper(py_name, parsed['parameters']) - return_type = parsed['return_type'] + if parsed.get('is_lambda'): + return f'def {py_name}(*args, **kwargs) -> Any: ...' + + sig_info = parsed.get('cpp_parsed_signature') + default_args = dict(parsed.get('default_args', {})) + schema_default_by_name = parsed.get('schema_default_args', {}) + + 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) + + # Build parameter list param_lines = [] - for param in parameters: - name = sanitize_param_name(param['name']) - if param['default'] is not None: - param_lines.append(f' {name}: {param["py_type"]} = {param["default"]}') + 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 (py::arg defaults take precedence over schema defaults). + py_default = None + if i in default_args: + py_default = cpp_default_to_python_default(default_args[i]) + elif param_name in schema_default_by_name: + py_default = schema_default_by_name[param_name] + if param_type == 'torch.dtype' and py_default == '6': + py_default = 'torch.float32' + + if py_default is not None: + param_str = f' {param_name}: {param_type} = {py_default}' else: - param_lines.append(f' {name}: {param["py_type"]}') + param_str = f' {param_name}: {param_type}' + + param_lines.append(param_str) 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}: ...' - + func_def = f'def {py_name}(\n{params_block}\n) -> {return_type}: ...' + else: + func_def = f'def {py_name}() -> {return_type}: ...' -def _alias_pyi_decl(decl: str, alias_name: str, source_name: str) -> str: - return decl.replace(f'def {source_name}(', f'def {alias_name}(', 1) + return func_def def generate_pyi_file_content(enhanced_results, module_name: str = 'my_module'): - by_name = {} - for item in enhanced_results: - for stmt in item['m_def_statements']: - name = stmt['parsed']['python_function_name'] - by_name[name] = stmt - - decl_by_name = {} + function_decls = [] has_optional = False has_torch = False + has_numpy = False - for name in sorted(by_name): - stmt = by_name[name] - try: - decl = generate_pyi_function(stmt) - decl_by_name[name] = decl - if 'Optional[' in decl: - has_optional = True - if 'torch.' in decl: - has_torch = True - except Exception as e: - decl_by_name[name] = f'# ERROR: failed to generate stub for {name}: {e}' - - for alias_name, source_name in _PYI_ALIASES.items(): - if source_name in decl_by_name and alias_name not in decl_by_name: - decl_by_name[alias_name] = _alias_pyi_decl(decl_by_name[source_name], alias_name, source_name) - if 'Optional[' in decl_by_name[alias_name]: - has_optional = True - if 'torch.' in decl_by_name[alias_name]: - has_torch = True - - lines = [ - f'# Stubs for module: {module_name}', - '', - 'from typing import Any', - ] + 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: - lines[2] += ', Optional' + imports[0] += ', Optional' + if has_torch: - lines.append('import torch') - lines.extend(['', '']) + 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 name in sorted(decl_by_name): - lines.extend([decl_by_name[name], '', '']) + for decl in function_decls: + lines.append(decl) + lines.append('') + lines.append('') 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) - enhanced_results = [] + cpp_results = [] for item in results: - statements = [] - for stmt in item['m_def_statements']: - statements.append({ - 'raw': stmt, - 'parsed': parse_m_def_statement(stmt), - }) - enhanced_results.append({'m_def_statements': statements}) + enhanced_item = parse_mdef_and_attach_cpp_signatures(item, func_index) + cpp_item = extract_cpp_signature_details(enhanced_item) + cpp_results.append(cpp_item) - pyi_content = generate_pyi_file_content(enhanced_results, module_name=name) + pyi_content = generate_pyi_file_content(cpp_results, module_name=name) output_path = Path(output_dir) / f'{name}.pyi' output_path.parent.mkdir(parents=True, exist_ok=True) @@ -557,7 +1064,7 @@ def main(argv=None) -> int: import sys parser = argparse.ArgumentParser( - description='Generate deep_gemm/_C.pyi stubs from TORCH_LIBRARY schemas.', + description='Generate deep_gemm/_C.pyi stubs from C++ signatures and m.def registrations.', ) parser.add_argument('--name', default='_C', help='Module name for the .pyi file (default: _C)') parser.add_argument('--root', default='./csrc', help='Root to scan for m.def(...) (default: ./csrc)') @@ -582,14 +1089,14 @@ def main(argv=None) -> int: pyi_path = output_dir / f'{args.name}.pyi' if args.check: content = pyi_path.read_text(encoding='utf-8') - if '*args, **kwargs' in content: - print(f'CHECK FAILED: generic stubs found in {pyi_path}', file=sys.stderr) - return 1 + generic_count = content.count('*args, **kwargs') stub_count = content.count('def ') if stub_count == 0: print(f'CHECK FAILED: no function stubs in {pyi_path}', file=sys.stderr) return 1 - print(f'CHECK PASSED: {stub_count} typed stubs in {pyi_path}') + print(f'CHECK PASSED: {stub_count} stubs in {pyi_path} ({stub_count - generic_count} typed, {generic_count} generic)') + if generic_count > 0: + return 1 return 0 From 796a48e0e424b7e64a57bc58d53d1b4a7b583c73 Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Wed, 8 Jul 2026 19:55:48 +0000 Subject: [PATCH 06/28] used _C.py to generate default argumements in gererate_pyi.py since pybind default arguments are no longer there Signed-off-by: Chris Leonard --- scripts/generate_pyi.py | 119 +++++++++++++++++++++++++++++++++++++--- 1 file changed, 112 insertions(+), 7 deletions(-) diff --git a/scripts/generate_pyi.py b/scripts/generate_pyi.py index 6408496eac..608afff0b9 100644 --- a/scripts/generate_pyi.py +++ b/scripts/generate_pyi.py @@ -1,3 +1,4 @@ +import ast import re from pathlib import Path @@ -940,7 +941,82 @@ def cpp_default_to_python_default(cpp_default: str): return 'None' -def generate_pyi_function(item_entry): +def format_ast_default(node: ast.AST) -> str: + """Convert an AST default value node to a Python expression string for stubs.""" + 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 parse_wrapper_defaults(c_py_path: Path) -> dict[str, dict[str, str]]: + """Parse deep_gemm/_C.py wrapper function defaults keyed by exported name.""" + 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 ( + isinstance(node.func, ast.Attribute) + and node.func.attr == 'update' + and isinstance(node.func.value, ast.Name) + and node.func.value.id == 'globals' + ): + 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) and value_node.id in func_defaults: + func_defaults[alias_name] = func_defaults[value_node.id] + + return func_defaults + + +def generate_pyi_function(item_entry, wrapper_defaults=None): parsed = item_entry['parsed'] py_name = parsed['python_function_name'] @@ -950,6 +1026,7 @@ def generate_pyi_function(item_entry): sig_info = parsed.get('cpp_parsed_signature') default_args = dict(parsed.get('default_args', {})) schema_default_by_name = parsed.get('schema_default_args', {}) + wrapper_default_by_name = (wrapper_defaults or {}).get(py_name, {}) if not sig_info: return f'def {py_name}(*args, **kwargs) -> Any: ...' @@ -969,10 +1046,12 @@ def generate_pyi_function(item_entry): if param_name in {'def', 'class', 'from', 'import', 'None', 'True', 'False'}: param_name = f'{param_name}_' - # Check for default value (py::arg defaults take precedence over schema defaults). + # Defaults: py::arg > _C.py wrapper > TORCH schema. py_default = None if i in default_args: py_default = cpp_default_to_python_default(default_args[i]) + elif param_name in wrapper_default_by_name: + py_default = wrapper_default_by_name[param_name] elif param_name in schema_default_by_name: py_default = schema_default_by_name[param_name] if param_type == 'torch.dtype' and py_default == '6': @@ -994,7 +1073,7 @@ def generate_pyi_function(item_entry): return func_def -def generate_pyi_file_content(enhanced_results, module_name: str = 'my_module'): +def generate_pyi_file_content(enhanced_results, module_name: str = 'my_module', wrapper_defaults=None): function_decls = [] has_optional = False has_torch = False @@ -1003,7 +1082,7 @@ def generate_pyi_file_content(enhanced_results, module_name: str = 'my_module'): for item in enhanced_results: for stmt in item['m_def_statements']: try: - decl = generate_pyi_function(stmt) + decl = generate_pyi_function(stmt, wrapper_defaults=wrapper_defaults) function_decls.append(decl) if 'Optional[' in decl: @@ -1038,7 +1117,7 @@ def generate_pyi_file_content(enhanced_results, module_name: str = 'my_module'): return '\n'.join(lines) -def generate_pyi_file(name, root, output_dir='.'): +def generate_pyi_file(name, root, output_dir='.', c_py_path=None): func_index = build_cpp_function_index(root) results = extract_m_def_statements(root) @@ -1048,7 +1127,19 @@ def generate_pyi_file(name, root, output_dir='.'): cpp_item = extract_cpp_signature_details(enhanced_item) cpp_results.append(cpp_item) - pyi_content = generate_pyi_file_content(cpp_results, module_name=name) + 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_wrapper_defaults(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, + wrapper_defaults=wrapper_defaults, + ) output_path = Path(output_dir) / f'{name}.pyi' output_path.parent.mkdir(parents=True, exist_ok=True) @@ -1069,6 +1160,11 @@ def main(argv=None) -> int: parser.add_argument('--name', default='_C', help='Module name for the .pyi file (default: _C)') parser.add_argument('--root', default='./csrc', help='Root to scan for m.def(...) (default: ./csrc)') parser.add_argument('--output-dir', default='./stubs', help='Output directory (default: ./stubs)') + parser.add_argument( + '--c-py', + default='./deep_gemm/_C.py', + help='Python wrapper module to read public API defaults from (default: ./deep_gemm/_C.py)', + ) parser.add_argument( '--check', action='store_true', @@ -1084,7 +1180,16 @@ def main(argv=None) -> int: if not output_dir.is_absolute(): output_dir = repo_root / output_dir - generate_pyi_file(name=args.name, root=str(root), output_dir=str(output_dir)) + c_py_path = Path(args.c_py) + if not c_py_path.is_absolute(): + c_py_path = repo_root / c_py_path + + generate_pyi_file( + name=args.name, + root=str(root), + output_dir=str(output_dir), + c_py_path=str(c_py_path), + ) pyi_path = output_dir / f'{args.name}.pyi' if args.check: From 074956920d55ffa98368acc8ddc863663f487bb9 Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Fri, 10 Jul 2026 13:59:49 +0000 Subject: [PATCH 07/28] reverted back to using TORCH_LIBRARY schema to generate pyi file Signed-off-by: Chris Leonard --- scripts/generate_pyi.py | 1390 +++++++++++++-------------------------- 1 file changed, 467 insertions(+), 923 deletions(-) diff --git a/scripts/generate_pyi.py b/scripts/generate_pyi.py index 608afff0b9..b75701330c 100644 --- a/scripts/generate_pyi.py +++ b/scripts/generate_pyi.py @@ -2,227 +2,8 @@ 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 - - -def extract_torch_op_name(schema: str) -> str: - """Extract the operator name from a TORCH schema string or plain name.""" - paren_pos = schema.find('(') - if paren_pos == -1: - return schema.strip() - return schema[:paren_pos].strip() - - -def split_schema_args(args_str: str) -> list[str]: - """Split a TORCH schema argument list by top-level commas.""" - if not args_str.strip(): - return [] - return split_cpp_parameters(args_str) - - -def parse_schema_arg_default(spec: str) -> str | None: - """Extract the default value from a single TORCH schema argument, if present.""" - spec = spec.strip() - if not spec: - return None - - in_quote = None - paren = bracket = angle = 0 - for i, ch in enumerate(spec): - if in_quote: - if ch == in_quote and (i == 0 or spec[i - 1] != '\\'): - in_quote = None - continue - if ch in ('"', "'"): - in_quote = ch - continue - if ch == '(': - paren += 1 - elif ch == ')': - paren -= 1 - elif ch == '[': - bracket += 1 - elif ch == ']': - bracket -= 1 - elif ch == '<': - angle += 1 - elif ch == '>': - angle -= 1 - elif ch == '=' and paren == bracket == angle == 0: - return schema_default_to_python(spec[i + 1:].strip()) - return None - - -def schema_default_to_python(val: str) -> str: - """Convert a TORCH schema default literal to a Python expression string.""" - val = val.strip() - if not val: - return 'None' - if val == 'None': - return 'None' - if val in ('True', 'true'): - return 'True' - if val in ('False', 'false'): - return 'False' - if (val.startswith("'") and val.endswith("'")) or (val.startswith('"') and val.endswith('"')): - return f'"{val[1:-1]}"' - if re.match(r'^[+-]?\d+$', val): - return val - if re.match(r'^[+-]?\d*\.\d+([eE][+-]?\d+)?$', val): - return val - print(f'Warning: Unrecognized schema default value: {val}') - return val - - -def parse_schema_parameter_name(spec: str) -> str | None: - """Extract the parameter name from a TORCH schema argument.""" - spec = spec.strip() - if not spec: - return None - - in_quote = None - paren = bracket = angle = 0 - eq_pos = -1 - for i, ch in enumerate(spec): - if in_quote: - if ch == in_quote and (i == 0 or spec[i - 1] != '\\'): - in_quote = None - continue - if ch in ('"', "'"): - in_quote = ch - continue - if ch == '(': - paren += 1 - elif ch == ')': - paren -= 1 - elif ch == '[': - bracket += 1 - elif ch == ']': - bracket -= 1 - elif ch == '<': - angle += 1 - elif ch == '>': - angle -= 1 - elif ch == '=' and paren == bracket == angle == 0: - eq_pos = i - break - - left = spec[:eq_pos].strip() if eq_pos != -1 else spec - name_match = re.search(r'([a-zA-Z_][a-zA-Z0-9_]*)\s*$', left) - return name_match.group(1) if name_match else None - - -def parse_schema_parameter_defaults(schema: str) -> dict[str, str]: - """Parse parameter defaults from a TORCH schema string, keyed by parameter name.""" - arrow = schema.rfind(' -> ') - if arrow == -1: - return {} - - sig_part = schema[:arrow].strip() - open_paren = sig_part.find('(') - close_paren = sig_part.rfind(')') - if open_paren == -1 or close_paren == -1 or close_paren <= open_paren: - return {} - - defaults = {} - for spec in split_schema_args(sig_part[open_paren + 1:close_paren]): - default_val = parse_schema_arg_default(spec) - if default_val is None: - continue - param_name = parse_schema_parameter_name(spec) - if param_name: - defaults[param_name] = default_val - return defaults +_TENSOR_PAIR = 'tuple[torch.Tensor, torch.Tensor]' +_Q_TYPE = 'torch.Tensor | tuple[torch.Tensor, Optional[torch.Tensor]]' class BracketTracker: @@ -279,666 +60,299 @@ def is_top_level(self): 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'} - - # Regex: match m.def( ... ), supports multi-line - pattern = re.compile(r'm\.def\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 - - m_def_list = [] - lines = content.splitlines(keepends=True) - i = 0 - 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 - while j < len(lines): - current_line = lines[j] - for k, char in enumerate(current_line): - if char == '(': - if not found_start and re.search(r'm\.def\s*\(', current_line[:k+1]): - found_start = True - if found_start: - paren_count += 1 - elif char == ')': - 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 - break - 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 - }) - - return results - - -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 - start = m_def_str.find('m.def(') - if start == -1: - raise ValueError(f'[{m_def_str}] Could not find m.def start position') - - paren_count = 0 - content_start = start + len('m.def(') - content_end = -1 - for i in range(content_start, len(m_def_str)): - ch = m_def_str[i] - if ch == '(': - paren_count += 1 - elif ch == ')': - if paren_count == 0: - content_end = i - break - else: - paren_count -= 1 - if content_end == -1: - raise ValueError(f'[{m_def_str}] m.def parentheses not closed') - - args_content = m_def_str[content_start:content_end] - - # Split arguments using BracketTracker - args_list = [] +def split_top_level_commas(value: str) -> list[str]: + """Split a string on top-level commas.""" + parts = [] current = [] tracker = BracketTracker() - - for ch in args_content: + for ch in value: if ch in '()[]{}<>': tracker.update(ch) if ch == ',' and tracker.is_top_level(): - args_list.append(''.join(current).strip()) + parts.append(''.join(current).strip()) current = [] else: current.append(ch) - if current: - args_list.append(''.join(current).strip()) - - if not args_list: - raise ValueError(f'[{m_def_str}] m.def has no arguments') - - # Extract Python function name from the first string literal (plain name or schema). - first = args_list[0].strip() - str_match = re.match(r'^"([^"\\]*(?:\\.[^"\\]*)*)"', first) - if str_match: - schema_or_name = str_match.group(1) - result['python_function_name'] = extract_torch_op_name(schema_or_name) - if '(' in schema_or_name: - result['schema_default_args'] = parse_schema_parameter_defaults(schema_or_name) - else: - raise ValueError(f'[{m_def_str}] m.def first argument should be a string literal') + parts.append(''.join(current).strip()) + return parts - if len(args_list) == 1: - result['cpp_function_name'] = result['python_function_name'] - else: - cpp_func_part = args_list[1].strip() - if cpp_func_part.startswith('&'): - cpp_func_part = cpp_func_part[1:].strip() - - if cpp_func_part.startswith('['): - result['is_lambda'] = True - result['cpp_function_name'] = None - elif cpp_func_part.startswith(('DEEP_GEMM_IMPL(', 'TORCH_FN(')): - result['cpp_function_name'] = result['python_function_name'] - 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 - - # Parse py::arg arguments (legacy pybind registrations only). - py_args = args_list[2:] - result['num_args'] = len(py_args) - - 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 +def find_top_level_equals(value: str) -> int: + """Return index of top-level '=' in a schema argument, 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 - 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_torch_op_name(schema: str) -> str: + """Extract the operator name from a TORCH schema string.""" + paren_pos = schema.find('(') + if paren_pos == -1: + return schema.strip() + return schema[:paren_pos].strip() -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 schema_type_to_python(type_str: str) -> str: + """Map a TORCH_LIBRARY schema type to a Python type annotation string.""" + 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]' + else: + print(f'Warning: unrecognized schema type {type_str!r}, using Any') + py_type = 'Any' - 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 + if optional: + return f'Optional[{py_type}]' + return py_type - # 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 schema_return_to_python(return_str: str) -> str: + """Map a TORCH_LIBRARY return type to a Python annotation.""" + return_str = return_str.strip() + 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' -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 = [] +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 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' - 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') +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') - parsed['cpp_signature'] = cpp_sig - statements_with_parsed_signatures.append({ - 'raw': stmt, - 'parsed': parsed - }) + 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 { - 'm_def_statements': statements_with_parsed_signatures + 'name': match.group(2), + 'py_type': schema_type_to_python(match.group(1)), + 'default': default, } -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() - if return_type.startswith('static '): - return_type = return_type[len('static '):].strip() +def parse_torch_schema(schema: str) -> dict: + """Parse a TORCH_LIBRARY 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}') - # Now extract parameter list content - param_list_str = cpp_sig[paren_pos+1:cpp_sig.rfind(')')].strip() + args_blob = signature[open_paren + 1:close_paren].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) + if args_blob: + for arg in split_top_level_commas(args_blob): + parameters.append(parse_schema_arg(arg)) return { - 'return_type': return_type, + 'python_function_name': name, 'parameters': parameters, - 'num_parameters': len(parameters) + 'return_type': return_type, + 'schema': schema, } -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() +def _merge_named_pairs(parameters: list[dict], pairs: tuple[tuple[str, str], ...]) -> list[dict]: + """Replace (left, right) arg pairs with a single tuple-typed parameter.""" + drop = {right for left, right in pairs} + merged_left = {left for left, _ in pairs} + out = [] + for param in parameters: + if param['name'] in drop: + continue + if param['name'] in merged_left: + out.append({ + 'name': param['name'], + 'py_type': _TENSOR_PAIR, + 'default': None, + }) + continue + out.append(dict(param)) + return out - # 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 - } +def _is_tensor_schema_param(param: dict) -> bool: + py_type = param['py_type'] + return py_type in {'torch.Tensor', 'Optional[torch.Tensor]'} - # 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] +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' - 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() +def detect_tensor_sf_pairs(parameters: list[dict]) -> list[tuple[str, str]]: + """Detect consecutive (tensor, scale_factor) arg pairs in a TORCH schema.""" + 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 - # Clean up the type - remove trailing &, * and whitespace - param_type = param_type.rstrip('&* \t') - return { - 'type': param_type, - 'name': param_name - } +def _apply_q_qsf_merge(parameters: list[dict]) -> list[dict]: + """Merge optional q_sf into q for attention wrappers that accept either form.""" + if not any(param['name'] == 'q_sf' for param in parameters): + return [dict(param) for param in parameters] - # Fallback: if we can't find a clear parameter name, just return the type - return { - 'type': decl, - 'name': None - } + out = [] + for param in parameters: + if param['name'] == 'q_sf': + continue + param = dict(param) + if param['name'] == 'q': + param['py_type'] = _Q_TYPE + out.append(param) + return out + + +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 _infer_recipe_tuple_type( + op_name: str, + parameters: list[dict], + wrapper_defaults: dict[str, dict[str, str]], +) -> None: + """Promote int[] recipe args to tuple[int, int, int] when the wrapper uses tuples.""" + wrapper_default = wrapper_defaults.get(op_name, {}).get('recipe') + has_weight_tuples = any( + param['name'] in {'l1_weights', 'l2_weights'} and param['py_type'] == _TENSOR_PAIR + for param in parameters + ) + for param in parameters: + if param['name'] != 'recipe': + continue + if param['py_type'] not in {'list[int]', 'Optional[list[int]]'}: + continue + default_expr = wrapper_default if wrapper_default is not None else param.get('default') + if (default_expr and default_expr.startswith('(')) or has_weight_tuples: + param['py_type'] = 'tuple[int, int, int]' -def extract_cpp_signature_details(item): - """ - For each m.def entry in item, parse cpp_signature to extract return type and parameter details. +def adjust_for_c_py_wrapper( + name: str, + parameters: list[dict], + wrapper_defaults: dict[str, dict[str, str]] | None = None, +) -> list[dict]: """ - 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 - } - + Adjust parsed schema parameters to match deep_gemm._C Python wrappers. -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 / c10::optional - if cleaned.startswith('std::optional<') or cleaned.startswith('c10::optional<'): - inner = cleaned[cleaned.index('<') + 1:-1].strip() - 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: optional with unexpected args: {cleaned}') - return 'Any' - - # c10::List - if cleaned.startswith('c10::List<'): - inner = cleaned[10:-1].strip() - args = split_template_args(inner) - if len(args) == 1: - inner_py = cpp_type_to_python_type(args[0]) - return f'list[{inner_py}]' - print(f'Warning: c10::List 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' - - # at::ScalarType - if re.search(r'\bat::ScalarType\b', original): - return 'torch.dtype' - - # mega.hpp type alias - if re.search(r'\bSymmBufferSlice\b', original): - tensor = 'torch.Tensor' - return f'tuple[{", ".join([tensor] * 8)}]' - - # 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'] + TORCH_LIBRARY registers flat tensor/scales args; _C.py preserves the legacy + pybind API by accepting (tensor, scale_factor) tuples for many kernels. """ - if not template_args.strip(): - return [] - args = [] - current = [] - tracker = BracketTracker() + parameters = _apply_q_qsf_merge(parameters) - for ch in template_args: - if ch in '()[]{}<>': - tracker.update(ch) - if ch == ',' and tracker.is_top_level(): - args.append(''.join(current).strip()) - current = [] - else: - current.append(ch) - - if current: - args.append(''.join(current).strip()) - return args + pairs = detect_tensor_sf_pairs(parameters) + if pairs: + parameters = _merge_named_pairs(parameters, tuple(pairs)) + for param in parameters: + if param['name'] == 'logits_dtype': + param['py_type'] = 'torch.dtype' -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' + _maybe_widen_int_list_value_param(parameters) + _infer_recipe_tuple_type(name, parameters, wrapper_defaults or {}) - s = cpp_default.strip() + return parameters - # 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 - - if s == 'torch::kFloat32': - return 'torch.float32' - - # Fallback: unrecognized → warn and return None - print(f'Warning: Unrecognized default value: {s}') - return 'None' +def sanitize_param_name(name: str) -> str: + if name in {'def', 'class', 'from', 'import', 'None', 'True', 'False'}: + return f'{name}_' + return name def format_ast_default(node: ast.AST) -> str: @@ -982,12 +396,24 @@ def extract_function_defaults(func_def: ast.FunctionDef) -> dict[str, str]: return defaults -def parse_wrapper_defaults(c_py_path: Path) -> dict[str, dict[str, str]]: - """Parse deep_gemm/_C.py wrapper function defaults keyed by exported name.""" +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) -> tuple[dict[str, dict[str, str]], dict[str, str]]: + """Parse wrapper defaults and legacy alias exports from deep_gemm/_C.py.""" 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]] = {} + aliases: dict[str, str] = {} for node in ast.walk(module): if isinstance(node, ast.FunctionDef): @@ -996,12 +422,7 @@ def parse_wrapper_defaults(c_py_path: Path) -> dict[str, dict[str, str]]: for node in ast.walk(module): if not isinstance(node, ast.Call): continue - if not ( - isinstance(node.func, ast.Attribute) - and node.func.attr == 'update' - and isinstance(node.func.value, ast.Name) - and node.func.value.id == 'globals' - ): + if not _is_globals_update_call(node): continue if not node.args or not isinstance(node.args[0], ast.Dict): continue @@ -1010,135 +431,255 @@ def parse_wrapper_defaults(c_py_path: Path) -> dict[str, dict[str, str]]: 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) and value_node.id in func_defaults: - func_defaults[alias_name] = func_defaults[value_node.id] + if isinstance(value_node, ast.Name): + if value_node.id in func_defaults: + func_defaults[alias_name] = func_defaults[value_node.id] + if alias_name != value_node.id: + aliases[alias_name] = value_node.id - return func_defaults + return func_defaults, aliases -def generate_pyi_function(item_entry, wrapper_defaults=None): - parsed = item_entry['parsed'] - py_name = parsed['python_function_name'] +def extract_m_def_statements(root_path): + """ + Scan all C++ files under root_path and extract all m.def(...) statements. + Supports multi-line m.def(...) calls. + """ + results = [] + extensions = {'.hpp', '.cpp', '.h', '.cc'} + + 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 + + m_def_list = [] + lines = content.splitlines(keepends=True) + i = 0 + while i < len(lines): + line = lines[i] + if 'm.def(' in line: + # Found a potential starting line + # 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 + while j < len(lines): + current_line = lines[j] + for k, char in enumerate(current_line): + if char == '(': + if not found_start and re.search(r'm\.def\s*\(', current_line[:k+1]): + found_start = True + if found_start: + paren_count += 1 + elif char == ')': + 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 + break + if paren_count <= 0 and found_start: + break + j += 1 + i += 1 + + if m_def_list: + results.append({ + 'file': str(file_path), + 'm_def_statements': m_def_list + }) + + return results + + +def parse_m_def_statement(m_def_str): + """ + Parse a TORCH_LIBRARY m.def(...) statement. + + DeepGEMM registers ops via TORCH_LIBRARY_FRAGMENT, so the first argument is + always a schema string such as "fp8_fp4_gemm_nt(Tensor a, ...) -> ()". + """ + # Extract top-level arguments + start = m_def_str.find('m.def(') + if start == -1: + raise ValueError(f'[{m_def_str}] Could not find m.def start position') - if parsed.get('is_lambda'): - return f'def {py_name}(*args, **kwargs) -> Any: ...' + paren_count = 0 + content_start = start + len('m.def(') + content_end = -1 + for i in range(content_start, len(m_def_str)): + ch = m_def_str[i] + if ch == '(': + paren_count += 1 + elif ch == ')': + if paren_count == 0: + content_end = i + break + else: + paren_count -= 1 + if content_end == -1: + raise ValueError(f'[{m_def_str}] m.def parentheses not closed') - sig_info = parsed.get('cpp_parsed_signature') - default_args = dict(parsed.get('default_args', {})) - schema_default_by_name = parsed.get('schema_default_args', {}) - wrapper_default_by_name = (wrapper_defaults or {}).get(py_name, {}) + args_content = m_def_str[content_start:content_end] - if not sig_info: - return f'def {py_name}(*args, **kwargs) -> Any: ...' + # Split arguments using BracketTracker + args_list = split_top_level_commas(args_content) - return_type = cpp_type_to_python_type(sig_info['return_type']) - params = sig_info['parameters'] - num_params = len(params) + if not args_list: + raise ValueError(f'[{m_def_str}] m.def has no arguments') + + # Extract operator schema from the first string literal + first = args_list[0].strip() + str_match = re.match(r'^"([^"\\]*(?:\\.[^"\\]*)*)"', first) + if not str_match: + raise ValueError(f'[{m_def_str}] m.def first argument should be a string literal') + + return parse_torch_schema(str_match.group(1)) + + +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 + + 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 + + +def generate_pyi_function(item_entry, wrapper_defaults=None): + """Generate a typed .pyi stub for one registered op.""" + parsed = item_entry['parsed'] + 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'] - # 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}_' - - # Defaults: py::arg > _C.py wrapper > TORCH schema. - py_default = None - if i in default_args: - py_default = cpp_default_to_python_default(default_args[i]) - elif param_name in wrapper_default_by_name: - py_default = wrapper_default_by_name[param_name] - elif param_name in schema_default_by_name: - py_default = schema_default_by_name[param_name] - if param_type == 'torch.dtype' and py_default == '6': - py_default = 'torch.float32' - - if py_default is not None: - param_str = f' {param_name}: {param_type} = {py_default}' + for param in parameters: + name = sanitize_param_name(param['name']) + if param['default'] is not None: + param_lines.append(f' {name}: {param["py_type"]} = {param["default"]}') else: - param_str = f' {param_name}: {param_type}' - - param_lines.append(param_str) + param_lines.append(f' {name}: {param["py_type"]}') 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 f'def {py_name}(\n{params_block}\n) -> {return_type}: ...' + return f'def {py_name}() -> {return_type}: ...' - return func_def +def _alias_pyi_decl(decl: str, alias_name: str, source_name: str) -> str: + return decl.replace(f'def {source_name}(', f'def {alias_name}(', 1) -def generate_pyi_file_content(enhanced_results, module_name: str = 'my_module', wrapper_defaults=None): - function_decls = [] - has_optional = False - has_torch = False - has_numpy = False +def generate_pyi_file_content( + enhanced_results, + module_name: str = 'my_module', + wrapper_defaults=None, + pyi_aliases=None, +): + by_name = {} for item in enhanced_results: for stmt in item['m_def_statements']: - try: - decl = generate_pyi_function(stmt, wrapper_defaults=wrapper_defaults) - 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' + name = stmt['parsed']['python_function_name'] + by_name[name] = stmt - if has_torch: - imports.append('import torch') - if has_numpy: - imports.append('import numpy') + decl_by_name = {} + has_optional = False + has_torch = False - lines = [f'# Stubs for module: {module_name}', ''] - lines.extend(imports) - lines.append('') - lines.append('') + for name in sorted(by_name): + stmt = by_name[name] + try: + decl = generate_pyi_function(stmt, wrapper_defaults=wrapper_defaults) + decl_by_name[name] = decl + if 'Optional[' in decl: + has_optional = True + if 'torch.' in decl: + has_torch = True + except Exception as e: + decl_by_name[name] = f'# ERROR: failed to generate stub for {name}: {e}' + + for alias_name, source_name in sorted((pyi_aliases or {}).items()): + if source_name in decl_by_name and alias_name not in decl_by_name: + decl_by_name[alias_name] = _alias_pyi_decl(decl_by_name[source_name], alias_name, source_name) + if 'Optional[' in decl_by_name[alias_name]: + has_optional = True + if 'torch.' in decl_by_name[alias_name]: + has_torch = True + + lines = [ + f'# Stubs for module: {module_name}', + '', + 'from typing import Any', + ] + if has_optional: + lines[2] += ', Optional' + if has_torch: + lines.append('import torch') + lines.extend(['', '']) - for decl in function_decls: - lines.append(decl) - lines.append('') - lines.append('') + for name in sorted(decl_by_name): + lines.extend([decl_by_name[name], '', '']) return '\n'.join(lines) def generate_pyi_file(name, root, output_dir='.', c_py_path=None): - func_index = build_cpp_function_index(root) results = extract_m_def_statements(root) - cpp_results = [] + enhanced_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) + statements = [] + for stmt in item['m_def_statements']: + statements.append({ + 'raw': stmt, + 'parsed': parse_m_def_statement(stmt), + }) + enhanced_results.append({'m_def_statements': statements}) wrapper_defaults = {} + pyi_aliases = {} if c_py_path is not None: c_py_path = Path(c_py_path) if c_py_path.is_file(): - wrapper_defaults = parse_wrapper_defaults(c_py_path) + wrapper_defaults, pyi_aliases = parse_c_py_metadata(c_py_path) else: print(f'Warning: wrapper file not found: {c_py_path}') pyi_content = generate_pyi_file_content( - cpp_results, + enhanced_results, module_name=name, wrapper_defaults=wrapper_defaults, + pyi_aliases=pyi_aliases, ) output_path = Path(output_dir) / f'{name}.pyi' @@ -1155,7 +696,7 @@ def main(argv=None) -> int: import sys parser = argparse.ArgumentParser( - description='Generate deep_gemm/_C.pyi stubs from C++ signatures and m.def registrations.', + description='Generate deep_gemm/_C.pyi stubs from TORCH_LIBRARY schemas.', ) parser.add_argument('--name', default='_C', help='Module name for the .pyi file (default: _C)') parser.add_argument('--root', default='./csrc', help='Root to scan for m.def(...) (default: ./csrc)') @@ -1199,7 +740,10 @@ def main(argv=None) -> int: if stub_count == 0: print(f'CHECK FAILED: no function stubs in {pyi_path}', file=sys.stderr) return 1 - print(f'CHECK PASSED: {stub_count} stubs in {pyi_path} ({stub_count - generic_count} typed, {generic_count} generic)') + print( + f'CHECK PASSED: {stub_count} stubs in {pyi_path} ' + f'({stub_count - generic_count} typed, {generic_count} generic)', + ) if generic_count > 0: return 1 return 0 From fd57d1cf4e220ed06c8136f2feb6112732f9cd79 Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Fri, 10 Jul 2026 15:09:34 +0000 Subject: [PATCH 08/28] made some updates to the generate_pyi.py file to remove aliases and other minor formatting issues. Also added comments with examples to help describe each step. These will be removed (as well as main()) in a followup commit, but wanted them here for reference Signed-off-by: Chris Leonard --- scripts/generate_pyi.py | 321 ++++++++++++++++++++++++++-------------- 1 file changed, 208 insertions(+), 113 deletions(-) diff --git a/scripts/generate_pyi.py b/scripts/generate_pyi.py index b75701330c..a9d7294bd5 100644 --- a/scripts/generate_pyi.py +++ b/scripts/generate_pyi.py @@ -1,9 +1,49 @@ +""" +Generate deep_gemm/_C.pyi from TORCH_LIBRARY schemas and deep_gemm/_C.py wrappers. + +Pipeline (one op at a time, e.g. fp8_fp4_gemm_nt): + + 1. extract_m_def_statements(csrc/) + Scan csrc/apis/*.hpp for m.def("...schema...") registrations. + Returns a flat list of raw m.def(...) statement strings. + + 2. parse_m_def_statement(stmt) -> parse_torch_schema(schema) + Split the schema string into structured params (flat, as registered in C++). + Example output: + name='fp8_fp4_gemm_nt' + parameters=[ + {'name': 'a', 'py_type': 'torch.Tensor', 'default': None}, + {'name': 'sfa', 'py_type': 'torch.Tensor', 'default': None}, + {'name': 'b', 'py_type': 'torch.Tensor', 'default': None}, + ... + ] + + 3. parse_c_py_metadata(deep_gemm/_C.py) + Read wrapper defaults from the Python shim (including names exported via + globals().update aliases, so fp8_gemm_nt picks up fp8_fp4_gemm_nt defaults). + Example: + defaults['fp8_einsum']['recipe'] = '(1, 128, 128)' + + 4. adjust_for_c_py_wrapper(name, parameters, wrapper_defaults) + Reshape flat schema params to match the public _C.py API. + Example: (a, sfa), (b, sfb) -> a: tuple[Tensor, Tensor], b: tuple[...] + + 5. apply_wrapper_defaults(name, parameters, wrapper_defaults) + Overlay defaults from _C.py where they differ from the TORCH schema. + Example: recipe=None in schema -> recipe=(1, 128, 128) from fp8_einsum() + + 6. generate_pyi_function(...) -> str + Render one stub def line block for the .pyi file. + + 7. generate_pyi_file_content(...) + Emit one stub per TORCH_LIBRARY op found in csrc/. +""" import ast import re from pathlib import Path _TENSOR_PAIR = 'tuple[torch.Tensor, torch.Tensor]' -_Q_TYPE = 'torch.Tensor | tuple[torch.Tensor, Optional[torch.Tensor]]' +_Q_TUPLE = 'tuple[torch.Tensor, Optional[torch.Tensor]]' class BracketTracker: @@ -61,7 +101,13 @@ def is_top_level(self): def split_top_level_commas(value: str) -> list[str]: - """Split a string on top-level commas.""" + """Split a string on top-level commas. + + Example: + "Tensor a, Tensor? c=None, int[]? recipe=None" + -> ["Tensor a", "Tensor? c=None", "int[]? recipe=None"] + Commas inside brackets/parens (e.g. Tensor(d!) d) are ignored. + """ parts = [] current = [] tracker = BracketTracker() @@ -89,16 +135,15 @@ def find_top_level_equals(value: str) -> int: return -1 -def extract_torch_op_name(schema: str) -> str: - """Extract the operator name from a TORCH schema string.""" - paren_pos = schema.find('(') - if paren_pos == -1: - return schema.strip() - return schema[:paren_pos].strip() - - def schema_type_to_python(type_str: str) -> str: - """Map a TORCH_LIBRARY schema type to a Python type annotation string.""" + """Map a TORCH_LIBRARY schema type to a Python type annotation string. + + Examples: + "Tensor" -> "torch.Tensor" + "Tensor?" -> "Optional[torch.Tensor]" + "int[]" -> "list[int]" + "int[]?" -> "Optional[list[int]]" + """ type_str = type_str.strip() optional = type_str.endswith('?') if optional: @@ -166,7 +211,12 @@ def schema_default_to_python(default_str: str) -> str: def parse_schema_arg(arg_str: str) -> dict: - """Parse one TORCH schema argument such as 'Tensor? c=None'.""" + """Parse one TORCH schema argument such as 'Tensor? c=None'. + + Example: + "str compiled_dims='nk'" + -> {'name': 'compiled_dims', 'py_type': 'str', 'default': "'nk'"} + """ arg_str = arg_str.strip() if not arg_str: raise ValueError('empty schema argument') @@ -188,7 +238,23 @@ def parse_schema_arg(arg_str: str) -> dict: def parse_torch_schema(schema: str) -> dict: - """Parse a TORCH_LIBRARY schema into name, parameters, and return type.""" + """Parse a TORCH_LIBRARY schema into name, parameters, and return type. + + Example input: + "fp8_fp4_gemm_nt(Tensor a, Tensor sfa, Tensor b, Tensor sfb, " + "Tensor(d!) d, Tensor? c=None, str compiled_dims='nk') -> ()" + + Example output (abbreviated): + { + 'python_function_name': 'fp8_fp4_gemm_nt', + 'return_type': 'None', + 'parameters': [ + {'name': 'a', 'py_type': 'torch.Tensor', 'default': None}, + {'name': 'sfa', 'py_type': 'torch.Tensor', 'default': None}, + ... + ], + } + """ arrow = schema.rfind(' -> ') if arrow == -1: raise ValueError(f'schema missing return type: {schema!r}') @@ -224,12 +290,15 @@ def parse_torch_schema(schema: str) -> dict: 'python_function_name': name, 'parameters': parameters, 'return_type': return_type, - 'schema': schema, } def _merge_named_pairs(parameters: list[dict], pairs: tuple[tuple[str, str], ...]) -> list[dict]: - """Replace (left, right) arg pairs with a single tuple-typed parameter.""" + """Replace (left, right) arg pairs with a single tuple-typed parameter. + + Example: pairs=(('a', 'sfa'), ('b', 'sfb')) + [a, sfa, b, sfb, d, ...] -> [a: tuple[Tensor, Tensor], b: tuple[...], d, ...] + """ drop = {right for left, right in pairs} merged_left = {left for left, _ in pairs} out = [] @@ -261,7 +330,13 @@ def _is_tensor_scale_factor_pair(base_name: str, sf_name: str) -> bool: def detect_tensor_sf_pairs(parameters: list[dict]) -> list[tuple[str, str]]: - """Detect consecutive (tensor, scale_factor) arg pairs in a TORCH schema.""" + """Detect consecutive (tensor, scale_factor) arg pairs in a TORCH schema. + + Examples (flat schema params from step 2): + [a, sfa, b, sfb, ...] -> [('a', 'sfa'), ('b', 'sfb')] + [kv, kv_sf, weights, ...] -> [('kv', 'kv_sf')] + [l1_weights, l1_weights_sf, ...] -> [('l1_weights', 'l1_weights_sf')] + """ pairs = [] i = 0 while i < len(parameters) - 1: @@ -279,7 +354,12 @@ def detect_tensor_sf_pairs(parameters: list[dict]) -> list[tuple[str, str]]: def _apply_q_qsf_merge(parameters: list[dict]) -> list[dict]: - """Merge optional q_sf into q for attention wrappers that accept either form.""" + """Merge optional q_sf into q for attention wrappers that accept either form. + + Example schema: q, q_sf, kv, kv_sf, ... + -> q: tuple[Tensor, Optional[Tensor]], kv, kv_sf, ... + (q_sf is dropped; kv/kv_sf merging happens separately via detect_tensor_sf_pairs.) + """ if not any(param['name'] == 'q_sf' for param in parameters): return [dict(param) for param in parameters] @@ -289,7 +369,7 @@ def _apply_q_qsf_merge(parameters: list[dict]) -> list[dict]: continue param = dict(param) if param['name'] == 'q': - param['py_type'] = _Q_TYPE + param['py_type'] = _Q_TUPLE out.append(param) return out @@ -301,25 +381,29 @@ def _maybe_widen_int_list_value_param(parameters: list[dict]) -> None: parameters[0]['py_type'] = 'int | list[int]' -def _infer_recipe_tuple_type( - op_name: str, - parameters: list[dict], - wrapper_defaults: dict[str, dict[str, str]], -) -> None: - """Promote int[] recipe args to tuple[int, int, int] when the wrapper uses tuples.""" - wrapper_default = wrapper_defaults.get(op_name, {}).get('recipe') - has_weight_tuples = any( - param['name'] in {'l1_weights', 'l2_weights'} and param['py_type'] == _TENSOR_PAIR - for param in parameters - ) +def _promote_int_list_tuple_types(op_name: str, parameters: list[dict]) -> None: + """Promote int[] schema params to fixed-size tuples matching the public _C.py API. + + TORCH schemas use int[] for C++ list/variant conversions; callers pass tuples. + """ for param in parameters: - if param['name'] != 'recipe': - continue - if param['py_type'] not in {'list[int]', 'Optional[list[int]]'}: - continue - default_expr = wrapper_default if wrapper_default is not None else param.get('default') - if (default_expr and default_expr.startswith('(')) or has_weight_tuples: + name = param['name'] + py_type = param['py_type'] + + if name == 'head_splits' and py_type == 'list[int]': param['py_type'] = 'tuple[int, int, int]' + elif name == 'recipe_a' and py_type == 'Optional[list[int]]': + param['py_type'] = 'Optional[tuple[int, int]]' + elif name == 'recipe_b' and py_type == 'Optional[list[int]]': + param['py_type'] = 'Optional[tuple[int, int]]' + elif name == 'recipe' and op_name == 'transform_sf_into_required_layout': + if py_type == 'list[int]': + param['py_type'] = 'tuple[int, int] | tuple[int, int, int]' + elif name == 'recipe': + if py_type == 'list[int]': + param['py_type'] = 'tuple[int, int, int]' + elif py_type == 'Optional[list[int]]': + param['py_type'] = 'Optional[tuple[int, int, int]]' def adjust_for_c_py_wrapper( @@ -332,6 +416,10 @@ def adjust_for_c_py_wrapper( TORCH_LIBRARY registers flat tensor/scales args; _C.py preserves the legacy pybind API by accepting (tensor, scale_factor) tuples for many kernels. + + Example transformation for fp8_fp4_gemm_nt: + schema: a, sfa, b, sfb, d, c=None, recipe=None, compiled_dims='nk', ... + stub: a: tuple[Tensor, Tensor], b: tuple[Tensor, Tensor], d, c=None, ... """ parameters = _apply_q_qsf_merge(parameters) @@ -344,7 +432,7 @@ def adjust_for_c_py_wrapper( param['py_type'] = 'torch.dtype' _maybe_widen_int_list_value_param(parameters) - _infer_recipe_tuple_type(name, parameters, wrapper_defaults or {}) + _promote_int_list_tuple_types(name, parameters) return parameters @@ -407,13 +495,23 @@ def _is_globals_update_call(node: ast.Call) -> bool: return False -def parse_c_py_metadata(c_py_path: Path) -> tuple[dict[str, dict[str, str]], dict[str, str]]: - """Parse wrapper defaults and legacy alias exports from deep_gemm/_C.py.""" +def parse_c_py_metadata(c_py_path: Path) -> dict[str, dict[str, str]]: + """Parse wrapper defaults from deep_gemm/_C.py. + + Walks the AST (does not import or execute _C.py). + + Defaults example — from: + def fp8_einsum(..., recipe=(1, 128, 128)): + produces: + defaults['fp8_einsum']['recipe'] = '(1, 128, 128)' + + Also copies defaults onto legacy alias names from globals().update(...) so + wrapper defaults apply when the public name differs from the TORCH op name. + """ 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]] = {} - aliases: dict[str, str] = {} for node in ast.walk(module): if isinstance(node, ast.FunctionDef): @@ -434,18 +532,22 @@ def parse_c_py_metadata(c_py_path: Path) -> tuple[dict[str, dict[str, str]], dic if isinstance(value_node, ast.Name): if value_node.id in func_defaults: func_defaults[alias_name] = func_defaults[value_node.id] - if alias_name != value_node.id: - aliases[alias_name] = value_node.id - return func_defaults, aliases + return func_defaults -def extract_m_def_statements(root_path): +def extract_m_def_statements(root_path) -> list[str]: """ Scan all C++ files under root_path and extract all m.def(...) statements. - Supports multi-line m.def(...) calls. + + Returns a flat list of raw statement strings (one per registration found). + Supports multi-line m.def(...) calls. This is pipeline step 1. + + Example match in gemm.hpp: + m.def( + "fp8_fp4_gemm_nt(Tensor a, Tensor sfa, ...) -> ()"); """ - results = [] + statements = [] extensions = {'.hpp', '.cpp', '.h', '.cc'} for file_path in Path(root_path).rglob('*'): @@ -501,20 +603,24 @@ def extract_m_def_statements(root_path): 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): """ - Parse a TORCH_LIBRARY m.def(...) statement. + Parse a TORCH_LIBRARY m.def(...) statement (pipeline step 2). + + DeepGEMM registers ops via TORCH_LIBRARY_FRAGMENT; the first m.def argument + is always a schema string. Extra args like DEEP_GEMM_IMPL(...) are ignored. + + Example input: + m.def("bf16_gemm_nt(Tensor a, Tensor b, Tensor(d!) d, " + "Tensor? c=None, str compiled_dims='nk') -> ()", + DEEP_GEMM_IMPL(bf16_gemm_nt)); - DeepGEMM registers ops via TORCH_LIBRARY_FRAGMENT, so the first argument is - always a schema string such as "fp8_fp4_gemm_nt(Tensor a, ...) -> ()". + Delegates to parse_torch_schema() on the first string literal. """ # Extract top-level arguments start = m_def_str.find('m.def(') @@ -555,7 +661,16 @@ def parse_m_def_statement(m_def_str): 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.""" + """Overlay public API defaults from deep_gemm/_C.py onto schema-derived parameters. + + Schema defaults come from the TORCH registration string; wrapper defaults reflect + what callers actually get from _C.py. + + Example for fp8_einsum: + schema default: recipe=None + wrapper default: recipe=(1, 128, 128) # from def fp8_einsum(..., recipe=(1, 128, 128)) + stub result: recipe: tuple[int, int, int] = (1, 128, 128) + """ by_name = wrapper_defaults.get(name, {}) if not by_name: return parameters @@ -569,9 +684,18 @@ def apply_wrapper_defaults(name: str, parameters: list[dict], wrapper_defaults: return out -def generate_pyi_function(item_entry, wrapper_defaults=None): - """Generate a typed .pyi stub for one registered op.""" - parsed = item_entry['parsed'] +def generate_pyi_function(parsed, wrapper_defaults=None): + """Generate a typed .pyi stub for one registered op (pipeline steps 4-6). + + Example final output for fp8_fp4_gemm_nt: + def fp8_fp4_gemm_nt( + a: tuple[torch.Tensor, torch.Tensor], + b: tuple[torch.Tensor, torch.Tensor], + d: torch.Tensor, + c: Optional[torch.Tensor] = None, + ... + ) -> None: ... + """ py_name = parsed['python_function_name'] parameters = adjust_for_c_py_wrapper( py_name, @@ -596,90 +720,61 @@ def generate_pyi_function(item_entry, wrapper_defaults=None): return f'def {py_name}() -> {return_type}: ...' -def _alias_pyi_decl(decl: str, alias_name: str, source_name: str) -> str: - return decl.replace(f'def {source_name}(', f'def {alias_name}(', 1) - - def generate_pyi_file_content( - enhanced_results, + parsed_ops, module_name: str = 'my_module', wrapper_defaults=None, - pyi_aliases=None, ): - by_name = {} - for item in enhanced_results: - for stmt in item['m_def_statements']: - name = stmt['parsed']['python_function_name'] - by_name[name] = stmt - - decl_by_name = {} - has_optional = False - has_torch = False - - for name in sorted(by_name): - stmt = by_name[name] + """Assemble the full .pyi file from all parsed ops (pipeline step 7). + + parsed_ops: list of dicts returned by parse_m_def_statement / parse_torch_schema. + + - One stub per TORCH_LIBRARY op found in csrc/ + """ + decls = [] + + for parsed in parsed_ops: + name = parsed['python_function_name'] try: - decl = generate_pyi_function(stmt, wrapper_defaults=wrapper_defaults) - decl_by_name[name] = decl - if 'Optional[' in decl: - has_optional = True - if 'torch.' in decl: - has_torch = True + decl = generate_pyi_function(parsed, wrapper_defaults=wrapper_defaults) except Exception as e: - decl_by_name[name] = f'# ERROR: failed to generate stub for {name}: {e}' - - for alias_name, source_name in sorted((pyi_aliases or {}).items()): - if source_name in decl_by_name and alias_name not in decl_by_name: - decl_by_name[alias_name] = _alias_pyi_decl(decl_by_name[source_name], alias_name, source_name) - if 'Optional[' in decl_by_name[alias_name]: - has_optional = True - if 'torch.' in decl_by_name[alias_name]: - has_torch = True + decl = f'# ERROR: failed to generate stub for {name}: {e}' + decls.append(decl) lines = [ f'# Stubs for module: {module_name}', '', - 'from typing import Any', + 'from typing import Any, Optional', + 'import torch', + '', ] - if has_optional: - lines[2] += ', Optional' - if has_torch: - lines.append('import torch') - lines.extend(['', '']) - for name in sorted(decl_by_name): - lines.extend([decl_by_name[name], '', '']) + for decl in decls: + lines.extend([decl, '', '']) return '\n'.join(lines) def generate_pyi_file(name, root, output_dir='.', c_py_path=None): - results = extract_m_def_statements(root) - - enhanced_results = [] - for item in results: - statements = [] - for stmt in item['m_def_statements']: - statements.append({ - 'raw': stmt, - 'parsed': parse_m_def_statement(stmt), - }) - enhanced_results.append({'m_def_statements': statements}) + """Orchestrate the full pipeline and write stubs/.pyi.""" + # Step 1-2: scan csrc/ for m.def(...) and parse each TORCH schema. + m_def_statements = extract_m_def_statements(root) + parsed_ops = [parse_m_def_statement(stmt) for stmt in m_def_statements] + # Step 3: read wrapper defaults from deep_gemm/_C.py. wrapper_defaults = {} - pyi_aliases = {} if c_py_path is not None: c_py_path = Path(c_py_path) if c_py_path.is_file(): - wrapper_defaults, pyi_aliases = parse_c_py_metadata(c_py_path) + wrapper_defaults = parse_c_py_metadata(c_py_path) else: print(f'Warning: wrapper file not found: {c_py_path}') + # Steps 4-7: adjust params, apply defaults, render stubs, write file. pyi_content = generate_pyi_file_content( - enhanced_results, + parsed_ops, module_name=name, wrapper_defaults=wrapper_defaults, - pyi_aliases=pyi_aliases, ) output_path = Path(output_dir) / f'{name}.pyi' From 2a1da8ba79aad454e050800e8d154fbe01a87e64 Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Fri, 10 Jul 2026 15:17:57 +0000 Subject: [PATCH 09/28] removed verbose comments and main function from generate_pyi.py Signed-off-by: Chris Leonard --- scripts/generate_pyi.py | 296 ++++------------------------------------ 1 file changed, 24 insertions(+), 272 deletions(-) diff --git a/scripts/generate_pyi.py b/scripts/generate_pyi.py index a9d7294bd5..2ba6ef7e65 100644 --- a/scripts/generate_pyi.py +++ b/scripts/generate_pyi.py @@ -1,43 +1,4 @@ -""" -Generate deep_gemm/_C.pyi from TORCH_LIBRARY schemas and deep_gemm/_C.py wrappers. - -Pipeline (one op at a time, e.g. fp8_fp4_gemm_nt): - - 1. extract_m_def_statements(csrc/) - Scan csrc/apis/*.hpp for m.def("...schema...") registrations. - Returns a flat list of raw m.def(...) statement strings. - - 2. parse_m_def_statement(stmt) -> parse_torch_schema(schema) - Split the schema string into structured params (flat, as registered in C++). - Example output: - name='fp8_fp4_gemm_nt' - parameters=[ - {'name': 'a', 'py_type': 'torch.Tensor', 'default': None}, - {'name': 'sfa', 'py_type': 'torch.Tensor', 'default': None}, - {'name': 'b', 'py_type': 'torch.Tensor', 'default': None}, - ... - ] - - 3. parse_c_py_metadata(deep_gemm/_C.py) - Read wrapper defaults from the Python shim (including names exported via - globals().update aliases, so fp8_gemm_nt picks up fp8_fp4_gemm_nt defaults). - Example: - defaults['fp8_einsum']['recipe'] = '(1, 128, 128)' - - 4. adjust_for_c_py_wrapper(name, parameters, wrapper_defaults) - Reshape flat schema params to match the public _C.py API. - Example: (a, sfa), (b, sfb) -> a: tuple[Tensor, Tensor], b: tuple[...] - - 5. apply_wrapper_defaults(name, parameters, wrapper_defaults) - Overlay defaults from _C.py where they differ from the TORCH schema. - Example: recipe=None in schema -> recipe=(1, 128, 128) from fp8_einsum() - - 6. generate_pyi_function(...) -> str - Render one stub def line block for the .pyi file. - - 7. generate_pyi_file_content(...) - Emit one stub per TORCH_LIBRARY op found in csrc/. -""" +"""Generate deep_gemm/_C.pyi from TORCH_LIBRARY schemas and deep_gemm/_C.py wrappers.""" import ast import re from pathlib import Path @@ -47,14 +8,8 @@ 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 # [] @@ -62,9 +17,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 == ')': @@ -85,29 +37,14 @@ 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) + 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 a string on top-level commas. - - Example: - "Tensor a, Tensor? c=None, int[]? recipe=None" - -> ["Tensor a", "Tensor? c=None", "int[]? recipe=None"] - Commas inside brackets/parens (e.g. Tensor(d!) d) are ignored. - """ + """Split on commas not nested inside brackets.""" parts = [] current = [] tracker = BracketTracker() @@ -125,7 +62,7 @@ def split_top_level_commas(value: str) -> list[str]: def find_top_level_equals(value: str) -> int: - """Return index of top-level '=' in a schema argument, or -1.""" + """Return index of top-level '=', or -1.""" tracker = BracketTracker() for i, ch in enumerate(value): if ch in '()[]{}<>': @@ -136,14 +73,7 @@ def find_top_level_equals(value: str) -> int: def schema_type_to_python(type_str: str) -> str: - """Map a TORCH_LIBRARY schema type to a Python type annotation string. - - Examples: - "Tensor" -> "torch.Tensor" - "Tensor?" -> "Optional[torch.Tensor]" - "int[]" -> "list[int]" - "int[]?" -> "Optional[list[int]]" - """ + """Map a TORCH schema type to a Python annotation.""" type_str = type_str.strip() optional = type_str.endswith('?') if optional: @@ -171,7 +101,7 @@ def schema_type_to_python(type_str: str) -> str: def schema_return_to_python(return_str: str) -> str: - """Map a TORCH_LIBRARY return type to a Python annotation.""" + """Map a TORCH schema return type to a Python annotation.""" return_str = return_str.strip() if return_str == '()': return 'None' @@ -211,12 +141,7 @@ def schema_default_to_python(default_str: str) -> str: def parse_schema_arg(arg_str: str) -> dict: - """Parse one TORCH schema argument such as 'Tensor? c=None'. - - Example: - "str compiled_dims='nk'" - -> {'name': 'compiled_dims', 'py_type': 'str', 'default': "'nk'"} - """ + """Parse one TORCH schema argument such as 'Tensor? c=None'.""" arg_str = arg_str.strip() if not arg_str: raise ValueError('empty schema argument') @@ -238,23 +163,7 @@ def parse_schema_arg(arg_str: str) -> dict: def parse_torch_schema(schema: str) -> dict: - """Parse a TORCH_LIBRARY schema into name, parameters, and return type. - - Example input: - "fp8_fp4_gemm_nt(Tensor a, Tensor sfa, Tensor b, Tensor sfb, " - "Tensor(d!) d, Tensor? c=None, str compiled_dims='nk') -> ()" - - Example output (abbreviated): - { - 'python_function_name': 'fp8_fp4_gemm_nt', - 'return_type': 'None', - 'parameters': [ - {'name': 'a', 'py_type': 'torch.Tensor', 'default': None}, - {'name': 'sfa', 'py_type': 'torch.Tensor', 'default': None}, - ... - ], - } - """ + """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}') @@ -294,11 +203,7 @@ def parse_torch_schema(schema: str) -> dict: def _merge_named_pairs(parameters: list[dict], pairs: tuple[tuple[str, str], ...]) -> list[dict]: - """Replace (left, right) arg pairs with a single tuple-typed parameter. - - Example: pairs=(('a', 'sfa'), ('b', 'sfb')) - [a, sfa, b, sfb, d, ...] -> [a: tuple[Tensor, Tensor], b: tuple[...], d, ...] - """ + """Replace (tensor, scale_factor) arg pairs with one tuple-typed parameter.""" drop = {right for left, right in pairs} merged_left = {left for left, _ in pairs} out = [] @@ -330,13 +235,7 @@ def _is_tensor_scale_factor_pair(base_name: str, sf_name: str) -> bool: def detect_tensor_sf_pairs(parameters: list[dict]) -> list[tuple[str, str]]: - """Detect consecutive (tensor, scale_factor) arg pairs in a TORCH schema. - - Examples (flat schema params from step 2): - [a, sfa, b, sfb, ...] -> [('a', 'sfa'), ('b', 'sfb')] - [kv, kv_sf, weights, ...] -> [('kv', 'kv_sf')] - [l1_weights, l1_weights_sf, ...] -> [('l1_weights', 'l1_weights_sf')] - """ + """Detect consecutive (tensor, scale_factor) arg pairs.""" pairs = [] i = 0 while i < len(parameters) - 1: @@ -354,12 +253,7 @@ def detect_tensor_sf_pairs(parameters: list[dict]) -> list[tuple[str, str]]: def _apply_q_qsf_merge(parameters: list[dict]) -> list[dict]: - """Merge optional q_sf into q for attention wrappers that accept either form. - - Example schema: q, q_sf, kv, kv_sf, ... - -> q: tuple[Tensor, Optional[Tensor]], kv, kv_sf, ... - (q_sf is dropped; kv/kv_sf merging happens separately via detect_tensor_sf_pairs.) - """ + """Merge optional q_sf into q for attention wrappers.""" if not any(param['name'] == 'q_sf' for param in parameters): return [dict(param) for param in parameters] @@ -382,10 +276,7 @@ def _maybe_widen_int_list_value_param(parameters: list[dict]) -> None: def _promote_int_list_tuple_types(op_name: str, parameters: list[dict]) -> None: - """Promote int[] schema params to fixed-size tuples matching the public _C.py API. - - TORCH schemas use int[] for C++ list/variant conversions; callers pass tuples. - """ + """Promote int[] schema params to fixed-size tuples matching the public _C.py API.""" for param in parameters: name = param['name'] py_type = param['py_type'] @@ -411,16 +302,7 @@ def adjust_for_c_py_wrapper( parameters: list[dict], wrapper_defaults: dict[str, dict[str, str]] | None = None, ) -> list[dict]: - """ - Adjust parsed schema parameters to match deep_gemm._C Python wrappers. - - TORCH_LIBRARY registers flat tensor/scales args; _C.py preserves the legacy - pybind API by accepting (tensor, scale_factor) tuples for many kernels. - - Example transformation for fp8_fp4_gemm_nt: - schema: a, sfa, b, sfb, d, c=None, recipe=None, compiled_dims='nk', ... - stub: a: tuple[Tensor, Tensor], b: tuple[Tensor, Tensor], d, c=None, ... - """ + """Adjust flat schema params to match deep_gemm._C Python wrappers.""" parameters = _apply_q_qsf_merge(parameters) pairs = detect_tensor_sf_pairs(parameters) @@ -444,7 +326,7 @@ def sanitize_param_name(name: str) -> str: def format_ast_default(node: ast.AST) -> str: - """Convert an AST default value node to a Python expression string for stubs.""" + """Convert an AST default value node to a Python expression string.""" if isinstance(node, ast.Constant): if node.value is None: return 'None' @@ -496,18 +378,7 @@ def _is_globals_update_call(node: ast.Call) -> bool: def parse_c_py_metadata(c_py_path: Path) -> dict[str, dict[str, str]]: - """Parse wrapper defaults from deep_gemm/_C.py. - - Walks the AST (does not import or execute _C.py). - - Defaults example — from: - def fp8_einsum(..., recipe=(1, 128, 128)): - produces: - defaults['fp8_einsum']['recipe'] = '(1, 128, 128)' - - Also copies defaults onto legacy alias names from globals().update(...) so - wrapper defaults apply when the public name differs from the TORCH op name. - """ + """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)) @@ -537,16 +408,7 @@ def fp8_einsum(..., recipe=(1, 128, 128)): def extract_m_def_statements(root_path) -> list[str]: - """ - Scan all C++ files under root_path and extract all m.def(...) statements. - - Returns a flat list of raw statement strings (one per registration found). - Supports multi-line m.def(...) calls. This is pipeline step 1. - - Example match in gemm.hpp: - m.def( - "fp8_fp4_gemm_nt(Tensor a, Tensor sfa, ...) -> ()"); - """ + """Scan C++ sources under root_path for m.def(...) registrations.""" statements = [] extensions = {'.hpp', '.cpp', '.h', '.cc'} @@ -569,14 +431,11 @@ def extract_m_def_statements(root_path) -> list[str]: while i < len(lines): line = lines[i] if 'm.def(' in line: - # Found a potential starting line - # 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 @@ -592,7 +451,6 @@ def extract_m_def_statements(root_path) -> list[str]: 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 @@ -609,20 +467,7 @@ def extract_m_def_statements(root_path) -> list[str]: def parse_m_def_statement(m_def_str): - """ - Parse a TORCH_LIBRARY m.def(...) statement (pipeline step 2). - - DeepGEMM registers ops via TORCH_LIBRARY_FRAGMENT; the first m.def argument - is always a schema string. Extra args like DEEP_GEMM_IMPL(...) are ignored. - - Example input: - m.def("bf16_gemm_nt(Tensor a, Tensor b, Tensor(d!) d, " - "Tensor? c=None, str compiled_dims='nk') -> ()", - DEEP_GEMM_IMPL(bf16_gemm_nt)); - - Delegates to parse_torch_schema() on the first string literal. - """ - # 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') @@ -644,14 +489,10 @@ 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] - - # Split arguments using BracketTracker args_list = split_top_level_commas(args_content) - if not args_list: raise ValueError(f'[{m_def_str}] m.def has no arguments') - # Extract operator schema from the first string literal first = args_list[0].strip() str_match = re.match(r'^"([^"\\]*(?:\\.[^"\\]*)*)"', first) if not str_match: @@ -661,16 +502,7 @@ def parse_m_def_statement(m_def_str): 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. - - Schema defaults come from the TORCH registration string; wrapper defaults reflect - what callers actually get from _C.py. - - Example for fp8_einsum: - schema default: recipe=None - wrapper default: recipe=(1, 128, 128) # from def fp8_einsum(..., recipe=(1, 128, 128)) - stub result: recipe: tuple[int, int, int] = (1, 128, 128) - """ + """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 @@ -685,17 +517,7 @@ def apply_wrapper_defaults(name: str, parameters: list[dict], wrapper_defaults: def generate_pyi_function(parsed, wrapper_defaults=None): - """Generate a typed .pyi stub for one registered op (pipeline steps 4-6). - - Example final output for fp8_fp4_gemm_nt: - def fp8_fp4_gemm_nt( - a: tuple[torch.Tensor, torch.Tensor], - b: tuple[torch.Tensor, torch.Tensor], - d: torch.Tensor, - c: Optional[torch.Tensor] = None, - ... - ) -> None: ... - """ + """Generate a typed .pyi stub for one registered op.""" py_name = parsed['python_function_name'] parameters = adjust_for_c_py_wrapper( py_name, @@ -725,12 +547,7 @@ def generate_pyi_file_content( module_name: str = 'my_module', wrapper_defaults=None, ): - """Assemble the full .pyi file from all parsed ops (pipeline step 7). - - parsed_ops: list of dicts returned by parse_m_def_statement / parse_torch_schema. - - - One stub per TORCH_LIBRARY op found in csrc/ - """ + """Assemble the full .pyi file from parsed TORCH ops.""" decls = [] for parsed in parsed_ops: @@ -747,6 +564,7 @@ def generate_pyi_file_content( 'from typing import Any, Optional', 'import torch', '', + '', ] for decl in decls: @@ -756,12 +574,10 @@ def generate_pyi_file_content( def generate_pyi_file(name, root, output_dir='.', c_py_path=None): - """Orchestrate the full pipeline and write stubs/.pyi.""" - # Step 1-2: scan csrc/ for m.def(...) and parse each TORCH schema. + """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] - # Step 3: read wrapper defaults from deep_gemm/_C.py. wrapper_defaults = {} if c_py_path is not None: c_py_path = Path(c_py_path) @@ -770,7 +586,6 @@ def generate_pyi_file(name, root, output_dir='.', c_py_path=None): else: print(f'Warning: wrapper file not found: {c_py_path}') - # Steps 4-7: adjust params, apply defaults, render stubs, write file. pyi_content = generate_pyi_file_content( parsed_ops, module_name=name, @@ -784,66 +599,3 @@ def generate_pyi_file(name, root, output_dir='.', c_py_path=None): f.write(pyi_content) print(f'.pyi file generated: {output_path}') - - -def main(argv=None) -> int: - import argparse - import sys - - parser = argparse.ArgumentParser( - description='Generate deep_gemm/_C.pyi stubs from TORCH_LIBRARY schemas.', - ) - parser.add_argument('--name', default='_C', help='Module name for the .pyi file (default: _C)') - parser.add_argument('--root', default='./csrc', help='Root to scan for m.def(...) (default: ./csrc)') - parser.add_argument('--output-dir', default='./stubs', help='Output directory (default: ./stubs)') - parser.add_argument( - '--c-py', - default='./deep_gemm/_C.py', - help='Python wrapper module to read public API defaults from (default: ./deep_gemm/_C.py)', - ) - parser.add_argument( - '--check', - action='store_true', - help='Verify the output has typed stubs (no generic *args, **kwargs)', - ) - args = parser.parse_args(argv) - - repo_root = Path(__file__).resolve().parent.parent - root = Path(args.root) - output_dir = Path(args.output_dir) - if not root.is_absolute(): - root = repo_root / root - if not output_dir.is_absolute(): - output_dir = repo_root / output_dir - - c_py_path = Path(args.c_py) - if not c_py_path.is_absolute(): - c_py_path = repo_root / c_py_path - - generate_pyi_file( - name=args.name, - root=str(root), - output_dir=str(output_dir), - c_py_path=str(c_py_path), - ) - - pyi_path = output_dir / f'{args.name}.pyi' - if args.check: - content = pyi_path.read_text(encoding='utf-8') - generic_count = content.count('*args, **kwargs') - stub_count = content.count('def ') - if stub_count == 0: - print(f'CHECK FAILED: no function stubs in {pyi_path}', file=sys.stderr) - return 1 - print( - f'CHECK PASSED: {stub_count} stubs in {pyi_path} ' - f'({stub_count - generic_count} typed, {generic_count} generic)', - ) - if generic_count > 0: - return 1 - return 0 - - -if __name__ == '__main__': - import sys - raise SystemExit(main()) From 4aff592809d924d382649034e3df10ffc85301d5 Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Fri, 10 Jul 2026 18:50:07 +0000 Subject: [PATCH 10/28] added c_py_path for generate_pyi_file in setup.py Signed-off-by: Chris Leonard --- setup.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/setup.py b/setup.py index 48482efc9c..67f762c694 100644 --- a/setup.py +++ b/setup.py @@ -128,7 +128,12 @@ def run(self): build_py.run(self) 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') From ee762439363c4e15d33699aaec79fd1f64061702 Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Fri, 10 Jul 2026 19:26:34 +0000 Subject: [PATCH 11/28] Refactor mega_moe to build symm buffer layout once per kernel launch Signed-off-by: Chris Leonard --- csrc/apis/mega.hpp | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/csrc/apis/mega.hpp b/csrc/apis/mega.hpp index adc98750c0..09dbad3366 100644 --- a/csrc/apis/mega.hpp +++ b/csrc/apis/mega.hpp @@ -260,6 +260,19 @@ static SymmBufferSlice slice_symm_buffer_for_mega_moe( hidden, intermediate_hidden, mma_type, activation, num_shared_experts)); } +static SymmBufferSlice slice_symm_buffer_for_mega_moe( + const torch::Tensor& buffer, + 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_ring_tokens) { + 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_ring_tokens); + return slice_symm_buffer_from_layout(buffer, layout_info); +} + static void fp8_fp4_mega_moe( const torch::Tensor& y, const std::tuple& l1_weights_tuple, From 75d4a62bff1a463c4ee57bc59ef954ada04bdbcd Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Fri, 17 Jul 2026 14:48:24 +0000 Subject: [PATCH 12/28] Match the FA/TORCH_LIBRARY packaging pattern: ops still register via TORCH_LIBRARY; deep_gemm/_C.py can keep using torch.ops.load_library. Signed-off-by: Chris Leonard --- csrc/python_api.cpp | 3 +++ csrc/utils/registration.h | 17 +++++++++++++++++ 2 files changed, 20 insertions(+) create mode 100644 csrc/utils/registration.h diff --git a/csrc/python_api.cpp b/csrc/python_api.cpp index efd5a622d5..f6ca80b1e8 100644 --- a/csrc/python_api.cpp +++ b/csrc/python_api.cpp @@ -1,4 +1,5 @@ #include "utils/torch_compat.hpp" +#include "utils/registration.h" #include "apis/attention.hpp" #include "apis/einsum.hpp" @@ -8,3 +9,5 @@ #include "apis/mega.hpp" #include "apis/sm90_mega.hpp" #include "apis/runtime.hpp" + +REGISTER_EXTENSION(_C_extension) 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); \ + } From 927c5a74790f6a28c5ec7dc4fb44636dfd555841 Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Tue, 21 Jul 2026 16:35:53 +0000 Subject: [PATCH 13/28] rebased onto nv_dev, removed duplicate slice_symm_buffer_for_mega_moe function that was leftover from migration, replaces torch_compat with torch/all.h, and updated the sm120 files to use torch/all.h instead of torch/python.h --- csrc/apis/einsum.hpp | 2 +- csrc/apis/mega.hpp | 15 +------------- csrc/jit/device_runtime.hpp | 2 +- csrc/jit_kernels/impls/runtime_utils.hpp | 2 +- csrc/jit_kernels/impls/sm100_bf16_gemm.hpp | 2 +- .../jit_kernels/impls/sm100_bf16_mega_moe.hpp | 2 +- csrc/jit_kernels/impls/sm100_bmk_bnk_mn.hpp | 2 +- .../impls/sm100_fp8_fp4_gemm_1d1d.hpp | 2 +- .../impls/sm100_fp8_fp4_mega_moe.hpp | 2 +- .../impls/sm100_tf32_hc_prenorm_gemm.hpp | 2 +- csrc/jit_kernels/impls/sm120_bf16_gemm.hpp | 2 +- csrc/jit_kernels/impls/sm120_bmk_bnk_mn.hpp | 2 +- .../impls/sm120_fp8_fp4_gemm_1d1d.hpp | 2 +- .../impls/sm120_tf32_hc_prenorm_gemm.hpp | 2 +- csrc/jit_kernels/impls/sm90_bf16_gemm.hpp | 2 +- csrc/jit_kernels/impls/sm90_bmk_bnk_mn.hpp | 2 +- csrc/jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp | 2 +- csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp | 2 +- .../impls/sm90_tf32_hc_prenorm_gemm.hpp | 2 +- csrc/jit_kernels/impls/smxx_layout.hpp | 2 +- csrc/python_api.cpp | 2 +- csrc/utils/layout.hpp | 2 +- csrc/utils/math.hpp | 2 +- csrc/utils/torch_compat.hpp | 20 ------------------- 24 files changed, 23 insertions(+), 56 deletions(-) delete mode 100644 csrc/utils/torch_compat.hpp diff --git a/csrc/apis/einsum.hpp b/csrc/apis/einsum.hpp index a2bd1f84cd..40e2a05079 100644 --- a/csrc/apis/einsum.hpp +++ b/csrc/apis/einsum.hpp @@ -1,6 +1,6 @@ #pragma once -#include "../utils/torch_compat.hpp" +#include #include "../utils/exception.hpp" #include "../utils/format.hpp" diff --git a/csrc/apis/mega.hpp b/csrc/apis/mega.hpp index 09dbad3366..cacf4e3372 100644 --- a/csrc/apis/mega.hpp +++ b/csrc/apis/mega.hpp @@ -4,7 +4,7 @@ #include #include -#include "../utils/torch_compat.hpp" +#include #include #include #include "../utils/math.hpp" @@ -260,19 +260,6 @@ static SymmBufferSlice slice_symm_buffer_for_mega_moe( hidden, intermediate_hidden, mma_type, activation, num_shared_experts)); } -static SymmBufferSlice slice_symm_buffer_for_mega_moe( - const torch::Tensor& buffer, - 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_ring_tokens) { - 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_ring_tokens); - return slice_symm_buffer_from_layout(buffer, layout_info); -} - static void fp8_fp4_mega_moe( const torch::Tensor& y, const std::tuple& l1_weights_tuple, diff --git a/csrc/jit/device_runtime.hpp b/csrc/jit/device_runtime.hpp index 14a3f79442..31b08b7a8f 100644 --- a/csrc/jit/device_runtime.hpp +++ b/csrc/jit/device_runtime.hpp @@ -4,7 +4,7 @@ #include #include -#include "../utils/torch_compat.hpp" +#include #include "../utils/exception.hpp" #include "../utils/lazy_init.hpp" diff --git a/csrc/jit_kernels/impls/runtime_utils.hpp b/csrc/jit_kernels/impls/runtime_utils.hpp index 739543d6f9..839d702710 100644 --- a/csrc/jit_kernels/impls/runtime_utils.hpp +++ b/csrc/jit_kernels/impls/runtime_utils.hpp @@ -1,7 +1,7 @@ #pragma once #include -#include "../../utils/torch_compat.hpp" +#include #include "../heuristics/sm90.hpp" #include "../../jit/handle.hpp" diff --git a/csrc/jit_kernels/impls/sm100_bf16_gemm.hpp b/csrc/jit_kernels/impls/sm100_bf16_gemm.hpp index 5c2d8b8295..9129b7bcc8 100644 --- a/csrc/jit_kernels/impls/sm100_bf16_gemm.hpp +++ b/csrc/jit_kernels/impls/sm100_bf16_gemm.hpp @@ -1,6 +1,6 @@ #pragma once -#include "../../utils/torch_compat.hpp" +#include #include "../../jit/compiler.hpp" #include "../../jit/device_runtime.hpp" diff --git a/csrc/jit_kernels/impls/sm100_bf16_mega_moe.hpp b/csrc/jit_kernels/impls/sm100_bf16_mega_moe.hpp index 3be0ef94f3..47ad292bc2 100644 --- a/csrc/jit_kernels/impls/sm100_bf16_mega_moe.hpp +++ b/csrc/jit_kernels/impls/sm100_bf16_mega_moe.hpp @@ -1,6 +1,6 @@ #pragma once -#include "../../utils/torch_compat.hpp" +#include #include "../../jit/compiler.hpp" #include "../../jit/kernel_runtime.hpp" diff --git a/csrc/jit_kernels/impls/sm100_bmk_bnk_mn.hpp b/csrc/jit_kernels/impls/sm100_bmk_bnk_mn.hpp index 2ec5d12a7f..292bd3903e 100644 --- a/csrc/jit_kernels/impls/sm100_bmk_bnk_mn.hpp +++ b/csrc/jit_kernels/impls/sm100_bmk_bnk_mn.hpp @@ -1,6 +1,6 @@ #pragma once -#include "../../utils/torch_compat.hpp" +#include #include "../../jit/compiler.hpp" #include "../../jit/device_runtime.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 d2fb50e951..95f855298c 100644 --- a/csrc/jit_kernels/impls/sm100_fp8_fp4_gemm_1d1d.hpp +++ b/csrc/jit_kernels/impls/sm100_fp8_fp4_gemm_1d1d.hpp @@ -1,6 +1,6 @@ #pragma once -#include "../../utils/torch_compat.hpp" +#include #include "../../jit/compiler.hpp" #include "../../jit/device_runtime.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 c11d2dd6c7..2746d41b47 100644 --- a/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp +++ b/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp @@ -1,6 +1,6 @@ #pragma once -#include "../../utils/torch_compat.hpp" +#include #include "../../jit/compiler.hpp" #include "../../jit/kernel_runtime.hpp" 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 e91a5d41fe..4ee309e67c 100644 --- a/csrc/jit_kernels/impls/sm100_tf32_hc_prenorm_gemm.hpp +++ b/csrc/jit_kernels/impls/sm100_tf32_hc_prenorm_gemm.hpp @@ -1,6 +1,6 @@ #pragma once -#include "../../utils/torch_compat.hpp" +#include #include "../../jit/compiler.hpp" #include "../../jit/device_runtime.hpp" diff --git a/csrc/jit_kernels/impls/sm120_bf16_gemm.hpp b/csrc/jit_kernels/impls/sm120_bf16_gemm.hpp index e272caf346..94be80b79b 100644 --- a/csrc/jit_kernels/impls/sm120_bf16_gemm.hpp +++ b/csrc/jit_kernels/impls/sm120_bf16_gemm.hpp @@ -1,6 +1,6 @@ #pragma once -#include +#include #include "../../jit/compiler.hpp" #include "../../jit/device_runtime.hpp" diff --git a/csrc/jit_kernels/impls/sm120_bmk_bnk_mn.hpp b/csrc/jit_kernels/impls/sm120_bmk_bnk_mn.hpp index 55365324c3..1472b26de9 100644 --- a/csrc/jit_kernels/impls/sm120_bmk_bnk_mn.hpp +++ b/csrc/jit_kernels/impls/sm120_bmk_bnk_mn.hpp @@ -1,6 +1,6 @@ #pragma once -#include +#include #include "../../jit/compiler.hpp" #include "../../jit/device_runtime.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 c8fab5839e..47176295e5 100644 --- a/csrc/jit_kernels/impls/sm120_fp8_fp4_gemm_1d1d.hpp +++ b/csrc/jit_kernels/impls/sm120_fp8_fp4_gemm_1d1d.hpp @@ -1,6 +1,6 @@ #pragma once -#include +#include #include "../../jit/compiler.hpp" #include "../../jit/device_runtime.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 3067baca5a..f375ca75d6 100644 --- a/csrc/jit_kernels/impls/sm120_tf32_hc_prenorm_gemm.hpp +++ b/csrc/jit_kernels/impls/sm120_tf32_hc_prenorm_gemm.hpp @@ -1,6 +1,6 @@ #pragma once -#include +#include #include "../../jit/compiler.hpp" #include "../../jit/device_runtime.hpp" diff --git a/csrc/jit_kernels/impls/sm90_bf16_gemm.hpp b/csrc/jit_kernels/impls/sm90_bf16_gemm.hpp index e901351b66..130dcc220b 100644 --- a/csrc/jit_kernels/impls/sm90_bf16_gemm.hpp +++ b/csrc/jit_kernels/impls/sm90_bf16_gemm.hpp @@ -1,6 +1,6 @@ #pragma once -#include "../../utils/torch_compat.hpp" +#include #include "../../jit/compiler.hpp" #include "../../jit/kernel_runtime.hpp" diff --git a/csrc/jit_kernels/impls/sm90_bmk_bnk_mn.hpp b/csrc/jit_kernels/impls/sm90_bmk_bnk_mn.hpp index 19a1556e6c..7aed97598b 100644 --- a/csrc/jit_kernels/impls/sm90_bmk_bnk_mn.hpp +++ b/csrc/jit_kernels/impls/sm90_bmk_bnk_mn.hpp @@ -1,6 +1,6 @@ #pragma once -#include "../../utils/torch_compat.hpp" +#include #include "../../jit/compiler.hpp" #include "../../jit/device_runtime.hpp" diff --git a/csrc/jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp b/csrc/jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp index 960578047a..1350f32f9a 100644 --- a/csrc/jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp +++ b/csrc/jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp @@ -1,6 +1,6 @@ #pragma once -#include "../../utils/torch_compat.hpp" +#include #include "../../jit/compiler.hpp" #include "../../jit/device_runtime.hpp" diff --git a/csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp b/csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp index c296524b60..daeece0326 100644 --- a/csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp +++ b/csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp @@ -1,6 +1,6 @@ #pragma once -#include "../../utils/torch_compat.hpp" +#include #include "../../jit/compiler.hpp" #include "../../jit/device_runtime.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 4a10d69775..7fbedf813f 100644 --- a/csrc/jit_kernels/impls/sm90_tf32_hc_prenorm_gemm.hpp +++ b/csrc/jit_kernels/impls/sm90_tf32_hc_prenorm_gemm.hpp @@ -1,6 +1,6 @@ #pragma once -#include "../../utils/torch_compat.hpp" +#include #include "../../jit/compiler.hpp" #include "../../jit/device_runtime.hpp" diff --git a/csrc/jit_kernels/impls/smxx_layout.hpp b/csrc/jit_kernels/impls/smxx_layout.hpp index ef5b6d4080..23dfb19c4c 100644 --- a/csrc/jit_kernels/impls/smxx_layout.hpp +++ b/csrc/jit_kernels/impls/smxx_layout.hpp @@ -1,6 +1,6 @@ #pragma once -#include "../../utils/torch_compat.hpp" +#include #include "../../jit/kernel_runtime.hpp" #include "../../jit/compiler.hpp" diff --git a/csrc/python_api.cpp b/csrc/python_api.cpp index f6ca80b1e8..6a95de9c24 100644 --- a/csrc/python_api.cpp +++ b/csrc/python_api.cpp @@ -1,4 +1,4 @@ -#include "utils/torch_compat.hpp" +#include #include "utils/registration.h" #include "apis/attention.hpp" diff --git a/csrc/utils/layout.hpp b/csrc/utils/layout.hpp index 09a9126d68..ae030571d3 100644 --- a/csrc/utils/layout.hpp +++ b/csrc/utils/layout.hpp @@ -1,7 +1,7 @@ #pragma once #include -#include "torch_compat.hpp" +#include #include "math.hpp" #include "exception.hpp" diff --git a/csrc/utils/math.hpp b/csrc/utils/math.hpp index f77049584e..81ecd427c5 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 "torch_compat.hpp" +#include #include "exception.hpp" diff --git a/csrc/utils/torch_compat.hpp b/csrc/utils/torch_compat.hpp deleted file mode 100644 index 9bc017ac14..0000000000 --- a/csrc/utils/torch_compat.hpp +++ /dev/null @@ -1,20 +0,0 @@ -#pragma once - -#include - -// torch/library.h declares `namespace torch` for op registration but does not -// re-export ATen types. DeepGEMM csrc uses torch::Tensor throughout; under -// Py_LIMITED_API we cannot include torch/python.h or torch/types.h (autograd -// pulls the full Python C-API). Re-export at:: into torch:: instead. -namespace torch { -using namespace at; - -constexpr auto kUInt8 = at::kByte; -constexpr auto kInt8 = at::kChar; -constexpr auto kInt16 = at::kShort; -constexpr auto kInt32 = at::kInt; -constexpr auto kInt64 = at::kLong; -constexpr auto kFloat16 = at::kHalf; -constexpr auto kFloat32 = at::kFloat; -constexpr auto kFloat64 = at::kDouble; -} // namespace torch From a71e0a8b5ba1bf40d7584bcf55dd516ee4488a2c Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Tue, 21 Jul 2026 19:08:57 +0000 Subject: [PATCH 14/28] Rename TORCH_LIBRARY schemas to match C++ parameter names; drop torch_library_macros - Drop torch_library_macros.hpp; register ops via TORCH_FN directly and rename torch_library_utils namespace to torch_utils. - Use at::ScalarType/torch.dtype for logits_dtype instead of int + _SCALAR_TYPE dict. - Align schema/wrapper param names with C++ impl (fused_kv_cache, activation_clamp_opt, *_tuple/*_tuple_opt), propagated to _C.py, mega/__init__.py, and tests. - Prefix slice_symm_buffer_for_mega_moe with _ to mark it private. Signed-off-by: Chris Leonard --- csrc/apis/attention.hpp | 23 ++++++------ csrc/apis/einsum.hpp | 5 ++- csrc/apis/gemm.hpp | 5 ++- csrc/apis/hyperconnection.hpp | 2 +- csrc/apis/layout.hpp | 13 ++++--- csrc/apis/mega.hpp | 67 ++++++++++++++++++----------------- csrc/apis/runtime.hpp | 20 +++++------ csrc/torch_library_macros.hpp | 17 --------- csrc/torch_library_utils.hpp | 4 +-- deep_gemm/_C.py | 53 +++++++++++---------------- deep_gemm/mega/__init__.py | 30 ++++++++-------- scripts/generate_pyi.py | 24 ++++++++++--- tests/test_mega_moe.py | 14 ++++---- 13 files changed, 140 insertions(+), 137 deletions(-) delete mode 100644 csrc/torch_library_macros.hpp diff --git a/csrc/apis/attention.hpp b/csrc/apis/attention.hpp index 001235fac0..799fe770cb 100644 --- a/csrc/apis/attention.hpp +++ b/csrc/apis/attention.hpp @@ -15,7 +15,8 @@ #endif #include "layout.hpp" -#include "../torch_library_macros.hpp" +#include +#include "../torch_library_utils.hpp" namespace deep_gemm::attention { @@ -468,6 +469,8 @@ static torch::Tensor fp8_paged_mqa_logits(const torch::Tensor& q, namespace deep_gemm::torch_registration { +using namespace deep_gemm::torch_utils; + #if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE static void fp8_gemm_nt_skip_head_mid( const torch::Tensor& a, const torch::Tensor& sfa, @@ -492,13 +495,13 @@ static torch::Tensor fp8_fp4_mqa_logits( const torch::Tensor& cu_seq_len_k_end, const bool& clean_logits, const int64_t& max_seqlen_k, - const int64_t& logits_dtype) { + at::ScalarType logits_dtype) { return attention::fp8_fp4_mqa_logits( std::make_tuple(q, q_sf), std::make_tuple(kv, kv_sf), weights, cu_seq_len_k_start, cu_seq_len_k_end, clean_logits, static_cast(max_seqlen_k), - static_cast(logits_dtype)); + logits_dtype); } static torch::Tensor get_paged_mqa_logits_metadata( @@ -511,20 +514,20 @@ static torch::Tensor get_paged_mqa_logits_metadata( 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& fused_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 int64_t& logits_dtype, + 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, + fused_kv_cache, weights, context_lens, block_table, schedule_meta, static_cast(max_context_len), clean_logits, - static_cast(logits_dtype), indices); + logits_dtype, indices); } static torch::Tensor fp8_mqa_logits( @@ -565,15 +568,15 @@ 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[] head_splits, int[]? 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, int logits_dtype=6) -> Tensor"); + "fp8_fp4_mqa_logits(Tensor q, Tensor? q_sf, Tensor kv, Tensor kv_sf, Tensor weights, Tensor cu_seq_len_k_start, Tensor cu_seq_len_k_end, bool clean_logits=True, int max_seqlen_k=0, ScalarType logits_dtype=float) -> Tensor"); m.def( "get_paged_mqa_logits_metadata(Tensor context_lens, int block_kv, int num_sms, Tensor? indices=None) -> Tensor"); m.def( - "fp8_fp4_paged_mqa_logits(Tensor q, Tensor? q_sf, Tensor kv_cache, Tensor weights, Tensor context_lens, Tensor block_table, Tensor schedule_meta, int max_context_len, bool clean_logits=False, int logits_dtype=6, Tensor? indices=None) -> Tensor"); + "fp8_fp4_paged_mqa_logits(Tensor q, Tensor? q_sf, Tensor fused_kv_cache, Tensor weights, Tensor context_lens, Tensor block_table, Tensor schedule_meta, int max_context_len, bool clean_logits=False, ScalarType logits_dtype=float, Tensor? indices=None) -> Tensor"); m.def( "fp8_mqa_logits(Tensor q, Tensor kv, Tensor kv_sf, Tensor weights, Tensor cu_seq_len_k_start, Tensor cu_seq_len_k_end, bool clean_logits=True, int max_seqlen_k=0) -> Tensor"); m.def( - "fp8_paged_mqa_logits(Tensor q, Tensor kv_cache, Tensor weights, Tensor context_lens, Tensor block_table, Tensor schedule_meta, int max_context_len, bool clean_logits=False, Tensor? indices=None) -> Tensor"); + "fp8_paged_mqa_logits(Tensor q, Tensor fused_kv_cache, Tensor weights, Tensor context_lens, Tensor block_table, Tensor schedule_meta, int max_context_len, bool clean_logits=False, Tensor? indices=None) -> Tensor"); #endif } diff --git a/csrc/apis/einsum.hpp b/csrc/apis/einsum.hpp index 40e2a05079..7cfec7a8f4 100644 --- a/csrc/apis/einsum.hpp +++ b/csrc/apis/einsum.hpp @@ -18,7 +18,8 @@ #include "../jit_kernels/impls/sm120_fp8_fp4_gemm_1d1d.hpp" #include "../jit_kernels/impls/smxx_cublaslt.hpp" #endif -#include "../torch_library_macros.hpp" +#include +#include "../torch_library_utils.hpp" namespace deep_gemm::einsum { @@ -272,6 +273,8 @@ static void fp8_einsum(const std::string& expr, namespace deep_gemm::torch_registration { +using namespace deep_gemm::torch_utils; + #if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE static void einsum(const std::string& expr, const torch::Tensor& a, const torch::Tensor& b, diff --git a/csrc/apis/gemm.hpp b/csrc/apis/gemm.hpp index da158a90d9..dd3be44a28 100644 --- a/csrc/apis/gemm.hpp +++ b/csrc/apis/gemm.hpp @@ -15,7 +15,8 @@ #include "../jit_kernels/impls/smxx_cublaslt.hpp" #include "layout.hpp" -#include "../torch_library_macros.hpp" +#include +#include "../torch_library_utils.hpp" namespace deep_gemm::gemm { @@ -787,6 +788,8 @@ static void cublaslt_gemm_tt(const torch::Tensor& a, const torch::Tensor& b, namespace deep_gemm::torch_registration { +using namespace deep_gemm::torch_utils; + #if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE static void fp8_fp4_gemm_nt( const torch::Tensor& a, const torch::Tensor& sfa, diff --git a/csrc/apis/hyperconnection.hpp b/csrc/apis/hyperconnection.hpp index ffe04c0672..8b37f65bb6 100644 --- a/csrc/apis/hyperconnection.hpp +++ b/csrc/apis/hyperconnection.hpp @@ -7,7 +7,7 @@ #include "../jit_kernels/impls/sm100_tf32_hc_prenorm_gemm.hpp" #include "../jit_kernels/impls/sm120_tf32_hc_prenorm_gemm.hpp" #endif -#include "../torch_library_macros.hpp" +#include namespace deep_gemm::hyperconnection { diff --git a/csrc/apis/layout.hpp b/csrc/apis/layout.hpp index 512e67876c..a5e6d42acb 100644 --- a/csrc/apis/layout.hpp +++ b/csrc/apis/layout.hpp @@ -7,7 +7,8 @@ #if DG_TENSORMAP_COMPATIBLE #include "../jit_kernels/impls/smxx_layout.hpp" #endif -#include "../torch_library_macros.hpp" +#include +#include "../torch_library_utils.hpp" namespace deep_gemm::layout { @@ -142,6 +143,8 @@ static torch::Tensor transform_k_grouped_sf_into_required_layout(const torch::Te namespace deep_gemm::torch_registration { +using namespace deep_gemm::torch_utils; + #if DG_TENSORMAP_COMPATIBLE static torch::Tensor transform_sf_into_required_layout( const torch::Tensor& sf, const int64_t& mn, const int64_t& k, @@ -207,7 +210,7 @@ TORCH_LIBRARY_FRAGMENT(deep_gemm, m) { #if DG_TENSORMAP_COMPATIBLE m.def( "transform_sf_into_required_layout(Tensor 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"); - m.def("get_tma_aligned_size(int x, int element_size) -> int", DEEP_GEMM_IMPL(get_tma_aligned_size)); + 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 sf) -> Tensor"); m.def( "get_mn_major_tma_aligned_packed_ue8m0_tensor(Tensor sf, Tensor? psum_layout=None) -> Tensor"); @@ -216,11 +219,11 @@ TORCH_LIBRARY_FRAGMENT(deep_gemm, m) { #endif m.def("set_mk_alignment_for_contiguous_layout(int new_value) -> ()", - DEEP_GEMM_IMPL(set_mk_alignment_for_contiguous_layout)); + TORCH_FN(deep_gemm::torch_registration::set_mk_alignment_for_contiguous_layout)); m.def("get_mk_alignment_for_contiguous_layout() -> int", - DEEP_GEMM_IMPL(get_mk_alignment_for_contiguous_layout)); + 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? num_groups=None) -> int", - DEEP_GEMM_IMPL(get_theoretical_mk_alignment_for_contiguous_layout)); + TORCH_FN(deep_gemm::torch_registration::get_theoretical_mk_alignment_for_contiguous_layout)); } TORCH_LIBRARY_IMPL(deep_gemm, CUDA, m) { diff --git a/csrc/apis/mega.hpp b/csrc/apis/mega.hpp index cacf4e3372..0899f7816a 100644 --- a/csrc/apis/mega.hpp +++ b/csrc/apis/mega.hpp @@ -16,7 +16,8 @@ #include "../jit_kernels/impls/sm100_bf16_mega_moe.hpp" #include "../jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp" #include "../jit_kernels/impls/sm100_fp4_fp4_mega_moe.hpp" -#include "../torch_library_macros.hpp" +#include +#include "../torch_library_utils.hpp" namespace deep_gemm::mega { @@ -711,6 +712,8 @@ static void bf16_mega_moe( namespace deep_gemm::torch_registration { +using namespace deep_gemm::torch_utils; + static int64_t get_token_alignment_for_mega_moe() { return static_cast(mega::get_token_alignment_for_mega_moe()); } @@ -738,7 +741,7 @@ static int64_t get_symm_buffer_size_for_mega_moe( mma_type, activation, static_cast(num_shared_experts)); } -static mega::SymmBufferSlice slice_symm_buffer_for_mega_moe( +static mega::SymmBufferSlice _slice_symm_buffer_for_mega_moe( const torch::Tensor& buffer, const int64_t& num_ranks, const int64_t& num_experts, const int64_t& num_max_tokens_per_rank, const int64_t& num_topk, @@ -755,12 +758,12 @@ static mega::SymmBufferSlice slice_symm_buffer_for_mega_moe( static void fp8_fp4_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 c10::optional& shared_l1_weights, - const c10::optional& shared_l1_weights_sf, - const c10::optional& shared_l2_weights, - const c10::optional& shared_l2_weights_sf, + 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 c10::optional& shared_l1_weights_tuple_opt, + const c10::optional& shared_l1_weights_tuple_opt_sf, + const c10::optional& shared_l2_weights_tuple_opt, + const c10::optional& shared_l2_weights_tuple_opt_sf, const c10::optional& cumulative_local_expert_recv_stats, const torch::Tensor& sym_buffer, const c10::List& sym_buffer_ptrs, @@ -769,22 +772,22 @@ static void fp8_fp4_mega_moe( const int64_t& num_experts, const int64_t& num_topk, const c10::List& recipe, const std::string& activation, - const c10::optional& activation_clamp, + const c10::optional& activation_clamp_opt, const bool& fast_math) { std::optional> shared_l1_opt = std::nullopt; std::optional> shared_l2_opt = std::nullopt; - if (shared_l1_weights.has_value()) { - DG_HOST_ASSERT(shared_l1_weights_sf.has_value() and shared_l2_weights.has_value() and shared_l2_weights_sf.has_value()); - shared_l1_opt = std::make_tuple(shared_l1_weights.value(), shared_l1_weights_sf.value()); - shared_l2_opt = std::make_tuple(shared_l2_weights.value(), shared_l2_weights_sf.value()); + if (shared_l1_weights_tuple_opt.has_value()) { + DG_HOST_ASSERT(shared_l1_weights_tuple_opt_sf.has_value() and shared_l2_weights_tuple_opt.has_value() and shared_l2_weights_tuple_opt_sf.has_value()); + shared_l1_opt = std::make_tuple(shared_l1_weights_tuple_opt.value(), shared_l1_weights_tuple_opt_sf.value()); + shared_l2_opt = std::make_tuple(shared_l2_weights_tuple_opt.value(), shared_l2_weights_tuple_opt_sf.value()); } else { - DG_HOST_ASSERT(not shared_l1_weights_sf.has_value() and not shared_l2_weights.has_value() and not shared_l2_weights_sf.has_value()); + DG_HOST_ASSERT(not shared_l1_weights_tuple_opt_sf.has_value() and not shared_l2_weights_tuple_opt.has_value() and not shared_l2_weights_tuple_opt_sf.has_value()); } mega::fp8_fp4_mega_moe( y, - std::make_tuple(l1_weights, l1_weights_sf), - std::make_tuple(l2_weights, l2_weights_sf), + std::make_tuple(l1_weights_tuple, l1_weights_tuple_sf), + std::make_tuple(l2_weights_tuple, l2_weights_tuple_sf), shared_l1_opt, shared_l2_opt, cumulative_local_expert_recv_stats, @@ -795,8 +798,8 @@ static void fp8_fp4_mega_moe( static_cast(num_experts), static_cast(num_topk), list_to_tuple3(recipe), activation, - activation_clamp.has_value() - ? std::make_optional(static_cast(activation_clamp.value())) + activation_clamp_opt.has_value() + ? std::make_optional(static_cast(activation_clamp_opt.value())) : std::nullopt, fast_math); } @@ -805,8 +808,8 @@ static void bf16_mega_moe( const torch::Tensor& y, const torch::Tensor& l1_weights, const torch::Tensor& l2_weights, - const c10::optional& shared_l1_weights, - const c10::optional& shared_l2_weights, + const c10::optional& shared_l1_weights_opt, + const c10::optional& shared_l2_weights_opt, const c10::optional& cumulative_local_expert_recv_stats, const torch::Tensor& sym_buffer, const c10::List& sym_buffer_ptrs, @@ -814,12 +817,12 @@ static void bf16_mega_moe( const int64_t& num_max_tokens_per_rank, const int64_t& num_experts, const int64_t& num_topk, const std::string& activation, - const c10::optional& activation_clamp, + const c10::optional& activation_clamp_opt, const bool& fast_math) { mega::bf16_mega_moe( y, l1_weights, l2_weights, - shared_l1_weights, - shared_l2_weights, + shared_l1_weights_opt, + shared_l2_weights_opt, cumulative_local_expert_recv_stats, sym_buffer, std::vector(sym_buffer_ptrs.begin(), sym_buffer_ptrs.end()), @@ -827,8 +830,8 @@ static void bf16_mega_moe( static_cast(num_max_tokens_per_rank), static_cast(num_experts), static_cast(num_topk), activation, - activation_clamp.has_value() - ? std::make_optional(static_cast(activation_clamp.value())) + activation_clamp_opt.has_value() + ? std::make_optional(static_cast(activation_clamp_opt.value())) : std::nullopt, fast_math); } @@ -839,19 +842,19 @@ TORCH_LIBRARY_FRAGMENT(deep_gemm, m) { #if DG_TENSORMAP_COMPATIBLE m.def( "get_token_alignment_for_mega_moe() -> int", - DEEP_GEMM_IMPL(get_token_alignment_for_mega_moe)); + TORCH_FN(deep_gemm::torch_registration::get_token_alignment_for_mega_moe)); 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", - DEEP_GEMM_IMPL(get_block_m_for_mega_moe)); + TORCH_FN(deep_gemm::torch_registration::get_block_m_for_mega_moe)); 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=0) -> int", - DEEP_GEMM_IMPL(get_symm_buffer_size_for_mega_moe)); + TORCH_FN(deep_gemm::torch_registration::get_symm_buffer_size_for_mega_moe)); m.def( - "slice_symm_buffer_for_mega_moe(Tensor buffer, 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=0) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor)"); + "_slice_symm_buffer_for_mega_moe(Tensor buffer, 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=0) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor)"); m.def( - "fp8_fp4_mega_moe(Tensor(y!) y, Tensor l1_weights, Tensor l1_weights_sf, Tensor l2_weights, Tensor l2_weights_sf, Tensor? shared_l1_weights, Tensor? shared_l1_weights_sf, Tensor? shared_l2_weights, Tensor? shared_l2_weights_sf, Tensor? cumulative_local_expert_recv_stats, Tensor(sym_buffer!) sym_buffer, int[] sym_buffer_ptrs, int rank_idx, int num_max_tokens_per_rank, int num_experts, int num_topk, int[] recipe, str activation, float? activation_clamp, bool fast_math) -> ()"); + "fp8_fp4_mega_moe(Tensor(y!) 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? cumulative_local_expert_recv_stats, Tensor(sym_buffer!) sym_buffer, int[] sym_buffer_ptrs, int rank_idx, int num_max_tokens_per_rank, int num_experts, int num_topk, int[] recipe, str activation, float? activation_clamp_opt, bool fast_math) -> ()"); m.def( - "bf16_mega_moe(Tensor(y!) y, Tensor l1_weights, Tensor l2_weights, Tensor? shared_l1_weights, Tensor? shared_l2_weights, Tensor? cumulative_local_expert_recv_stats, Tensor(sym_buffer!) 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, bool fast_math) -> ()"); + "bf16_mega_moe(Tensor(y!) y, Tensor l1_weights, Tensor l2_weights, Tensor? shared_l1_weights_opt, Tensor? shared_l2_weights_opt, Tensor? cumulative_local_expert_recv_stats, Tensor(sym_buffer!) 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) -> ()"); #endif } @@ -859,7 +862,7 @@ TORCH_LIBRARY_IMPL(deep_gemm, CUDA, m) { using namespace deep_gemm::torch_registration; #if DG_TENSORMAP_COMPATIBLE - m.impl("slice_symm_buffer_for_mega_moe", TORCH_FN(slice_symm_buffer_for_mega_moe)); + m.impl("_slice_symm_buffer_for_mega_moe", TORCH_FN(_slice_symm_buffer_for_mega_moe)); m.impl("fp8_fp4_mega_moe", TORCH_FN(fp8_fp4_mega_moe)); m.impl("bf16_mega_moe", TORCH_FN(bf16_mega_moe)); #endif diff --git a/csrc/apis/runtime.hpp b/csrc/apis/runtime.hpp index c0173c5b6d..9d79d89d61 100644 --- a/csrc/apis/runtime.hpp +++ b/csrc/apis/runtime.hpp @@ -7,7 +7,7 @@ #include "../jit/device_runtime.hpp" #include "../jit_kernels/heuristics/runtime.hpp" -#include "../torch_library_macros.hpp" +#include namespace deep_gemm::torch_registration { @@ -62,13 +62,13 @@ static void init(const std::string& library_root_path, } // namespace deep_gemm::torch_registration TORCH_LIBRARY_FRAGMENT(deep_gemm, m) { - m.def("set_num_sms(int new_num_sms) -> ()", DEEP_GEMM_IMPL(set_num_sms)); - m.def("get_num_sms() -> int", DEEP_GEMM_IMPL(get_num_sms)); - m.def("set_tc_util(int new_tc_util) -> ()", DEEP_GEMM_IMPL(set_tc_util)); - m.def("get_tc_util() -> int", DEEP_GEMM_IMPL(get_tc_util)); - m.def("set_pdl(bool new_enable_pdl) -> ()", DEEP_GEMM_IMPL(set_pdl)); - m.def("get_pdl() -> bool", DEEP_GEMM_IMPL(get_pdl)); - m.def("set_ignore_compile_dims(bool new_value) -> ()", DEEP_GEMM_IMPL(set_ignore_compile_dims)); - m.def("set_block_size_multiple_of(int[] value) -> ()", DEEP_GEMM_IMPL(set_block_size_multiple_of)); - m.def("init(str library_root_path, str cuda_home_path_by_python) -> ()", DEEP_GEMM_IMPL(init)); + m.def("set_num_sms(int new_num_sms) -> ()", TORCH_FN(deep_gemm::torch_registration::set_num_sms)); + m.def("get_num_sms() -> int", TORCH_FN(deep_gemm::torch_registration::get_num_sms)); + m.def("set_tc_util(int new_tc_util) -> ()", TORCH_FN(deep_gemm::torch_registration::set_tc_util)); + m.def("get_tc_util() -> int", TORCH_FN(deep_gemm::torch_registration::get_tc_util)); + m.def("set_pdl(bool new_enable_pdl) -> ()", TORCH_FN(deep_gemm::torch_registration::set_pdl)); + m.def("get_pdl() -> bool", TORCH_FN(deep_gemm::torch_registration::get_pdl)); + m.def("set_ignore_compile_dims(bool new_value) -> ()", TORCH_FN(deep_gemm::torch_registration::set_ignore_compile_dims)); + m.def("set_block_size_multiple_of(int[] value) -> ()", TORCH_FN(deep_gemm::torch_registration::set_block_size_multiple_of)); + m.def("init(str library_root_path, str cuda_home_path_by_python) -> ()", TORCH_FN(deep_gemm::torch_registration::init)); } diff --git a/csrc/torch_library_macros.hpp b/csrc/torch_library_macros.hpp deleted file mode 100644 index 5ab5688185..0000000000 --- a/csrc/torch_library_macros.hpp +++ /dev/null @@ -1,17 +0,0 @@ -#pragma once - -#include - -#include "torch_library_utils.hpp" - -namespace deep_gemm::torch_registration { - -using deep_gemm::torch_library_utils::list_to_optional_vector_int; -using deep_gemm::torch_library_utils::list_to_recipe2; -using deep_gemm::torch_library_utils::list_to_recipe3; -using deep_gemm::torch_library_utils::list_to_recipe_variant; -using deep_gemm::torch_library_utils::list_to_tuple3; - -} // namespace deep_gemm::torch_registration - -#define DEEP_GEMM_IMPL(fn) TORCH_FN(deep_gemm::torch_registration::fn) diff --git a/csrc/torch_library_utils.hpp b/csrc/torch_library_utils.hpp index 77cfa0abcc..e84bc143ea 100644 --- a/csrc/torch_library_utils.hpp +++ b/csrc/torch_library_utils.hpp @@ -8,7 +8,7 @@ #include "utils/exception.hpp" -namespace deep_gemm::torch_library_utils { +namespace deep_gemm::torch_utils { inline std::optional> list_to_recipe3( const c10::optional>& recipe) { @@ -61,4 +61,4 @@ inline std::optional> list_to_optional_vector_int( return out; } -} // namespace deep_gemm::torch_library_utils +} // namespace deep_gemm::torch_utils diff --git a/deep_gemm/_C.py b/deep_gemm/_C.py index 5fbdbde1b4..b86eaffea1 100644 --- a/deep_gemm/_C.py +++ b/deep_gemm/_C.py @@ -1,17 +1,6 @@ import torch from pathlib import Path -_SCALAR_TYPE = { - torch.float32: 6, - torch.bfloat16: 15, -} - - -def _as_scalar_type(dtype): - if isinstance(dtype, int): - return dtype - return _SCALAR_TYPE.get(dtype, 6) - def _load_extension(): so_files = list(Path(__file__).parent.glob('_C_extension*.so')) @@ -171,25 +160,25 @@ def fp8_fp4_mqa_logits(q, kv, weights, cu_seq_len_k_start, cu_seq_len_k_end, cle 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, _as_scalar_type(logits_dtype), + clean_logits, max_seqlen_k, logits_dtype, ) - def fp8_fp4_paged_mqa_logits(q, kv_cache, weights, context_lens, block_table, schedule_meta, max_context_len, + def fp8_fp4_paged_mqa_logits(q, fused_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 = _unpack_q(q) 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, _as_scalar_type(logits_dtype), indices, + q_fp, q_sf, fused_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, + def fp8_paged_mqa_logits(q, fused_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, + q, fused_kv_cache, weights, context_lens, block_table, schedule_meta, max_context_len, clean_logits, indices, ) globals().update({ @@ -253,8 +242,8 @@ def get_symm_buffer_size_for_mega_moe(*args, **kwargs): return _torch_ops.get_symm_buffer_size_for_mega_moe(*args, **kwargs) -def slice_symm_buffer_for_mega_moe(buffer, *args, **kwargs): - return _torch_ops.slice_symm_buffer_for_mega_moe(buffer, *args, **kwargs) +def _slice_symm_buffer_for_mega_moe(buffer, *args, **kwargs): + return _torch_ops._slice_symm_buffer_for_mega_moe(buffer, *args, **kwargs) # DG_TENSORMAP_COMPATIBLE — mega.hpp (C++ impl conditional; matches legacy pybind export guard) @@ -264,32 +253,32 @@ def slice_symm_buffer_for_mega_moe(buffer, *args, **kwargs): ) -def fp8_fp4_mega_moe(y, l1_weights, l2_weights, shared_l1_weights, shared_l2_weights, - cumulative_local_expert_recv_stats, sym_buffer, +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, fast_math): + activation, activation_clamp_opt, fast_math): shared_l1_w = shared_l1_sf = shared_l2_w = shared_l2_sf = None - if shared_l1_weights is not None: - shared_l1_w, shared_l1_sf = shared_l1_weights - shared_l2_w, shared_l2_sf = shared_l2_weights + if shared_l1_weights_tuple_opt is not None: + shared_l1_w, shared_l1_sf = shared_l1_weights_tuple_opt + shared_l2_w, shared_l2_sf = shared_l2_weights_tuple_opt return _torch_ops.fp8_fp4_mega_moe( - y, l1_weights[0], l1_weights[1], l2_weights[0], l2_weights[1], + y, l1_weights_tuple[0], l1_weights_tuple[1], l2_weights_tuple[0], l2_weights_tuple[1], shared_l1_w, shared_l1_sf, shared_l2_w, shared_l2_sf, cumulative_local_expert_recv_stats, sym_buffer, list(sym_buffer_ptrs), rank_idx, num_max_tokens_per_rank, num_experts, num_topk, list(recipe), activation, - activation_clamp, fast_math, + activation_clamp_opt, fast_math, ) -def bf16_mega_moe(y, l1_weights, l2_weights, shared_l1_weights, shared_l2_weights, +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, fast_math): + activation, activation_clamp_opt, fast_math): return _torch_ops.bf16_mega_moe( - y, l1_weights, l2_weights, shared_l1_weights, shared_l2_weights, + y, l1_weights, l2_weights, shared_l1_weights_opt, shared_l2_weights_opt, cumulative_local_expert_recv_stats, sym_buffer, list(sym_buffer_ptrs), rank_idx, num_max_tokens_per_rank, num_experts, num_topk, - activation, activation_clamp, fast_math, + activation, activation_clamp_opt, fast_math, ) @@ -309,7 +298,7 @@ def bf16_mega_moe(y, l1_weights, l2_weights, shared_l1_weights, shared_l2_weight 'cublaslt_gemm_tn', 'cublaslt_gemm_tt', # Mega MoE (imported via deep_gemm.mega; always defined, fails at call if unregistered) 'get_symm_buffer_size_for_mega_moe', - 'slice_symm_buffer_for_mega_moe', + '_slice_symm_buffer_for_mega_moe', 'fp8_fp4_mega_moe', 'bf16_mega_moe', ) diff --git a/deep_gemm/mega/__init__.py b/deep_gemm/mega/__init__.py index e43a5ab409..cd93042408 100644 --- a/deep_gemm/mega/__init__.py +++ b/deep_gemm/mega/__init__.py @@ -43,7 +43,7 @@ def __init__(self, group: dist.ProcessGroup, mma_type, activation, num_shared_experts, ) - slice_input_buffers = lambda buffer: _C.slice_symm_buffer_for_mega_moe( + slice_input_buffers = lambda buffer: _C._slice_symm_buffer_for_mega_moe( buffer, group.size(), num_experts, num_max_tokens_per_rank, num_topk, @@ -282,29 +282,29 @@ def validate_pair(name, pair): def fp8_fp4_mega_moe(y: torch.Tensor, - l1_weights: Tuple[torch.Tensor, torch.Tensor], - l2_weights: Tuple[torch.Tensor, torch.Tensor], + l1_weights_tuple: Tuple[torch.Tensor, torch.Tensor], + l2_weights_tuple: Tuple[torch.Tensor, torch.Tensor], sym_buffer: SymmBuffer, - shared_l1_weights: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, - shared_l2_weights: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + shared_l1_weights_tuple_opt: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + shared_l2_weights_tuple_opt: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, cumulative_local_expert_recv_stats: Optional[torch.Tensor] = None, recipe: Tuple[int, int, int] = (1, 1, 32), activation: str = 'swiglu', - activation_clamp: Optional[float] = None, + activation_clamp_opt: Optional[float] = None, fast_math: bool = True, situ_beta: Optional[float] = None, situ_linear_beta: Optional[float] = None): _C.fp8_fp4_mega_moe( y, - l1_weights, l2_weights, - shared_l1_weights, shared_l2_weights, + l1_weights_tuple, l2_weights_tuple, + shared_l1_weights_tuple_opt, shared_l2_weights_tuple_opt, cumulative_local_expert_recv_stats, sym_buffer.buffer, sym_buffer.handle.buffer_ptrs, sym_buffer.group.rank(), sym_buffer.num_max_tokens_per_rank, sym_buffer.num_experts, sym_buffer.num_topk, recipe, - activation, activation_clamp, + activation, activation_clamp_opt, fast_math, situ_beta, situ_linear_beta ) @@ -370,18 +370,18 @@ def bf16_mega_moe(y: torch.Tensor, l1_weights: torch.Tensor, l2_weights: torch.Tensor, sym_buffer: SymmBuffer, - shared_l1_weights: Optional[torch.Tensor] = None, - shared_l2_weights: Optional[torch.Tensor] = None, + shared_l1_weights_opt: Optional[torch.Tensor] = None, + shared_l2_weights_opt: Optional[torch.Tensor] = None, cumulative_local_expert_recv_stats: Optional[torch.Tensor] = None, activation: str = 'swiglu', - activation_clamp: Optional[float] = None, + activation_clamp_opt: Optional[float] = None, fast_math: bool = True): _C.bf16_mega_moe( y, l1_weights, l2_weights, - shared_l1_weights, - shared_l2_weights, + shared_l1_weights_opt, + shared_l2_weights_opt, cumulative_local_expert_recv_stats, sym_buffer.buffer, sym_buffer.handle.buffer_ptrs, @@ -389,7 +389,7 @@ def bf16_mega_moe(y: torch.Tensor, sym_buffer.num_max_tokens_per_rank, sym_buffer.num_experts, sym_buffer.num_topk, - activation, activation_clamp, + activation, activation_clamp_opt, fast_math ) diff --git a/scripts/generate_pyi.py b/scripts/generate_pyi.py index 2ba6ef7e65..0704d797ac 100644 --- a/scripts/generate_pyi.py +++ b/scripts/generate_pyi.py @@ -91,6 +91,8 @@ def schema_type_to_python(type_str: str) -> str: py_type = 'str' elif type_str == 'int[]': py_type = 'list[int]' + elif type_str == 'ScalarType': + py_type = 'torch.dtype' else: print(f'Warning: unrecognized schema type {type_str!r}, using Any') py_type = 'Any' @@ -124,6 +126,22 @@ def schema_return_to_python(return_str: str) -> str: return 'Any' +_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() @@ -132,6 +150,8 @@ def schema_default_to_python(default_str: str) -> 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): @@ -309,10 +329,6 @@ def adjust_for_c_py_wrapper( if pairs: parameters = _merge_named_pairs(parameters, tuple(pairs)) - for param in parameters: - if param['name'] == 'logits_dtype': - param['py_type'] = 'torch.dtype' - _maybe_widen_int_list_value_param(parameters) _promote_int_list_tuple_types(name, parameters) diff --git a/tests/test_mega_moe.py b/tests/test_mega_moe.py index 9987e024a0..2958162f2a 100644 --- a/tests/test_mega_moe.py +++ b/tests/test_mega_moe.py @@ -183,17 +183,17 @@ def run_fused(): copy_inputs_to_buffer() y = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device='cuda') + weights_suffix = '' if is_bf16xbf16 else '_tuple' kernel_kwargs = dict( - y=y, l1_weights=transformed_l1_weights, l2_weights=transformed_l2_weights, - sym_buffer=buffer, + y=y, sym_buffer=buffer, cumulative_local_expert_recv_stats=cumulative_local_expert_recv_stats_fused, - activation_clamp=args.activation_clamp, + activation_clamp_opt=args.activation_clamp, fast_math=bool(args.fast_math)) + kernel_kwargs[f'l1_weights{weights_suffix}'] = transformed_l1_weights + kernel_kwargs[f'l2_weights{weights_suffix}'] = transformed_l2_weights if num_shared_experts > 0: - kernel_kwargs.update( - shared_l1_weights=transformed_shared_l1_weights, - shared_l2_weights=transformed_shared_l2_weights - ) + kernel_kwargs[f'shared_l1_weights{weights_suffix}_opt'] = transformed_shared_l1_weights + kernel_kwargs[f'shared_l2_weights{weights_suffix}_opt'] = transformed_shared_l2_weights (deep_gemm.bf16_mega_moe if is_bf16xbf16 else deep_gemm.fp8_fp4_mega_moe)(**kernel_kwargs) return y, cumulative_local_expert_recv_stats_fused From 0022d4543a698bb1caf8014711f7fda856443c5e Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Wed, 22 Jul 2026 17:19:00 +0000 Subject: [PATCH 15/28] Fix test_paged_mqa_logits kwarg to match fused_kv_cache rename Signed-off-by: Chris Leonard --- tests/test_attention.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_attention.py b/tests/test_attention.py index 7a47d478ee..41ee0150ee 100644 --- a/tests/test_attention.py +++ b/tests/test_attention.py @@ -523,7 +523,7 @@ def enumerate_paged_mqa_logits(): num_kv_multicast = 2 if get_arch_major() == 9 and next_n == 4 else 1 num_clusters = deep_gemm.get_num_sms() // num_kv_multicast kernel_kwargs = dict( - q=q_in, kv_cache=kv_in, weights=kernel_weights, + q=q_in, fused_kv_cache=kv_in, weights=kernel_weights, context_lens=context_lens_nextn, block_table=block_table, schedule_meta=deep_gemm.get_paged_mqa_logits_metadata(context_lens_nextn, block_kv, num_clusters, indices=indices), max_context_len=max_model_len, clean_logits=clean_logits, logits_dtype=logits_dtype, From 94d3bfbc79015ea0a8ae77eda2cb5950299438cb Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Wed, 22 Jul 2026 18:10:37 +0000 Subject: [PATCH 16/28] reverting pytorch apis back to the original names instead of trying to align it with the C++ names Signed-off-by: Chris Leonard --- csrc/apis/attention.hpp | 12 ++++----- csrc/apis/mega.hpp | 50 +++++++++++++++++++------------------- deep_gemm/_C.py | 32 ++++++++++++------------ deep_gemm/mega/__init__.py | 28 ++++++++++----------- tests/test_attention.py | 2 +- tests/test_mega_moe.py | 14 +++++------ 6 files changed, 69 insertions(+), 69 deletions(-) diff --git a/csrc/apis/attention.hpp b/csrc/apis/attention.hpp index 799fe770cb..187c16fea9 100644 --- a/csrc/apis/attention.hpp +++ b/csrc/apis/attention.hpp @@ -514,7 +514,7 @@ static torch::Tensor get_paged_mqa_logits_metadata( static torch::Tensor fp8_fp4_paged_mqa_logits( const torch::Tensor& q, const c10::optional& q_sf, - const torch::Tensor& fused_kv_cache, + const torch::Tensor& kv_cache, const torch::Tensor& weights, const torch::Tensor& context_lens, const torch::Tensor& block_table, @@ -525,7 +525,7 @@ static torch::Tensor fp8_fp4_paged_mqa_logits( const c10::optional& indices) { return attention::fp8_fp4_paged_mqa_logits( std::make_tuple(q, q_sf), - fused_kv_cache, weights, context_lens, block_table, schedule_meta, + kv_cache, weights, context_lens, block_table, schedule_meta, static_cast(max_context_len), clean_logits, logits_dtype, indices); } @@ -546,7 +546,7 @@ static torch::Tensor fp8_mqa_logits( static torch::Tensor fp8_paged_mqa_logits( const torch::Tensor& q, - const torch::Tensor& fused_kv_cache, + const torch::Tensor& kv_cache, const torch::Tensor& weights, const torch::Tensor& context_lens, const torch::Tensor& block_table, @@ -555,7 +555,7 @@ static torch::Tensor fp8_paged_mqa_logits( const bool& clean_logits, const c10::optional& indices) { return attention::fp8_paged_mqa_logits( - q, fused_kv_cache, weights, + q, kv_cache, weights, context_lens, block_table, schedule_meta, static_cast(max_context_len), clean_logits, indices); } @@ -572,11 +572,11 @@ TORCH_LIBRARY_FRAGMENT(deep_gemm, m) { 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 fused_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"); + "fp8_fp4_paged_mqa_logits(Tensor q, Tensor? q_sf, Tensor kv_cache, Tensor weights, Tensor context_lens, Tensor block_table, Tensor schedule_meta, int max_context_len, bool clean_logits=False, ScalarType logits_dtype=float, Tensor? indices=None) -> Tensor"); m.def( "fp8_mqa_logits(Tensor q, Tensor kv, Tensor kv_sf, Tensor weights, Tensor cu_seq_len_k_start, Tensor cu_seq_len_k_end, bool clean_logits=True, int max_seqlen_k=0) -> Tensor"); m.def( - "fp8_paged_mqa_logits(Tensor q, Tensor fused_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"); + "fp8_paged_mqa_logits(Tensor q, Tensor kv_cache, Tensor weights, Tensor context_lens, Tensor block_table, Tensor schedule_meta, int max_context_len, bool clean_logits=False, Tensor? indices=None) -> Tensor"); #endif } diff --git a/csrc/apis/mega.hpp b/csrc/apis/mega.hpp index 0899f7816a..0e0a34f6d0 100644 --- a/csrc/apis/mega.hpp +++ b/csrc/apis/mega.hpp @@ -758,12 +758,12 @@ static mega::SymmBufferSlice _slice_symm_buffer_for_mega_moe( 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 c10::optional& shared_l1_weights_tuple_opt, - const c10::optional& shared_l1_weights_tuple_opt_sf, - const c10::optional& shared_l2_weights_tuple_opt, - const c10::optional& shared_l2_weights_tuple_opt_sf, + const torch::Tensor& l1_weights, const torch::Tensor& l1_weights_sf, + const torch::Tensor& l2_weights, const torch::Tensor& l2_weights_sf, + const c10::optional& shared_l1_weights, + const c10::optional& shared_l1_weights_sf, + const c10::optional& shared_l2_weights, + const c10::optional& shared_l2_weights_sf, const c10::optional& cumulative_local_expert_recv_stats, const torch::Tensor& sym_buffer, const c10::List& sym_buffer_ptrs, @@ -772,22 +772,22 @@ static void fp8_fp4_mega_moe( const int64_t& num_experts, const int64_t& num_topk, const c10::List& recipe, const std::string& activation, - const c10::optional& activation_clamp_opt, + const c10::optional& activation_clamp, const bool& fast_math) { std::optional> shared_l1_opt = std::nullopt; std::optional> shared_l2_opt = std::nullopt; - if (shared_l1_weights_tuple_opt.has_value()) { - DG_HOST_ASSERT(shared_l1_weights_tuple_opt_sf.has_value() and shared_l2_weights_tuple_opt.has_value() and shared_l2_weights_tuple_opt_sf.has_value()); - shared_l1_opt = std::make_tuple(shared_l1_weights_tuple_opt.value(), shared_l1_weights_tuple_opt_sf.value()); - shared_l2_opt = std::make_tuple(shared_l2_weights_tuple_opt.value(), shared_l2_weights_tuple_opt_sf.value()); + if (shared_l1_weights.has_value()) { + DG_HOST_ASSERT(shared_l1_weights_sf.has_value() and shared_l2_weights.has_value() and shared_l2_weights_sf.has_value()); + shared_l1_opt = std::make_tuple(shared_l1_weights.value(), shared_l1_weights_sf.value()); + shared_l2_opt = std::make_tuple(shared_l2_weights.value(), shared_l2_weights_sf.value()); } else { - DG_HOST_ASSERT(not shared_l1_weights_tuple_opt_sf.has_value() and not shared_l2_weights_tuple_opt.has_value() and not shared_l2_weights_tuple_opt_sf.has_value()); + DG_HOST_ASSERT(not shared_l1_weights_sf.has_value() and not shared_l2_weights.has_value() and not shared_l2_weights_sf.has_value()); } mega::fp8_fp4_mega_moe( y, - std::make_tuple(l1_weights_tuple, l1_weights_tuple_sf), - std::make_tuple(l2_weights_tuple, l2_weights_tuple_sf), + std::make_tuple(l1_weights, l1_weights_sf), + std::make_tuple(l2_weights, l2_weights_sf), shared_l1_opt, shared_l2_opt, cumulative_local_expert_recv_stats, @@ -798,8 +798,8 @@ static void fp8_fp4_mega_moe( static_cast(num_experts), static_cast(num_topk), list_to_tuple3(recipe), activation, - activation_clamp_opt.has_value() - ? std::make_optional(static_cast(activation_clamp_opt.value())) + activation_clamp.has_value() + ? std::make_optional(static_cast(activation_clamp.value())) : std::nullopt, fast_math); } @@ -808,8 +808,8 @@ static void bf16_mega_moe( const torch::Tensor& y, const torch::Tensor& l1_weights, const torch::Tensor& l2_weights, - const c10::optional& shared_l1_weights_opt, - const c10::optional& shared_l2_weights_opt, + const c10::optional& shared_l1_weights, + const c10::optional& shared_l2_weights, const c10::optional& cumulative_local_expert_recv_stats, const torch::Tensor& sym_buffer, const c10::List& sym_buffer_ptrs, @@ -817,12 +817,12 @@ static void bf16_mega_moe( const int64_t& num_max_tokens_per_rank, const int64_t& num_experts, const int64_t& num_topk, const std::string& activation, - const c10::optional& activation_clamp_opt, + const c10::optional& activation_clamp, const bool& fast_math) { mega::bf16_mega_moe( y, l1_weights, l2_weights, - shared_l1_weights_opt, - shared_l2_weights_opt, + shared_l1_weights, + shared_l2_weights, cumulative_local_expert_recv_stats, sym_buffer, std::vector(sym_buffer_ptrs.begin(), sym_buffer_ptrs.end()), @@ -830,8 +830,8 @@ static void bf16_mega_moe( static_cast(num_max_tokens_per_rank), static_cast(num_experts), static_cast(num_topk), activation, - activation_clamp_opt.has_value() - ? std::make_optional(static_cast(activation_clamp_opt.value())) + activation_clamp.has_value() + ? std::make_optional(static_cast(activation_clamp.value())) : std::nullopt, fast_math); } @@ -852,9 +852,9 @@ TORCH_LIBRARY_FRAGMENT(deep_gemm, m) { m.def( "_slice_symm_buffer_for_mega_moe(Tensor buffer, 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=0) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor)"); m.def( - "fp8_fp4_mega_moe(Tensor(y!) 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? cumulative_local_expert_recv_stats, Tensor(sym_buffer!) sym_buffer, int[] sym_buffer_ptrs, int rank_idx, int num_max_tokens_per_rank, int num_experts, int num_topk, int[] recipe, str activation, float? activation_clamp_opt, bool fast_math) -> ()"); + "fp8_fp4_mega_moe(Tensor(y!) y, Tensor l1_weights, Tensor l1_weights_sf, Tensor l2_weights, Tensor l2_weights_sf, Tensor? shared_l1_weights, Tensor? shared_l1_weights_sf, Tensor? shared_l2_weights, Tensor? shared_l2_weights_sf, Tensor? cumulative_local_expert_recv_stats, Tensor(sym_buffer!) sym_buffer, int[] sym_buffer_ptrs, int rank_idx, int num_max_tokens_per_rank, int num_experts, int num_topk, int[] recipe, str activation, float? activation_clamp, bool fast_math) -> ()"); m.def( - "bf16_mega_moe(Tensor(y!) y, Tensor l1_weights, Tensor l2_weights, Tensor? shared_l1_weights_opt, Tensor? shared_l2_weights_opt, Tensor? cumulative_local_expert_recv_stats, Tensor(sym_buffer!) 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) -> ()"); + "bf16_mega_moe(Tensor(y!) y, Tensor l1_weights, Tensor l2_weights, Tensor? shared_l1_weights, Tensor? shared_l2_weights, Tensor? cumulative_local_expert_recv_stats, Tensor(sym_buffer!) 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, bool fast_math) -> ()"); #endif } diff --git a/deep_gemm/_C.py b/deep_gemm/_C.py index b86eaffea1..607c2b56b8 100644 --- a/deep_gemm/_C.py +++ b/deep_gemm/_C.py @@ -163,11 +163,11 @@ def fp8_fp4_mqa_logits(q, kv, weights, cu_seq_len_k_start, cu_seq_len_k_end, cle clean_logits, max_seqlen_k, logits_dtype, ) - def fp8_fp4_paged_mqa_logits(q, fused_kv_cache, weights, context_lens, block_table, schedule_meta, max_context_len, + 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 = _unpack_q(q) return _torch_ops.fp8_fp4_paged_mqa_logits( - q_fp, q_sf, fused_kv_cache, weights, context_lens, block_table, schedule_meta, max_context_len, + q_fp, q_sf, kv_cache, weights, context_lens, block_table, schedule_meta, max_context_len, clean_logits, logits_dtype, indices, ) @@ -175,10 +175,10 @@ def fp8_mqa_logits(q, kv, weights, cu_seq_len_k_start, cu_seq_len_k_end, clean_l 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, fused_kv_cache, weights, context_lens, block_table, schedule_meta, max_context_len, + 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, fused_kv_cache, weights, context_lens, block_table, schedule_meta, max_context_len, clean_logits, indices, + q, kv_cache, weights, context_lens, block_table, schedule_meta, max_context_len, clean_logits, indices, ) globals().update({ @@ -253,32 +253,32 @@ def _slice_symm_buffer_for_mega_moe(buffer, *args, **kwargs): ) -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, +def fp8_fp4_mega_moe(y, l1_weights, l2_weights, shared_l1_weights, shared_l2_weights, + 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, activation_clamp, fast_math): shared_l1_w = shared_l1_sf = shared_l2_w = shared_l2_sf = None - if shared_l1_weights_tuple_opt is not None: - shared_l1_w, shared_l1_sf = shared_l1_weights_tuple_opt - shared_l2_w, shared_l2_sf = shared_l2_weights_tuple_opt + if shared_l1_weights is not None: + shared_l1_w, shared_l1_sf = shared_l1_weights + shared_l2_w, shared_l2_sf = shared_l2_weights return _torch_ops.fp8_fp4_mega_moe( - y, l1_weights_tuple[0], l1_weights_tuple[1], l2_weights_tuple[0], l2_weights_tuple[1], + y, l1_weights[0], l1_weights[1], l2_weights[0], l2_weights[1], shared_l1_w, shared_l1_sf, shared_l2_w, shared_l2_sf, cumulative_local_expert_recv_stats, sym_buffer, list(sym_buffer_ptrs), rank_idx, num_max_tokens_per_rank, num_experts, num_topk, list(recipe), activation, - activation_clamp_opt, fast_math, + activation_clamp, fast_math, ) -def bf16_mega_moe(y, l1_weights, l2_weights, shared_l1_weights_opt, shared_l2_weights_opt, +def bf16_mega_moe(y, l1_weights, l2_weights, shared_l1_weights, shared_l2_weights, 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, activation_clamp, fast_math): return _torch_ops.bf16_mega_moe( - y, l1_weights, l2_weights, shared_l1_weights_opt, shared_l2_weights_opt, + y, l1_weights, l2_weights, shared_l1_weights, shared_l2_weights, cumulative_local_expert_recv_stats, sym_buffer, list(sym_buffer_ptrs), rank_idx, num_max_tokens_per_rank, num_experts, num_topk, - activation, activation_clamp_opt, fast_math, + activation, activation_clamp, fast_math, ) diff --git a/deep_gemm/mega/__init__.py b/deep_gemm/mega/__init__.py index cd93042408..997d05eb50 100644 --- a/deep_gemm/mega/__init__.py +++ b/deep_gemm/mega/__init__.py @@ -282,29 +282,29 @@ def validate_pair(name, pair): def fp8_fp4_mega_moe(y: torch.Tensor, - l1_weights_tuple: Tuple[torch.Tensor, torch.Tensor], - l2_weights_tuple: Tuple[torch.Tensor, torch.Tensor], + l1_weights: Tuple[torch.Tensor, torch.Tensor], + l2_weights: Tuple[torch.Tensor, torch.Tensor], sym_buffer: SymmBuffer, - shared_l1_weights_tuple_opt: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, - shared_l2_weights_tuple_opt: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + shared_l1_weights: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + shared_l2_weights: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, cumulative_local_expert_recv_stats: Optional[torch.Tensor] = None, recipe: Tuple[int, int, int] = (1, 1, 32), activation: str = 'swiglu', - activation_clamp_opt: Optional[float] = None, + activation_clamp: Optional[float] = None, fast_math: bool = True, situ_beta: Optional[float] = None, situ_linear_beta: Optional[float] = None): _C.fp8_fp4_mega_moe( y, - l1_weights_tuple, l2_weights_tuple, - shared_l1_weights_tuple_opt, shared_l2_weights_tuple_opt, + l1_weights, l2_weights, + shared_l1_weights, shared_l2_weights, cumulative_local_expert_recv_stats, sym_buffer.buffer, sym_buffer.handle.buffer_ptrs, sym_buffer.group.rank(), sym_buffer.num_max_tokens_per_rank, sym_buffer.num_experts, sym_buffer.num_topk, recipe, - activation, activation_clamp_opt, + activation, activation_clamp, fast_math, situ_beta, situ_linear_beta ) @@ -370,18 +370,18 @@ def bf16_mega_moe(y: torch.Tensor, l1_weights: torch.Tensor, l2_weights: torch.Tensor, sym_buffer: SymmBuffer, - shared_l1_weights_opt: Optional[torch.Tensor] = None, - shared_l2_weights_opt: Optional[torch.Tensor] = None, + shared_l1_weights: Optional[torch.Tensor] = None, + shared_l2_weights: Optional[torch.Tensor] = None, cumulative_local_expert_recv_stats: Optional[torch.Tensor] = None, activation: str = 'swiglu', - activation_clamp_opt: Optional[float] = None, + activation_clamp: Optional[float] = None, fast_math: bool = True): _C.bf16_mega_moe( y, l1_weights, l2_weights, - shared_l1_weights_opt, - shared_l2_weights_opt, + shared_l1_weights, + shared_l2_weights, cumulative_local_expert_recv_stats, sym_buffer.buffer, sym_buffer.handle.buffer_ptrs, @@ -389,7 +389,7 @@ def bf16_mega_moe(y: torch.Tensor, sym_buffer.num_max_tokens_per_rank, sym_buffer.num_experts, sym_buffer.num_topk, - activation, activation_clamp_opt, + activation, activation_clamp, fast_math ) diff --git a/tests/test_attention.py b/tests/test_attention.py index 41ee0150ee..7a47d478ee 100644 --- a/tests/test_attention.py +++ b/tests/test_attention.py @@ -523,7 +523,7 @@ def enumerate_paged_mqa_logits(): num_kv_multicast = 2 if get_arch_major() == 9 and next_n == 4 else 1 num_clusters = deep_gemm.get_num_sms() // num_kv_multicast kernel_kwargs = dict( - q=q_in, fused_kv_cache=kv_in, weights=kernel_weights, + q=q_in, kv_cache=kv_in, weights=kernel_weights, context_lens=context_lens_nextn, block_table=block_table, schedule_meta=deep_gemm.get_paged_mqa_logits_metadata(context_lens_nextn, block_kv, num_clusters, indices=indices), max_context_len=max_model_len, clean_logits=clean_logits, logits_dtype=logits_dtype, diff --git a/tests/test_mega_moe.py b/tests/test_mega_moe.py index 2958162f2a..9987e024a0 100644 --- a/tests/test_mega_moe.py +++ b/tests/test_mega_moe.py @@ -183,17 +183,17 @@ def run_fused(): copy_inputs_to_buffer() y = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device='cuda') - weights_suffix = '' if is_bf16xbf16 else '_tuple' kernel_kwargs = dict( - y=y, sym_buffer=buffer, + y=y, l1_weights=transformed_l1_weights, l2_weights=transformed_l2_weights, + sym_buffer=buffer, cumulative_local_expert_recv_stats=cumulative_local_expert_recv_stats_fused, - activation_clamp_opt=args.activation_clamp, + activation_clamp=args.activation_clamp, fast_math=bool(args.fast_math)) - kernel_kwargs[f'l1_weights{weights_suffix}'] = transformed_l1_weights - kernel_kwargs[f'l2_weights{weights_suffix}'] = transformed_l2_weights if num_shared_experts > 0: - kernel_kwargs[f'shared_l1_weights{weights_suffix}_opt'] = transformed_shared_l1_weights - kernel_kwargs[f'shared_l2_weights{weights_suffix}_opt'] = transformed_shared_l2_weights + kernel_kwargs.update( + shared_l1_weights=transformed_shared_l1_weights, + shared_l2_weights=transformed_shared_l2_weights + ) (deep_gemm.bf16_mega_moe if is_bf16xbf16 else deep_gemm.fp8_fp4_mega_moe)(**kernel_kwargs) return y, cumulative_local_expert_recv_stats_fused From 82c7c96e4a48b0aecb54d262722faa5f45e36ea5 Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Thu, 23 Jul 2026 19:29:08 +0000 Subject: [PATCH 17/28] Mark cumulative_local_expert_recv_stats as mutable in mega MoE schemas. fp8_fp4_mega_moe and bf16_mega_moe both write into this tensor via a device-side red.add reduction (accumulating per-expert recv counts across calls), but the TORCH_LIBRARY schema declared it as immutable (Tensor?). Fix by annotating it Tensor(cumulative_local_expert_recv_stats!)?. Signed-off-by: Chris Leonard --- csrc/apis/mega.hpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/csrc/apis/mega.hpp b/csrc/apis/mega.hpp index 0e0a34f6d0..334af52cc8 100644 --- a/csrc/apis/mega.hpp +++ b/csrc/apis/mega.hpp @@ -852,9 +852,9 @@ TORCH_LIBRARY_FRAGMENT(deep_gemm, m) { m.def( "_slice_symm_buffer_for_mega_moe(Tensor buffer, 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=0) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor)"); m.def( - "fp8_fp4_mega_moe(Tensor(y!) y, Tensor l1_weights, Tensor l1_weights_sf, Tensor l2_weights, Tensor l2_weights_sf, Tensor? shared_l1_weights, Tensor? shared_l1_weights_sf, Tensor? shared_l2_weights, Tensor? shared_l2_weights_sf, Tensor? cumulative_local_expert_recv_stats, Tensor(sym_buffer!) sym_buffer, int[] sym_buffer_ptrs, int rank_idx, int num_max_tokens_per_rank, int num_experts, int num_topk, int[] recipe, str activation, float? activation_clamp, bool fast_math) -> ()"); + "fp8_fp4_mega_moe(Tensor(y!) y, Tensor l1_weights, Tensor l1_weights_sf, Tensor l2_weights, Tensor l2_weights_sf, Tensor? shared_l1_weights, Tensor? shared_l1_weights_sf, Tensor? shared_l2_weights, Tensor? shared_l2_weights_sf, Tensor(cumulative_local_expert_recv_stats!)? cumulative_local_expert_recv_stats, Tensor(sym_buffer!) sym_buffer, int[] sym_buffer_ptrs, int rank_idx, int num_max_tokens_per_rank, int num_experts, int num_topk, int[] recipe, str activation, float? activation_clamp, bool fast_math) -> ()"); m.def( - "bf16_mega_moe(Tensor(y!) y, Tensor l1_weights, Tensor l2_weights, Tensor? shared_l1_weights, Tensor? shared_l2_weights, Tensor? cumulative_local_expert_recv_stats, Tensor(sym_buffer!) 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, bool fast_math) -> ()"); + "bf16_mega_moe(Tensor(y!) y, Tensor l1_weights, Tensor l2_weights, Tensor? shared_l1_weights, Tensor? shared_l2_weights, Tensor(cumulative_local_expert_recv_stats!)? cumulative_local_expert_recv_stats, Tensor(sym_buffer!) 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, bool fast_math) -> ()"); #endif } From 944a443f7e61bf13c8451e178593a7a60302ec8f Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Mon, 27 Jul 2026 20:42:51 +0000 Subject: [PATCH 18/28] Align the abi3 floor with the release matrix to stop colliding wheel filenames on release uploads. Bump the floor to cp310, which every torch version in the matrix (2.4-2.8) already supports. Signed-off-by: Chris Leonard --- .github/workflows/publish.yml | 9 +++------ setup.py | 4 ++-- 2 files changed, 5 insertions(+), 8 deletions(-) diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index a7b3e6b880..4b9d91ed56 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 a single abi3 wheel pinned to the cp310 floor, which every torch + # version here (2.4-2.8) supports, so one Python build covers the whole matrix. + python-version: ["3.10"] torch-version: ["2.4.0", "2.5.1", "2.6.0", "2.7.1", "2.8.0"] cuda-version: ["12.9.1"] # We need separate wheels that either uses C++11 ABI (-D_GLIBCXX_USE_CXX11_ABI) or not. @@ -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/setup.py b/setup.py index 67f762c694..50a5fd4202 100644 --- a/setup.py +++ b/setup.py @@ -27,7 +27,7 @@ # Compiler flags cxx_flags = ['-std=c++17', '-O3', '-fPIC', '-Wno-psabi', '-Wno-deprecated-declarations', f'-D_GLIBCXX_USE_CXX11_ABI={int(torch.compiled_with_cxx11_abi())}', - '-DPy_LIMITED_API=0x03090000'] + '-DPy_LIMITED_API=0x030a0000'] if DG_JIT_USE_RUNTIME_API: cxx_flags.append('-DDG_JIT_USE_RUNTIME_API') @@ -214,7 +214,7 @@ def run(self): }, ext_modules=get_ext_modules(), zip_safe=False, - options={'bdist_wheel': {'py_limited_api': 'cp39'}}, + options={'bdist_wheel': {'py_limited_api': 'cp310'}}, cmdclass={ 'build_py': CustomBuildPy, 'bdist_wheel': CachedWheelsCommand, From d689a67cd1d14f6cf531ca115b825de2abbc540d Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Tue, 28 Jul 2026 12:29:07 +0000 Subject: [PATCH 19/28] Make required in fp8_einsum and k_grouped_fp8_gemm_*_contiguous schemas. These were declared optional but the C++ impl always called unconditionally, so would have crashed anyway; the Python wrappers already default to a concrete tuple. Signed-off-by: Chris Leonard --- csrc/apis/einsum.hpp | 6 +++--- csrc/apis/gemm.hpp | 12 ++++++------ deep_gemm/_C.py | 2 +- 3 files changed, 10 insertions(+), 10 deletions(-) diff --git a/csrc/apis/einsum.hpp b/csrc/apis/einsum.hpp index 7cfec7a8f4..b999ecccb1 100644 --- a/csrc/apis/einsum.hpp +++ b/csrc/apis/einsum.hpp @@ -287,8 +287,8 @@ 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 c10::optional>& recipe) { - einsum::fp8_einsum(expr, {a, sfa}, {b, sfb}, d, c, list_to_tuple3(recipe.value())); + const c10::List& recipe) { + einsum::fp8_einsum(expr, {a, sfa}, {b, sfb}, d, c, list_to_tuple3(recipe)); } #endif @@ -299,7 +299,7 @@ TORCH_LIBRARY_FRAGMENT(deep_gemm, m) { m.def( "einsum(str expr, Tensor a, Tensor b, Tensor(d!) d, Tensor? c=None, bool use_cublaslt=False) -> ()"); m.def( - "fp8_einsum(str expr, Tensor a, Tensor sfa, Tensor b, Tensor sfb, Tensor(d!) d, Tensor? c=None, int[]? recipe=None) -> ()"); + "fp8_einsum(str expr, Tensor a, Tensor sfa, Tensor b, Tensor sfb, Tensor(d!) d, Tensor? c=None, int[] recipe) -> ()"); #endif } diff --git a/csrc/apis/gemm.hpp b/csrc/apis/gemm.hpp index dd3be44a28..2c63ce2b67 100644 --- a/csrc/apis/gemm.hpp +++ b/csrc/apis/gemm.hpp @@ -899,12 +899,12 @@ static void k_grouped_fp8_gemm_tn_contiguous( const c10::optional>& ks_cpu, const torch::Tensor& grouped_layout, const c10::optional& c, - const c10::optional>& recipe, + const c10::List& 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.value()), compiled_dims, use_psum_layout); + list_to_tuple3(recipe), compiled_dims, use_psum_layout); } static void k_grouped_fp8_gemm_nt_contiguous( @@ -914,12 +914,12 @@ static void k_grouped_fp8_gemm_nt_contiguous( const c10::optional>& ks_cpu, const torch::Tensor& grouped_layout, const c10::optional& c, - const c10::optional>& recipe, + const c10::List& 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.value()), compiled_dims, use_psum_layout); + list_to_tuple3(recipe), compiled_dims, use_psum_layout); } #endif @@ -1032,9 +1032,9 @@ TORCH_LIBRARY_FRAGMENT(deep_gemm, m) { 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[]? recipe=None, int[]? recipe_a=None, int[]? 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[]? recipe=None, str compiled_dims='mn', bool use_psum_layout=False) -> ()"); + "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[] 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[]? recipe=None, str compiled_dims='mn', bool use_psum_layout=False) -> ()"); + "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[] recipe, str compiled_dims='mn', bool use_psum_layout=False) -> ()"); #endif #if DG_TENSORMAP_COMPATIBLE diff --git a/deep_gemm/_C.py b/deep_gemm/_C.py index 607c2b56b8..392ef60fa1 100644 --- a/deep_gemm/_C.py +++ b/deep_gemm/_C.py @@ -152,7 +152,7 @@ def fp8_gemm_nt_skip_head_mid(a, b, d, head_splits, recipe=None, compiled_dims=' ) def fp8_einsum(expr, a, b, d, c=None, recipe=(1, 128, 128)): - return _torch_ops.fp8_einsum(expr, a[0], a[1], b[0], b[1], d, c, list(recipe) if recipe is not None else None) + return _torch_ops.fp8_einsum(expr, a[0], a[1], b[0], b[1], d, c, list(recipe)) 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): From 50997886bab9eb9396aee961ffdcd78d6fcab73c Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Tue, 28 Jul 2026 13:36:48 +0000 Subject: [PATCH 20/28] Avoid recomputing the mega MoE symm-buffer layout twice per allocation by having get_symm_buffer_size_for_mega_moe return the computed layout as an int[] alongside num_bytes, which _slice_symm_buffer_for_mega_moe now reuses instead of rederiving it (and its num_sms-dependent ring sizing) from the original args a second time. Signed-off-by: Chris Leonard --- csrc/apis/mega.hpp | 81 +++++++++++++++++++++++--------------- deep_gemm/mega/__init__.py | 15 ++----- 2 files changed, 54 insertions(+), 42 deletions(-) diff --git a/csrc/apis/mega.hpp b/csrc/apis/mega.hpp index 334af52cc8..e065b2e0c3 100644 --- a/csrc/apis/mega.hpp +++ b/csrc/apis/mega.hpp @@ -60,6 +60,46 @@ struct SymmBufferLayoutInfo { 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 + c10::List 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 c10::List& 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 = static_cast(values[13]); + info.num_topk = static_cast(values[14]); + info.hidden = static_cast(values[15]); + info.intermediate_hidden = static_cast(values[16]); + info.num_shared_experts = static_cast(values[17]); + info.shared_intermediate_hidden = static_cast(values[18]); + info.num_ring_tokens = static_cast(values[19]); + info.num_sf_ring_tokens = static_cast(values[20]); + return info; + } }; static SymmBufferLayoutInfo build_symm_buffer_layout( @@ -157,15 +197,16 @@ static SymmBufferLayoutInfo build_symm_buffer_layout( return layout_info; } -static int64_t get_symm_buffer_size_for_mega_moe( +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) { - return build_symm_buffer_layout( + 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).num_bytes; + 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& l1_weights_tuple, @@ -728,7 +755,7 @@ static int64_t get_block_m_for_mega_moe( static_cast(num_topk), mma_type)); } -static int64_t get_symm_buffer_size_for_mega_moe( +static std::tuple> get_symm_buffer_size_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_topk, const int64_t& hidden, const int64_t& intermediate_hidden, @@ -743,17 +770,9 @@ static int64_t get_symm_buffer_size_for_mega_moe( static mega::SymmBufferSlice _slice_symm_buffer_for_mega_moe( const torch::Tensor& buffer, - const int64_t& num_ranks, const int64_t& num_experts, - const int64_t& num_max_tokens_per_rank, const int64_t& num_topk, - const int64_t& hidden, const int64_t& intermediate_hidden, - const std::string& mma_type, const std::string& activation, - const int64_t& num_shared_experts) { - return mega::slice_symm_buffer_for_mega_moe( - buffer, - static_cast(num_ranks), static_cast(num_experts), - static_cast(num_max_tokens_per_rank), static_cast(num_topk), - static_cast(hidden), static_cast(intermediate_hidden), - mma_type, activation, static_cast(num_shared_experts)); + const c10::List& layout_info) { + return mega::slice_symm_buffer_from_layout( + buffer, mega::SymmBufferLayoutInfo::from_int_list(layout_info)); } static void fp8_fp4_mega_moe( @@ -847,10 +866,10 @@ TORCH_LIBRARY_FRAGMENT(deep_gemm, m) { "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", TORCH_FN(deep_gemm::torch_registration::get_block_m_for_mega_moe)); 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=0) -> int", + "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=0) -> (int, int[])", TORCH_FN(deep_gemm::torch_registration::get_symm_buffer_size_for_mega_moe)); m.def( - "_slice_symm_buffer_for_mega_moe(Tensor buffer, 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=0) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor)"); + "_slice_symm_buffer_for_mega_moe(Tensor buffer, int[] layout_info) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor)"); m.def( "fp8_fp4_mega_moe(Tensor(y!) y, Tensor l1_weights, Tensor l1_weights_sf, Tensor l2_weights, Tensor l2_weights_sf, Tensor? shared_l1_weights, Tensor? shared_l1_weights_sf, Tensor? shared_l2_weights, Tensor? shared_l2_weights_sf, Tensor(cumulative_local_expert_recv_stats!)? cumulative_local_expert_recv_stats, Tensor(sym_buffer!) sym_buffer, int[] sym_buffer_ptrs, int rank_idx, int num_max_tokens_per_rank, int num_experts, int num_topk, int[] recipe, str activation, float? activation_clamp, bool fast_math) -> ()"); m.def( diff --git a/deep_gemm/mega/__init__.py b/deep_gemm/mega/__init__.py index 997d05eb50..7e826a5745 100644 --- a/deep_gemm/mega/__init__.py +++ b/deep_gemm/mega/__init__.py @@ -35,16 +35,9 @@ def __init__(self, group: dist.ProcessGroup, self.mma_type = mma_type self.activation = activation - # Allocate a symmetric buffer - num_bytes = _C.get_symm_buffer_size_for_mega_moe( - group.size(), num_experts, - num_max_tokens_per_rank, num_topk, - hidden, intermediate_hidden, - mma_type, activation, - num_shared_experts, - ) - slice_input_buffers = lambda buffer: _C._slice_symm_buffer_for_mega_moe( - buffer, + # Allocate a symmetric buffer. The layout is computed once here and reused for + # slicing below. + num_bytes, layout_info = _C.get_symm_buffer_size_for_mega_moe( group.size(), num_experts, num_max_tokens_per_rank, num_topk, hidden, intermediate_hidden, @@ -68,7 +61,7 @@ def __init__(self, group: dist.ProcessGroup, self.shared_l1_acts, self.shared_l1_acts_sf, self.shared_l2_acts, self.shared_l2_acts_sf, self.l1_acts, self.l1_acts_sf, - self.l2_acts, self.l2_acts_sf) = slice_input_buffers(self.buffer) + self.l2_acts, self.l2_acts_sf) = _C._slice_symm_buffer_for_mega_moe(self.buffer, layout_info) def destroy(self): self.handle = None From 5f39e32908f2e7fe376664470a97d8077298c5b3 Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Tue, 28 Jul 2026 15:36:00 +0000 Subject: [PATCH 21/28] Use int[N] instead of int[] in TORCH_LIBRARY schemas for fixed-arity params like recipe/head_splits, so the schema itself is the source of truth for stub generation instead of a hand-maintained per-name promotion table in generate_pyi.py. Add an _as_int_list wrapper helper so int[N]'s scalar-broadcasting behavior can't silently convert a bad scalar argument into a repeated list. Signed-off-by: Chris Leonard --- csrc/apis/attention.hpp | 2 +- csrc/apis/einsum.hpp | 2 +- csrc/apis/gemm.hpp | 18 +++++++++--------- csrc/apis/mega.hpp | 2 +- deep_gemm/_C.py | 21 +++++++++++++-------- scripts/generate_pyi.py | 33 ++++++++++++--------------------- 6 files changed, 37 insertions(+), 41 deletions(-) diff --git a/csrc/apis/attention.hpp b/csrc/apis/attention.hpp index 187c16fea9..b9d4d8dd64 100644 --- a/csrc/apis/attention.hpp +++ b/csrc/apis/attention.hpp @@ -566,7 +566,7 @@ static torch::Tensor fp8_paged_mqa_logits( TORCH_LIBRARY_FRAGMENT(deep_gemm, m) { #if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE m.def( - "fp8_gemm_nt_skip_head_mid(Tensor a, Tensor sfa, Tensor b, Tensor sfb, Tensor(d!) d, int[] head_splits, int[]? recipe=None, str compiled_dims='nk', bool disable_ue8m0_cast=False) -> ()"); + "fp8_gemm_nt_skip_head_mid(Tensor a, Tensor sfa, Tensor b, Tensor sfb, Tensor(d!) d, int[3] head_splits, int[3]? recipe=None, str compiled_dims='nk', bool disable_ue8m0_cast=False) -> ()"); m.def( "fp8_fp4_mqa_logits(Tensor q, Tensor? q_sf, Tensor kv, Tensor kv_sf, Tensor weights, Tensor cu_seq_len_k_start, Tensor cu_seq_len_k_end, bool clean_logits=True, int max_seqlen_k=0, ScalarType logits_dtype=float) -> Tensor"); m.def( diff --git a/csrc/apis/einsum.hpp b/csrc/apis/einsum.hpp index b999ecccb1..4b4c3150bc 100644 --- a/csrc/apis/einsum.hpp +++ b/csrc/apis/einsum.hpp @@ -299,7 +299,7 @@ TORCH_LIBRARY_FRAGMENT(deep_gemm, m) { m.def( "einsum(str expr, Tensor a, Tensor b, Tensor(d!) d, Tensor? c=None, bool use_cublaslt=False) -> ()"); m.def( - "fp8_einsum(str expr, Tensor a, Tensor sfa, Tensor b, Tensor sfb, Tensor(d!) d, Tensor? c=None, int[] recipe) -> ()"); + "fp8_einsum(str expr, Tensor a, Tensor sfa, Tensor b, Tensor sfb, Tensor(d!) d, Tensor? c=None, int[3] recipe) -> ()"); #endif } diff --git a/csrc/apis/gemm.hpp b/csrc/apis/gemm.hpp index 2c63ce2b67..314d8da13a 100644 --- a/csrc/apis/gemm.hpp +++ b/csrc/apis/gemm.hpp @@ -1018,23 +1018,23 @@ TORCH_LIBRARY_FRAGMENT(deep_gemm, m) { #if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE // GEMM — FP8/FP4 m.def( - "fp8_fp4_gemm_nt(Tensor a, Tensor sfa, Tensor b, Tensor sfb, Tensor(d!) d, Tensor? c=None, int[]? recipe=None, int[]? recipe_a=None, int[]? recipe_b=None, str compiled_dims='nk', bool disable_ue8m0_cast=False) -> ()"); + "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) -> ()"); m.def( - "fp8_fp4_gemm_nn(Tensor a, Tensor sfa, Tensor b, Tensor sfb, Tensor(d!) d, Tensor? c=None, int[]? recipe=None, int[]? recipe_a=None, int[]? recipe_b=None, str compiled_dims='nk', bool disable_ue8m0_cast=False) -> ()"); + "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) -> ()"); m.def( - "fp8_fp4_gemm_tn(Tensor a, Tensor sfa, Tensor b, Tensor sfb, Tensor(d!) d, Tensor? c=None, int[]? recipe=None, int[]? recipe_a=None, int[]? recipe_b=None, str compiled_dims='mn', bool disable_ue8m0_cast=False) -> ()"); + "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) -> ()"); m.def( - "fp8_fp4_gemm_tt(Tensor a, Tensor sfa, Tensor b, Tensor sfb, Tensor(d!) d, Tensor? c=None, int[]? recipe=None, int[]? recipe_a=None, int[]? recipe_b=None, str compiled_dims='mn', bool disable_ue8m0_cast=False) -> ()"); + "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) -> ()"); m.def( - "m_grouped_fp8_fp4_gemm_nt_contiguous(Tensor a, Tensor sfa, Tensor b, Tensor sfb, Tensor(d!) d, Tensor grouped_layout, int[]? recipe=None, int[]? recipe_a=None, int[]? 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_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[]? recipe=None, int[]? recipe_a=None, int[]? recipe_b=None, str compiled_dims='nk', bool disable_ue8m0_cast=False, bool use_psum_layout=False, bool ensure_zero_padding=True) -> ()"); + "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[]? recipe=None, int[]? recipe_a=None, int[]? recipe_b=None, str compiled_dims='nk', bool disable_ue8m0_cast=False) -> ()"); + "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[] recipe, str compiled_dims='mn', bool use_psum_layout=False) -> ()"); + "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[] recipe, str compiled_dims='mn', bool use_psum_layout=False) -> ()"); + "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) -> ()"); #endif #if DG_TENSORMAP_COMPATIBLE diff --git a/csrc/apis/mega.hpp b/csrc/apis/mega.hpp index e065b2e0c3..11ca2d5fc9 100644 --- a/csrc/apis/mega.hpp +++ b/csrc/apis/mega.hpp @@ -871,7 +871,7 @@ TORCH_LIBRARY_FRAGMENT(deep_gemm, m) { m.def( "_slice_symm_buffer_for_mega_moe(Tensor buffer, int[] layout_info) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor)"); m.def( - "fp8_fp4_mega_moe(Tensor(y!) y, Tensor l1_weights, Tensor l1_weights_sf, Tensor l2_weights, Tensor l2_weights_sf, Tensor? shared_l1_weights, Tensor? shared_l1_weights_sf, Tensor? shared_l2_weights, Tensor? shared_l2_weights_sf, Tensor(cumulative_local_expert_recv_stats!)? cumulative_local_expert_recv_stats, Tensor(sym_buffer!) sym_buffer, int[] sym_buffer_ptrs, int rank_idx, int num_max_tokens_per_rank, int num_experts, int num_topk, int[] recipe, str activation, float? activation_clamp, bool fast_math) -> ()"); + "fp8_fp4_mega_moe(Tensor(y!) y, Tensor l1_weights, Tensor l1_weights_sf, Tensor l2_weights, Tensor l2_weights_sf, Tensor? shared_l1_weights, Tensor? shared_l1_weights_sf, Tensor? shared_l2_weights, Tensor? shared_l2_weights_sf, Tensor(cumulative_local_expert_recv_stats!)? cumulative_local_expert_recv_stats, Tensor(sym_buffer!) 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, bool fast_math) -> ()"); m.def( "bf16_mega_moe(Tensor(y!) y, Tensor l1_weights, Tensor l2_weights, Tensor? shared_l1_weights, Tensor? shared_l2_weights, Tensor(cumulative_local_expert_recv_stats!)? cumulative_local_expert_recv_stats, Tensor(sym_buffer!) 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, bool fast_math) -> ()"); #endif diff --git a/deep_gemm/_C.py b/deep_gemm/_C.py index 392ef60fa1..c10f41636b 100644 --- a/deep_gemm/_C.py +++ b/deep_gemm/_C.py @@ -68,13 +68,18 @@ 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 DeepGEMM kernels only when C++ ops are registered.""" 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): 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, recipe, recipe_a, recipe_b, + 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, ) @@ -82,7 +87,7 @@ 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): 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, recipe, recipe_a, recipe_b, + 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, ) @@ -90,7 +95,7 @@ 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): 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, recipe, recipe_a, recipe_b, + 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, ) @@ -98,7 +103,7 @@ 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): 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, recipe, recipe_a, recipe_b, + 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, ) @@ -107,7 +112,7 @@ def m_grouped_fp8_fp4_gemm_nt_contiguous(a, b, d, grouped_layout, recipe=None, r 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, recipe, recipe_a, recipe_b, + 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, ) @@ -117,7 +122,7 @@ def m_grouped_fp8_fp4_gemm_nn_contiguous(a, b, d, grouped_layout, recipe=None, r 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, recipe, recipe_a, recipe_b, + 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, ) @@ -125,7 +130,7 @@ def m_grouped_fp8_fp4_gemm_nt_masked(a, b, d, masked_m, expected_m, recipe=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, recipe, recipe_a, recipe_b, + 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, ) @@ -148,7 +153,7 @@ def k_grouped_fp8_gemm_nt_contiguous(a, b, d, ks_cpu, grouped_layout, c=None, re 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), recipe, compiled_dims, disable_ue8m0_cast, + 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)): diff --git a/scripts/generate_pyi.py b/scripts/generate_pyi.py index 0704d797ac..476a49a245 100644 --- a/scripts/generate_pyi.py +++ b/scripts/generate_pyi.py @@ -91,6 +91,10 @@ def schema_type_to_python(type_str: 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: @@ -295,26 +299,13 @@ def _maybe_widen_int_list_value_param(parameters: list[dict]) -> None: parameters[0]['py_type'] = 'int | list[int]' -def _promote_int_list_tuple_types(op_name: str, parameters: list[dict]) -> None: - """Promote int[] schema params to fixed-size tuples matching the public _C.py API.""" - for param in parameters: - name = param['name'] - py_type = param['py_type'] - - if name == 'head_splits' and py_type == 'list[int]': - param['py_type'] = 'tuple[int, int, int]' - elif name == 'recipe_a' and py_type == 'Optional[list[int]]': - param['py_type'] = 'Optional[tuple[int, int]]' - elif name == 'recipe_b' and py_type == 'Optional[list[int]]': - param['py_type'] = 'Optional[tuple[int, int]]' - elif name == 'recipe' and op_name == 'transform_sf_into_required_layout': - if py_type == 'list[int]': - param['py_type'] = 'tuple[int, int] | tuple[int, int, int]' - elif name == 'recipe': - if py_type == 'list[int]': - param['py_type'] = 'tuple[int, int, int]' - elif py_type == 'Optional[list[int]]': - param['py_type'] = 'Optional[tuple[int, int, 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( @@ -330,7 +321,7 @@ def adjust_for_c_py_wrapper( parameters = _merge_named_pairs(parameters, tuple(pairs)) _maybe_widen_int_list_value_param(parameters) - _promote_int_list_tuple_types(name, parameters) + _promote_transform_sf_recipe_type(name, parameters) return parameters From 5c74217b103f780df3e6f20593e9ae7ae6fa1532 Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Tue, 28 Jul 2026 19:01:45 +0000 Subject: [PATCH 22/28] Generalize _merge_named_pairs to cover asymmetric-optional tensor/scale-factor pairs like q/q_sf, replacing the special-cased _apply_q_qsf_merge, and sort the csrc/ file scan so op ordering in the generated stub no longer shuffles across machines/checkouts. Also drop sanitize_param_name, since it could only keep the .pyi syntactically valid and not guarantee the stub's keyword name actually matches the real wrapper at runtime. Signed-off-by: Chris Leonard --- scripts/generate_pyi.py | 61 +++++++++++++++-------------------------- 1 file changed, 22 insertions(+), 39 deletions(-) diff --git a/scripts/generate_pyi.py b/scripts/generate_pyi.py index 476a49a245..1f1dcbe9a0 100644 --- a/scripts/generate_pyi.py +++ b/scripts/generate_pyi.py @@ -3,10 +3,6 @@ import re from pathlib import Path -_TENSOR_PAIR = 'tuple[torch.Tensor, torch.Tensor]' -_Q_TUPLE = 'tuple[torch.Tensor, Optional[torch.Tensor]]' - - class BracketTracker: """Track () [] {} <> nesting for top-level comma/default splitting.""" @@ -228,20 +224,31 @@ def parse_torch_schema(schema: str) -> dict: 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.""" - drop = {right for left, right in pairs} - merged_left = {left for left, _ in pairs} + 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 - if param['name'] in merged_left: - out.append({ - 'name': param['name'], - 'py_type': _TENSOR_PAIR, - 'default': None, - }) + sf_name = sf_of.get(param['name']) + if sf_name is None: + out.append(dict(param)) continue - out.append(dict(param)) + 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 @@ -276,22 +283,6 @@ def detect_tensor_sf_pairs(parameters: list[dict]) -> list[tuple[str, str]]: return pairs -def _apply_q_qsf_merge(parameters: list[dict]) -> list[dict]: - """Merge optional q_sf into q for attention wrappers.""" - if not any(param['name'] == 'q_sf' for param in parameters): - return [dict(param) for param in parameters] - - out = [] - for param in parameters: - if param['name'] == 'q_sf': - continue - param = dict(param) - if param['name'] == 'q': - param['py_type'] = _Q_TUPLE - out.append(param) - return out - - 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': @@ -314,8 +305,6 @@ def adjust_for_c_py_wrapper( wrapper_defaults: dict[str, dict[str, str]] | None = None, ) -> list[dict]: """Adjust flat schema params to match deep_gemm._C Python wrappers.""" - parameters = _apply_q_qsf_merge(parameters) - pairs = detect_tensor_sf_pairs(parameters) if pairs: parameters = _merge_named_pairs(parameters, tuple(pairs)) @@ -326,12 +315,6 @@ def adjust_for_c_py_wrapper( return parameters -def sanitize_param_name(name: str) -> str: - if name in {'def', 'class', 'from', 'import', 'None', 'True', 'False'}: - return f'{name}_' - return name - - def format_ast_default(node: ast.AST) -> str: """Convert an AST default value node to a Python expression string.""" if isinstance(node, ast.Constant): @@ -419,7 +402,7 @@ def extract_m_def_statements(root_path) -> list[str]: statements = [] extensions = {'.hpp', '.cpp', '.h', '.cc'} - for file_path in Path(root_path).rglob('*'): + for file_path in sorted(Path(root_path).rglob('*')): if file_path.suffix.lower() not in extensions: continue if not file_path.is_file(): @@ -537,7 +520,7 @@ def generate_pyi_function(parsed, wrapper_defaults=None): param_lines = [] for param in parameters: - name = sanitize_param_name(param['name']) + name = param['name'] if param['default'] is not None: param_lines.append(f' {name}: {param["py_type"]} = {param["default"]}') else: From a9a9576fa9a5c55fcc182ac5dfa29787e793c363 Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Thu, 30 Jul 2026 12:39:47 +0000 Subject: [PATCH 23/28] replaced c10::List that I added with std::vector. It doesn't really matter for this migration but it makes the migration to torch abi stable easier Signed-off-by: Chris Leonard --- csrc/apis/attention.hpp | 4 +-- csrc/apis/einsum.hpp | 2 +- csrc/apis/gemm.hpp | 52 ++++++++++++++++++------------------ csrc/apis/layout.hpp | 4 +-- csrc/apis/mega.hpp | 20 +++++++------- csrc/apis/runtime.hpp | 2 +- csrc/torch_library_utils.hpp | 11 ++++---- 7 files changed, 47 insertions(+), 48 deletions(-) diff --git a/csrc/apis/attention.hpp b/csrc/apis/attention.hpp index b9d4d8dd64..c9686112a3 100644 --- a/csrc/apis/attention.hpp +++ b/csrc/apis/attention.hpp @@ -476,8 +476,8 @@ 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 c10::List& head_splits, - const c10::optional>& recipe, + 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( diff --git a/csrc/apis/einsum.hpp b/csrc/apis/einsum.hpp index 4b4c3150bc..d7a9cabe7f 100644 --- a/csrc/apis/einsum.hpp +++ b/csrc/apis/einsum.hpp @@ -287,7 +287,7 @@ 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 c10::List& recipe) { + const std::vector& recipe) { einsum::fp8_einsum(expr, {a, sfa}, {b, sfb}, d, c, list_to_tuple3(recipe)); } #endif diff --git a/csrc/apis/gemm.hpp b/csrc/apis/gemm.hpp index 314d8da13a..a00954cf65 100644 --- a/csrc/apis/gemm.hpp +++ b/csrc/apis/gemm.hpp @@ -795,9 +795,9 @@ 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 c10::optional>& recipe, + const c10::optional>& recipe_a, + const c10::optional>& recipe_b, const std::string& compiled_dims, const bool& disable_ue8m0_cast) { 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), @@ -808,9 +808,9 @@ 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 c10::optional>& recipe, + const c10::optional>& recipe_a, + const c10::optional>& recipe_b, const std::string& compiled_dims, const bool& disable_ue8m0_cast) { 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), @@ -821,9 +821,9 @@ 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 c10::optional>& recipe, + const c10::optional>& recipe_a, + const c10::optional>& recipe_b, const std::string& compiled_dims, const bool& disable_ue8m0_cast) { 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), @@ -834,9 +834,9 @@ 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 c10::optional>& recipe, + const c10::optional>& recipe_a, + const c10::optional>& recipe_b, const std::string& compiled_dims, const bool& disable_ue8m0_cast) { 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), @@ -847,9 +847,9 @@ 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 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) { @@ -866,9 +866,9 @@ 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 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( @@ -882,9 +882,9 @@ static void m_grouped_fp8_fp4_gemm_nt_masked( 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 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), @@ -896,10 +896,10 @@ 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 c10::optional>& ks_cpu, const torch::Tensor& grouped_layout, const c10::optional& c, - const c10::List& recipe, + 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, @@ -911,10 +911,10 @@ 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 c10::optional>& ks_cpu, const torch::Tensor& grouped_layout, const c10::optional& c, - const c10::List& recipe, + 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, @@ -978,7 +978,7 @@ static void m_grouped_bf16_gemm_nt_masked( 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 c10::optional>& ks_cpu, const torch::Tensor& grouped_layout, const c10::optional& c, const std::string& compiled_dims, const bool& use_psum_layout) { diff --git a/csrc/apis/layout.hpp b/csrc/apis/layout.hpp index a5e6d42acb..10bceabf93 100644 --- a/csrc/apis/layout.hpp +++ b/csrc/apis/layout.hpp @@ -148,7 +148,7 @@ using namespace deep_gemm::torch_utils; #if DG_TENSORMAP_COMPATIBLE static torch::Tensor transform_sf_into_required_layout( const torch::Tensor& sf, const int64_t& mn, const int64_t& k, - const c10::List& recipe, + const std::vector& recipe, const c10::optional& num_groups, const c10::optional& is_sfa, const bool& disable_ue8m0_cast, @@ -173,7 +173,7 @@ static torch::Tensor get_mn_major_tma_aligned_packed_ue8m0_tensor( 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 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( diff --git a/csrc/apis/mega.hpp b/csrc/apis/mega.hpp index 11ca2d5fc9..8921e7b517 100644 --- a/csrc/apis/mega.hpp +++ b/csrc/apis/mega.hpp @@ -62,7 +62,7 @@ struct SymmBufferLayoutInfo { int num_sf_ring_tokens = 0; // Flatten into a plain `int[]` so it can cross the TORCH_LIBRARY boundary - c10::List to_int_list() const { + 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, @@ -73,7 +73,7 @@ struct SymmBufferLayoutInfo { }; } - static SymmBufferLayoutInfo from_int_list(const c10::List& values) { + static SymmBufferLayoutInfo from_int_list(const std::vector& values) { DG_HOST_ASSERT(static_cast(values.size()) == 21); SymmBufferLayoutInfo info; info.num_bytes = values[0]; @@ -197,7 +197,7 @@ static SymmBufferLayoutInfo build_symm_buffer_layout( return layout_info; } -static std::tuple> get_symm_buffer_size_for_mega_moe( +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, @@ -755,7 +755,7 @@ static int64_t get_block_m_for_mega_moe( static_cast(num_topk), mma_type)); } -static std::tuple> get_symm_buffer_size_for_mega_moe( +static std::tuple> get_symm_buffer_size_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_topk, const int64_t& hidden, const int64_t& intermediate_hidden, @@ -770,7 +770,7 @@ static std::tuple> get_symm_buffer_size_for_mega_moe static mega::SymmBufferSlice _slice_symm_buffer_for_mega_moe( const torch::Tensor& buffer, - const c10::List& layout_info) { + const std::vector& layout_info) { return mega::slice_symm_buffer_from_layout( buffer, mega::SymmBufferLayoutInfo::from_int_list(layout_info)); } @@ -785,11 +785,11 @@ static void fp8_fp4_mega_moe( const c10::optional& shared_l2_weights_sf, const c10::optional& cumulative_local_expert_recv_stats, const torch::Tensor& sym_buffer, - const c10::List& sym_buffer_ptrs, + 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 c10::List& recipe, + const std::vector& recipe, const std::string& activation, const c10::optional& activation_clamp, const bool& fast_math) { @@ -811,7 +811,7 @@ static void fp8_fp4_mega_moe( shared_l2_opt, cumulative_local_expert_recv_stats, sym_buffer, - std::vector(sym_buffer_ptrs.begin(), sym_buffer_ptrs.end()), + sym_buffer_ptrs, static_cast(rank_idx), static_cast(num_max_tokens_per_rank), static_cast(num_experts), static_cast(num_topk), @@ -831,7 +831,7 @@ static void bf16_mega_moe( const c10::optional& shared_l2_weights, const c10::optional& cumulative_local_expert_recv_stats, const torch::Tensor& sym_buffer, - const c10::List& sym_buffer_ptrs, + 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, @@ -844,7 +844,7 @@ static void bf16_mega_moe( shared_l2_weights, cumulative_local_expert_recv_stats, sym_buffer, - std::vector(sym_buffer_ptrs.begin(), sym_buffer_ptrs.end()), + sym_buffer_ptrs, static_cast(rank_idx), static_cast(num_max_tokens_per_rank), static_cast(num_experts), static_cast(num_topk), diff --git a/csrc/apis/runtime.hpp b/csrc/apis/runtime.hpp index 9d79d89d61..c97bd3f6b4 100644 --- a/csrc/apis/runtime.hpp +++ b/csrc/apis/runtime.hpp @@ -39,7 +39,7 @@ static void set_ignore_compile_dims(const bool& new_value) { heuristics_runtime->set_ignore_compile_dims(new_value); } -static void set_block_size_multiple_of(const c10::List& value) { +static void set_block_size_multiple_of(const std::vector& value) { if (value.size() == 1) { const int v = static_cast(value[0]); heuristics_runtime->set_block_size_multiple_of(v, v); diff --git a/csrc/torch_library_utils.hpp b/csrc/torch_library_utils.hpp index e84bc143ea..2d80758ba3 100644 --- a/csrc/torch_library_utils.hpp +++ b/csrc/torch_library_utils.hpp @@ -1,6 +1,5 @@ #pragma once -#include #include #include #include @@ -11,7 +10,7 @@ namespace deep_gemm::torch_utils { inline std::optional> list_to_recipe3( - const c10::optional>& recipe) { + const c10::optional>& recipe) { if (not recipe.has_value() or recipe->empty()) { return std::nullopt; } @@ -22,7 +21,7 @@ inline std::optional> list_to_recipe3( } inline std::optional> list_to_recipe2( - const c10::optional>& recipe) { + const c10::optional>& recipe) { if (not recipe.has_value() or recipe->empty()) { return std::nullopt; } @@ -31,7 +30,7 @@ inline std::optional> list_to_recipe2( } inline std::variant, std::tuple> list_to_recipe_variant( - const c10::List& recipe) { + 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])); @@ -41,7 +40,7 @@ inline std::variant, std::tuple> list_to_rec static_cast(recipe[2])); } -inline std::tuple list_to_tuple3(const c10::List& values) { +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]), @@ -49,7 +48,7 @@ inline std::tuple list_to_tuple3(const c10::List& values } inline std::optional> list_to_optional_vector_int( - const c10::optional>& values) { + const c10::optional>& values) { if (not values.has_value()) { return std::nullopt; } From 3ea26ef2377d180049f5b2d26433244fc983be09 Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Fri, 31 Jul 2026 18:11:38 +0000 Subject: [PATCH 24/28] Make _bind_guarded_ops assert on partial guard-group matches instead of silently dropping the whole group, so a future C++ #if drift between grouped ops fails loudly at import with the exact op names involved rather than quietly vanishing from the API. Also document which file each op in the einsum/tf32_hc_prenorm_gemm/get_paged_mqa_logits_metadata group comes from, since they're only bundled by coincidence of a shared guard today. Signed-off-by: Chris Leonard --- deep_gemm/_C.py | 24 +++++++++++++----------- 1 file changed, 13 insertions(+), 11 deletions(-) diff --git a/deep_gemm/_C.py b/deep_gemm/_C.py index c10f41636b..fb8960bb86 100644 --- a/deep_gemm/_C.py +++ b/deep_gemm/_C.py @@ -16,13 +16,14 @@ def _load_extension(): def _bind_guarded_ops(*names): """Bind ops when all are registered (matches one C++ #if guard group).""" - bound = {} - for name in names: - op = getattr(_torch_ops, name, None) - if op is None: - return - bound[name] = op - globals().update(bound) + 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 @@ -223,11 +224,12 @@ def fp8_paged_mqa_logits(q, kv_cache, weights, context_lens, block_table, schedu 'k_grouped_bf16_gemm_tn_contiguous', ) - # DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE — einsum.hpp, attention.hpp, hyperconnection.hpp + # DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE — grouped only because these three + # happen to share the same guard today; if any one's guard changes, split it out. _bind_guarded_ops( - 'einsum', - 'tf32_hc_prenorm_gemm', - 'get_paged_mqa_logits_metadata', + 'einsum', # einsum.hpp + 'tf32_hc_prenorm_gemm', # hyperconnection.hpp + 'get_paged_mqa_logits_metadata', # attention.hpp ) # DG_TENSORMAP_COMPATIBLE — layout.hpp (schema and impl conditional) From e6e1dcfef77dbf3f4798b21b9d87d1cc1988f898 Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Fri, 4 Sep 2026 17:16:20 +0000 Subject: [PATCH 25/28] Preserve NVFP4 MegaMoE support after TORCH_LIBRARY refactor - serialize NVFP4-specific symmetric buffer layout fields - register fp4_fp4_mega_moe and preserve SiTU parameters - add matching Python operator wrappers Signed-off-by: Chris Leonard --- csrc/apis/mega.hpp | 81 ++++++++++++++++++++++++++++++++++++++-------- deep_gemm/_C.py | 24 ++++++++++++-- 2 files changed, 89 insertions(+), 16 deletions(-) diff --git a/csrc/apis/mega.hpp b/csrc/apis/mega.hpp index 8921e7b517..2761e70448 100644 --- a/csrc/apis/mega.hpp +++ b/csrc/apis/mega.hpp @@ -67,14 +67,15 @@ struct SymmBufferLayoutInfo { 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, + static_cast(mma_kind), static_cast(with_sf), sf_gran_k, + 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); + DG_HOST_ASSERT(static_cast(values.size()) == 23); SymmBufferLayoutInfo info; info.num_bytes = values[0]; info.input_token_base = values[1]; @@ -88,16 +89,18 @@ struct SymmBufferLayoutInfo { info.l1_sf_base = values[9]; info.l2_token_base = values[10]; info.l2_sf_base = values[11]; + info.mma_kind = static_cast(values[12]); // `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 = static_cast(values[13]); - info.num_topk = static_cast(values[14]); - info.hidden = static_cast(values[15]); - info.intermediate_hidden = static_cast(values[16]); - info.num_shared_experts = static_cast(values[17]); - info.shared_intermediate_hidden = static_cast(values[18]); - info.num_ring_tokens = static_cast(values[19]); - info.num_sf_ring_tokens = static_cast(values[20]); + info.with_sf = values[13] != 0; + info.sf_gran_k = static_cast(values[14]); + info.num_max_tokens_per_rank = static_cast(values[15]); + info.num_topk = static_cast(values[16]); + info.hidden = static_cast(values[17]); + info.intermediate_hidden = static_cast(values[18]); + info.num_shared_experts = static_cast(values[19]); + info.shared_intermediate_hidden = static_cast(values[20]); + info.num_ring_tokens = static_cast(values[21]); + info.num_sf_ring_tokens = static_cast(values[22]); return info; } }; @@ -792,7 +795,9 @@ static void fp8_fp4_mega_moe( const std::vector& recipe, const std::string& activation, const c10::optional& activation_clamp, - const bool& fast_math) { + const bool& fast_math, + const c10::optional& situ_beta, + const c10::optional& situ_linear_beta) { std::optional> shared_l1_opt = std::nullopt; std::optional> shared_l2_opt = std::nullopt; if (shared_l1_weights.has_value()) { @@ -820,7 +825,52 @@ static void fp8_fp4_mega_moe( activation_clamp.has_value() ? std::make_optional(static_cast(activation_clamp.value())) : std::nullopt, - fast_math); + fast_math, + situ_beta.has_value() + ? std::make_optional(static_cast(situ_beta.value())) + : std::nullopt, + situ_linear_beta.has_value() + ? std::make_optional(static_cast(situ_linear_beta.value())) + : std::nullopt); +} + +static void fp4_fp4_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 c10::optional& shared_l1_weights, + const c10::optional& shared_l2_weights, + const c10::optional& x_bf16, + const c10::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 c10::optional& activation_clamp, + const bool& fast_math, + const c10::optional& l1_alphas, + const c10::optional& l2_alphas, + const c10::optional& a2_scales, + const double& routed_scaling_factor) { + mega::fp4_fp4_mega_moe( + y, + std::make_tuple(l1_weights, l1_weights_sf), + std::make_tuple(l2_weights, l2_weights_sf), + shared_l1_weights, shared_l2_weights, x_bf16, + cumulative_local_expert_recv_stats, + sym_buffer, sym_buffer_ptrs, + static_cast(rank_idx), + static_cast(num_max_tokens_per_rank), + static_cast(num_experts), static_cast(num_topk), + list_to_tuple3(recipe), activation, + activation_clamp.has_value() + ? std::make_optional(static_cast(activation_clamp.value())) + : std::nullopt, + fast_math, l1_alphas, l2_alphas, a2_scales, + static_cast(routed_scaling_factor)); } static void bf16_mega_moe( @@ -871,7 +921,9 @@ TORCH_LIBRARY_FRAGMENT(deep_gemm, m) { m.def( "_slice_symm_buffer_for_mega_moe(Tensor buffer, int[] layout_info) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor)"); m.def( - "fp8_fp4_mega_moe(Tensor(y!) y, Tensor l1_weights, Tensor l1_weights_sf, Tensor l2_weights, Tensor l2_weights_sf, Tensor? shared_l1_weights, Tensor? shared_l1_weights_sf, Tensor? shared_l2_weights, Tensor? shared_l2_weights_sf, Tensor(cumulative_local_expert_recv_stats!)? cumulative_local_expert_recv_stats, Tensor(sym_buffer!) 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, bool fast_math) -> ()"); + "fp8_fp4_mega_moe(Tensor(y!) y, Tensor l1_weights, Tensor l1_weights_sf, Tensor l2_weights, Tensor l2_weights_sf, Tensor? shared_l1_weights, Tensor? shared_l1_weights_sf, Tensor? shared_l2_weights, Tensor? shared_l2_weights_sf, Tensor(cumulative_local_expert_recv_stats!)? cumulative_local_expert_recv_stats, Tensor(sym_buffer!) 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, bool fast_math, float? situ_beta=None, float? situ_linear_beta=None) -> ()"); + m.def( + "fp4_fp4_mega_moe(Tensor(y!) y, Tensor l1_weights, Tensor l1_weights_sf, Tensor l2_weights, Tensor l2_weights_sf, Tensor? shared_l1_weights, Tensor? shared_l2_weights, Tensor? x_bf16, Tensor(cumulative_local_expert_recv_stats!)? cumulative_local_expert_recv_stats, Tensor(sym_buffer!) 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, bool fast_math, Tensor? l1_alphas, Tensor? l2_alphas, Tensor? a2_scales, float routed_scaling_factor) -> ()"); m.def( "bf16_mega_moe(Tensor(y!) y, Tensor l1_weights, Tensor l2_weights, Tensor? shared_l1_weights, Tensor? shared_l2_weights, Tensor(cumulative_local_expert_recv_stats!)? cumulative_local_expert_recv_stats, Tensor(sym_buffer!) 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, bool fast_math) -> ()"); #endif @@ -883,6 +935,7 @@ TORCH_LIBRARY_IMPL(deep_gemm, CUDA, m) { #if DG_TENSORMAP_COMPATIBLE m.impl("_slice_symm_buffer_for_mega_moe", TORCH_FN(_slice_symm_buffer_for_mega_moe)); m.impl("fp8_fp4_mega_moe", TORCH_FN(fp8_fp4_mega_moe)); + m.impl("fp4_fp4_mega_moe", TORCH_FN(fp4_fp4_mega_moe)); m.impl("bf16_mega_moe", TORCH_FN(bf16_mega_moe)); #endif } diff --git a/deep_gemm/_C.py b/deep_gemm/_C.py index fb8960bb86..5561be95ca 100644 --- a/deep_gemm/_C.py +++ b/deep_gemm/_C.py @@ -263,7 +263,8 @@ def _slice_symm_buffer_for_mega_moe(buffer, *args, **kwargs): def fp8_fp4_mega_moe(y, l1_weights, l2_weights, shared_l1_weights, shared_l2_weights, 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, fast_math): + activation, activation_clamp, fast_math, + situ_beta=None, situ_linear_beta=None): shared_l1_w = shared_l1_sf = shared_l2_w = shared_l2_sf = None if shared_l1_weights is not None: shared_l1_w, shared_l1_sf = shared_l1_weights @@ -273,7 +274,25 @@ def fp8_fp4_mega_moe(y, l1_weights, l2_weights, shared_l1_weights, shared_l2_wei shared_l1_w, shared_l1_sf, shared_l2_w, shared_l2_sf, cumulative_local_expert_recv_stats, sym_buffer, list(sym_buffer_ptrs), rank_idx, num_max_tokens_per_rank, num_experts, num_topk, list(recipe), activation, - activation_clamp, fast_math, + activation_clamp, fast_math, situ_beta, situ_linear_beta, + ) + + +def fp4_fp4_mega_moe(y, l1_weights, l2_weights, + shared_l1_weights, shared_l2_weights, x_bf16, + 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, fast_math, l1_alphas, l2_alphas, + a2_scales, routed_scaling_factor): + return _torch_ops.fp4_fp4_mega_moe( + y, l1_weights[0], l1_weights[1], l2_weights[0], l2_weights[1], + shared_l1_weights, shared_l2_weights, x_bf16, + cumulative_local_expert_recv_stats, sym_buffer, + list(sym_buffer_ptrs), rank_idx, num_max_tokens_per_rank, + num_experts, num_topk, list(recipe), activation, + activation_clamp, fast_math, l1_alphas, l2_alphas, + a2_scales, routed_scaling_factor, ) @@ -307,6 +326,7 @@ def bf16_mega_moe(y, l1_weights, l2_weights, shared_l1_weights, shared_l2_weight 'get_symm_buffer_size_for_mega_moe', '_slice_symm_buffer_for_mega_moe', 'fp8_fp4_mega_moe', + 'fp4_fp4_mega_moe', 'bf16_mega_moe', ) From 9b82e2e58db9f8a76bf345fcfd6ceaa64ec96243 Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Fri, 4 Sep 2026 19:05:32 +0000 Subject: [PATCH 26/28] fixed comments and changed torch/python.h to torch/all.h in new csrc/jit_kernels/impls/sm100_fp4_fp4_mega_moe.hpp file Signed-off-by: Chris Leonard --- csrc/apis/mega.hpp | 5 ++--- csrc/jit_kernels/impls/sm100_fp4_fp4_mega_moe.hpp | 2 +- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/csrc/apis/mega.hpp b/csrc/apis/mega.hpp index 2761e70448..c44d283dbe 100644 --- a/csrc/apis/mega.hpp +++ b/csrc/apis/mega.hpp @@ -219,8 +219,8 @@ using SymmBufferSlice = std::tuple +#include #include "../../jit/compiler.hpp" #include "../../jit/kernel_runtime.hpp" From 343826cce1c6ea87ab49ca1daa545755e756fc8d Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Wed, 9 Sep 2026 14:51:52 +0000 Subject: [PATCH 27/28] Migrate SM90 MegaMoE APIs to TORCH_LIBRARY Signed-off-by: Chris Leonard --- csrc/apis/sm90_mega.hpp | 265 +++++++++++++++---- csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp | 2 +- deep_gemm/_C.py | 29 ++ deep_gemm/mega/__init__.py | 5 +- 4 files changed, 248 insertions(+), 53 deletions(-) diff --git a/csrc/apis/sm90_mega.hpp b/csrc/apis/sm90_mega.hpp index a7cee9469d..9ac74b5791 100644 --- a/csrc/apis/sm90_mega.hpp +++ b/csrc/apis/sm90_mega.hpp @@ -1,18 +1,20 @@ #pragma once -#include #include #include #include #include #include -#include + +#include +#include #if DG_TENSORMAP_COMPATIBLE #include "../jit/compiler.hpp" #endif #include "../jit/device_runtime.hpp" #include "../jit_kernels/impls/sm90_fp8_mega_moe.hpp" +#include "../torch_library_utils.hpp" #include "../utils/layout.hpp" #include "../utils/system.hpp" @@ -24,8 +26,63 @@ static int get_token_alignment_for_sm90_mega_moe() { return kSM90MegaMoETokenAlignment; } -static std::tuple(const torch::Tensor&)>> -get_symm_buffer_size_for_sm90_mega_moe( +struct SM90SymmBufferLayoutInfo { + 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 l1_token_base = 0; + int64_t l1_sf_base = 0; + int64_t l2_token_base = 0; + int64_t l2_sf_base = 0; + int num_max_tokens_per_rank = 0; + int num_topk = 0; + int hidden = 0; + int intermediate_hidden = 0; + int num_max_pool_tokens = 0; + int num_max_padded_sf_pool_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, + l1_token_base, l1_sf_base, + l2_token_base, l2_sf_base, + num_max_tokens_per_rank, num_topk, + hidden, intermediate_hidden, + num_max_pool_tokens, num_max_padded_sf_pool_tokens, + }; + } + + static SM90SymmBufferLayoutInfo from_int_list(const std::vector& values) { + DG_HOST_ASSERT(static_cast(values.size()) == 15); + SM90SymmBufferLayoutInfo 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.l1_token_base = values[5]; + info.l1_sf_base = values[6]; + info.l2_token_base = values[7]; + info.l2_sf_base = values[8]; + info.num_max_tokens_per_rank = static_cast(values[9]); + info.num_topk = static_cast(values[10]); + info.hidden = static_cast(values[11]); + info.intermediate_hidden = static_cast(values[12]); + info.num_max_pool_tokens = static_cast(values[13]); + info.num_max_padded_sf_pool_tokens = static_cast(values[14]); + return info; + } +}; + +using SM90SymmBufferSlice = std::tuple< + torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, + torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>; + +static SM90SymmBufferLayoutInfo build_sm90_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, @@ -104,45 +161,76 @@ get_symm_buffer_size_for_sm90_mega_moe( bf16_token_layout, num_topk, num_max_tokens_per_rank, l2_sf_buffer.get_end_ptr()); + SM90SymmBufferLayoutInfo info; + info.num_bytes = reinterpret_cast(combine_token_buffer.get_end_ptr()); + info.input_token_base = reinterpret_cast(input_token_buffer.base); + info.input_sf_base = reinterpret_cast(input_sf_buffer.base); + info.input_topk_idx_base = reinterpret_cast(input_topk_idx_buffer.base); + info.input_topk_weights_base = reinterpret_cast(input_topk_weights_buffer.base); + info.l1_token_base = reinterpret_cast(l1_token_buffer.base); + info.l1_sf_base = reinterpret_cast(l1_sf_buffer.base); + info.l2_token_base = reinterpret_cast(l2_token_buffer.base); + info.l2_sf_base = reinterpret_cast(l2_sf_buffer.base); + info.num_max_tokens_per_rank = num_max_tokens_per_rank; + info.num_topk = num_topk; + info.hidden = hidden; + info.intermediate_hidden = intermediate_hidden; + info.num_max_pool_tokens = num_max_pool_tokens; + info.num_max_padded_sf_pool_tokens = num_max_padded_sf_pool_tokens; + return info; +} + +static std::tuple> get_symm_buffer_size_for_sm90_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 bool& use_fp8_dispatch, const std::string& activation) { + const auto info = build_sm90_symm_buffer_layout( + num_ranks, num_experts, num_max_tokens_per_rank, num_topk, + hidden, intermediate_hidden, use_fp8_dispatch, activation); + return {info.num_bytes, info.to_int_list()}; +} + +static SM90SymmBufferSlice slice_sm90_symm_buffer_from_layout( + const torch::Tensor& buffer, const SM90SymmBufferLayoutInfo& info) { + DG_HOST_ASSERT(buffer.nbytes() >= static_cast(info.num_bytes)); + // `x_sf` is K-major; pool scale factors are M-major. - auto slice_input_buffers = [=](const torch::Tensor& buffer) { - auto x = torch::from_blob( - math::advance_ptr(buffer.data_ptr(), reinterpret_cast(input_token_buffer.base)), - {num_max_tokens_per_rank, hidden}, - torch::TensorOptions().dtype(torch::kFloat8_e4m3fn).device(buffer.device())); - auto x_sf = torch::from_blob( - math::advance_ptr(buffer.data_ptr(), reinterpret_cast(input_sf_buffer.base)), - {num_max_tokens_per_rank, hidden / 128}, - torch::TensorOptions().dtype(torch::kFloat32).device(buffer.device())); - auto topk_idx = torch::from_blob( - math::advance_ptr(buffer.data_ptr(), reinterpret_cast(input_topk_idx_buffer.base)), - {num_max_tokens_per_rank, num_topk}, - torch::TensorOptions().dtype(torch::kInt64).device(buffer.device())); - auto topk_weights = torch::from_blob( - math::advance_ptr(buffer.data_ptr(), reinterpret_cast(input_topk_weights_buffer.base)), - {num_max_tokens_per_rank, num_topk}, - torch::TensorOptions().dtype(torch::kFloat32).device(buffer.device())); - auto l1_acts = torch::from_blob( - math::advance_ptr(buffer.data_ptr(), reinterpret_cast(l1_token_buffer.base)), - {num_max_pool_tokens, hidden}, - torch::TensorOptions().dtype(torch::kFloat8_e4m3fn).device(buffer.device())); - auto l1_acts_sf = torch::from_blob( - math::advance_ptr(buffer.data_ptr(), reinterpret_cast(l1_sf_buffer.base)), - {num_max_padded_sf_pool_tokens, hidden / 128}, - {1, num_max_padded_sf_pool_tokens}, - torch::TensorOptions().dtype(torch::kFloat32).device(buffer.device())); - auto l2_acts = torch::from_blob( - math::advance_ptr(buffer.data_ptr(), reinterpret_cast(l2_token_buffer.base)), - {num_max_pool_tokens, intermediate_hidden}, - torch::TensorOptions().dtype(torch::kFloat8_e4m3fn).device(buffer.device())); - auto l2_acts_sf = torch::from_blob( - math::advance_ptr(buffer.data_ptr(), reinterpret_cast(l2_sf_buffer.base)), - {num_max_padded_sf_pool_tokens, intermediate_hidden / 64}, - {1, num_max_padded_sf_pool_tokens}, - torch::TensorOptions().dtype(torch::kFloat32).device(buffer.device())); - return std::make_tuple(x, x_sf, topk_idx, topk_weights, l1_acts, l1_acts_sf, l2_acts, l2_acts_sf); - }; - return {reinterpret_cast(combine_token_buffer.get_end_ptr()), slice_input_buffers}; + auto x = torch::from_blob( + math::advance_ptr(buffer.data_ptr(), info.input_token_base), + {info.num_max_tokens_per_rank, info.hidden}, + torch::TensorOptions().dtype(torch::kFloat8_e4m3fn).device(buffer.device())); + auto x_sf = torch::from_blob( + math::advance_ptr(buffer.data_ptr(), info.input_sf_base), + {info.num_max_tokens_per_rank, info.hidden / 128}, + torch::TensorOptions().dtype(torch::kFloat32).device(buffer.device())); + auto topk_idx = torch::from_blob( + math::advance_ptr(buffer.data_ptr(), info.input_topk_idx_base), + {info.num_max_tokens_per_rank, info.num_topk}, + torch::TensorOptions().dtype(torch::kInt64).device(buffer.device())); + auto topk_weights = torch::from_blob( + math::advance_ptr(buffer.data_ptr(), info.input_topk_weights_base), + {info.num_max_tokens_per_rank, info.num_topk}, + torch::TensorOptions().dtype(torch::kFloat32).device(buffer.device())); + auto l1_acts = torch::from_blob( + math::advance_ptr(buffer.data_ptr(), info.l1_token_base), + {info.num_max_pool_tokens, info.hidden}, + torch::TensorOptions().dtype(torch::kFloat8_e4m3fn).device(buffer.device())); + auto l1_acts_sf = torch::from_blob( + math::advance_ptr(buffer.data_ptr(), info.l1_sf_base), + {info.num_max_padded_sf_pool_tokens, info.hidden / 128}, + {1, info.num_max_padded_sf_pool_tokens}, + torch::TensorOptions().dtype(torch::kFloat32).device(buffer.device())); + auto l2_acts = torch::from_blob( + math::advance_ptr(buffer.data_ptr(), info.l2_token_base), + {info.num_max_pool_tokens, info.intermediate_hidden}, + torch::TensorOptions().dtype(torch::kFloat8_e4m3fn).device(buffer.device())); + auto l2_acts_sf = torch::from_blob( + math::advance_ptr(buffer.data_ptr(), info.l2_sf_base), + {info.num_max_padded_sf_pool_tokens, info.intermediate_hidden / 64}, + {1, info.num_max_padded_sf_pool_tokens}, + torch::TensorOptions().dtype(torch::kFloat32).device(buffer.device())); + return {x, x_sf, topk_idx, topk_weights, l1_acts, l1_acts_sf, l2_acts, l2_acts_sf}; } // SM90 (Hopper) FP8 MegaMoE entry point. @@ -228,16 +316,17 @@ static void fp8_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_sm90_mega_moe( + const auto layout_info = build_sm90_symm_buffer_layout( num_ranks, num_experts, num_max_tokens_per_rank, num_topk, hidden, intermediate_hidden, true, activation); - 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, l1_acts, l1_acts_sf, l2_acts, l2_acts_sf] = slice(sym_buffer); + const auto [x, x_sf, topk_idx, topk_weights, l1_acts, l1_acts_sf, l2_acts, l2_acts_sf] = + slice_sm90_symm_buffer_from_layout(sym_buffer, layout_info); sm90_fp8_mega_moe(y, l1_acts, l1_acts_sf, @@ -256,12 +345,88 @@ static void fp8_mega_moe( sym_buffer.zero_(); } -static void register_sm90_apis(pybind11::module_& m) { +} // namespace deep_gemm::mega + +namespace deep_gemm::torch_registration { + +using namespace deep_gemm::torch_utils; + +static int64_t get_token_alignment_for_sm90_mega_moe() { + return static_cast(mega::get_token_alignment_for_sm90_mega_moe()); +} + +static std::tuple> get_symm_buffer_size_for_sm90_mega_moe( + const int64_t& num_ranks, const int64_t& num_experts, + const int64_t& num_max_tokens_per_rank, const int64_t& num_topk, + const int64_t& hidden, const int64_t& intermediate_hidden, + const bool& use_fp8_dispatch, const std::string& activation) { + return mega::get_symm_buffer_size_for_sm90_mega_moe( + static_cast(num_ranks), static_cast(num_experts), + static_cast(num_max_tokens_per_rank), static_cast(num_topk), + static_cast(hidden), static_cast(intermediate_hidden), + use_fp8_dispatch, activation); +} + +static mega::SM90SymmBufferSlice _slice_symm_buffer_for_sm90_mega_moe( + const torch::Tensor& buffer, const std::vector& layout_info) { + return mega::slice_sm90_symm_buffer_from_layout( + buffer, mega::SM90SymmBufferLayoutInfo::from_int_list(layout_info)); +} + +static void fp8_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 c10::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 c10::optional& activation_clamp, + const bool& fast_math) { + mega::fp8_mega_moe( + y, + std::make_tuple(l1_weights, l1_weights_sf), + std::make_tuple(l2_weights, l2_weights_sf), + cumulative_local_expert_recv_stats, + sym_buffer, sym_buffer_ptrs, + static_cast(rank_idx), + static_cast(num_max_tokens_per_rank), + static_cast(num_experts), static_cast(num_topk), + list_to_tuple3(recipe), activation, + activation_clamp.has_value() + ? std::make_optional(static_cast(activation_clamp.value())) + : std::nullopt, + fast_math); +} + +} // namespace deep_gemm::torch_registration + +TORCH_LIBRARY_FRAGMENT(deep_gemm, m) { #if DG_TENSORMAP_COMPATIBLE - m.def("get_token_alignment_for_sm90_mega_moe", &get_token_alignment_for_sm90_mega_moe); - m.def("get_symm_buffer_size_for_sm90_mega_moe", &get_symm_buffer_size_for_sm90_mega_moe); - m.def("fp8_mega_moe", &fp8_mega_moe); + m.def( + "get_token_alignment_for_sm90_mega_moe() -> int", + TORCH_FN(deep_gemm::torch_registration::get_token_alignment_for_sm90_mega_moe)); + m.def( + "get_symm_buffer_size_for_sm90_mega_moe(int num_ranks, int num_experts, int num_max_tokens_per_rank, int num_topk, int hidden, int intermediate_hidden, bool use_fp8_dispatch, str activation) -> (int, int[])", + TORCH_FN(deep_gemm::torch_registration::get_symm_buffer_size_for_sm90_mega_moe)); + m.def( + "_slice_symm_buffer_for_sm90_mega_moe(Tensor buffer, int[] layout_info) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor)"); + m.def( + "fp8_mega_moe(Tensor(y!) y, Tensor l1_weights, Tensor l1_weights_sf, Tensor l2_weights, Tensor l2_weights_sf, Tensor(cumulative_local_expert_recv_stats!)? cumulative_local_expert_recv_stats, Tensor(sym_buffer!) 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, bool fast_math) -> ()"); #endif } -} // namespace deep_gemm::mega +TORCH_LIBRARY_IMPL(deep_gemm, CUDA, m) { + using namespace deep_gemm::torch_registration; + +#if DG_TENSORMAP_COMPATIBLE + m.impl( + "_slice_symm_buffer_for_sm90_mega_moe", + TORCH_FN(_slice_symm_buffer_for_sm90_mega_moe)); + m.impl("fp8_mega_moe", TORCH_FN(fp8_mega_moe)); +#endif +} diff --git a/csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp b/csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp index b1b408dc18..0d4129f6cc 100644 --- a/csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp +++ b/csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp @@ -1,6 +1,6 @@ #pragma once -#include +#include #include "../../jit/compiler.hpp" #include "../../jit/kernel_runtime.hpp" #include "../../utils/exception.hpp" diff --git a/deep_gemm/_C.py b/deep_gemm/_C.py index 5561be95ca..2c8bbd6746 100644 --- a/deep_gemm/_C.py +++ b/deep_gemm/_C.py @@ -253,12 +253,23 @@ def _slice_symm_buffer_for_mega_moe(buffer, *args, **kwargs): return _torch_ops._slice_symm_buffer_for_mega_moe(buffer, *args, **kwargs) +def get_symm_buffer_size_for_sm90_mega_moe(*args, **kwargs): + return _torch_ops.get_symm_buffer_size_for_sm90_mega_moe(*args, **kwargs) + + +def _slice_symm_buffer_for_sm90_mega_moe(buffer, *args, **kwargs): + return _torch_ops._slice_symm_buffer_for_sm90_mega_moe(buffer, *args, **kwargs) + + # DG_TENSORMAP_COMPATIBLE — mega.hpp (C++ impl conditional; matches legacy pybind export guard) _bind_guarded_ops( 'get_token_alignment_for_mega_moe', 'get_block_m_for_mega_moe', ) +# DG_TENSORMAP_COMPATIBLE — sm90_mega.hpp +_bind_guarded_ops('get_token_alignment_for_sm90_mega_moe') + def fp8_fp4_mega_moe(y, l1_weights, l2_weights, shared_l1_weights, shared_l2_weights, cumulative_local_expert_recv_stats, sym_buffer, @@ -308,6 +319,20 @@ def bf16_mega_moe(y, l1_weights, l2_weights, shared_l1_weights, shared_l2_weight ) +def fp8_mega_moe(y, l1_weights, l2_weights, + 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, fast_math): + return _torch_ops.fp8_mega_moe( + y, l1_weights[0], l1_weights[1], l2_weights[0], l2_weights[1], + cumulative_local_expert_recv_stats, sym_buffer, + list(sym_buffer_ptrs), rank_idx, num_max_tokens_per_rank, + num_experts, num_topk, list(recipe), activation, + activation_clamp, fast_math, + ) + + _UNCONDITIONAL_API = ( # Runtime 'init', @@ -325,9 +350,12 @@ def bf16_mega_moe(y, l1_weights, l2_weights, shared_l1_weights, shared_l2_weight # Mega MoE (imported via deep_gemm.mega; always defined, fails at call if unregistered) 'get_symm_buffer_size_for_mega_moe', '_slice_symm_buffer_for_mega_moe', + 'get_symm_buffer_size_for_sm90_mega_moe', + '_slice_symm_buffer_for_sm90_mega_moe', 'fp8_fp4_mega_moe', 'fp4_fp4_mega_moe', 'bf16_mega_moe', + 'fp8_mega_moe', ) _DEEP_GEMM_API = ( @@ -372,6 +400,7 @@ def bf16_mega_moe(y, l1_weights, l2_weights, shared_l1_weights, shared_l2_weight # Mega helpers (guarded) 'get_token_alignment_for_mega_moe', 'get_block_m_for_mega_moe', + 'get_token_alignment_for_sm90_mega_moe', ) __all__ = list(_UNCONDITIONAL_API) + [name for name in _DEEP_GEMM_API if name in globals()] diff --git a/deep_gemm/mega/__init__.py b/deep_gemm/mega/__init__.py index 7e826a5745..171cd15cec 100644 --- a/deep_gemm/mega/__init__.py +++ b/deep_gemm/mega/__init__.py @@ -85,7 +85,7 @@ def __init__(self, group: dist.ProcessGroup, self.hidden = hidden self.intermediate_hidden = intermediate_hidden - num_bytes, slice_input_buffers = _C.get_symm_buffer_size_for_sm90_mega_moe( + num_bytes, layout_info = _C.get_symm_buffer_size_for_sm90_mega_moe( group.size(), num_experts, num_max_tokens_per_rank, num_topk, hidden, intermediate_hidden, @@ -105,7 +105,8 @@ def __init__(self, group: dist.ProcessGroup, (self.x, self.x_sf, self.topk_idx, self.topk_weights, self.l1_acts, self.l1_acts_sf, - self.l2_acts, self.l2_acts_sf) = slice_input_buffers(self.buffer) + self.l2_acts, self.l2_acts_sf) = _C._slice_symm_buffer_for_sm90_mega_moe( + self.buffer, layout_info) def destroy(self): self.handle = None From 043dbdab90685351185774bd74cfb4c24c716f4e Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Wed, 9 Sep 2026 15:05:58 +0000 Subject: [PATCH 28/28] removed make_tuple and replaced it with curly brackets {}. The return type is explicit so the make_tuple was unnessesary, and this keeps it consistent. with sm90_mega Signed-off-by: Chris Leonard --- csrc/apis/mega.hpp | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/csrc/apis/mega.hpp b/csrc/apis/mega.hpp index c44d283dbe..46a3f9be33 100644 --- a/csrc/apis/mega.hpp +++ b/csrc/apis/mega.hpp @@ -286,9 +286,9 @@ static SymmBufferSlice slice_symm_buffer_from_layout( {layout_info.num_sf_ring_tokens, intermediate_sf_cols}, {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 {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}; } static void fp8_fp4_mega_moe(