From 16e9552778f9cf68e7bc40d64fdd6168943e241e Mon Sep 17 00:00:00 2001 From: William Lin Date: Fri, 26 Dec 2025 21:42:55 -0600 Subject: [PATCH] [kernel] Fix docker release build for kernel (#965) --- .../workflows/fastvideo-kernel-publish.yml | 2 + docker/Dockerfile.python3.10 | 2 +- docker/Dockerfile.python3.11 | 2 +- docker/Dockerfile.python3.12 | 2 +- docker/Dockerfile.python3.12.cuda12.9.1 | 2 +- docker/rocm7_1.Dockerfile | 2 +- fastvideo-kernel/CMakeLists.txt | 70 +++++++++++++------ fastvideo-kernel/README.md | 26 ++++++- fastvideo-kernel/build.sh | 19 ++++- fastvideo/training/training_pipeline.py | 14 +++- fastvideo/utils.py | 4 +- 11 files changed, 110 insertions(+), 35 deletions(-) diff --git a/.github/workflows/fastvideo-kernel-publish.yml b/.github/workflows/fastvideo-kernel-publish.yml index 973a900e8..1cfa480f3 100644 --- a/.github/workflows/fastvideo-kernel-publish.yml +++ b/.github/workflows/fastvideo-kernel-publish.yml @@ -147,7 +147,9 @@ jobs: cd fastvideo-kernel git submodule update --init --recursive # Ensure ThunderKittens submodule is initialized + # Release builds are produced on GPU-less runners, so force-enable TK and target Hopper. export TORCH_CUDA_ARCH_LIST="9.0a" + export CMAKE_ARGS="${CMAKE_ARGS:-} -DFASTVIDEO_KERNEL_BUILD_TK=ON -DCMAKE_CUDA_ARCHITECTURES=90a" # Build standard wheel (no local version suffix) for PyPI python -m build --wheel --outdir dist diff --git a/docker/Dockerfile.python3.10 b/docker/Dockerfile.python3.10 index 98a1e0bd7..39e3600a6 100644 --- a/docker/Dockerfile.python3.10 +++ b/docker/Dockerfile.python3.10 @@ -60,7 +60,7 @@ RUN source $HOME/.local/bin/env && \ source /opt/venv/bin/activate && \ cd fastvideo-kernel && \ git submodule update --init --recursive && \ - ./build.sh + ./build.sh --release EXPOSE 22 diff --git a/docker/Dockerfile.python3.11 b/docker/Dockerfile.python3.11 index fe3639a60..d8c0d2b35 100644 --- a/docker/Dockerfile.python3.11 +++ b/docker/Dockerfile.python3.11 @@ -60,7 +60,7 @@ RUN source $HOME/.local/bin/env && \ source /opt/venv/bin/activate && \ cd fastvideo-kernel && \ git submodule update --init --recursive && \ - ./build.sh + ./build.sh --release EXPOSE 22 diff --git a/docker/Dockerfile.python3.12 b/docker/Dockerfile.python3.12 index 64c0103c8..baa983423 100644 --- a/docker/Dockerfile.python3.12 +++ b/docker/Dockerfile.python3.12 @@ -60,6 +60,6 @@ RUN source $HOME/.local/bin/env && \ source /opt/venv/bin/activate && \ cd fastvideo-kernel && \ git submodule update --init --recursive && \ - ./build.sh + ./build.sh --release EXPOSE 22 diff --git a/docker/Dockerfile.python3.12.cuda12.9.1 b/docker/Dockerfile.python3.12.cuda12.9.1 index ce90131d8..764995b05 100644 --- a/docker/Dockerfile.python3.12.cuda12.9.1 +++ b/docker/Dockerfile.python3.12.cuda12.9.1 @@ -60,7 +60,7 @@ RUN source $HOME/.local/bin/env && \ source /opt/venv/bin/activate && \ cd fastvideo-kernel && \ git submodule update --init --recursive && \ - ./build.sh + ./build.sh --release EXPOSE 22 diff --git a/docker/rocm7_1.Dockerfile b/docker/rocm7_1.Dockerfile index 0ceba9651..f15ccc6bb 100644 --- a/docker/rocm7_1.Dockerfile +++ b/docker/rocm7_1.Dockerfile @@ -53,6 +53,6 @@ RUN source $HOME/.local/bin/env && \ source /opt/venv/bin/activate && \ cd fastvideo-kernel && \ git submodule update --init --recursive && \ - ./build.sh + ./build.sh --release EXPOSE 22 diff --git a/fastvideo-kernel/CMakeLists.txt b/fastvideo-kernel/CMakeLists.txt index ba6e5d3b8..84e2eaa8c 100644 --- a/fastvideo-kernel/CMakeLists.txt +++ b/fastvideo-kernel/CMakeLists.txt @@ -26,31 +26,54 @@ include_directories( ${TORCH_INCLUDE_DIRS} ) -# CUDA architecture flags -# If not defined, try to detect from PyTorch -if(NOT DEFINED TORCH_CUDA_ARCH_LIST) - execute_process( - COMMAND "${Python_EXECUTABLE}" -c "import torch; print(';'.join([x.replace('sm_', '') for x in torch.cuda.get_arch_list()]))" - OUTPUT_VARIABLE TORCH_CUDA_ARCH_LIST - OUTPUT_STRIP_TRAILING_WHITESPACE - ) +# --------------------------- +# ThunderKittens (TK) toggles +# --------------------------- +# AUTO: enable TK only when we can confidently target Hopper (sm_90a). +# ON: force-enable TK kernels (intended for release wheels/images; does NOT require a GPU). +# OFF: never build TK kernels. +set(FASTVIDEO_KERNEL_BUILD_TK "AUTO" CACHE STRING "Build ThunderKittens kernels: AUTO/ON/OFF") +set_property(CACHE FASTVIDEO_KERNEL_BUILD_TK PROPERTY STRINGS AUTO ON OFF) + +# Prefer environment variable (used by CI) if CMake var is not explicitly set. +if(NOT DEFINED TORCH_CUDA_ARCH_LIST AND DEFINED ENV{TORCH_CUDA_ARCH_LIST}) + set(TORCH_CUDA_ARCH_LIST "$ENV{TORCH_CUDA_ARCH_LIST}") endif() -if(NOT TORCH_CUDA_ARCH_LIST) - # Fallback default if detection failed or no GPU - set(TORCH_CUDA_ARCH_LIST "9.0a") -endif() +message(STATUS "TORCH_CUDA_ARCH_LIST (cmake/env): ${TORCH_CUDA_ARCH_LIST}") +message(STATUS "FASTVIDEO_KERNEL_BUILD_TK: ${FASTVIDEO_KERNEL_BUILD_TK}") -message(STATUS "TORCH_CUDA_ARCH_LIST: ${TORCH_CUDA_ARCH_LIST}") - -# Check if we are targeting Hopper (9.0a) for ThunderKittens -string(FIND "${TORCH_CUDA_ARCH_LIST}" "9.0a" HAS_HOPPER) -if(HAS_HOPPER EQUAL -1) - message(STATUS "TORCH_CUDA_ARCH_LIST (${TORCH_CUDA_ARCH_LIST}) does not contain 9.0a. ThunderKittens kernels will be disabled.") +set(ENABLE_TK_KERNELS OFF) +if(FASTVIDEO_KERNEL_BUILD_TK STREQUAL "ON") + set(ENABLE_TK_KERNELS ON) +elseif(FASTVIDEO_KERNEL_BUILD_TK STREQUAL "OFF") set(ENABLE_TK_KERNELS OFF) else() - message(STATUS "Enabling ThunderKittens kernels (Hopper supported).") - set(ENABLE_TK_KERNELS ON) + # AUTO: detect Hopper if possible. + if(TORCH_CUDA_ARCH_LIST) + # Accept common spellings: 9.0a, 90a, sm_90a. + string(REGEX MATCH "(^|[; ,])((9\\.0a)|(90a)|(sm_90a))([; ,]|$)" _HAS_90A "${TORCH_CUDA_ARCH_LIST}") + if(_HAS_90A) + set(ENABLE_TK_KERNELS ON) + endif() + else() + # Best-effort local detection (works when a CUDA device is visible). + execute_process( + COMMAND "${Python_EXECUTABLE}" -c "import torch; import sys; \nprint('1' if (torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 9) else '0')" + OUTPUT_VARIABLE _LOCAL_HAS_HOPPER + OUTPUT_STRIP_TRAILING_WHITESPACE + ERROR_QUIET + ) + if(_LOCAL_HAS_HOPPER STREQUAL "1") + set(ENABLE_TK_KERNELS ON) + endif() + endif() +endif() + +if(ENABLE_TK_KERNELS) + message(STATUS "ThunderKittens kernels: ENABLED") +else() + message(STATUS "ThunderKittens kernels: DISABLED (will use Triton fallbacks at runtime)") endif() # Always try to build the extension if CUDA is available, but conditionally add sources/flags @@ -72,10 +95,13 @@ set(CUDA_FLAGS "-Xptxas=--warn-on-spills" ) -# Only add -arch=sm_90a and TK flags if we are enabling TK kernels +# If TK is enabled, ensure we target Hopper. This is required even on GPU-less builders (CI). if(ENABLE_TK_KERNELS) - list(APPEND CUDA_FLAGS "-arch=sm_90a") + if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES OR CMAKE_CUDA_ARCHITECTURES STREQUAL "") + set(CMAKE_CUDA_ARCHITECTURES "90a" CACHE STRING "CUDA architectures" FORCE) + endif() list(APPEND CUDA_FLAGS "-DKITTENS_HOPPER") + message(STATUS "CMAKE_CUDA_ARCHITECTURES: ${CMAKE_CUDA_ARCHITECTURES}") endif() if(BUILD_CXX_KERNELS) diff --git a/fastvideo-kernel/README.md b/fastvideo-kernel/README.md index 95bcd3016..1dd982dc2 100644 --- a/fastvideo-kernel/README.md +++ b/fastvideo-kernel/README.md @@ -4,12 +4,24 @@ CUDA kernels for FastVideo video generation. ## Installation +### Standard Installation (Local Development) +This will automatically detect your GPU architecture. If an NVIDIA Hopper (H100/sm_90a) GPU is detected, ThunderKittens kernels will be enabled. Otherwise, they will be skipped, and the package will use Triton fallbacks at runtime. + ```bash git submodule update --init --recursive cd fastvideo-kernel -pip install . +./build.sh ``` +### Release Build (Force Enable Kernels) +If you are building a release wheel or docker image on a machine without a GPU (e.g., CI/CD), you can force-enable the compilation of Hopper-specific ThunderKittens kernels. + +```bash +cd fastvideo-kernel +./build.sh --release +``` +*Note: The resulting wheel will contain kernels that require an H100 GPU to run, but can be built on any machine with CUDA 12.3+ toolchain.* + ## Usage ```python @@ -27,5 +39,13 @@ out = moba_attn_varlen(q, k, v, cu_seqlens_q, cu_seqlens_k, ...) ## Requirements -- H100 GPU (sm_90a) for CUDA kernels -- Triton for non-H100 fallback +- **Runtime**: + - NVIDIA H100 (sm_90a) for C++ optimized kernels. + - Any CUDA GPU for Triton-based fallbacks. +- **Build**: + - CUDA Toolkit 12.3+ + - C++20 compatible compiler (GCC 10+, Clang 11+) + +## Acknowledgement + +This package structure and build system are based on [sgl-kernel](https://github.com/sgl-project/sglang/tree/main/sgl-kernel) from the SGLang project. diff --git a/fastvideo-kernel/build.sh b/fastvideo-kernel/build.sh index 5b7a65e80..8a7c9f993 100755 --- a/fastvideo-kernel/build.sh +++ b/fastvideo-kernel/build.sh @@ -2,15 +2,32 @@ set -ex # Simple build script wrapping uv/pip +# Usage: +# ./build.sh # local dev build (auto-detect / skip TK kernels when not available) +# ./build.sh --release # force-enable Hopper/TK kernels for release builds (no GPU required) echo "Building fastvideo-kernel..." # Ensure submodules are initialized if needed (tk) -# git submodule update --init --recursive +git submodule update --init --recursive # Install build dependencies pip install scikit-build-core cmake ninja +RELEASE=0 +if [ "${1:-}" = "--release" ] || [ "${1:-}" = "-r" ]; then + RELEASE=1 +fi + +if [ "$RELEASE" -eq 1 ]; then + # Force-enable ThunderKittens kernels and compile for Hopper. + # Intended for producing release wheels/images on machines without a GPU. + export TORCH_CUDA_ARCH_LIST="9.0a" + export CMAKE_ARGS="${CMAKE_ARGS:-} -DFASTVIDEO_KERNEL_BUILD_TK=ON -DCMAKE_CUDA_ARCHITECTURES=90a" +fi + +echo "TORCH_CUDA_ARCH_LIST: ${TORCH_CUDA_ARCH_LIST:-}" +echo "CMAKE_ARGS: ${CMAKE_ARGS:-}" # Build and install # Use -v for verbose output pip install . -v --no-build-isolation diff --git a/fastvideo/training/training_pipeline.py b/fastvideo/training/training_pipeline.py index 53d690ba2..8b0c8d262 100644 --- a/fastvideo/training/training_pipeline.py +++ b/fastvideo/training/training_pipeline.py @@ -360,7 +360,12 @@ class TrainingPipeline(LoRAPipeline, ABC): current_vsa_sparsity = training_batch.current_vsa_sparsity assert latents_shape is not None assert training_batch.timesteps is not None - if vsa_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VIDEO_SPARSE_ATTN": + if envs.FASTVIDEO_ATTENTION_BACKEND == "VIDEO_SPARSE_ATTN": + if not vsa_available: + raise ImportError( + "FASTVIDEO_ATTENTION_BACKEND is set to VIDEO_SPARSE_ATTN, " + "but fastvideo_kernel is not correctly installed or detected. " + "Please ensure fastvideo-kernel is installed.") training_batch.attn_metadata = VideoSparseAttentionMetadataBuilder( # type: ignore ).build( # type: ignore raw_latent_shape=latents_shape[2:5], @@ -368,7 +373,12 @@ class TrainingPipeline(LoRAPipeline, ABC): patch_size=patch_size, VSA_sparsity=current_vsa_sparsity, device=get_local_torch_device()) - elif vmoba_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VMOBA_ATTN": + elif envs.FASTVIDEO_ATTENTION_BACKEND == "VMOBA_ATTN": + if not vmoba_available: + raise ImportError( + "FASTVIDEO_ATTENTION_BACKEND is set to VMOBA_ATTN, " + "but fastvideo_kernel (or flash_attn>=2.7.4) is not correctly installed." + ) moba_params = self.training_args.moba_config.copy() moba_params.update({ "current_timestep": diff --git a/fastvideo/utils.py b/fastvideo/utils.py index db36af5e9..6c5074bce 100644 --- a/fastvideo/utils.py +++ b/fastvideo/utils.py @@ -851,12 +851,12 @@ def set_random_seed(seed: int) -> None: @lru_cache(maxsize=1) def is_vsa_available() -> bool: - return importlib.util.find_spec("vsa") is not None + return importlib.util.find_spec("fastvideo_kernel.ops") is not None @lru_cache(maxsize=1) def is_vmoba_available() -> bool: - if importlib.util.find_spec("csrc.attn.vmoba_attn.vmoba") is None: + if importlib.util.find_spec("fastvideo_kernel.vmoba") is None: return False try: import flash_attn