Update Dockerfile

Dockerfile

@@ -91,6 +91,11 @@ RUN curl https://sh.rustup.rs -sSf | bash -s -- -y
 ENV PATH="$HOME/.cargo/bin:${PATH}"
 RUN chmod -R 777 $HOME/.cargo/
 
+# Install ninja, which might help speed up flash attn build time
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        ninja-build \
+        && rm -rf /var/lib/apt/lists/*
+
 # Python packages
 RUN --mount=target=requirements.txt,source=requirements.txt \
     pip install --no-cache-dir --upgrade -r requirements.txt