Report CUTLASS setup and CUDA launch-configuration failures

MANIFEST.in

@@ -0,0 +1 @@
+include build_utils.py

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

csrc/smxx/fwd_launch.cu

@@ -2,6 +2,8 @@
 #include "fwd_kernel1.cuh"
 #include "fwd_kernel2.cuh"
 
+#include <c10/cuda/CUDAException.h>
+
 // ==================== launch_fwd ====================
 template <int D, bool HasStateIn, bool HasStateOut, bool StateFP32, bool IsVarlen>
 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);

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:

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))