fix

Signed-off-by: oliver könig <okoenig@nvidia.com>

Author: ko3n1g

.github/scripts/build.sh

@@ -20,16 +20,10 @@ EXIT_CODE=0
 timeout 5h python setup.py bdist_wheel --dist-dir=dist || EXIT_CODE=$?
 
 if [ $EXIT_CODE -eq 0 ]; then
-    MATRIX_TORCH_VERSION=$(echo "$MATRIX_TORCH_VERSION" | tr '+' '_')
-    tmpname=cu${WHEEL_CUDA_VERSION}torch${MATRIX_TORCH_VERSION}
-
-    if [ -n "$CXX11_ABI" ]; then
-        tmpname="${tmpname}_cxx11abi${CXX11_ABI}"
-    fi
-
-    wheel_name=$(ls dist/*whl | xargs -n 1 basename | sed "s/-/+$tmpname-/2")
-    ls dist/*whl |xargs -I {} mv {} dist/${wheel_name}
-    echo "wheel_name=${wheel_name}" >> $GITHUB_ENV
+tmpname=cu${WHEEL_CUDA_VERSION}torch${MATRIX_TORCH_VERSION}cxx11abi$CXX11_ABI
+wheel_name=$(ls dist/*whl | xargs -n 1 basename | sed "s/-/+$tmpname-/2")
+ls dist/*whl |xargs -I {} mv {} dist/${wheel_name}
+echo "wheel_name=${wheel_name}" >> $GITHUB_ENV
 fi
 
 echo $EXIT_CODE

.github/workflows/_build_in_container.yml

@@ -75,9 +75,10 @@ jobs:
           set -eoxu pipefail
 
           echo "MATRIX_CUDA_VERSION=$(echo $CUDA_VERSION | awk -F \. {'print $1 $2'})" >> $GITHUB_ENV
-          echo "MATRIX_TORCH_VERSION=$PYTORCH_VERSION" >> $GITHUB_ENV
+          echo "MATRIX_TORCH_VERSION=$NVIDIA_PYTORCH_VERSION" >> $GITHUB_ENV
           echo "WHEEL_CUDA_VERSION=$(echo $CUDA_VERSION | awk -F \. {'print $1'})" >> $GITHUB_ENV
           echo "MATRIX_PYTHON_VERSION=$(python -c "import sys; print('{}.{}'.format(sys.version_info[0], sys.version_info[1]))" | awk -F \. {'print $1 $2'})" >> $GITHUB_ENV
+          echo "CXX11_ABI=$(python -c 'import torch; print(str(torch._C._GLIBCXX_USE_CXX11_ABI).upper())')" >> $GITHUB_ENV
 
           cat $GITHUB_ENV
           EOF
@@ -94,7 +95,7 @@ jobs:
       - name: Build wheel
         id: build_wheel
         env: 
-          CXX11_ABI: ${{ inputs.cxx11_abi }}
+          CXX11_ABI: ${{ env.CXX11_ABI }}
           MATRIX_TORCH_VERSION: ${{ env.MATRIX_TORCH_VERSION}}
           WHEEL_CUDA_VERSION: ${{ env.WHEEL_CUDA_VERSION }}
           MATRIX_PYTHON_VERSION: ${{ env.MATRIX_PYTHON_VERSION }}

setup.py

@@ -303,17 +303,13 @@ def get_wheel_url():
     platform_name = get_platform()
     mamba_ssm_version = get_package_version()
     if os.environ.get("NVIDIA_PRODUCT_NAME", "") == "PyTorch":
-        torch_version = os.environ.get("PYTORCH_BUILD_VERSION").replace("+", "_")
+        torch_version = str(os.environ.get("NVIDIA_PYTORCH_VERSION"))
     else:
         torch_version = f"{torch_version_raw.major}.{torch_version_raw.minor}"
     cxx11_abi = str(torch._C._GLIBCXX_USE_CXX11_ABI).upper()
 
     # Determine wheel URL based on CUDA version, torch version, python version and OS
-    if os.environ.get("NVIDIA_PRODUCT_NAME", "") == "PyTorch": 
-        wheel_filename = f"{PACKAGE_NAME}-{mamba_ssm_version}+{cuda_or_hip}{gpu_compute_version}torch{torch_version}-{python_version}-{python_version}-{platform_name}.whl"
-    else:
-        wheel_filename = f"{PACKAGE_NAME}-{mamba_ssm_version}+{cuda_or_hip}{gpu_compute_version}torch{torch_version}cxx11abi{cxx11_abi}-{python_version}-{python_version}-{platform_name}.whl"
-    
+    wheel_filename = f"{PACKAGE_NAME}-{mamba_ssm_version}+{cuda_or_hip}{gpu_compute_version}torch{torch_version}cxx11abi{cxx11_abi}-{python_version}-{python_version}-{platform_name}.whl"
     wheel_url = BASE_WHEEL_URL.format(
         tag_name=f"v{mamba_ssm_version}", wheel_name=wheel_filename
     )