biondizzle/DeepGEMM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Multiple updates and refactorings (#280)

Browse Source
This commit is contained in:
Zhean Xu
2026-01-16 17:06:52 +08:00
committed by GitHub
parent 3ccf40c53a
commit 0f5f266202
55 changed files with 2706 additions and 891 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
csrc/utils/compatibility.hpp
View File

@@ -2,9 +2,16 @@
#include <torch/version.h>
#include <cuda.h>
#include <cuda_runtime.h>
// `torch::kFloat8_e4m3fn` is supported since PyTorch 2.1
#define DG_FP8_COMPATIBLE (TORCH_VERSION_MAJOR > 2 or (TORCH_VERSION_MAJOR == 2 and TORCH_VERSION_MINOR >= 1))
// `cuTensorMapEncodeTiled` is supported since CUDA Driver API 12.1
#define DG_TENSORMAP_COMPATIBLE (CUDA_VERSION >= 12010)
#define DG_TENSORMAP_COMPATIBLE (CUDA_VERSION >= 12010)
// `cublasGetErrorString` is supported since CUDA Runtime API 11.4.2
#define DG_CUBLAS_GET_ERROR_STRING_COMPATIBLE (CUDART_VERSION >= 11042)
// `CUBLASLT_MATMUL_DESC_FAST_ACCUM` and `CUBLASLT_MATMUL_DESC_SM_COUNT_TARGET` are supported since CUDA Runtime API 11.8
#define DG_CUBLASLT_ADVANCED_FEATURES_COMPATIBLE (CUDART_VERSION >= 11080)

21
csrc/utils/exception.hpp
View File

@@ -5,6 +5,8 @@
#include <string>
#include <sstream>
#include "compatibility.hpp"
namespace deep_gemm {
class DGException final : public std::exception {
@@ -74,6 +76,25 @@ do { \
#endif
#ifndef DG_CUBLASLT_CHECK
#if !DG_CUBLAS_GET_ERROR_STRING_COMPATIBLE
inline const char* cublasGetStatusString(cublasStatus_t status) {
switch(status) {
case CUBLAS_STATUS_SUCCESS: return "CUBLAS_STATUS_SUCCESS";
case CUBLAS_STATUS_NOT_INITIALIZED: return "CUBLAS_STATUS_NOT_INITIALIZED";
case CUBLAS_STATUS_ALLOC_FAILED: return "CUBLAS_STATUS_ALLOC_FAILED";
case CUBLAS_STATUS_INVALID_VALUE: return "CUBLAS_STATUS_INVALID_VALUE";
case CUBLAS_STATUS_ARCH_MISMATCH: return "CUBLAS_STATUS_ARCH_MISMATCH";
case CUBLAS_STATUS_MAPPING_ERROR: return "CUBLAS_STATUS_MAPPING_ERROR";
case CUBLAS_STATUS_EXECUTION_FAILED: return "CUBLAS_STATUS_EXECUTION_FAILED";
case CUBLAS_STATUS_INTERNAL_ERROR: return "CUBLAS_STATUS_INTERNAL_ERROR";
case CUBLAS_STATUS_NOT_SUPPORTED: return "CUBLAS_STATUS_NOT_SUPPORTED";
case CUBLAS_STATUS_LICENSE_ERROR: return "CUBLAS_STATUS_LICENSE_ERROR";
default: return "Unknown cuBLAS error";
}
}
#endif
#define DG_CUBLASLT_CHECK(cmd) \
do { \
const auto& e = (cmd); \

23
csrc/utils/layout.hpp
View File

@@ -36,15 +36,34 @@ static bool fp8_requires_k_major() {
// Tensor utils
template <int N>
static auto get_shape(const torch::Tensor& t) {
DG_HOST_ASSERT(t.dim() == N);
return [&t] <size_t... Is> (std::index_sequence<Is...>) {
return std::make_tuple(static_cast<int>(t.sizes()[Is])...);
}(std::make_index_sequence<N>());
}
static std::tuple<int, int> check_ab_fp8_fp4(const torch::Tensor& ab, const cute::UMMA::Major& major, const int& arch_major) {
auto [mn, k] = get_shape<2>(ab);
if (ab.scalar_type() != torch::kFloat8_e4m3fn) {
DG_HOST_ASSERT(ab.scalar_type() == kPackedFP4 and arch_major == 10);
major == cute::UMMA::Major::K ? (k *= 2) : (mn *= 2);
}
return std::make_tuple(mn, k);
}
static std::tuple<int, int, int> check_grouped_ab_fp8_fp4(const torch::Tensor& ab, const cute::UMMA::Major& major, const int& arch_major) {
auto [num_groups, mn, k] = get_shape<3>(ab);
if (ab.scalar_type() != torch::kFloat8_e4m3fn) {
DG_HOST_ASSERT(ab.scalar_type() == kPackedFP4 and arch_major == 10);
major == cute::UMMA::Major::K ? (k *= 2) : (mn *= 2);
}
return std::make_tuple(num_groups, mn, k);
}
// Recipe
static std::tuple<int, int, int>
get_default_recipe(const torch::ScalarType& sfa_dtype, const torch::ScalarType& sfb_dtype) {
const auto& arch_major = device_runtime->get_arch_major();
const auto arch_major = device_runtime->get_arch_major();
if (arch_major == 9) {
DG_HOST_ASSERT(sfa_dtype == torch::kFloat and sfb_dtype == torch::kFloat);
return {1, 128, 128};
@@ -70,7 +89,7 @@ static torch::Tensor check_sf_layout(const torch::Tensor& sf,
DG_HOST_ASSERT(sf.scalar_type() == type_check.value());
// Always do shape checks
const auto& sf_dtype = sf.scalar_type();
const auto sf_dtype = sf.scalar_type();
DG_HOST_ASSERT(sf_dtype == torch::kFloat or sf_dtype == torch::kInt);
DG_HOST_ASSERT(sf.dim() == static_cast<int>(num_groups.has_value()) + 2);
if (num_groups.has_value())

3
csrc/utils/math.hpp
View File

@@ -6,6 +6,9 @@
namespace deep_gemm {
// TODO: Use `torch::kFloat4_e2m1fn_x2`
constexpr auto kPackedFP4 = torch::kUInt8;
template <typename T>
static T ceil_div(const T& a, const T& b) {
return (a + b - 1) / b;