From f3be33a313608e09469b8d614ae0c962d3f135c9 Mon Sep 17 00:00:00 2001
From: Alisehen <814073252@qq.com>
Date: Thu, 15 May 2025 06:39:02 +0000
Subject: [PATCH] add xpu parameters to install.sh

---
 doc/en/xpu.md |  4 +---
 install.sh    | 22 ++++++++++++++++++++++
 2 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/doc/en/xpu.md b/doc/en/xpu.md
index 0691ce17..e120168d 100644
--- a/doc/en/xpu.md
+++ b/doc/en/xpu.md
@@ -62,9 +62,7 @@ cd ktransformers
 git submodule update --init
 
 # Install dependencies
-bash install.sh
-pip uninstall triton pytorch-triton-xpu
-pip install pytorch-triton-xpu==3.3.0 --extra-index-url https://download.pytorch.org/whl/xpu # to avoid potential triton import error
+bash install.sh --dev xpu
 ```
 
 ## Running DeepSeek-R1 Models
diff --git a/install.sh b/install.sh
index 79f31661..260ae465 100644
--- a/install.sh
+++ b/install.sh
@@ -1,6 +1,20 @@
 #!/bin/bash
 set -e  
 
+# default backend
+DEV="cuda"
+
+# parse --dev argument
+while [[ "$#" -gt 0 ]]; do
+    case $1 in
+        --dev) DEV="$2"; shift ;;
+        *) echo "Unknown parameter passed: $1"; exit 1 ;;
+    esac
+    shift
+done
+export DEV_BACKEND="$DEV"
+echo "Selected backend: $DEV_BACKEND"
+
 # clear build dirs
 rm -rf build
 rm -rf *.egg-info
@@ -13,6 +27,14 @@ rm -rf ~/.ktransformers
 echo "Installing python dependencies from requirements.txt"
 pip install -r requirements-local_chat.txt
 pip install -r ktransformers/server/requirements.txt
+
+# XPU-specific fix for triton
+if [[ "$DEV_BACKEND" == "xpu" ]]; then
+    echo "Replacing triton for XPU backend"
+    pip uninstall -y triton pytorch-triton-xpu || true
+    pip install pytorch-triton-xpu==3.3.0 --extra-index-url https://download.pytorch.org/whl/xpu
+fi
+
 echo "Installing ktransformers"
 KTRANSFORMERS_FORCE_BUILD=TRUE pip install -v . --no-build-isolation