leavelet/deepgemm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add signal for SBO in SM90 masked gemm.

Browse Source 
Co-authored-by: Zqy11 <841971412@qq.com>
Co-authored-by: AniZpZ <aniz1905@gmail.com>
This commit is contained in:
Eric Wong
2025-11-09 14:53:05 +08:00
parent f4adba8a66
commit 6635dd2ffd
6 changed files with 71 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
csrc/python_api.cpp
View File

@@ -342,17 +342,20 @@ TORCH_LIBRARY(deep_gemm, m) {
deep_gemm_wrappers::m_grouped_fp8_gemm_nn_contiguous_wrapper(a_val, a_scale, b_val, b_scale, d, m_indices, recipe, compiled_dims, disable_ue8m0_cast);
});
m.def(R"(m_grouped_fp8_gemm_nt_masked(Any a, Any b, Tensor d, Tensor masked_m, int expected_m, int[]? recipe=None, str compiled_dims="nk", bool disable_ue8m0_cast=False) -> ())");
m.def(R"(m_grouped_fp8_gemm_nt_masked(Any a, Any b, Tensor d, Tensor masked_m, int expected_m, int[]? recipe=None, str compiled_dims="nk", bool disable_ue8m0_cast=False, int max_block_n=256, bool enable_overlap=False, Tensor? signal=None) -> (int?, int?))");
m.impl("m_grouped_fp8_gemm_nt_masked", torch::kCUDA, [](const c10::IValue& a_input, const c10::IValue& b_input,
const torch::Tensor& d,
const torch::Tensor& masked_m,
int64_t expected_m,
const c10::optional<c10::IntArrayRef>& recipe,
const std::string& compiled_dims,
bool disable_ue8m0_cast) {
bool disable_ue8m0_cast,
int64_t max_block_n,
bool enable_overlap,
const c10::optional<torch::Tensor>& signal) {
auto [a_val, a_scale] = parse_tensor_or_tuple(a_input);
auto [b_val, b_scale] = parse_tensor_or_tuple(b_input);
deep_gemm_wrappers::m_grouped_fp8_gemm_nt_masked_wrapper(a_val, a_scale, b_val, b_scale, d, masked_m, expected_m, recipe, compiled_dims, disable_ue8m0_cast);
return deep_gemm_wrappers::m_grouped_fp8_gemm_nt_masked_wrapper(a_val, a_scale, b_val, b_scale, d, masked_m, expected_m, recipe, compiled_dims, disable_ue8m0_cast, max_block_n, enable_overlap, signal);
});
m.def(R"(k_grouped_fp8_gemm_nt_contiguous(Any a, Any b, Tensor d, int[] ks, Tensor ks_tensor, Tensor? c=None, int[] recipe=[1, 1, 128], str compiled_dims="mn") -> ())");