Install on ROCm without CUDA wheels, and run the sLSTM kernel on ROCm

CMakeLists.txt

@@ -52,6 +52,26 @@ find_package(pybind11 CONFIG REQUIRED)
 
 # Find CUDA and set up the CUDA language
 find_package(CUDAToolkit QUIET)
+if (CUDAToolkit_FOUND)
+  # torch is only required at build time for the CUDA extension. It is
+  # deliberately not part of build-system.requires (that would pull the CUDA
+  # torch wheel and its nvidia-* dependencies onto ROCm/CPU-only machines), so
+  # probe for it and skip the extension gracefully if it is missing.
+  execute_process(
+      COMMAND "${PYTHON_EXECUTABLE}"
+      -c "import torch;print(torch.utils.cmake_prefix_path)"
+      OUTPUT_VARIABLE TORCH_CMAKE_PREFIX_PATH
+      OUTPUT_STRIP_TRAILING_WHITESPACE
+      RESULT_VARIABLE TORCH_IMPORT_RESULT
+  )
+  if (NOT TORCH_IMPORT_RESULT EQUAL 0)
+    message(WARNING
+      "CUDAToolkit found but torch is not importable in the build environment; "
+      "skipping the sLSTM CUDA extension. Install torch in the build "
+      "environment and build with --no-build-isolation to enable it.")
+    set(CUDAToolkit_FOUND FALSE)
+  endif()
+endif()
 if (CUDAToolkit_FOUND)
   message(STATUS "CUDAToolkit found: ${CUDAToolkit_VERSION}, building with CUDA support")
   if (DEFINED CMAKE_CUDA_ARCHITECTURES)
@@ -83,14 +103,7 @@ if (CUDAToolkit_FOUND)
     set(ENV{TORCH_CUDA_ARCH_LIST} "7.5;8.0;8.6;9.0+PTX")
   endif()
 
-  # Get Torch's CMake package from the build environment used by scikit-build.
-  execute_process(
-      COMMAND "${PYTHON_EXECUTABLE}"
-      -c "import torch;print(torch.utils.cmake_prefix_path)"
-      OUTPUT_VARIABLE TORCH_CMAKE_PREFIX_PATH
-      OUTPUT_STRIP_TRAILING_WHITESPACE
-      COMMAND_ERROR_IS_FATAL ANY
-  )
+  # Torch's CMake package location was probed above.
   list(PREPEND CMAKE_PREFIX_PATH "${TORCH_CMAKE_PREFIX_PATH}")
   find_package(Torch REQUIRED)
   target_include_directories(_slstm PRIVATE ${TORCH_INCLUDE_DIRS})

pyproject.toml

@@ -3,8 +3,13 @@ requires = [
   "scikit-build-core>=0.11",
   "pybind11>=3.0",
   "numpy>=2.0",
-  "torch>=2.0",
 ]
+# NOTE: torch is intentionally NOT a build requirement. It is only needed at
+# build time when the CUDA sLSTM extension is compiled (CUDAToolkit present).
+# Listing it here forces isolated builds (pip/uv default) to download the CUDA
+# torch wheel plus all nvidia-* dependencies, which is wrong on ROCm/CPU-only
+# machines. When building with a CUDA toolchain, install torch into the build
+# environment and use --no-build-isolation.
 build-backend = "scikit_build_core.build"
 
 [project]
@@ -24,18 +29,25 @@ classifiers = [
     "Operating System :: OS Independent",
 ]
 keywords = ["LSTM", "Transformer", "Machine Learning", "Deep Learning", "State Space Models"]
