Add cutlass_4_0 as a dependency

.github/workflows/rust.yaml

@@ -9,6 +9,8 @@ jobs:
     steps:
       - uses: actions/checkout@v4
       - uses: dtolnay/rust-toolchain@stable
+        with:
+          components: rustfmt
       - name: Cargo fmt (kernel-abi-check)
         run: |
           ( cd kernel-abi-check/kernel-abi-check && cargo fmt --all -- --check )

build2cmake/src/config/v2.rs

@@ -217,6 +217,8 @@ pub enum Dependencies {
     Cutlass3_8,
     #[serde(rename = "cutlass_3_9")]
     Cutlass3_9,
+    #[serde(rename = "cutlass_4_0")]
+    Cutlass4_0,
     #[serde(rename = "cutlass_sycl")]
     CutlassSycl,
     Torch,

build2cmake/src/torch/cuda.rs

@@ -276,6 +276,17 @@ fn render_deps(env: &Environment, build: &Build, write: &mut impl Write) -> Resu
                     )
                     .wrap_err("Cannot render CUTLASS dependency template")?;
             }
+            Dependencies::Cutlass4_0 => {
+                env.get_template("cuda/dep-cutlass.cmake")
+                    .wrap_err("Cannot get CUTLASS dependency template")?
+                    .render_to_write(
+                        context! {
+                            version => "4.0.0",
+                        },
+                        &mut *write,
+                    )
+                    .wrap_err("Cannot render CUTLASS dependency template")?;
+            }
             Dependencies::Torch => (),
             _ => {
                 eprintln!("Warning: CUDA backend doesn't need/support dependency: {dep:?}");

lib/deps.nix

@@ -25,6 +25,9 @@ let
     "cutlass_3_9" = [
       pkgs.cutlass_3_9
     ];
+    "cutlass_4_0" = [
+      pkgs.cutlass_4_0
+    ];
     "torch" = [
       torch
       torch.cxxdev