From 5a19a6c6705fe83db2e3517a2d2f473586901743 Mon Sep 17 00:00:00 2001
From: elvischenv <219235043+elvischenv@users.noreply.github.com>
Date: Thu, 24 Jul 2025 18:25:41 +0800
Subject: [PATCH] [Fix] Update mamba_ssm to 2.2.5 (#21421)

Signed-off-by: elvischenv <219235043+elvischenv@users.noreply.github.com>
---
 docker/Dockerfile                              | 8 --------
 docs/contributing/ci/update_pytorch_version.md | 2 +-
 requirements/test.in                           | 2 +-
 requirements/test.txt                          | 6 ++++--
 4 files changed, 6 insertions(+), 12 deletions(-)

diff --git a/docker/Dockerfile b/docker/Dockerfile
index 3c2bdc2066e83..11991829968a4 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -276,10 +276,6 @@ ARG PYTORCH_CUDA_INDEX_BASE_URL
 ENV UV_HTTP_TIMEOUT=500
 ENV UV_INDEX_STRATEGY="unsafe-best-match"
 
-# Workaround for #17068
-RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --system --no-build-isolation "git+https://github.com/state-spaces/mamba@v2.2.4"
-
 COPY requirements/lint.txt requirements/lint.txt
 COPY requirements/test.txt requirements/test.txt
 COPY requirements/dev.txt requirements/dev.txt
@@ -452,10 +448,6 @@ ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
 ENV UV_HTTP_TIMEOUT=500
 ENV UV_INDEX_STRATEGY="unsafe-best-match"
 
-# Workaround for #17068
-RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --system --no-build-isolation "git+https://github.com/state-spaces/mamba@v2.2.4"
-
 # install development dependencies (for testing)
 RUN --mount=type=cache,target=/root/.cache/uv \
     CUDA_MAJOR="${CUDA_VERSION%%.*}"; \
diff --git a/docs/contributing/ci/update_pytorch_version.md b/docs/contributing/ci/update_pytorch_version.md
index 1fe18d5d88565..5046db11a4715 100644
--- a/docs/contributing/ci/update_pytorch_version.md
+++ b/docs/contributing/ci/update_pytorch_version.md
@@ -134,7 +134,7 @@ MAX_JOBS=16 uv pip install --system \
 
 ```bash
 uv pip install --system \
-    --no-build-isolation "git+https://github.com/state-spaces/mamba@v2.2.4"
+    --no-build-isolation "git+https://github.com/state-spaces/mamba@v2.2.5"
 ```
 
 ### causal-conv1d
diff --git a/requirements/test.in b/requirements/test.in
index 429d1a50422f7..c794d1b3cb8d4 100644
--- a/requirements/test.in
+++ b/requirements/test.in
@@ -26,7 +26,7 @@ torch==2.7.1
 torchaudio==2.7.1
 torchvision==0.22.1
 transformers_stream_generator # required for qwen-vl test
-mamba_ssm # required for plamo2 test
+mamba_ssm==2.2.5 # required for plamo2 test
 matplotlib # required for qwen-vl test
 mistral_common[image,audio] >= 1.8.2 # required for voxtral test
 num2words # required for smolvlm test
diff --git a/requirements/test.txt b/requirements/test.txt
index 8e5af8d74bad7..c4e3c33f373fc 100644
--- a/requirements/test.txt
+++ b/requirements/test.txt
@@ -421,7 +421,7 @@ lxml==5.3.0
     #   sacrebleu
 mako==1.3.10
     # via alembic
-mamba-ssm==2.2.4
+mamba-ssm==2.2.5
     # via -r requirements/test.in
 markdown==3.8.2
     # via mlflow
@@ -1152,7 +1152,9 @@ transformers==4.53.2
 transformers-stream-generator==0.0.5
     # via -r requirements/test.in
 triton==3.3.1
-    # via torch
+    # via
+    #   mamba-ssm
+    #   torch
 tritonclient==2.51.0
     # via
     #   -r requirements/test.in