[ci]: cache FastVideo kernel builds in Modal (#1562)

Author: macthecadillac

.github/workflows/infra-build-image.yml

@@ -13,16 +13,23 @@ on:
         required: false
         default: false
         type: boolean
-  # Auto-rebuild the CUDA images when their Dockerfile changes on main. The CUDA
-  # matrix is the only lane that builds from docker/Dockerfile, so a path-scoped
-  # push trigger is a sufficient change detector on its own -- no separate
-  # detect-changes/paths-filter job is needed now that there is a single
-  # in-scope Dockerfile. Dreamverse (apps/dreamverse/docker/Dockerfile) and the
-  # rocm Dockerfile stay manual-dispatch only.
+  # Auto-rebuild the CUDA images when a repository-controlled image input
+  # changes on main. This includes the trusted SM89 kernel artifact's source,
+  # metadata/key helper, ABI dependency metadata, and build orchestration.
+  # Dreamverse (apps/dreamverse/docker/Dockerfile) and the ROCm Dockerfile stay
+  # manual-dispatch only.
   push:
     branches: [main]
     paths:
+      - '.dockerignore'
+      - '.github/workflows/_template-build-image.yml'
+      - '.github/workflows/infra-build-image.yml'
+      - '.gitmodules'
       - 'docker/Dockerfile'
+      - 'docker/uv-excludes'
+      - 'fastvideo-kernel/**'
+      - 'fastvideo/tests/modal/kernel_build_cache.py'
+      - 'pyproject.toml'
 
 
 permissions:
@@ -50,7 +57,7 @@ jobs:
   # 2.8.3 comes from the architecture-specific prebuilt releases.
   build-cuda-images:
     # Runs on a manual dispatch when build_cuda_matrix is set, or automatically
-    # on a push that changed docker/Dockerfile (inputs are null on push). The
+    # on an in-scope main push (inputs are null on push). The
     # repository guard keeps fork syncs from auto-building; manual dispatch
     # still works in forks.
     if: ${{ (github.event_name == 'push' && github.repository == 'hao-ai-lab/FastVideo') || github.event.inputs.build_cuda_matrix == 'true' }}

docker/Dockerfile

@@ -186,15 +186,40 @@ RUN --mount=type=cache,target=/opt/uv/cache \
 
 COPY . .
 
-# Install FastVideo Unified Kernel exactly once. build.sh initializes only its
-# CUTLASS/ThunderKittens submodules, then compiles for TORCH_CUDA_ARCH_LIST
-# (default Hopper sm_90a) without requiring a live GPU.
+# Build immutable FastVideo kernel wheels for the published image. The requested
+# architecture remains installed for normal image users; amd64 images also carry
+# an SM89 artifact so the predominant L40S Modal lanes can reuse it exactly.
+ARG FASTVIDEO_KERNEL_PREBUILT_DIR=/opt/fastvideo-kernel-prebuilt
 RUN --mount=type=cache,target=/opt/uv/cache \
     source $HOME/.local/bin/env && \
     source /opt/venv/bin/activate && \
+    default_arch="${TORCH_CUDA_ARCH_LIST}" && \
+    default_wheel_dir="${FASTVIDEO_KERNEL_PREBUILT_DIR}/${default_arch}" && \
+    export TORCH_CUDA_ARCH_LIST="${default_arch}" && \
     cd fastvideo-kernel && \
-    CMAKE_BUILD_PARALLEL_LEVEL=${CMAKE_BUILD_PARALLEL_LEVEL} \
-    TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST} ./build.sh
+    CMAKE_ARGS= CMAKE_BUILD_PARALLEL_LEVEL=${CMAKE_BUILD_PARALLEL_LEVEL} \
+    ./build.sh --wheel-dir "${default_wheel_dir}" && \
+    cd /FastVideo && \
+    CMAKE_ARGS= python fastvideo/tests/modal/kernel_build_cache.py write-build-info \
+        --wheel-dir "${default_wheel_dir}" \
+        --output "${default_wheel_dir}/metadata.json" && \
+    if [[ "${TARGETARCH:-amd64}" == "amd64" && "${default_arch}" != "8.9" ]]; then \
+        export TORCH_CUDA_ARCH_LIST=8.9 && \
+        l40s_wheel_dir="${FASTVIDEO_KERNEL_PREBUILT_DIR}/8.9" && \
+        cd /FastVideo/fastvideo-kernel && \
+        CMAKE_ARGS= CMAKE_BUILD_PARALLEL_LEVEL=${CMAKE_BUILD_PARALLEL_LEVEL} \
+        ./build.sh --wheel-dir "${l40s_wheel_dir}" && \
+        cd /FastVideo && \
+        CMAKE_ARGS= python fastvideo/tests/modal/kernel_build_cache.py write-build-info \
+            --wheel-dir "${l40s_wheel_dir}" \
+            --output "${l40s_wheel_dir}/metadata.json"; \
+    fi && \
+    export TORCH_CUDA_ARCH_LIST="${default_arch}" && \
+    default_wheel="$(find "${default_wheel_dir}" -maxdepth 1 -type f \
+        \( -name 'fastvideo_kernel-*.whl' -o -name 'fastvideo-kernel-*.whl' \) \
+        | sort | tail -n 1)" && \
+    uv pip install "${default_wheel}" \
+        --reinstall-package fastvideo-kernel --no-deps
 
 # Install FastVideo itself (editable) now that the source is present, and set up
 # shell configuration. Dependencies and the local kernel are already installed,

