nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[kernel] Fix docker release build for kernel (#965)

Browse Source
This commit is contained in:
William Lin
2025-12-26 21:42:55 -06:00
committed by GitHub
parent 87f8a2782d
commit 16e9552778
11 changed files with 110 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/fastvideo-kernel-publish.yml
+2
View File
@@ -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
docker/Dockerfile.python3.10
+1 -1
View File
@@ -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
docker/Dockerfile.python3.11
+1 -1
View File
@@ -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
docker/Dockerfile.python3.12
+1 -1
View File
@@ -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
docker/Dockerfile.python3.12.cuda12.9.1
+1 -1
View File
@@ -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
docker/rocm7_1.Dockerfile
+1 -1
View File
@@ -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
fastvideo-kernel/CMakeLists.txt
+48 -22
View File
@@ -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)
fastvideo-kernel/README.md
+23 -3
View File
@@ -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.
fastvideo-kernel/build.sh
+18 -1
View File
@@ -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:-<unset>}"
echo "CMAKE_ARGS: ${CMAKE_ARGS:-<unset>}"
# Build and install
# Use -v for verbose output
pip install . -v --no-build-isolation
fastvideo/training/training_pipeline.py
+12 -2
View File
@@ -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":
fastvideo/utils.py
+2 -2
View File
@@ -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