Add missing dependencies in bzlmod.

PiperOrigin-RevId: 872956195

Author: ethanluoyc

MODULE.bazel

@@ -91,6 +91,7 @@ use_repo(
     "XNNPACK",
     "cpuinfo",
     "cudnn_frontend_archive",
+    "cutlass_archive",
     "dlpack",
     "ducc",
     "eigen_archive",
@@ -108,6 +109,7 @@ use_repo(
     "nvshmem",
     "onednn",
     "pthreadpool",
+    "raft",
     "riegeli",
     "rocm_device_libs",
     "shardy",

third_party/cutlass/BUILD.bazel

@@ -0,0 +1,10 @@
+load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls")
+
+def repo():
+    tf_http_archive(
+        name = "cutlass_archive",
+        build_file = "//third_party:cutlass.BUILD",
+        sha256 = "a7739ca3dc74e3a5cb57f93fc95224c5e2a3c2dff2c16bb09a5e459463604c08",
+        strip_prefix = "cutlass-3.8.0",
+        urls = tf_mirror_urls("https://github.com/NVIDIA/cutlass/archive/refs/tags/v3.8.0.zip"),
+    )

third_party/cutlass/workspace.bzl

@@ -4,6 +4,7 @@ load("//third_party/benchmark:workspace.bzl", benchmark = "repo")
 load("//third_party/brotli:workspace.bzl", brotli = "repo")
 load("//third_party/cpuinfo:workspace.bzl", cpuinfo = "repo")
 load("//third_party/cudnn_frontend:workspace.bzl", cudnn_frontend = "repo")
+load("//third_party/cutlass:workspace.bzl", cutlass = "repo")
 load("//third_party/dlpack:workspace.bzl", dlpack = "repo")
 load("//third_party/ducc:workspace.bzl", ducc = "repo")
 load("//third_party/eigen3:workspace.bzl", eigen3 = "repo")
@@ -84,6 +85,7 @@ def _third_party_ext_impl(mctx):  # @unused
     uv()
     xnnpack()
     xxd()
+    cutlass()
 
 third_party_ext = module_extension(
     implementation = _third_party_ext_impl,

third_party/extensions/third_party.bzl

@@ -14,6 +14,7 @@ load("//third_party/brotli:workspace.bzl", brotli = "repo")
 load("//third_party/clang_toolchain:cc_configure_clang.bzl", "cc_download_clang_toolchain")
 load("//third_party/cpuinfo:workspace.bzl", cpuinfo = "repo")
 load("//third_party/cudnn_frontend:workspace.bzl", cudnn_frontend = "repo")
+load("//third_party/cutlass:workspace.bzl", cutlass = "repo")
 load("//third_party/dlpack:workspace.bzl", dlpack = "repo")
 load("//third_party/ducc:workspace.bzl", ducc = "repo")
 load("//third_party/eigen3:workspace.bzl", eigen3 = "repo")
@@ -113,6 +114,7 @@ def _initialize_third_party():
     uv()
     xnnpack()
     xxd()
+    cutlass()
 
     # copybara: tsl vendor
 
@@ -405,14 +407,6 @@ def _tf_repositories():
         urls = tf_mirror_urls("https://github.com/google/snappy/archive/refs/tags/1.2.1.tar.gz"),
     )
 
-    tf_http_archive(
-        name = "cutlass_archive",
-        build_file = "//third_party:cutlass.BUILD",
-        sha256 = "a7739ca3dc74e3a5cb57f93fc95224c5e2a3c2dff2c16bb09a5e459463604c08",
-        strip_prefix = "cutlass-3.8.0",
-        urls = tf_mirror_urls("https://github.com/NVIDIA/cutlass/archive/refs/tags/v3.8.0.zip"),
-    )
-
     tf_http_archive(
         name = "nccl_archive",
         build_file = "//third_party/nccl:archive.BUILD",