Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 7 additions & 4 deletions docker/Dockerfile
Original file line number Diff line number Diff line change
Expand Up @@ -183,9 +183,12 @@ RUN mkdir -p /wheels && \
rm -rf /tmp/torchao

# FA3 (Hopper-only, SM90a-locked, ~25 min nvcc) + FA4 cute (Blackwell, lightweight). arm64 -> skip FA3: GB200 target, QEMU cross-build impractical.
ARG FLASH_ATTN_REF=002cce0a1
# Pin >= #2762/#2507: older refs rebuilt the CuTe compile key every call when max_seqlen
# is a tensor (varlen/packed), so HF models recompiled each step. Also picks up #2745
# (fwd dynamic-shape correctness) and #2819 (faster SM100 mask compile).
ARG FLASH_ATTN_REF=ce088ab9ce0fc0434dcd8afa0a791da9fcc3a820
ARG INSTALL_FA3=true
ARG INSTALL_FA4=false
ARG INSTALL_FA4=true
ARG TARGETARCH
RUN if [ "$TARGETARCH" = "arm64" ]; then INSTALL_FA3=false; fi && \
if [ "$INSTALL_FA3" = "true" ] || [ "$INSTALL_FA4" = "true" ]; then \
Expand All @@ -210,10 +213,10 @@ WORKDIR /opt/Automodel
# torchao MXFP8 + FA3/FA4: prebuilt wheels from the wheel_builder stage, installed into system site-packages
# before uv sync — pip targets system deps, uv the venv (torchao/FA gated `never`), so uv sync won't prune or shadow them.
COPY --from=wheel_builder /wheels /tmp/wheels/
ARG INSTALL_FA4=false
ARG INSTALL_FA4=true
RUN pip install --no-cache-dir --no-deps /tmp/wheels/*.whl && \
if [ "$INSTALL_FA4" = "true" ]; then \
pip install --no-cache-dir "nvidia-cutlass-dsl[cu13]==4.6.0.dev0" && \
pip install --no-cache-dir "nvidia-cutlass-dsl[cu13]==4.6.2" && \
FA2_DIR=$(python -c "import flash_attn, os; print(os.path.dirname(flash_attn.__file__))") && \
VENV_CUTE=/opt/venv/lib/python3.12/site-packages/flash_attn/cute && \
{ [ -e "$FA2_DIR/cute" ] || ln -s "$VENV_CUTE" "$FA2_DIR/cute"; }; \
Expand Down
Loading
Loading