From 3c83bbb03ed419d9551a540e93524acad29e61c7 Mon Sep 17 00:00:00 2001
From: hailin <hailin@gdzx.xyz>
Date: Fri, 8 Aug 2025 21:50:56 +0800
Subject: [PATCH] .

---
 .../install_deepspeed_src_mamba_single_gpu.sh | 27 ++++++++++++++++---
 1 file changed, 24 insertions(+), 3 deletions(-)

diff --git a/dspeed/install_deepspeed_src_mamba_single_gpu.sh b/dspeed/install_deepspeed_src_mamba_single_gpu.sh
index 476a951..f6919d1 100755
--- a/dspeed/install_deepspeed_src_mamba_single_gpu.sh
+++ b/dspeed/install_deepspeed_src_mamba_single_gpu.sh
@@ -149,18 +149,39 @@ if git -C "${REPO_PATH}" rev-parse --is-inside-work-tree >/dev/null 2>&1; then
   if [[ "${cur_remote}" != "${REPO_URL}" ]]; then
     log "remote 不匹配，重建仓库"
     rm -rf "${REPO_PATH}"
-    git clone --branch "${REPO_BRANCH}" --depth 1 "${REPO_URL}" "${REPO_PATH}"
+    git clone --branch "${REPO_BRANCH}" --depth 1 --recurse-submodules --shallow-submodules "${REPO_URL}" "${REPO_PATH}"
   else
     log "更新已有仓库到 ${REPO_BRANCH}"
     git -C "${REPO_PATH}" fetch --depth 1 origin "${REPO_BRANCH}"
     git -C "${REPO_PATH}" checkout -f "${REPO_BRANCH}"
     git -C "${REPO_PATH}" reset --hard "origin/${REPO_BRANCH}"
     git -C "${REPO_PATH}" clean -fdx
+    # 同步并更新子模块到正确提交
+    git -C "${REPO_PATH}" submodule sync --recursive
+    git -C "${REPO_PATH}" submodule update --init --recursive --depth 1
   fi
 else
-  # 不存在或不是仓库：重拉
+  # 不存在或不是仓库：重拉（连带子模块）
   rm -rf "${REPO_PATH}"
-  git clone --branch "${REPO_BRANCH}" --depth 1 "${REPO_URL}" "${REPO_PATH}"
+  git clone --branch "${REPO_BRANCH}" --depth 1 --recurse-submodules --shallow-submodules "${REPO_URL}" "${REPO_PATH}"
+fi
+
+# 再保险：确保子模块已到位
+git -C "${REPO_PATH}" submodule sync --recursive
+git -C "${REPO_PATH}" submodule update --init --recursive --depth 1
+
+# 安装 ds-kernels（DeepSpeed Inference CUTLASS 依赖）
+DK_DIR=""
+# 从 .gitmodules 找子模块路径（不同版本命名可能是 ds-kernels 或 dskernels）
+if [[ -f "${REPO_PATH}/.gitmodules" ]]; then
+  DK_DIR=$(git -C "${REPO_PATH}" config --file .gitmodules --get-regexp path | awk '/ds[-_]?kernels/{print $2; exit}' || true)
+fi
+if [[ -n "${DK_DIR}" && -f "${REPO_PATH}/${DK_DIR}/setup.py" ]]; then
+  log "Install ds-kernels from submodule: ${DK_DIR}"
+  pip install -v "${REPO_PATH}/${DK_DIR}"
+else
+  log "Submodule ds-kernels 不存在或未包含构建脚本，直接从仓库安装"
+  pip install -v "git+https://github.com/microsoft/ds-kernels.git"
 fi
 
 # 基本健检