Modify vendor setup to handle triton and flagtree

Uninstall triton and install flagtree or triton based on USE_TRITON variable.

Signed-off-by: Qiming Teng <tengqm@outlook.com>

Author: tengqm

tools/setup_vendor.sh

@@ -118,14 +118,18 @@ case $VENDOR in
         "torch==2.7.1+musa.4.0.0" \
         "torch_musa==2.7.1" \
         "numpy==1.26.4" \
-        "flagtree==0.5.0+mthreads3.1" \
         "mkl==2024.0.0"
 
     # Replace flagtree with Triton if requested
     if [ -n "${USE_TRITON}" ]; then
       uv pip uninstall flagtree
+      uv pip uninstall triton
       uv pip install --index $FLAGOS_PYPI \
         "triton==3.1.0+musa1.4.6"
+    else
+      uv pip uninstall triton
+      uv pip install --index $FLAGOS_PYPI \
+        "flagtree==0.5.0+mthreads3.1"
     fi
     ;;