+# torch is NOT a base dependency: its CUDA, ROCm and CPU builds live on
+# separate PyTorch wheel indexes and there is no way to auto-select the right
+# one here, so pinning it would force a CUDA build (and the nvidia-* stack)
+# onto ROCm/CPU machines. Install a backend extra plus the matching index, e.g.
+#   pip install "xlstm[rocm]" --extra-index-url https://download.pytorch.org/whl/rocm6.2
+#   pip install "xlstm[cuda]" --extra-index-url https://download.pytorch.org/whl/cu126
+#   pip install "xlstm[cpu]"  --extra-index-url https://download.pytorch.org/whl/cpu
+# or just have torch already installed for your platform.
 dependencies = [
-    "torch>=2.0",
     "einops",
     "numpy",
-    "opt_einsum", 
+    "opt_einsum",
     "omegaconf",
     "transformers",
     "reportlab",
     "joypy",
     "ipykernel",
     "dacite",
-    "ftfy", 
+    "ftfy",
     "ninja",
     "huggingface-hub",
     "rich",
@@ -45,6 +57,11 @@ dependencies = [
     "mlstm_kernels; python_version >= '3.11'",
 ]
 
+[project.optional-dependencies]
+cuda = ["torch>=2.0"]
+rocm = ["torch>=2.0"]
+cpu  = ["torch>=2.0"]
+
 [tool.uv]
 cache-keys = [
     {file = "pyproject.toml"},
@@ -61,21 +78,6 @@ cache-keys = [
     {file = "blocks/slstm/src/util/*.cuh"},
 ]
 
-[[tool.uv.index]]
-name = "pytorch-cu126"
-url = "https://download.pytorch.org/whl/cu126"
-explicit = true
-
-[[tool.uv.index]]
-name = "pytorch-cu130"
-url = "https://download.pytorch.org/whl/cu130"
-explicit = true
-
-[tool.uv.sources]
-torch = [
-  { index = "pytorch-cu126", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
-  # { index = "pytorch-cu130", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
-]
 
 
 [tool.scikit-build]

xlstm/blocks/slstm/src/cuda_init.py

@@ -26,10 +26,19 @@ def defines_to_cflags(defines=Union[dict[str, Union[int, str]], Sequence[tuple[s
 
 curdir = os.path.dirname(__file__)
 
+# ROCm builds of torch expose the CUDA API through HIP: torch.cuda works, .cu
+# sources are hipified transparently, but the toolchain is hipcc and cuBLAS is
+# hipBLAS, so nvcc-only flags and cublas linkage must be swapped below.
+IS_HIP = torch.version.hip is not None
+
 if torch.cuda.is_available():
     from packaging import version
 
-    if version.parse(torch.__version__) >= version.parse("2.6.0"):
+    if IS_HIP:
+        from torch.utils.cpp_extension import ROCM_HOME
+
+        os.environ["CUDA_LIB"] = os.path.join(ROCM_HOME or "/opt/rocm", "lib")
+    elif version.parse(torch.__version__) >= version.parse("2.6.0"):
         os.environ["CUDA_LIB"] = os.path.join(
             os.path.split(torch.utils.cpp_extension.include_paths(device_type="cuda")[-1])[0], "lib"
         )
@@ -57,15 +66,137 @@ def defines_to_cflags(defines=Union[dict[str, Union[int, str]], Sequence[tuple[s
         )
 
 
+def _hipify_sources(sources):
+    """Translate the sLSTM CUDA sources — and the headers they include — to HIP.
+
+    torch's JIT builder only hipifies the files passed as ``sources``, not the
+    headers they pull in (blas.h, inline_ops*.cuh, ...), so those would keep
+    their cuBLAS / ``__nv_bfloat16`` spellings while the hipified .cu bodies use
+    the HIP ones. Instead copy the whole src tree into a cache dir, hipify
+    everything there, and compile from that copy. The repo sources are never
+    touched.
+    """
+    import shutil
+    from torch.utils.hipify import hipify_python
+    from torch.utils.file_baton import FileBaton
+
+    src_root = os.path.abspath(curdir)
+    out_root = os.environ.get(
+        "XLSTM_HIP_SRC_DIR",
+        os.path.join(
+            os.environ.get("XDG_CACHE_HOME", os.path.expanduser("~/.cache")),
+            "xlstm",
+            "hip_src",
+        ),
+    )
+
+    def _map_source(source):
+        rel = os.path.relpath(os.path.abspath(source), src_root)
+        hip_rel = hipify_python.get_hip_file_path(rel, is_pytorch_extension=True)
+        hip_path = os.path.join(out_root, hip_rel)
+        if not os.path.exists(hip_path):
+            hip_path = os.path.join(out_root, rel)
+        # The pybind glue (.cc) references the HIP fp16/bf16 types for its dtype
+        # dispatch. Those headers only compile under clang, so route the glue
+        # through hipcc by giving it a .hip extension (torch selects the compiler
+        # by suffix); a plain-C++ host compile would fail.
+        if hip_path.endswith((".cc", ".cpp")):
+            hip_path = os.path.splitext(hip_path)[0] + "_glue.hip"
+        return hip_path
+
+    def _produce():
+        for sub in ("cuda", "util"):
+            shutil.copytree(
+                os.path.join(src_root, sub), os.path.join(out_root, sub), dirs_exist_ok=True
+            )
+        hipify_python.hipify(
+            project_directory=out_root,
+            output_directory=out_root,
+            includes=[os.path.join(out_root, "*")],
+            is_pytorch_extension=True,
+            show_detailed=False,
+        )
+        # hipify writes renamed copies (cuda/foo.cu -> hip/foo.hip, blas.h ->
+        # blas_hip.h) and leaves the untranslated originals; overwrite the
+        # originals with the hipified text so stale relative includes resolve.
+        for dirpath, _, filenames in os.walk(out_root):
+            for filename in filenames:
+                path = os.path.join(dirpath, filename)
+                rel = os.path.relpath(path, out_root)
+                hip_rel = hipify_python.get_hip_file_path(rel, is_pytorch_extension=True)
+                hip_path = os.path.join(out_root, hip_rel)
+                if hip_path != path and os.path.exists(hip_path):
+                    shutil.copyfile(hip_path, path)
+
+        # Residual fixups hipify's substitution map does not cover: CUDA-only
+        # driver headers with no HIP counterpart, and the fp16 gemm pointer
+        # (hipify rewrites &cublasHgemm -> &hipblasHgemm, whose hipblasHalf
+        # signature is incompatible with the __half-typed wrapper; point it at
+        # the local cublasHgemm2 wrapper instead, matching the strided path).
+        import re
+
+        _dead_includes = re.compile(
+            r'^\s*#\s*include\s*[<"](?:cuda|cuda_runtime_api|cuda_device_runtime_api)\.h[>"]\s*$',
+            re.MULTILINE,
+        )
+        _text_exts = (".cu", ".cuh", ".cc", ".cpp", ".c", ".h", ".hpp", ".hip")
+        for dirpath, _, filenames in os.walk(out_root):
+            if "__pycache__" in dirpath:
+                continue
+            for filename in filenames:
+                if not filename.endswith(_text_exts):
+                    continue  # skip .pyc and other binaries copied alongside sources
+                path = os.path.join(dirpath, filename)
+                with open(path, "r") as fh:
+                    text = fh.read()
+                new_text = _dead_includes.sub("// [hip] removed CUDA-only include", text)
+                new_text = re.sub(r"&\s*hipblasHgemm\b", "&cublasHgemm2", new_text)
+                # bf16 blas support is gated on CUDART_VERSION, which HIP lacks;
+                # enable the same block on ROCm.
+                new_text = new_text.replace(
+                    "CUDART_VERSION >= 11020",
+                    "(CUDART_VERSION >= 11020 || defined(__HIP_PLATFORM_AMD__))",
+                )
+                if new_text != text:
+                    with open(path, "w") as fh:
+                        fh.write(new_text)
+
+        for source in sources:
+            if source.endswith((".cc", ".cpp")):
+                rel = os.path.relpath(os.path.abspath(source), src_root)
+                hip_rel = hipify_python.get_hip_file_path(rel, is_pytorch_extension=True)
+                orig = os.path.join(out_root, hip_rel)
+                if not os.path.exists(orig):
+                    orig = os.path.join(out_root, rel)
+                shutil.copyfile(orig, _map_source(source))
+
+    # Serialize the shared cache tree across concurrent workers (e.g. several
+    # dataloader / distributed processes initializing the backend at once): the
+    # first to arrive builds it, the rest wait for that build to finish.
+    os.makedirs(os.path.dirname(out_root), exist_ok=True)
+    baton = FileBaton(out_root.rstrip("/") + ".lock")
+    if baton.try_acquire():
+        try:
+            _produce()
+        finally:
+            baton.release()
+    else:
+        baton.wait()
+
+    return [_map_source(s) for s in sources]
+
+
 def load(*, name, sources, extra_cflags=(), extra_cuda_cflags=(), **kwargs):
+    if IS_HIP:
+        sources = _hipify_sources(sources)
     suffix = ""
     for flag in extra_cflags:
         pref = [st[0] for st in flag[2:].split("=")[0].split("_")]
         if len(pref) > 1:
             pref = pref[1:]
         suffix += "".join(pref)
         value = flag[2:].split("=")[1].replace("-", "m").replace(".", "d")
-        value_map = {"float": "f", "__half": "h", "__nv_bfloat16": "b", "true": "1", "false": "0"}
+        value_map = {"float": "f", "__half": "h", "__nv_bfloat16": "b", "__hip_bfloat16": "b", "true": "1", "false": "0"}
         if value in value_map:
             value = value_map[value]
         suffix += value
@@ -86,27 +217,50 @@ def load(*, name, sources, extra_cflags=(), extra_cuda_cflags=(), **kwargs):
         extra_cflags.append("-isystem")
         extra_cflags.append(eip)
 
-    myargs = {
-        "verbose": True,
-        "with_cuda": True,
-        "extra_ldflags": [f"-L{os.environ['CUDA_LIB']}", "-lcublas"],
-        "extra_cflags": [*extra_cflags],
-        "extra_cuda_cflags": [
-            # "-gencode",
-            # "arch=compute_70,code=compute_70",
-            # "-dbg=1",
-            '-Xptxas="-v"',
-            "-gencode",
-            "arch=compute_80,code=compute_80",
-            "-res-usage",
-            "--use_fast_math",
-            "-O3",
-            "-Xptxas -O3",
-            "--extra-device-vectorization",
-            *extra_cflags,
-            *extra_cuda_cflags,
-        ],
-    }
+    if IS_HIP:
+        # hipcc rejects nvcc-only flags (-Xptxas, -gencode, -res-usage, ...).
+        # cuBLAS calls hipify to hipBLAS, so link hipblas; force-include the
+        # compat shim for the enums/intrinsics hipify does not translate.
+        # Force-include the compat shim ONLY on the device (hipcc) pass: it
+        # pulls in hip_bf16.h/hip_fp16.h, which rely on clang builtins and do
+        # not compile under the g++ host compiler used for the pybind glue.
+        compat_header = os.path.join(curdir, "util", "hip_compat.h")
+        myargs = {
+            "verbose": True,
+            "with_cuda": True,
+            "extra_ldflags": [f"-L{os.environ['CUDA_LIB']}", "-lhipblas"],
+            "extra_cflags": [*extra_cflags],
+            "extra_cuda_cflags": [
+                "-O3",
+                "-ffast-math",
+                "-include",
+                compat_header,
+                *extra_cflags,
+                *extra_cuda_cflags,
+            ],
+        }
+    else:
+        myargs = {
+            "verbose": True,
+            "with_cuda": True,
+            "extra_ldflags": [f"-L{os.environ['CUDA_LIB']}", "-lcublas"],
+            "extra_cflags": [*extra_cflags],
+            "extra_cuda_cflags": [
+                # "-gencode",
+                # "arch=compute_70,code=compute_70",
+                # "-dbg=1",
+                '-Xptxas="-v"',
+                "-gencode",
+                "arch=compute_80,code=compute_80",
+                "-res-usage",
+                "--use_fast_math",
+                "-O3",
+                "-Xptxas -O3",
+                "--extra-device-vectorization",
+                *extra_cflags,
+                *extra_cuda_cflags,
+            ],
+        }
     print(myargs)
     myargs.update(**kwargs)
     # add random waiting time to minimize deadlocks because of badly managed multicompile of pytorch ext