[ci]: prebuild L40S kernel artifact (#1562)

Author: macthecadillac

docker/Dockerfile

@@ -186,22 +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.
-ARG FASTVIDEO_KERNEL_WHEEL_DIR=/opt/fastvideo-kernel-wheels
-ARG FASTVIDEO_KERNEL_BUILD_INFO=/opt/fastvideo-kernel-build-info.json
+# 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 && \
-    export TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST} && \
+    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} \
-    ./build.sh --wheel-dir ${FASTVIDEO_KERNEL_WHEEL_DIR} && \
+    CMAKE_ARGS= CMAKE_BUILD_PARALLEL_LEVEL=${CMAKE_BUILD_PARALLEL_LEVEL} \
+    ./build.sh --wheel-dir "${default_wheel_dir}" && \
     cd /FastVideo && \
-    python fastvideo/tests/modal/kernel_build_cache.py write-build-info \
-        --wheel-dir ${FASTVIDEO_KERNEL_WHEEL_DIR} \
-        --output ${FASTVIDEO_KERNEL_BUILD_INFO}
+    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

@@ -281,6 +281,12 @@ then publishes one multi-platform manifest per CUDA version. CUDA 12.6 owns 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.
+
 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
 stack is not yet validated on ARM64.

fastvideo/tests/modal/kernel_build_cache.py

@@ -24,8 +24,9 @@
 
 
 CACHE_SCHEMA_VERSION = 3
-DEFAULT_PREBUILT_INFO_PATH = "/opt/fastvideo-kernel-build-info.json"
+DEFAULT_PREBUILT_INFO_PATH = "/opt/fastvideo-kernel-prebuilt"
 KERNEL_RELATIVE_DIR = "fastvideo-kernel"
+DEFAULT_BUILD_INFO_OUTPUT = "/opt/fastvideo-kernel-prebuilt/default/metadata.json"
 METADATA_FILE = "metadata.json"
 EXPECTED_DISTRIBUTION = "fastvideo-kernel"
 
@@ -372,26 +373,35 @@ def _cache_entry_wheel(cache_entry: Path, cache_key: str) -> tuple[Path | None,
     return _validated_payload_wheel(metadata, wheel, cache_key)
 
 
+def _prebuilt_metadata_paths(prebuilt_info_path: Path) -> list[Path]:
+    if prebuilt_info_path.is_dir() and not prebuilt_info_path.is_symlink():
+        return sorted(prebuilt_info_path.glob(f"*/{METADATA_FILE}"))
+    if prebuilt_info_path.exists() or prebuilt_info_path.is_symlink():
+        return [prebuilt_info_path]
+    return []
+
+
 def _try_install_prebuilt(metadata: dict[str, object], prebuilt_info_path: Path) -> bool:
     cache_key = str(metadata["cache_key"])
-    if not prebuilt_info_path.exists():
-        return False
-    prebuilt, reason = _load_metadata(prebuilt_info_path)
-    if prebuilt is None:
-        _log(f"Docker-prebuilt kernel metadata is invalid at {prebuilt_info_path}: {reason}")
-        return False
-    wheel_path = Path(str(prebuilt.get("wheel_path", "")))
-    wheel, reason = _validated_payload_wheel(prebuilt, wheel_path, cache_key)
-    if wheel is None:
-        _log(f"Docker-prebuilt kernel artifact rejected at {prebuilt_info_path}: {reason}")
-        return False
-    try:
-        _log(f"using Docker-prebuilt kernel wheel for cache key {cache_key}")
-        _install_wheel(wheel)
-    except (FileNotFoundError, OSError, subprocess.CalledProcessError) as error:
-        _log(f"Docker-prebuilt kernel installation failed: {error}; falling back")
-        return False
-    return True
+    for metadata_path in _prebuilt_metadata_paths(prebuilt_info_path):
+        prebuilt, reason = _load_metadata(metadata_path)
+        if prebuilt is None:
+            _log(f"Docker-prebuilt kernel metadata is invalid at {metadata_path}: {reason}")
+            continue
+        wheel_path = Path(str(prebuilt.get("wheel_path", "")))
+        wheel, reason = _validated_payload_wheel(prebuilt, wheel_path, cache_key)
+        if wheel is None:
+            _log(f"Docker-prebuilt kernel artifact rejected at {metadata_path}: {reason}")
+            continue
+        try:
+            _log(f"Docker-prebuilt cache hit for key {cache_key}: {wheel}")
+            _install_wheel(wheel)
+        except (FileNotFoundError, OSError, subprocess.CalledProcessError) as error:
+            _log(f"Docker-prebuilt kernel installation failed: {error}; falling back")
+            return False
+        return True
+    _log(f"Docker-prebuilt cache miss for key {cache_key}")
+    return False
 
 
 def _store_cache_entry(cache_root: Path, metadata: dict[str, object], wheel: Path) -> Path:
@@ -535,7 +545,7 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--prebuilt-info-path",
                         default=os.environ.get("FASTVIDEO_KERNEL_PREBUILT_INFO", DEFAULT_PREBUILT_INFO_PATH))
     parser.add_argument("--wheel-dir", default="")
