Update Dockerfile

Dockerfile

@@ -26,6 +26,20 @@ RUN rm -f /opt/conda/.condarc && \
     conda config --add channels conda-forge && \
     conda config --set channel_priority strict
 
+# Ensure build tools are present
+RUN apt-get update && apt-get install -y --no-install-recommends \
+      ninja-build build-essential cmake git python3-dev && \
+      rm -rf /var/lib/apt/lists/*
+
+# Install packaging helpers
+RUN pip install --upgrade pip setuptools wheel packaging ninja
+
+# Uninstall any prebuilt flash-attn
+RUN pip uninstall -y flash-attn || true
+
+# Build from source so it links correctly
+RUN pip install flash-attn --no-build-isolation --no-cache-dir
+
 ENV CONDA_FORGE=1 \
     CONDA_AUTO_UPDATE_CONDA=false