Fix mamba cicd installation

Signed-off-by: Keval Morabia <[email protected]>

Author: kevalmorabia97

.github/workflows/gpu_tests.yml

@@ -76,7 +76,7 @@ jobs:
           echo "LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:/usr/include:/usr/lib/x86_64-linux-gnu:/usr/local/tensorrt/targets/x86_64-linux-gnu/lib" >> $GITHUB_ENV
           echo "PATH=${PATH}:/usr/local/tensorrt/targets/x86_64-linux-gnu/bin" >> $GITHUB_ENV
       - name: Run gpu tests
-        run: pip install tox-current-env && tox -e py312-cuda13-gpu --current-env
+        run: pip install tox-current-env && tox -e gpu --current-env
   gpu-tests-non-pr:
     if: ${{ !startsWith(github.ref, 'refs/heads/pull-request/') }}
     # Runner list at https://github.com/nv-gha-runners/enterprise-runner-configuration/blob/main/docs/runner-groups.md

.gitlab/tests.yml

@@ -47,7 +47,7 @@ multi-gpu:
   script:
     # Use pre-installed packages without a new venv with tox-current-env
     - pip install tox-current-env
-    - tox -e py312-cuda13-gpu --current-env
+    - tox -e gpu --current-env
 
 ##### Example Tests #####
 example-torch:

tox.ini

@@ -60,10 +60,10 @@ commands =
     torch_deploy: python -m pytest tests/unit/torch/deploy
 
 
-########################################################
-# GPU test environments (Can be used with --current-env)
-########################################################
-[testenv:{py310,py311,py312}-cuda13-gpu]
+###########################################################
+# GPU test environments (Should be used with --current-env)
+###########################################################
+[testenv:gpu]
 setenv =
     MAMBA_FORCE_BUILD=TRUE
 commands_pre =
@@ -72,8 +72,9 @@ commands_pre =
     pip install git+https://github.com/Dao-AILab/fast-hadamard-transform.git
 
     # Install Mamba model dependencies (takes 8-10mins!)
-    # Install same triton as pytorch-triton in the NGC PyTorch 25.08 docker otherwise Mamba may install an incompatible version
-    pip install triton==3.3.1
+    # Skip triton because pytorch-triton is installed in the NGC PyTorch containers
+    pip install pip-mark-installed
+    pip install triton
     pip install --no-build-isolation git+https://github.com/state-spaces/mamba.git
 
     # Install Eagle-3 test dependencies