diff --git a/python/pyproject.toml b/python/pyproject.toml
index 1a0a498f0..a4e83bfea 100644
--- a/python/pyproject.toml
+++ b/python/pyproject.toml
@@ -103,7 +103,7 @@ test = [
     "accelerate",
     "peft",
 ]
-all = ["sglang[srt]", "sglang[openai]", "sglang[anthropic]", "sglang[litellm]"]
+all = ["sglang[srt]", "sglang[openai]", "sglang[anthropic]", "sglang[litellm]", "sglang[torch_memory_saver]"]
 all_hip = ["sglang[srt_hip]", "sglang[openai]", "sglang[anthropic]", "sglang[litellm]"]
 all_xpu = ["sglang[srt_xpu]", "sglang[openai]", "sglang[anthropic]", "sglang[litellm]"]
 all_hpu = ["sglang[srt_hpu]", "sglang[openai]", "sglang[anthropic]", "sglang[litellm]"]
diff --git a/scripts/ci_install_dependency.sh b/scripts/ci_install_dependency.sh
index ae99fe4e0..e8599e8bd 100755
--- a/scripts/ci_install_dependency.sh
+++ b/scripts/ci_install_dependency.sh
@@ -19,7 +19,6 @@ rm -rf /usr/local/lib/python3.10/dist-packages/sgl_kernel*
 pip install -e "python[all]"
 
 # Install additional dependencies
-pip install torch_memory_saver
 pip install transformers==4.51.0 timm torchaudio==2.6.0 sentence_transformers accelerate peft pandas datasets mooncake-transfer-engine==0.3.0
 
 # For compiling xgrammar kernels