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

[Public release 26/04] Introducing Mega MoE, FP4 Indexer and other features/fixes (#304)

Browse Source 
* Merge with private repo

* Update README

* Update README

* Update README

* Add PyTorch requirements

* Fix sync scopes for MQA logits (#256)

* Update README
This commit is contained in:
Chenggang Zhao
2026-04-17 09:45:14 +08:00
committed by GitHub
parent d30fc36c8f
commit 7f2a703ed5
109 changed files with 12101 additions and 3219 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
csrc/utils/exception.hpp
View File

@@ -42,7 +42,7 @@ do { \
#ifndef DG_NVRTC_CHECK
#define DG_NVRTC_CHECK(cmd) \
do { \
const auto& e = (cmd); \
const auto e = (cmd); \
if (e != NVRTC_SUCCESS) { \
throw DGException("NVRTC", __FILE__, __LINE__, nvrtcGetErrorString(e)); \
} \
@@ -52,7 +52,7 @@ do { \
#ifndef DG_CUDA_DRIVER_CHECK
#define DG_CUDA_DRIVER_CHECK(cmd) \
do { \
const auto& e = (cmd); \
const auto e = (cmd); \
if (e != CUDA_SUCCESS) { \
std::stringstream ss; \
const char *name, *info; \
@@ -66,7 +66,7 @@ do { \
#ifndef DG_CUDA_RUNTIME_CHECK
#define DG_CUDA_RUNTIME_CHECK(cmd) \
do { \
const auto& e = (cmd); \
const auto e = (cmd); \
if (e != cudaSuccess) { \
std::stringstream ss; \
ss << static_cast<int>(e) << " (" << cudaGetErrorName(e) << ", " << cudaGetErrorString(e) << ")"; \
@@ -97,7 +97,7 @@ inline const char* cublasGetStatusString(cublasStatus_t status) {
#define DG_CUBLASLT_CHECK(cmd) \
do { \
const auto& e = (cmd); \
const auto e = (cmd); \
if (e != CUBLAS_STATUS_SUCCESS) { \
std::ostringstream ss; \
ss << static_cast<int>(e) << " (" << cublasGetStatusString(e) << ")"; \

8
csrc/utils/hash.hpp
View File

@@ -6,7 +6,7 @@ namespace deep_gemm {
static uint64_t fnv1a(const std::vector<char>& data, const uint64_t& seed) {
uint64_t h = seed;
const uint64_t& prime = 0x100000001b3ull;
const uint64_t prime = 0x100000001b3ull;
for (const char& c: data) {
h ^= static_cast<uint8_t>(c);
h *= prime;
@@ -15,11 +15,11 @@ static uint64_t fnv1a(const std::vector<char>& data, const uint64_t& seed) {
}
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);
const auto state_0 = fnv1a(data, 0xc6a4a7935bd1e995ull);
const auto state_1 = fnv1a(data, 0x9e3779b97f4a7c15ull);
// Split-mix 64
const auto& split_mix = [](uint64_t z) {
const auto split_mix = [](uint64_t z) {
z = (z ^ (z >> 30)) * 0xbf58476d1ce4e5b9ull;
z = (z ^ (z >> 27)) * 0x94d049bb133111ebull;
return z ^ (z >> 31);

5
csrc/utils/layout.hpp
View File

@@ -116,9 +116,4 @@ static torch::Tensor check_sf_layout(const torch::Tensor& sf,
return sf;
}
// Value matrix layout
static int get_mk_alignment_for_contiguous_layout() {
return 128;
}
} // namespace deep_gemm

5
csrc/utils/math.hpp
View File

@@ -1,3 +1,4 @@
// TODO: merge this file with `math.cuh` (the device part)
#pragma once
#include <torch/python.h>
@@ -6,8 +7,8 @@
namespace deep_gemm {
// TODO: Use `torch::kFloat4_e2m1fn_x2`
constexpr auto kPackedFP4 = torch::kUInt8;
// TODO: use `torch::kFloat4_e2m1fn_x2`
constexpr auto kPackedFP4 = torch::kInt8;
template <typename T>
static T ceil_div(const T& a, const T& b) {

39
csrc/utils/system.hpp
View File

@@ -16,7 +16,7 @@ namespace deep_gemm {
// ReSharper disable once CppNotAllPathsReturnValue
template <typename dtype_t>
static dtype_t get_env(const std::string& name, const dtype_t& default_value = dtype_t()) {
const auto& c_str = std::getenv(name.c_str());
const auto c_str = std::getenv(name.c_str());
if (c_str == nullptr)
return default_value;
@@ -34,7 +34,7 @@ static dtype_t get_env(const std::string& name, const dtype_t& default_value = d
static std::tuple<int, std::string> call_external_command(std::string command) {
command = command + " 2>&1";
const auto& deleter = [](FILE* f) { if (f) pclose(f); };
const auto deleter = [](FILE* f) { if (f) pclose(f); };
std::unique_ptr<FILE, decltype(deleter)> pipe(popen(command.c_str(), "r"), deleter);
DG_HOST_ASSERT(pipe != nullptr);
@@ -42,7 +42,10 @@ static std::tuple<int, std::string> call_external_command(std::string command) {
std::string output;
while (fgets(buffer.data(), buffer.size(), pipe.get()))
output += buffer.data();
const auto& exit_code = WEXITSTATUS(pclose(pipe.release()));
const auto status = pclose(pipe.release());
// NOTES: if the child was killed by a signal (e.g., SIGINT from Ctrl+C),
// WEXITSTATUS would incorrectly return 0. Treat signal death as failure.
const auto exit_code = WIFEXITED(status) ? WEXITSTATUS(status) : 128 + WTERMSIG(status);
return {exit_code, output};
}
@@ -68,7 +71,7 @@ static std::vector<std::filesystem::path> collect_files(const std::filesystem::p
static std::filesystem::path make_dirs(const std::filesystem::path& path) {
// OK if existed
std::error_code capture;
const bool& created = std::filesystem::create_directories(path, capture);
const bool created = std::filesystem::create_directories(path, capture);
if (not (created or capture.value() == 0)) {
DG_HOST_UNREACHABLE(fmt::format("Failed to make directory: {}, created: {}, value: {}",
path.c_str(), created, capture.value()));
@@ -94,4 +97,32 @@ static std::string get_uuid() {
return ss.str();
}
static void safe_remove_all(const std::filesystem::path& path) {
std::error_code ec;
if (not std::filesystem::exists(path, ec) or ec)
return;
// A single file
if (not std::filesystem::is_directory(path, ec) or ec) {
std::filesystem::remove(path, ec);
return;
}
// Remove directory
auto it = std::filesystem::directory_iterator(path,
std::filesystem::directory_options::skip_permission_denied, ec);
for (auto end = std::filesystem::directory_iterator(); it != end and not ec;) {
const auto entry_path = it->path();
// Increase firstly to avoid failures
it.increment(ec);
if (ec)
break;
// Recursively clean
safe_remove_all(entry_path);
}
std::filesystem::remove(path, ec);
}
} // deep_gemm