-    parser.add_argument("--output", default=DEFAULT_PREBUILT_INFO_PATH)
+    parser.add_argument("--output", default=DEFAULT_BUILD_INFO_OUTPUT)
     return parser.parse_args()
 
 

fastvideo/tests/modal/test_kernel_build_cache.py

@@ -60,6 +60,19 @@ def _write_cache_entry(cache_root: Path, cache_key: str, *, payload: str = "payl
     return cache_entry
 
 
+def _write_prebuilt_artifact(prebuilt_root: Path, name: str, cache_key: str) -> Path:
+    artifact_dir = prebuilt_root / name
+    artifact_dir.mkdir(parents=True)
+    wheel = _write_wheel(artifact_dir / WHEEL_NAME)
+    payload = {
+        **_metadata(cache_key),
+        "artifact": kernel_build_cache._wheel_artifact(wheel),
+        "wheel_path": str(wheel),
+    }
+    (artifact_dir / kernel_build_cache.METADATA_FILE).write_text(json.dumps(payload), encoding="utf-8")
+    return wheel
+
+
 def _patch_stable_metadata(monkeypatch) -> None:
     for name in (
         "GPU_BACKEND",
@@ -318,6 +331,31 @@ def test_consumer_falls_back_when_cached_install_fails(monkeypatch, tmp_path) ->
     assert built == [True]
 
 
+def test_normal_l40s_pr_ssim_install_uses_docker_prebuilt_without_build(monkeypatch, tmp_path, capsys) -> None:
+    prebuilt_root = tmp_path / "prebuilt"
+    _write_prebuilt_artifact(prebuilt_root, "9.0a", "hopper-key")
+    l40s_wheel = _write_prebuilt_artifact(prebuilt_root, "8.9", "l40s-key")
+    installed = []
+    monkeypatch.setattr(kernel_build_cache, "_build_metadata", lambda repo_root: _metadata("l40s-key"))
+    monkeypatch.setattr(kernel_build_cache, "_install_wheel", installed.append)
+    monkeypatch.setattr(
+        kernel_build_cache,
+        "_build_and_install_local",
+        lambda repo_root, metadata: pytest.fail("L40S prebuilt hit unexpectedly called build.sh"),
+    )
+
+    kernel_build_cache.install_cached_or_build(
+        tmp_path,
+        cache_root=None,
+        prebuilt_info_path=prebuilt_root,
+    )
+
+    output = capsys.readouterr().out
+    assert installed == [l40s_wheel]
+    assert "Docker-prebuilt cache hit for key l40s-key" in output
+    assert "build.sh" not in output
+
+
 def test_prebuilt_requires_validated_artifact(monkeypatch, tmp_path) -> None:
     wheel = _write_wheel(tmp_path / WHEEL_NAME)
     prebuilt_info = tmp_path / "prebuilt.json"