packaging: build SageAttention into the OCI image (non-fatal, arch-gated)

The diffusers video path uses SageAttention (INT8 attention) when available
for faster Wan2.2 rendering. Like flash-attn it is CUDA-arch-sensitive, so
it is built from source in the devel/builder stage against the just-installed
torch, gated by BUILD_SAGEATTENTION (default 1) and SAGEATTENTION_ARCH
(default 8.6 = RTX 3090). The build is non-fatal: on failure the image still
works and the runtime attention-backend resolver falls back to flash/SDPA.

build_oci_image.sh passes the three new args through (overridable via env:
BUILD_SAGEATTENTION / SAGEATTENTION_REF / SAGEATTENTION_ARCH).

Note: the fast update_oci_image.sh overlay is based on the runtime image (no
nvcc), so it cannot build SageAttention — a full build_oci_image.sh is needed
to bake it in.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01RdMufYvtTbtGDWsiZVoXce

packaging/linux/Dockerfile.oci

@@ -12,6 +12,14 @@ ARG UV_VERSION=0.7.13
 ARG WHISPERCPP_REF=master
 ARG LLAMA_CPP_PYTHON_VERSION=
 ARG SD_CPP_PYTHON_VERSION=
+# SageAttention (INT8 attention for the diffusers video transformer). Like
+# flash-attn it is CUDA-arch-sensitive, so it is built from source and gated by
+# an ARG; the build is non-fatal (video falls back to flash/SDPA automatically
+# via the attention-backend resolver). Default arch covers RTX 3090 (8.6);
+# override SAGEATTENTION_ARCH for other GPUs (e.g. "8.0;8.6;8.9;9.0").
+ARG BUILD_SAGEATTENTION=1
+ARG SAGEATTENTION_REF=main
+ARG SAGEATTENTION_ARCH=8.6
 
 ENV DEBIAN_FRONTEND=noninteractive \
     PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin \
@@ -92,6 +100,21 @@ RUN uv pip install --python /opt/coderai/python/bin/python3 \
       -r /tmp/requirements-oci.txt \
     && rm -rf /root/.cache/uv /tmp/requirements*.txt
 
+# SageAttention (optional, non-fatal): built against the torch just installed.
+# Needs no-build-isolation so it sees the installed torch; arch list keeps the
+# build small. On any failure the image still works — the video path resolves to
+# flash/SDPA. Set --build-arg BUILD_SAGEATTENTION=0 to skip entirely.
+RUN set -eux; \
+    if [ "${BUILD_SAGEATTENTION}" = "1" ]; then \
+      TORCH_CUDA_ARCH_LIST="${SAGEATTENTION_ARCH}" MAX_JOBS="${MAX_JOBS:-4}" \
+        uv pip install --python /opt/coderai/python/bin/python3 --no-build-isolation \
+        "git+https://github.com/thu-ml/SageAttention.git@${SAGEATTENTION_REF}" \
+      || echo "WARN: SageAttention build failed — video falls back to flash/SDPA"; \
+    else \
+      echo "SageAttention build skipped (BUILD_SAGEATTENTION=${BUILD_SAGEATTENTION})"; \
+    fi; \
+    rm -rf /root/.cache/uv
+
 RUN set -eux; \
     chmod +x /opt/coderai/app/coderai; \
     mkdir -p /config /models /cache /opt/coderai/app/models; \

packaging/linux/build_oci_image.sh

@@ -246,6 +246,12 @@ UBUNTU_VERSION="${UBUNTU_VERSION:-22.04}"
 WHISPERCPP_REF="${WHISPERCPP_REF:-master}"
 LLAMA_CPP_PYTHON_VERSION="${LLAMA_CPP_PYTHON_VERSION:-}"
 SD_CPP_PYTHON_VERSION="${SD_CPP_PYTHON_VERSION:-}"
+# SageAttention (INT8 video attention): built from source (CUDA-arch-sensitive).
+# Default ON for the RTX 3090 arch (8.6); non-fatal in the Dockerfile so a build
+# on another arch/CUDA just falls back to flash/SDPA at runtime.
+BUILD_SAGEATTENTION="${BUILD_SAGEATTENTION:-1}"
+SAGEATTENTION_REF="${SAGEATTENTION_REF:-main}"
+SAGEATTENTION_ARCH="${SAGEATTENTION_ARCH:-8.6}"
 
 write_build_manifest() {
   local mode="$1"
@@ -572,6 +578,9 @@ else
     --build-arg WHISPERCPP_REF="$WHISPERCPP_REF" \
     --build-arg LLAMA_CPP_PYTHON_VERSION="$LLAMA_CPP_PYTHON_VERSION" \
     --build-arg SD_CPP_PYTHON_VERSION="$SD_CPP_PYTHON_VERSION" \
+    --build-arg BUILD_SAGEATTENTION="$BUILD_SAGEATTENTION" \
+    --build-arg SAGEATTENTION_REF="$SAGEATTENTION_REF" \
+    --build-arg SAGEATTENTION_ARCH="$SAGEATTENTION_ARCH" \
     -t "$IMAGE_TAG" \
     "$ROOT_DIR"
 fi