diff --git a/MANIFEST.in b/MANIFEST.in new file mode 100644 index 0000000..71c9c73 --- /dev/null +++ b/MANIFEST.in @@ -0,0 +1 @@ +include build_utils.py diff --git a/build_utils.py b/build_utils.py new file mode 100644 index 0000000..ff7dcfb --- /dev/null +++ b/build_utils.py @@ -0,0 +1,23 @@ +import subprocess + + +def update_cutlass_submodule(repo_dir: str) -> None: + command = ["git", "submodule", "update", "--init", "cutlass"] + try: + subprocess.run(command, cwd=repo_dir, check=True) + except FileNotFoundError as exc: + raise RuntimeError( + "Failed to initialize the CUTLASS submodule because Git was not found. " + "Install Git and run `git submodule update --init --recursive`." + ) from exc + except (OSError, subprocess.CalledProcessError) as exc: + detail = ( + f" (Git exited with status {exc.returncode})" + if isinstance(exc, subprocess.CalledProcessError) + else "" + ) + raise RuntimeError( + "Failed to initialize the CUTLASS submodule" + f"{detail}. Run `git submodule update --init --recursive` " + "from the FlashKDA repository and retry the installation." + ) from exc diff --git a/csrc/smxx/fwd_launch.cu b/csrc/smxx/fwd_launch.cu index 74f7ee4..ba96b0f 100644 --- a/csrc/smxx/fwd_launch.cu +++ b/csrc/smxx/fwd_launch.cu @@ -2,6 +2,8 @@ #include "fwd_kernel1.cuh" #include "fwd_kernel2.cuh" +#include + // ==================== launch_fwd ==================== template void launch_fwd( @@ -157,7 +159,8 @@ void launch_fwd( CHUNK, D, kK1Threads, IsVarlen >; - cudaFuncSetAttribute(kernel1, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size_k1); + C10_CUDA_CHECK(cudaFuncSetAttribute( + kernel1, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size_k1)); dim3 grid_k1(total_tiles, H); dim3 block_k1(kK1Threads); @@ -191,7 +194,8 @@ void launch_fwd( HasStateIn, HasStateOut, StateFP32, IsVarlen >; - cudaFuncSetAttribute(kernel2, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size_k2); + C10_CUDA_CHECK(cudaFuncSetAttribute( + kernel2, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size_k2)); dim3 grid_k2(N, H); dim3 block_k2(kK2Threads); diff --git a/setup.py b/setup.py index 76e44ed..756640a 100644 --- a/setup.py +++ b/setup.py @@ -3,8 +3,10 @@ from setuptools import setup from torch.utils.cpp_extension import CUDAExtension, BuildExtension, CUDA_HOME +from build_utils import update_cutlass_submodule + this_dir = os.path.dirname(os.path.abspath(__file__)) -subprocess.run(["git", "submodule", "update", "--init", "cutlass"]) +update_cutlass_submodule(this_dir) def is_flag_set(flag: str) -> bool: diff --git a/tests/test_build_utils.py b/tests/test_build_utils.py new file mode 100644 index 0000000..3cd7917 --- /dev/null +++ b/tests/test_build_utils.py @@ -0,0 +1,41 @@ +import subprocess +from unittest import mock + +import pytest + +from build_utils import update_cutlass_submodule + + +def test_update_cutlass_submodule_runs_from_repository(tmp_path): + with mock.patch("build_utils.subprocess.run") as run: + update_cutlass_submodule(str(tmp_path)) + + run.assert_called_once_with( + ["git", "submodule", "update", "--init", "cutlass"], + cwd=str(tmp_path), + check=True, + ) + + +def test_update_cutlass_submodule_reports_git_failure(tmp_path): + error = subprocess.CalledProcessError(128, ["git", "submodule"]) + + with mock.patch("build_utils.subprocess.run", side_effect=error): + with pytest.raises( + RuntimeError, + match=r"Failed to initialize the CUTLASS submodule " + r"\(Git exited with status 128\).*git submodule update", + ): + update_cutlass_submodule(str(tmp_path)) + + +def test_update_cutlass_submodule_reports_missing_git(tmp_path): + with mock.patch( + "build_utils.subprocess.run", + side_effect=FileNotFoundError("git"), + ): + with pytest.raises( + RuntimeError, + match=r"CUTLASS submodule because Git was not found", + ): + update_cutlass_submodule(str(tmp_path))