let's try this again

Dockerfile

@@ -97,6 +97,22 @@ RUN --mount=type=cache,target=/root/.cache/pip \
 
 #################### DEV IMAGE ####################
 
+#################### MAMBA Build IMAGE ####################
+FROM dev as mamba-builder
+# max jobs used for build
+ARG max_jobs=2
+ENV MAX_JOBS=${max_jobs}
+
+WORKDIR /usr/src/mamba
+
+COPY requirements-mamba.txt requirements-mamba.txt
+
+# Download the wheel or build it if a pre-compiled release doesn't exist
+RUN pip --verbose wheel -r requirements-mamba.txt \
+    --no-build-isolation --no-deps --no-cache-dir
+
+#################### MAMBA Build IMAGE ####################
+
 #################### Aphrodite installation IMAGE ####################
 # image with Aphrodite installed
 FROM nvidia/cuda:${CUDA_VERSION}-base-ubuntu22.04 AS aphrodite-base

requirements-cuda.txt

@@ -7,6 +7,4 @@ torch == 2.3.1
 torchvision == 0.18.1  # for phi3v
 xformers == 0.0.27  # Requires torch 2.3.1
 triton >= 2.2.1
-vllm-flash-attn == 2.5.9.post1 # Requires PyTorch 2.3.1
-causal-conv1d >= 1.2.1
-mamba-ssm >= 1.2.2
+vllm-flash-attn == 2.5.9.post1 # Requires PyTorch 2.3.1

requirements-mamba.txt

@@ -0,0 +1,2 @@
+causal-conv1d >= 1.2.1
+mamba-ssm >= 1.2.2