feat: support misc kernel launch

This commit is contained in:
Rain Jiang
2025-10-22 06:23:08 +00:00
committed by FlamingoPg
parent 301cbc1d75
commit f4adba8a66
9 changed files with 52 additions and 24 deletions
+3 -3
View File
@@ -115,7 +115,7 @@ static void sm90_bf16_gemm(const torch::Tensor& a,
};
const auto& code = SM90BF16GemmRuntime::generate(args);
const auto& runtime = compiler->build("sm90_bf16_gemm", code);
SM90BF16GemmRuntime::launch(runtime, args);
MAYBE_LAUNCH(SM90BF16GemmRuntime::launch(runtime, args));
}
static void sm90_m_grouped_bf16_gemm_contiguous(const torch::Tensor& a,
@@ -168,7 +168,7 @@ static void sm90_m_grouped_bf16_gemm_contiguous(const torch::Tensor& a,
};
const auto& code = SM90BF16GemmRuntime::generate(args);
const auto& runtime = compiler->build("sm90_m_grouped_bf16_gemm_contiguous", code);
SM90BF16GemmRuntime::launch(runtime, args);
MAYBE_LAUNCH(SM90BF16GemmRuntime::launch(runtime, args));
}
static void sm90_bf16_m_grouped_gemm_masked(const torch::Tensor& a,
@@ -222,7 +222,7 @@ static void sm90_bf16_m_grouped_gemm_masked(const torch::Tensor& a,
};
const auto& code = SM90BF16GemmRuntime::generate(args);
const auto& runtime = compiler->build("sm90_bf16_m_grouped_gemm_masked", code);
SM90BF16GemmRuntime::launch(runtime, args);
MAYBE_LAUNCH(SM90BF16GemmRuntime::launch(runtime, args));
}
} // namespace deep_gemm