leavelet/deepgemm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: support libtorch

Browse Source
This commit is contained in:
FlamingoPg
2025-08-26 19:17:08 -07:00
parent 239112cb4c
commit 301cbc1d75
22 changed files with 652 additions and 245 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
csrc/apis/attention.hpp
View File

@@ -219,22 +219,4 @@ static torch::Tensor fp8_paged_mqa_logits(const torch::Tensor& q,
return logits;
}
static void register_apis(pybind11::module_& m) {
m.def("fp8_gemm_nt_skip_head_mid", &fp8_gemm_nt_skip_head_mid,
py::arg("a"), py::arg("b"), py::arg("d"), py::arg("head_splits"),
py::arg("recipe") = std::nullopt,
py::arg("compiled_dims") = "nk",
py::arg("disable_ue8m0_cast") = false);
m.def("fp8_mqa_logits", &fp8_mqa_logits,
py::arg("q"), py::arg("kv"), py::arg("weights"),
py::arg("cu_seq_len_k_start"), py::arg("cu_seq_len_k_end"),
py::arg("clean_logits") = true);
m.def("get_paged_mqa_logits_metadata", &get_paged_mqa_logits_metadata,
py::arg("context_lens"), py::arg("block_kv"), py::arg("num_sms"));
m.def("fp8_paged_mqa_logits", &fp8_paged_mqa_logits,
py::arg("q"), py::arg("kv_cache"), py::arg("weights"),
py::arg("context_lens"), py::arg("block_table"), py::arg("schedule_meta"),
py::arg("max_context_len"), py::arg("clean_logits") = false);
}
} // namespace deep_gemm::attention

9
csrc/apis/einsum.hpp
View File

@@ -1,8 +1,5 @@
#pragma once
#include <pybind11/pybind11.h>
#include <torch/python.h>
#include "../utils/exception.hpp"
#include "../utils/format.hpp"
#include "../utils/layout.hpp"
@@ -106,10 +103,4 @@ static void einsum(const std::string& expr,
}
}
static void register_apis(pybind11::module_& m) {
m.def("einsum", &einsum,
py::arg("expr"), py::arg("a"), py::arg("b"),
py::arg("d"), py::arg("c") = std::nullopt);
}
} // namespace deep_gemm::einsum

80
csrc/apis/gemm.hpp
View File

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

13
csrc/apis/layout.hpp
View File

@@ -69,17 +69,4 @@ static torch::Tensor transform_k_grouped_sf_into_required_layout(const torch::Te
DG_HOST_UNREACHABLE("Unknown cases");
}
static void register_apis(pybind11::module_& m) {
m.def("transform_sf_into_required_layout", &transform_sf_into_required_layout,
py::arg("sf"), py::arg("mn"), py::arg("k"), py::arg("recipe"),
py::arg("num_groups") = std::nullopt, py::arg("is_sfa") = false,
py::arg("disable_ue8m0_cast") = false);
m.def("get_tma_aligned_size", &get_tma_aligned_size);
m.def("get_mk_alignment_for_contiguous_layout", &get_mk_alignment_for_contiguous_layout);
m.def("get_mn_major_tma_aligned_tensor", &get_mn_major_tma_aligned_tensor);
m.def("get_mn_major_tma_aligned_packed_ue8m0_tensor", &get_mn_major_tma_aligned_packed_ue8m0_tensor);
m.def("get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor", &get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor);
}
} // namespace deep_gemm::layout

20
csrc/apis/runtime.hpp
View File

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

5
csrc/jit/device_runtime.hpp
View File

