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-10-21 09:55:33 +00: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
csrc/jit_kernels/impls/runtime_utils.hpp
-5
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;
csrc/jit_kernels/impls/sm100_bf16_gemm.hpp
-2
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"
csrc/jit_kernels/impls/sm100_bmk_bnk_mn.hpp
+1 -3
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"
@@ -134,4 +132,4 @@ static void sm100_bmn_bnk_mn_gemm(const torch::Tensor &a,
SM100BmkBnkMnRuntime::launch(runtime, args);
}
} // namespace deep_gemm
} // namespace deep_gemm
csrc/jit_kernels/impls/sm100_fp8_gemm_1d1d.hpp
-2
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"
csrc/jit_kernels/impls/sm100_fp8_gemm_1d2d.hpp
-2
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"
csrc/jit_kernels/impls/sm90_bf16_gemm.hpp
-2
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"
csrc/jit_kernels/impls/sm90_bmk_bnk_mn.hpp
-2
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"
csrc/jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp
+2 -4
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
csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp
-2
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"
csrc/jit_kernels/impls/smxx_layout.hpp
-2
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"