Fix CUTLASS submodule setup

.gitmodules

-[submodule "csrc/cutlass"]
-	path = csrc/cutlass
-	url = https://github.com/NVIDIA/cutlass.git
+[submodule "csrc/cutlass/cutlass_3.2.1"]
+	path = csrc/cutlass/cutlass_3.2.1
+	url = ssh://git@10.16.1.204:10022/dcutoolkit/deeplearing/cutlass_3.2.1.git
+	branch = feature/16x64-mmac

README.md

@@ -78,6 +78,10 @@ git submodule update --init --recursive
 pip install -v .
 ```
 
+The CUTLASS dependency is pinned as `csrc/cutlass/cutlass_3.2.1` on branch
+`feature/16x64-mmac`. If the submodule is missing, `setup.py` will try to
+initialize it before compiling.
+
 ## Usage
 
 ### MLA Decoding

cutlass @ 147f5673

-Subproject commit 147f5673d0c1c3dcf66f78d677fd647e4a020219

cutlass_3.2.1 @ a46390fe

+Subproject commit a46390fe2220776773f08eaeca8a13f9a76055b6

setup.py

@@ -38,11 +38,22 @@ def get_arch_flags():
 #     # nvcc_threads = os.getenv("NVCC_THREADS") or "32"
 #     return ["--threads", nvcc_threads]
 
-# subprocess.run(["git", "submodule", "update", "--init", "csrc/cutlass"])
-
 this_dir = os.path.dirname(os.path.abspath(__file__))
 os.environ['PYTORCH_NVCC'] = 'aicc'
 
+cutlass_dir = Path(this_dir) / "csrc" / "cutlass" / "cutlass_3.2.1"
+if not (cutlass_dir / "include").exists():
+    subprocess.run(
+        ["git", "submodule", "update", "--init", "--recursive", "csrc/cutlass/cutlass_3.2.1"],
+        cwd=this_dir,
+        check=True,
+    )
+if not (cutlass_dir / "include").exists():
+    raise RuntimeError(
+        "CUTLASS dependency is missing. Run "
+        "`git submodule update --init --recursive csrc/cutlass/cutlass_3.2.1` first."
+    )
+
 if False:
     cxx_args = ["/O2", "/std:c++20", "/DNDEBUG", "/W0"]
 else: