fix: auto-detect CUDA version for nvidia-cutlass-dsl extra and clean install

- requirements.txt: remove [cu13] extra (keep generic for Docker builds)
- setup_test_env.sh: detect CUDA major version from torch to select
  [cu12] or [cu13] extra, and clean uninstall old packages before
  installing (per NVIDIA docs recommendation)

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: limin2021

requirements.txt

@@ -5,7 +5,7 @@ einops
 ninja
 numpy
 nvidia-cudnn-frontend>=1.13.0
-nvidia-cutlass-dsl[cu13]>=4.4.2
+nvidia-cutlass-dsl>=4.4.2
 nvidia-ml-py
 packaging>=24.2
 requests

scripts/setup_test_env.sh

@@ -26,10 +26,19 @@ fi
 
 # Override nvidia-cutlass-dsl if specified
 if [ -n "${CUTLASS_DSL_VERSION:-}" ]; then
+  # Detect CUDA major version to select the correct extra (cu12 or cu13)
+  CUDA_MAJOR=$(python -c "import torch; print(torch.version.cuda.split('.')[0])" 2>/dev/null || echo "12")
+  if [ "$CUDA_MAJOR" = "13" ]; then
+    CUTLASS_DSL_EXTRA="cu13"
+  else
+    CUTLASS_DSL_EXTRA="cu12"
+  fi
   echo "========================================"
-  echo "Overriding nvidia-cutlass-dsl with version: ${CUTLASS_DSL_VERSION}"
+  echo "Overriding nvidia-cutlass-dsl with version: ${CUTLASS_DSL_VERSION} [${CUTLASS_DSL_EXTRA}]"
   echo "========================================"
-  pip install --force-reinstall "nvidia-cutlass-dsl[cu13]==${CUTLASS_DSL_VERSION}"
+  # Clean uninstall old packages first (recommended by NVIDIA docs)
+  pip uninstall nvidia-cutlass-dsl nvidia-cutlass-dsl-libs-base nvidia-cutlass-dsl-libs-cu12 nvidia-cutlass-dsl-libs-cu13 -y 2>/dev/null || true
+  pip install "nvidia-cutlass-dsl[${CUTLASS_DSL_EXTRA}]==${CUTLASS_DSL_VERSION}"
   echo "nvidia-cutlass-dsl override complete."
   echo ""
 fi