1.0.5

[danilyef]

Fixes

1. Fixes sLSTM kernel compilation

Before: compilation of the sLSTM kernel failed on the CUDA version >= 13.0, because of the -static-global-template-stub flag, which defaults to true. For the CUDA version 12.8 the exact same flag defaults to false (which is what we need)

Fix: this flag was introduced only starting from the CUDA version 12.8 and simply inserting it would lead to an error (see example)

In order to fix the issue the _find_cuda_home (from pytorch cpp_extensions) as well as get_cuda_version were introduced. The flag -static-global-template-stub=false is inserted to the args, only if CUDA version is >= 12.8 and CUDA path is found.

Important: Major CUDA Toolkit version used for compilation should match with the CUDA version used by the PyTorch wheel. Mismatched versions (e.g., PyTorch cu128 with CUDA Toolkit 13.0) lead to build issues,

2. Triton installation

Before: flashrnn attempts to install triton unconditionally even for windows and macos.

**Fix:**installs triton only for linux.

3. Home path:

Before: constrint.py path failed on windows

Fix: Path.home() should fix home path for windows (still to be tested)

How to test:

main.py script for testing:

import os

#for cuda home
# os.environ["CUDA_HOME"] = "/home/dan_yef/.local/cuda-nvcc-12.8"
os.environ["CUDA_HOME"] = "/usr/local/cuda-13.0"
# os.environ["CUDA_HOME"] = "/usr/local/cuda-13.3"

#for nvcc
os.environ["PATH"] = f"{os.environ['CUDA_HOME']}/bin:" + os.environ["PATH"]

#cuda libraries
ld_library_path = os.environ.get("LD_LIBRARY_PATH", "")
os.environ["LD_LIBRARY_PATH"] = (
    f"{os.environ['CUDA_HOME']}/lib64"
    + (":" + ld_library_path if ld_library_path else "")
)


import torch
from tirex2 import TimeseriesType, load_model

# load model
model = load_model("NX-AI/TiRex-2", device="cuda")


context = torch.sin(torch.arange(128).float() / 8)
ts = TimeseriesType(target=context.unsqueeze(0), past_covariates=None, future_covariates=None)

forecast = model.forecast([ts], prediction_length=32, output_type="numpy")[0]
print(forecast)

MacOS:

# 1. Go to the root
cd /your/root/folder

# 2. venv on Python 3.11
python3.11 -m venv .venv
source .venv/bin/activate
python -m pip install -U pip

# 3. Install flashrnn from source
git -C flashrnn checkout 1.0.5   
pip install -e ./flashrnn

# 4. Install tirex-2 from source
pip install -e ./tirex-2

# 5. HF_TOKEN
export HF_TOKEN=hf_xxxxxxxxxxxxxxxxx     

# 6. Run main.py from a subfolder 
mkdir -p run
cd run
python main.py

Linux:

# 0. Go to the repo root

# 1. Initialize conda for the current shell
source /home/dan_yef/miniconda3/etc/profile.d/conda.sh

# 2. Conda Env
conda create -y -n tirex311 python=3.11
conda activate tirex311

# 3. Install flashrnn
git -C flashrnn checkout 1.0.5   
pip install -e ./flashrnn

# 4. Install tirex-2 from source
# for cuda 13.0 and above relax torch requirement in pyproject.toml
#  "torch>=2.8"
pip install -e ./tirex-2

# 5. Check PyTorch version (especially cuda wheel):
python -c "import torch; print(torch.__version__)"

# 6. HF_TOKEN
export HF_TOKEN=hf_xxxxxxxxxxxxxxxxx   

# 7. Run main.py from a subfolder 
mkdir -p run
cd run
python main.py

# 8. Deactivate
conda deactivate

# 9. Remove Env
conda env remove -n tirex311

[martinloretzzz]

torch.version.cuda is your friend ;), otherwise looks good.

[danilyef]

torch.version.cuda is your friend ;), otherwise looks good.

The question is how reliable is torch.version.cuda ?
For example, pytorch does a check, if torch.version.cuda is equal the cuda version found by --nvcc

 if cuda_ver != torch_cuda_version:
        # major/minor attributes are only available in setuptools>=49.4.0
        if getattr(cuda_ver, "major", None) is None:
            raise ValueError("setuptools>=49.4.0 is required")
        if cuda_ver.major != torch_cuda_version.major:
            raise RuntimeError(CUDA_MISMATCH_MESSAGE, cuda_str_version, torch.version.cuda)
        logger.warning(CUDA_MISMATCH_WARN, cuda_str_version, torch.version.cuda)

I will leave it like that, I think