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

Multiple updates and refactorings (#231)

Browse Source
This commit is contained in:
Ray Wang
2025-11-21 17:49:47 +08:00
committed by GitHub
parent bb4424aad4
commit 38f8ef73a4
80 changed files with 3767 additions and 2103 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/utils/compatibility.hpp
+10
View File
@@ -0,0 +1,10 @@
#pragma once
#include <torch/version.h>
#include <cuda.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)
csrc/utils/exception.hpp
+1 -1
View File
@@ -54,7 +54,7 @@ do { \
if (e != CUDA_SUCCESS) { \
std::stringstream ss; \
const char *name, *info; \
cuGetErrorName(e, &name), cuGetErrorString(e, &info); \
lazy_cuGetErrorName(e, &name), lazy_cuGetErrorString(e, &info); \
ss << static_cast<int>(e) << " (" << name << ", " << info << ")"; \
throw DGException("CUDA driver", __FILE__, __LINE__, ss.str()); \
} \
csrc/utils/hash.hpp
+6 -2
View File
@@ -4,7 +4,7 @@
namespace deep_gemm {
static uint64_t fnv1a(const std::string& data, const uint64_t& seed) {
static uint64_t fnv1a(const std::vector<char>& data, const uint64_t& seed) {
uint64_t h = seed;
const uint64_t& prime = 0x100000001b3ull;
for (const char& c: data) {
@@ -14,7 +14,7 @@ static uint64_t fnv1a(const std::string& data, const uint64_t& seed) {
return h;
}
static std::string get_hex_digest(const std::string& data) {
static std::string get_hex_digest(const std::vector<char>& data) {
const auto& state_0 = fnv1a(data, 0xc6a4a7935bd1e995ull);
const auto& state_1 = fnv1a(data, 0x9e3779b97f4a7c15ull);
@@ -32,4 +32,8 @@ static std::string get_hex_digest(const std::string& data) {
return oss.str();
}
static std::string get_hex_digest(const std::string& data) {
return get_hex_digest(std::vector<char>{data.begin(), data.end()});
}
} // namespace deep_gemm
csrc/utils/layout.hpp
+11 -6
View File
@@ -51,7 +51,7 @@ get_default_recipe(const torch::ScalarType& sfa_dtype, const torch::ScalarType&
} else if (arch_major == 10) {
DG_HOST_ASSERT(sfb_dtype == torch::kFloat or sfb_dtype == torch::kInt);
return sfb_dtype == torch::kFloat ?
std::make_tuple(1, 128, 128): // Legacy format or 1D2D kernels
std::make_tuple(1, 128, 128): // Legacy format
std::make_tuple(1, 1, 128); // 1D1D kernels
}
DG_HOST_UNREACHABLE("Unknown recipe");
@@ -63,7 +63,7 @@ static torch::Tensor check_sf_layout(const torch::Tensor& sf,
const int& gran_mn, const int& gran_k,
const std::optional<int>& num_groups,
const bool& tma_stride_check = false,
const bool& contiguous_check = false,
const bool& sm90_sfb_check = false,
const std::optional<torch::ScalarType>& type_check = std::nullopt) {
// Type check
if (type_check.has_value())
@@ -82,13 +82,18 @@ static torch::Tensor check_sf_layout(const torch::Tensor& sf,
if (tma_stride_check) {
if (num_groups.has_value())
DG_HOST_ASSERT(sf.stride(-3) == sf.stride(-1) * sf.size(-1));
DG_HOST_ASSERT(sf.stride(-2) == 1);
// Check contiguity in the MN direction
DG_HOST_ASSERT(sf.stride(-2) == 1 or mn == 1);
DG_HOST_ASSERT(sf.stride(-1) == get_tma_aligned_size(mn, sf.element_size()));
}
// Hopper SFB must be contiguous
if (contiguous_check)
DG_HOST_ASSERT(sf.is_contiguous());
// SM90 SFB must be contiguous, or contiguous after transposing the last two dimensions
if (sm90_sfb_check) {
if (num_groups.has_value())
DG_HOST_ASSERT(sf.stride(-3) == sf.size(-2) * sf.size(-1));
DG_HOST_ASSERT((sf.stride(-1) == 1 and sf.stride(-2) == sf.size(-1)) or
(sf.stride(-1) == sf.size(-2) and sf.stride(-2) == 1));
}
return sf;
}