docs/contributing/ci_architecture.md

@@ -269,16 +269,26 @@ The docs job:
 ### Docker Images
 
 `.github/workflows/infra-build-image.yml` supports manual `workflow_dispatch`
-runs and automatically rebuilds the CUDA matrix when `docker/Dockerfile`
-changes on `main` in the canonical repository. Manual runs let maintainers
-choose which image families to build. The `fastvideo-dev` matrix builds Python
-3.12 images for CUDA 12.6 and CUDA 13 on native `amd64` and `arm64` runners,
-then publishes one multi-platform manifest per CUDA version. CUDA 12.6 owns the
-`py3.12-latest` and global `latest` tags, as well as the explicit
-`py3.12-cuda12.6.3-latest` alias. CUDA 13 is published under the explicit
-`py3.12-cuda13.0.0-latest` tag. This publication policy does not change the
-unparameterized `docker/Dockerfile` build defaults, which remain CUDA 13 and
-`cu130`.
+runs and automatically rebuilds the CUDA matrix when a repository-controlled
+image input changes on `main` in the canonical repository. Those inputs include
+the CUDA Dockerfile and reusable workflow, dependency metadata, Docker context
+policy, `fastvideo-kernel/**`, and the kernel artifact metadata/key helper.
+Manual runs let maintainers choose which image families to build. The
+`fastvideo-dev` matrix builds Python 3.12 images for CUDA 12.6 and CUDA 13 on
+native `amd64` and `arm64` runners, then publishes one multi-platform manifest
+per CUDA version. CUDA 12.6 owns the `py3.12-latest` and global `latest` tags, as
+well as the explicit `py3.12-cuda12.6.3-latest` alias. CUDA 13 is published under
+the explicit `py3.12-cuda13.0.0-latest` tag. This publication policy does not
+change the unparameterized `docker/Dockerfile` build defaults, which remain CUDA
+13 and `cu130`.
+
+Published amd64 development images keep their configured Hopper kernel wheel
+installed and also carry an immutable SM89 wheel under
+`/opt/fastvideo-kernel-prebuilt`. Modal PR and SSIM jobs select the exact
+source, ABI, and GPU-architecture match from that directory, so L40S jobs reuse
+the trusted image artifact while kernel-changing PRs still build locally. Once
+a kernel or artifact-key change reaches `main`, the image workflow republishes
+the matching trusted artifact before later jobs consume the updated image tag.
 
 The optional Dreamverse matrix builds backend and UI images for CUDA 12.6 and
 CUDA 13 on `amd64`. Dreamverse remains `amd64`-only because its FA4 dependency

fastvideo-kernel/build.sh

@@ -4,6 +4,7 @@ set -ex
 # Simple build script wrapping uv/pip
 # Usage:
 #   ./build.sh  # local build (torch-based arch detection, TK only on SM90)
+#   ./build.sh --wheel-dir /tmp/wheels  # build + install a reusable wheel
 # Environment overrides (if set, they win over auto-detection):
 #   TORCH_CUDA_ARCH_LIST
 #   CMAKE_ARGS (for FASTVIDEO_KERNEL_BUILD_TK / CMAKE_CUDA_ARCHITECTURES / GPU_BACKEND)
@@ -67,12 +68,25 @@ fi
 # Install build dependencies
 uv pip install scikit-build-core cmake ninja
 
-RELEASE=0
 GPU_BACKEND=CUDA
-for arg in "$@"; do
-    case "$arg" in
+WHEEL_DIR=""
+while [[ $# -gt 0 ]]; do
+    case "$1" in
         --rocm)
             GPU_BACKEND=ROCM
+            shift
+            ;;
+        --wheel-dir)
+            if [[ $# -lt 2 ]]; then
+                echo "ERROR: --wheel-dir requires a directory path" >&2
+                exit 2
+            fi
+            WHEEL_DIR="$2"
+            shift 2
+            ;;
+        *)
+            echo "ERROR: unknown argument: $1" >&2
+            exit 2
             ;;
     esac
 done
@@ -165,4 +179,15 @@ echo "CMAKE_ARGS: ${CMAKE_ARGS:-<unset>}"
 echo "GPU_BACKEND: ${GPU_BACKEND:-<unset>}"
 # Build and install
 # Use -v for verbose output
-uv pip install . -v --no-build-isolation
+if [ -n "${WHEEL_DIR}" ]; then
+    mkdir -p "${WHEEL_DIR}"
+    uv build --wheel -v --no-build-isolation --out-dir "${WHEEL_DIR}" .
+    wheel_path="$(find "${WHEEL_DIR}" -maxdepth 1 -type f \( -name 'fastvideo_kernel-*.whl' -o -name 'fastvideo-kernel-*.whl' \) | sort | tail -n 1)"
+    if [ -z "${wheel_path}" ]; then
+        echo "ERROR: no fastvideo-kernel wheel was built in ${WHEEL_DIR}" >&2
+        exit 1
+    fi
+    uv pip install "${wheel_path}" --reinstall-package fastvideo-kernel --no-deps
+else
+    uv pip install . -v --no-build-isolation
+fi