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
@@ -128,7 +128,7 @@ static void sm90_fp8_gemm_1d2d(const torch::Tensor& a, const torch::Tensor& sfa,
};
const auto& code = SM90FP8Gemm1D2DRuntime::generate(args);
const auto& runtime = compiler->build("sm90_fp8_gemm_1d2d", code);
SM90FP8Gemm1D2DRuntime::launch(runtime, args);
MAYBE_LAUNCH(SM90FP8Gemm1D2DRuntime::launch(runtime, args));
}
static void sm90_m_grouped_fp8_gemm_contiguous_1d2d(const torch::Tensor& a, const torch::Tensor& sfa,
@@ -188,7 +188,7 @@ static void sm90_m_grouped_fp8_gemm_contiguous_1d2d(const torch::Tensor& a, cons
};
const auto& code = SM90FP8Gemm1D2DRuntime::generate(args);
const auto& runtime = compiler->build("sm90_m_grouped_fp8_gemm_contiguous_1d2d", code);
SM90FP8Gemm1D2DRuntime::launch(runtime, args);
MAYBE_LAUNCH(SM90FP8Gemm1D2DRuntime::launch(runtime, args));
}
static void sm90_m_grouped_fp8_gemm_masked_1d2d(const torch::Tensor& a, const torch::Tensor& sfa,
@@ -249,7 +249,7 @@ static void sm90_m_grouped_fp8_gemm_masked_1d2d(const torch::Tensor& a, const to
};
const auto& code = SM90FP8Gemm1D2DRuntime::generate(args);
const auto& runtime = compiler->build("sm90_fp8_m_grouped_gemm_masked_1d2d", code);
SM90FP8Gemm1D2DRuntime::launch(runtime, args);
MAYBE_LAUNCH(SM90FP8Gemm1D2DRuntime::launch(runtime, args));
}
} // namespace deep_gemm