diff --git a/Dockerfile b/Dockerfile
index a7c920b9f..b01dd7065 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -20,6 +20,7 @@ ENV PATH="/opt/venv/bin:$PATH"
 # Upgrade pip
 RUN pip install --upgrade pip
 
+RUN pip install uv
 RUN pip install git+https://github.com/ayaka14732/jax-smi.git
 # If you encounter a checkpoint issue, try using following old version of pathways-utils.
 # RUN pip install git+https://github.com/AI-Hypercomputer/pathways-utils.git@b72729bb152b7b3426299405950b3af300d765a9#egg=pathwaysutils
@@ -37,6 +38,27 @@ RUN pip install -e .
 
 RUN bash /app/scripts/install_tunix_vllm_requirement.sh
 
+# Build argument to conditionally install MaxText dependencies
+ARG INSTALL_MAXTEXT=false
+
+# Install MaxText specific dependencies conditionally
+RUN if [ "$INSTALL_MAXTEXT" = "true" ]; then \
+      uv pip install -r /app/requirements/maxtext_requirements.txt --torch-backend=cpu; \
+    fi
+
+# Build argument to conditionally install Raiden weight sync dependencies
+ARG INSTALL_RAIDEN=false
+
+# Install Raiden specific dependencies conditionally
+RUN if [ "$INSTALL_RAIDEN" = "true" ]; then \
+      if [ -d "/app/raiden_wheels" ] && ls /app/raiden_wheels/*.whl 1>/dev/null 2>&1; then \
+        pip install --force-reinstall --no-deps /app/raiden_wheels/*.whl; \
+      else \
+        pip install keyrings.google-artifactregistry-auth && \
+        pip install tpu-raiden-jax --extra-index-url https://us-python.pkg.dev/cloud-tpu-inference-test/tpu-raiden/simple/; \
+      fi; \
+    fi
+
 # Build argument to conditionally install DeepSWE evaluation dependencies
 ARG INSTALL_DEEPSWE_DEPS=false
 
diff --git a/requirements/maxtext_requirements.txt b/requirements/maxtext_requirements.txt
new file mode 100644
index 000000000..38f0af3a1
--- /dev/null
+++ b/requirements/maxtext_requirements.txt
@@ -0,0 +1,14 @@
+maxtext @ git+https://github.com/AI-Hypercomputer/maxtext.git@8c2e29218c01b6287ba85e8d0dc9555561f7634c
+maxtext-vllm-adapter @ git+https://github.com/AI-Hypercomputer/maxtext.git@8c2e29218c01b6287ba85e8d0dc9555561f7634c#subdirectory=src/maxtext/integration/vllm
+aqtp
+tokamax>=0.0.4
+drjax>=0.1.4
+ml-goodput-measurement
+cloud-accelerator-diagnostics
+cloud-tpu-diagnostics
+google-cloud-mldiagnostics>=0.5.10
+google-cloud-monitoring
+einshape
+typeguard
+xprof
+cheroot
diff --git a/requirements/requirements.txt b/requirements/requirements.txt
index 7f3f9303d..92c9fd6c6 100644
--- a/requirements/requirements.txt
+++ b/requirements/requirements.txt
@@ -1,2 +1,2 @@
-vllm @ git+https://github.com/vllm-project/vllm.git@9842d701450214d4b78cd9aefb8eee0c616bce33
+vllm @ git+https://github.com/vllm-project/vllm.git@d626108b1841888ec90aced33367149a6bbc7e4b
 
diff --git a/requirements/special_requirements.txt b/requirements/special_requirements.txt
index 7f8bf1645..97937fb67 100644
--- a/requirements/special_requirements.txt
+++ b/requirements/special_requirements.txt
@@ -2,5 +2,4 @@
 # --find-links https://storage.googleapis.com/jax-releases/libtpu_releases.html
 # --pre
 
-tpu-inference @ git+https://github.com/vllm-project/tpu-inference.git@12eed6e661a1ffacbce9d107a9894c1a471e7479
-
+tpu-inference @ git+https://github.com/vllm-project/tpu-inference.git@4a8951606026d5c0307554bc4ed93f4690c4640a