Install on ROCm without CUDA wheels, and run the sLSTM kernel on ROCm - #126
Install on ROCm without CUDA wheels, and run the sLSTM kernel on ROCm#126kashif wants to merge 4 commits into
Conversation
Signed-off-by: Kashif Rasul <kashif.rasul@gmail.com>
Signed-off-by: Kashif Rasul <kashif.rasul@gmail.com>
|
All contributors have signed the CLA ✍️ ✅ |
|
I have read the CLA Document and I hereby sign the CLA |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: a442e6e6b7
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 5893467b72
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
Signed-off-by: Kashif Rasul <kashif.rasul@gmail.com>
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 2affde56f3
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| # 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 = [ |
There was a problem hiding this comment.
With torch removed from the base dependency list, the documented/default pip install xlstm path can now install a package that cannot even be imported in a fresh environment: xlstm/__init__.py imports modules such as xlstm/blocks/mlstm/block.py and xlstm/blocks/slstm/block.py that import torch at module load time, so import xlstm raises ModuleNotFoundError. If avoiding a forced CUDA wheel is the goal, the default install path still needs to require an appropriate torch extra or the docs/metadata need to make the base package non-default/unsupported.
Useful? React with 👍 / 👎.
Signed-off-by: Kashif Rasul <kashif.rasul@gmail.com>
1 participant
Two things for AMD/ROCm users.
Installing no longer drags in the CUDA torch wheel and the whole nvidia-* stack on ROCm or CPU only machines. torch is dropped from the build requirements (only the optional CUDA extension needs it at build time) and the forced cu126 index pin is commented out. When there is no CUDA toolkit the extension is skipped gracefully, so the wheel builds fine.
The compiled sLSTM cuda kernel now builds and runs on ROCm through HIP: hipify the sources, force include a small compat header for the cuBLAS enums and intrinsics hipify does not translate, compile the pybind glue with hipcc, and link hipBLAS. The CUDA path is unchanged.
Tested on ROCm 7.2 (gfx1150, torch 2.13+rocm7.2): an editable install pulls no nvidia wheels, and sLSTM with backend=cuda matches the vanilla backend (float32 forward and backward exact, bf16 within rounding).