diff --git a/csrc/apis/runtime.hpp b/csrc/apis/runtime.hpp index 94cbfd674..64a2df271 100644 --- a/csrc/apis/runtime.hpp +++ b/csrc/apis/runtime.hpp @@ -1,13 +1,133 @@ #pragma once +#include +#include +#include + #if DG_TENSORMAP_COMPATIBLE #include "../jit/compiler.hpp" #endif #include "../jit/device_runtime.hpp" #include "../jit_kernels/heuristics/runtime.hpp" +#include "../utils/compatibility.hpp" +#include "../utils/exception.hpp" + +#if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE +#include "../jit_kernels/heuristics/sm90.hpp" +#include "../jit_kernels/heuristics/sm100.hpp" +#include "../jit_kernels/heuristics/common.hpp" +#endif namespace deep_gemm::runtime { +#if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE + +static GemmType parse_gemm_type(const std::string& gemm_type) { + if (gemm_type == "Normal") + return GemmType::Normal; + if (gemm_type == "MGroupedContiguous") + return GemmType::MGroupedContiguous; + if (gemm_type == "MGroupedMasked") + return GemmType::MGroupedMasked; + DG_HOST_UNREACHABLE("Unsupported gemm_type for get_best_gemm_config_key"); +} + +// Fields that determine the JIT kernel identity (see *Runtime::generate_impl). +static std::vector gemm_config_to_key_vec(const GemmDesc& desc, const GemmConfig& config) { + return { + static_cast(desc.gemm_type), + static_cast(desc.kernel_type), + static_cast(desc.get_mma_kind()), + static_cast(desc.a_dtype), + static_cast(desc.b_dtype), + static_cast(desc.cd_dtype), + static_cast(desc.major_a), + static_cast(desc.major_b), + static_cast(desc.with_accumulation), + static_cast(config.layout.swap_ab), + static_cast(config.layout.block_m), + static_cast(config.layout.block_n), + static_cast(config.layout.block_k), + static_cast(config.layout.cluster_m), + static_cast(config.layout.cluster_n), + static_cast(config.pipeline_config.num_stages), + static_cast(config.pipeline_config.smem_size), + static_cast(config.launch_config.num_sms), + static_cast(config.launch_config.num_sms_per_cluster), + static_cast(config.launch_config.num_threads), + static_cast(config.launch_config.num_tma_threads), + static_cast(config.launch_config.num_math_threads), + static_cast(config.launch_config.num_non_epilogue_threads), + static_cast(config.launch_config.num_epilogue_threads), + static_cast(config.storage_config.swizzle_a_mode), + static_cast(config.storage_config.swizzle_b_mode), + static_cast(config.storage_config.swizzle_cd_mode), + static_cast(config.storage_config.load_block_m), + static_cast(config.storage_config.load_block_n), + static_cast(config.storage_config.store_block_m), + static_cast(config.storage_config.store_block_n), + }; +} + +template +static std::vector get_best_gemm_config_key_impl(const GemmDesc& desc) { + const auto& config = get_best_config(desc); + return gemm_config_to_key_vec(desc, config); +} + +static std::vector get_best_gemm_config_key_vec( + const std::string& gemm_type_str, + const std::string& mma, + const int& m, const int& n, const int& k, const int& num_groups) { + const auto gemm_type = parse_gemm_type(gemm_type_str); + const auto arch_major = device_runtime->get_arch_major(); + + KernelType kernel_type; + at::ScalarType a_dtype, b_dtype, cd_dtype; + if (mma == "fp8") { + a_dtype = torch::kFloat8_e4m3fn; + b_dtype = torch::kFloat8_e4m3fn; + cd_dtype = torch::kBFloat16; + // SM90 FP8 NT uses 1D2D; SM100 FP8 uses 1D1D (matches gemm.hpp dispatch). + kernel_type = (arch_major == 9) ? KernelType::Kernel1D2D : KernelType::Kernel1D1D; + } else if (mma == "bf16") { + a_dtype = torch::kBFloat16; + b_dtype = torch::kBFloat16; + cd_dtype = torch::kBFloat16; + kernel_type = KernelType::KernelNoSF; + } else if (mma == "bf16_fp32") { + a_dtype = torch::kBFloat16; + b_dtype = torch::kBFloat16; + cd_dtype = torch::kFloat; + kernel_type = KernelType::KernelNoSF; + } else { + DG_HOST_UNREACHABLE("Unsupported mma for get_best_gemm_config_key; expected fp8|bf16|bf16_fp32"); + } + + // For masked GEMM, callers pass expected_m as `m` (same as runtime). + const auto desc = GemmDesc { + .gemm_type = gemm_type, + .kernel_type = kernel_type, + .m = m, .n = n, .k = k, .num_groups = num_groups, + .a_dtype = a_dtype, .b_dtype = b_dtype, .cd_dtype = cd_dtype, + .major_a = cute::UMMA::Major::K, .major_b = cute::UMMA::Major::K, + .with_accumulation = false, + .num_sms = device_runtime->get_num_sms(), + .tc_util = device_runtime->get_tc_util(), + .compiled_dims = "", + }; + + if (arch_major == 9) { + return get_best_gemm_config_key_impl(desc); + } + if (arch_major == 10) { + return get_best_gemm_config_key_impl(desc); + } + DG_HOST_UNREACHABLE("Unsupported architecture for get_best_gemm_config_key"); +} + +#endif + #if 0 static void register_apis(pybind11::module_& m) { @@ -41,6 +161,24 @@ static void register_apis(pybind11::module_& m) { heuristics_runtime->set_block_size_multiple_of(x, y); } }); +#if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE + m.def("get_best_gemm_config_key", + [](const std::string& gemm_type, const std::string& mma, + const int& m, const int& n, const int& k, const int& num_groups) { + const auto key = get_best_gemm_config_key_vec(gemm_type, mma, m, n, k, num_groups); + pybind11::tuple result(key.size()); + for (size_t i = 0; i < key.size(); ++i) + result[i] = key[i]; + return result; + }, + pybind11::arg("gemm_type"), + pybind11::arg("mma"), + pybind11::arg("m"), + pybind11::arg("n"), + pybind11::arg("k"), + pybind11::arg("num_groups") = 1, + "Return a hashable tuple identifying the best JIT GEMM config for (gemm_type, mma, m, n, k, num_groups)."); +#endif m.def("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); diff --git a/csrc/tvm_ffi_api.cpp b/csrc/tvm_ffi_api.cpp index bee2310a7..0a1ef2811 100644 --- a/csrc/tvm_ffi_api.cpp +++ b/csrc/tvm_ffi_api.cpp @@ -18,6 +18,7 @@ #include "apis/gemm.hpp" #include "apis/layout.hpp" #include "apis/mega.hpp" +#include "apis/runtime.hpp" #include "apis/sm90_mega.hpp" #include "utils/torch_compat.hpp" @@ -65,6 +66,20 @@ TVM_FFI_DLL_EXPORT_TYPED_FUNC(set_tc_util, dg_set_tc_util); TVM_FFI_DLL_EXPORT_TYPED_FUNC(get_pdl, dg_get_pdl); TVM_FFI_DLL_EXPORT_TYPED_FUNC(set_pdl, dg_set_pdl); +#if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE +// Return a hashable Array identifying the best JIT GEMM config. +tvm::ffi::Array dg_get_best_gemm_config_key( + std::string gemm_type, std::string mma, + int64_t m, int64_t n, int64_t k, int64_t num_groups) { + auto key = deep_gemm::runtime::get_best_gemm_config_key_vec( + gemm_type, mma, + static_cast(m), static_cast(n), static_cast(k), + static_cast(num_groups)); + return tvm::ffi::Array(key); +} +TVM_FFI_DLL_EXPORT_TYPED_FUNC(get_best_gemm_config_key, dg_get_best_gemm_config_key); +#endif + // --------------------------------------------------------------------------- // Layout utilities // --------------------------------------------------------------------------- diff --git a/deep_gemm/__init__.py b/deep_gemm/__init__.py index c782e09e1..1402d38ff 100644 --- a/deep_gemm/__init__.py +++ b/deep_gemm/__init__.py @@ -25,6 +25,17 @@ get_pdl, ) +# Optional: present when built with FP8 + tensormap support (CUDA >= 12.1). +# Wrap as tuple: tvm_ffi.Array hashes by identity, so dict/set dedup would fail. +try: + _get_best_gemm_config_key = _C.get_best_gemm_config_key + + def get_best_gemm_config_key(*args, **kwargs): + return tuple(_get_best_gemm_config_key(*args, **kwargs)) +except AttributeError as e: + print(f"Could not import get_best_gemm_config_key. Error: {e}") + pass + # cuBLASLt Kernels from ._C import ( cublaslt_gemm_nt, cublaslt_gemm_nn, diff --git a/sgl_deep_gemm/__init__.py b/sgl_deep_gemm/__init__.py index c3e73931e..658819cf3 100644 --- a/sgl_deep_gemm/__init__.py +++ b/sgl_deep_gemm/__init__.py @@ -119,6 +119,17 @@ def _load_module() -> Module: set_pdl = _C.set_pdl get_pdl = _C.get_pdl +# Optional: present when built with FP8 + tensormap support (CUDA >= 12.1). +# Wrap as tuple: tvm_ffi.Array hashes by identity, so dict/set dedup would fail. +try: + _get_best_gemm_config_key = _C.get_best_gemm_config_key + + def get_best_gemm_config_key(*args, **kwargs): + return tuple(_get_best_gemm_config_key(*args, **kwargs)) +except AttributeError as e: + print(f"Could not import get_best_gemm_config_key. Error: {e}") + pass + # cuBLASLt Kernels def cublaslt_gemm_nt(a, b, d, c=None): _C.cublaslt_gemm_nt(a, b, d, c)