@@ -17,6 +17,11 @@ class DeviceRuntime {
cublasLtHandle_t cublaslt_handle{};
std::shared_ptr<torch::Tensor> cublaslt_workspace;
// cuBLASLt utils
static constexpr size_t kCublasLtWorkspaceSize = 32 * 1024 * 1024;
cublasLtHandle_t cublaslt_handle{};
std::shared_ptr<torch::Tensor> cublaslt_workspace;
public:
explicit DeviceRuntime() {
cublaslt_workspace = std::make_shared<torch::Tensor>(torch::empty({kCublasLtWorkspaceSize}, dtype(torch::kByte).device(at::kCUDA)));

5
csrc/jit_kernels/impls/runtime_utils.hpp
View File

@@ -1,7 +1,6 @@
#pragma once
#include <cuda.h>
#include <torch/python.h>
#include "../../utils/math.hpp"
#include "../heuristics/sm90.hpp"
@@ -75,10 +74,6 @@ static CUtensorMapSwizzle mode_into_tensor_map_swizzle(const int& mode, const in
}
#endif
DG_HOST_ASSERT(base == 0);
switch (mode) {
case 0:
case 16: return CU_TENSOR_MAP_SWIZZLE_NONE;
case 32: return CU_TENSOR_MAP_SWIZZLE_32B;
case 64: return CU_TENSOR_MAP_SWIZZLE_64B;
case 128: return CU_TENSOR_MAP_SWIZZLE_128B;

2
csrc/jit_kernels/impls/sm100_bf16_gemm.hpp
View File

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

2
csrc/jit_kernels/impls/sm100_bmk_bnk_mn.hpp
View File

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

2
csrc/jit_kernels/impls/sm100_fp8_gemm_1d1d.hpp
View File

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

2
csrc/jit_kernels/impls/sm100_fp8_gemm_1d2d.hpp
View File

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

2
csrc/jit_kernels/impls/sm90_bf16_gemm.hpp
View File

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

2
csrc/jit_kernels/impls/sm90_bmk_bnk_mn.hpp
View File

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

6
csrc/jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp
View File

@@ -1,7 +1,5 @@
#pragma once
#include <torch/python.h>
#include "../../jit/compiler.hpp"
#include "../../jit/device_runtime.hpp"
#include "../../jit/kernel_runtime.hpp"
@@ -133,7 +131,7 @@ static void sm90_fp8_gemm_1d1d(const torch::Tensor& a, const torch::Tensor& sfa,
const auto& code = SM90FP8Gemm1D1DRuntime::generate(args);
const auto& runtime = compiler->build("sm90_fp8_gemm_1d1d", code);
SM90FP8Gemm1D1DRuntime::launch(runtime, args);
MAYBE_LAUNCH(SM90FP8Gemm1D1DRuntime::launch(runtime, args));
}
static void sm90_fp8_k_grouped_gemm_1d1d(const torch::Tensor& a, const torch::Tensor& sfa,
@@ -208,7 +206,7 @@ static void sm90_fp8_k_grouped_gemm_1d1d(const torch::Tensor& a, const torch::Te
const auto& code = SM90FP8Gemm1D1DRuntime::generate(args);
const auto& runtime = compiler->build("sm90_fp8_gemm_1d1d", code);
SM90FP8Gemm1D1DRuntime::launch(runtime, args);
MAYBE_LAUNCH(SM90FP8Gemm1D1DRuntime::launch(runtime, args));
}
} // namespace deep_gemm

2
csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp
View File

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

2
csrc/jit_kernels/impls/smxx_layout.hpp
View File

@@ -1,7 +1,5 @@
#pragma once
#include <torch/python.h>
#include "../../jit/kernel_runtime.hpp"
#include "../../utils/exception.hpp"
#include "../../utils/format.hpp"

521
csrc/python_api.cpp
View File

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

1
csrc/utils/layout.hpp
View File

@@ -1,7 +1,6 @@
#pragma once
#include <cute/arch/mma_sm100_umma.hpp>
#include <torch/python.h>
#include "math.hpp"
#include "exception.hpp"

2
csrc/utils/math.hpp
View File

@@ -1,6 +1,6 @@
#pragma once
#include <torch/python.h>
#include <cstdint>
#include "exception.hpp"

161
deep_gemm/__init__.py
View File

@@ -1,5 +1,7 @@
import os
import subprocess
import torch
import torch.utils.cpp_extension
# Set some default environment provided at setup
try:
@@ -11,62 +13,12 @@ try:
except ImportError:
pass
# Configs
import deep_gemm_cpp
from deep_gemm_cpp import (
set_num_sms,
get_num_sms,
set_tc_util,
get_tc_util,
)
# Kernels
from deep_gemm_cpp import (
# FP8 GEMMs
fp8_gemm_nt, fp8_gemm_nn,
fp8_gemm_tn, fp8_gemm_tt,
fp8_gemm_nt_skip_head_mid,
m_grouped_fp8_gemm_nt_contiguous,
m_grouped_fp8_gemm_nn_contiguous,
m_grouped_fp8_gemm_nt_masked,
k_grouped_fp8_gemm_nt_contiguous,
k_grouped_fp8_gemm_tn_contiguous,
# BF16 GEMMs
bf16_gemm_nt, bf16_gemm_nn,
bf16_gemm_tn, bf16_gemm_tt,
m_grouped_bf16_gemm_nt_contiguous,
m_grouped_bf16_gemm_nt_masked,
# cuBLASLt GEMMs
cublaslt_gemm_nt, cublaslt_gemm_nn,
cublaslt_gemm_tn, cublaslt_gemm_tt,
# Einsum kernels
einsum,
# Attention kernels
fp8_mqa_logits,
get_paged_mqa_logits_metadata,
fp8_paged_mqa_logits,
# Layout kernels
transform_sf_into_required_layout
)
# Some alias for legacy supports
# TODO: remove these later
fp8_m_grouped_gemm_nt_masked = m_grouped_fp8_gemm_nt_masked
bf16_m_grouped_gemm_nt_masked = m_grouped_bf16_gemm_nt_masked
# Some utils
from . import testing
from . import utils
from .utils import *
from . import deep_gemm_cpp # noqa: F401 # Registers ops into torch.ops without touching CUDA
# Initialize CPP modules
def _find_cuda_home() -> str:
# TODO: reuse PyTorch API later
# For some PyTorch versions, the original `_find_cuda_home` will initialize CUDA, which is incompatible with process forks
cuda_home = os.environ.get('CUDA_HOME') or os.environ.get('CUDA_PATH')
if cuda_home is None:
# noinspection PyBroadException
try:
with open(os.devnull, 'w') as devnull:
nvcc = subprocess.check_output(['which', 'nvcc'], stderr=devnull).decode().rstrip('\r\n')
@@ -79,7 +31,106 @@ def _find_cuda_home() -> str:
return cuda_home
deep_gemm_cpp.init(
os.path.dirname(os.path.abspath(__file__)), # Library root directory path
_find_cuda_home() # CUDA home
)
# Lazy runtime init to be fork-safe on Linux (avoid initializing CUDA before fork)
_dg_initialized = False
def _ensure_initialized() -> None:
global _dg_initialized
if _dg_initialized:
return
library_root = os.path.dirname(os.path.abspath(__file__))
torch.ops.deep_gemm.init(library_root, _find_cuda_home())
_dg_initialized = True
def _wrap_op(name: str):
def _fn(*args, **kwargs):
_ensure_initialized()
return getattr(torch.ops.deep_gemm, name)(*args, **kwargs)
return _fn
set_num_sms = _wrap_op('set_num_sms')
get_num_sms = _wrap_op('get_num_sms')
set_tc_util = _wrap_op('set_tc_util')
get_tc_util = _wrap_op('get_tc_util')
fp8_gemm_nt = _wrap_op('fp8_gemm_nt')
fp8_gemm_nn = _wrap_op('fp8_gemm_nn')
fp8_gemm_tn = _wrap_op('fp8_gemm_tn')
fp8_gemm_tt = _wrap_op('fp8_gemm_tt')
m_grouped_fp8_gemm_nt_contiguous = _wrap_op('m_grouped_fp8_gemm_nt_contiguous')
m_grouped_fp8_gemm_nn_contiguous = _wrap_op('m_grouped_fp8_gemm_nn_contiguous')
# Export both canonical name and backward-compat alias
m_grouped_fp8_gemm_nt_masked = _wrap_op('m_grouped_fp8_gemm_nt_masked')
fp8_m_grouped_gemm_nt_masked = m_grouped_fp8_gemm_nt_masked
k_grouped_fp8_gemm_nt_contiguous = _wrap_op('k_grouped_fp8_gemm_nt_contiguous')
k_grouped_fp8_gemm_tn_contiguous = _wrap_op('k_grouped_fp8_gemm_tn_contiguous')
# BF16 GEMMs
bf16_gemm_nt = _wrap_op('bf16_gemm_nt')
bf16_gemm_nn = _wrap_op('bf16_gemm_nn')
bf16_gemm_tn = _wrap_op('bf16_gemm_tn')
bf16_gemm_tt = _wrap_op('bf16_gemm_tt')
m_grouped_bf16_gemm_nt_contiguous = _wrap_op('m_grouped_bf16_gemm_nt_contiguous')
m_grouped_bf16_gemm_nt_masked = _wrap_op('m_grouped_bf16_gemm_nt_masked')
# cuBLASLt GEMMs
cublaslt_gemm_nt = _wrap_op('cublaslt_gemm_nt')
cublaslt_gemm_nn = _wrap_op('cublaslt_gemm_nn')
cublaslt_gemm_tn = _wrap_op('cublaslt_gemm_tn')
cublaslt_gemm_tt = _wrap_op('cublaslt_gemm_tt')
# Attention kernel
fp8_gemm_nt_skip_head_mid = _wrap_op('fp8_gemm_nt_skip_head_mid')
fp8_mqa_logits = _wrap_op('fp8_mqa_logits')
get_paged_mqa_logits_metadata = _wrap_op('get_paged_mqa_logits_metadata')
fp8_paged_mqa_logits = _wrap_op('fp8_paged_mqa_logits')
# Einsum kernel
einsum = _wrap_op('einsum')
# Layout kernels
transform_sf_into_required_layout = _wrap_op('transform_sf_into_required_layout')
# Utility functions
get_tma_aligned_size = _wrap_op('get_tma_aligned_size')
get_mk_alignment_for_contiguous_layout = _wrap_op('get_mk_alignment_for_contiguous_layout')
get_mn_major_tma_aligned_tensor = _wrap_op('get_mn_major_tma_aligned_tensor')
get_mn_major_tma_aligned_packed_ue8m0_tensor = _wrap_op('get_mn_major_tma_aligned_packed_ue8m0_tensor')
get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor = _wrap_op('get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor')
# Some utils
from . import testing
from . import utils
from .utils import *
def _verify_ops_loaded():
expected_ops = [
'init', 'set_num_sms', 'get_num_sms', 'set_tc_util', 'get_tc_util',
'fp8_gemm_nt', 'fp8_gemm_nn', 'fp8_gemm_tn', 'fp8_gemm_tt',
'm_grouped_fp8_gemm_nt_contiguous', 'm_grouped_fp8_gemm_nn_contiguous',
'm_grouped_fp8_gemm_nt_masked', 'k_grouped_fp8_gemm_nt_contiguous',
'k_grouped_fp8_gemm_tn_contiguous',
'transform_sf_into_required_layout', 'get_tma_aligned_size',
'get_mk_alignment_for_contiguous_layout', 'get_mn_major_tma_aligned_tensor',
'get_mn_major_tma_aligned_packed_ue8m0_tensor',
'get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor',
'fp8_gemm_nt_skip_head_mid', 'fp8_mqa_logits',
'get_paged_mqa_logits_metadata', 'fp8_paged_mqa_logits',
'einsum',
'cublaslt_gemm_nt', 'cublaslt_gemm_nn',
'cublaslt_gemm_tn', 'cublaslt_gemm_tt',
]
available_ops = list(torch.ops.deep_gemm.__dict__.keys())
missing_ops = [op for op in expected_ops if op not in available_ops]
if missing_ops:
print(f"Warning: Missing operations: {missing_ops}")
_ensure_initialized()
if __debug__:
_verify_ops_loaded()

22
deep_gemm/utils/layout.py
View File

@@ -1,11 +1,13 @@
from deep_gemm_cpp import (
get_tma_aligned_size,
get_mk_alignment_for_contiguous_layout,
get_mn_major_tma_aligned_tensor,
get_mn_major_tma_aligned_packed_ue8m0_tensor,
get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor
)
import torch
from .. import _ensure_initialized
# Some alias
get_m_alignment_for_contiguous_layout = get_mk_alignment_for_contiguous_layout
get_k_alignment_for_contiguous_layout = get_mk_alignment_for_contiguous_layout
_ensure_initialized()
get_tma_aligned_size = torch.ops.deep_gemm.get_tma_aligned_size
get_mk_alignment_for_contiguous_layout = torch.ops.deep_gemm.get_mk_alignment_for_contiguous_layout
get_mn_major_tma_aligned_tensor = torch.ops.deep_gemm.get_mn_major_tma_aligned_tensor
get_mn_major_tma_aligned_packed_ue8m0_tensor = torch.ops.deep_gemm.get_mn_major_tma_aligned_packed_ue8m0_tensor
get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor = torch.ops.deep_gemm.get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor
get_m_alignment_for_contiguous_layout = torch.ops.deep_gemm.get_mk_alignment_for_contiguous_layout
get_k_alignment_for_contiguous_layout = torch.ops.deep_gemm.get_mk_alignment_for_contiguous_layout

16
tests/test_einsum.py
View File

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