diff --git a/fastvideo-kernel/CMakeLists.txt b/fastvideo-kernel/CMakeLists.txt index b84b32327..8849c0a96 100644 --- a/fastvideo-kernel/CMakeLists.txt +++ b/fastvideo-kernel/CMakeLists.txt @@ -303,6 +303,12 @@ endif() # there, and tk's base_types.cuh brace-initializes signed-char vector members # from char (narrowing error). Skip TK until upstream is aarch64-clean. if(ENABLE_TK_KERNELS AND CMAKE_SYSTEM_PROCESSOR MATCHES "^(aarch64|arm64)$") + if(FASTVIDEO_KERNEL_BUILD_TK STREQUAL "ON") + message(FATAL_ERROR + "fastvideo-kernel: FASTVIDEO_KERNEL_BUILD_TK=ON was requested, but the " + "ThunderKittens headers do not compile on ${CMAKE_SYSTEM_PROCESSOR}. " + "Use FASTVIDEO_KERNEL_BUILD_TK=AUTO or OFF to build with Triton fallbacks.") + endif() message(STATUS "ThunderKittens kernels: forced OFF on ${CMAKE_SYSTEM_PROCESSOR} (tk headers are not aarch64-clean)") set(ENABLE_TK_KERNELS OFF) endif() diff --git a/fastvideo-kernel/README.md b/fastvideo-kernel/README.md index b5f08734f..f332e2fed 100644 --- a/fastvideo-kernel/README.md +++ b/fastvideo-kernel/README.md @@ -9,7 +9,7 @@ Compiled CUDA extensions (CMake, see the build summary printed at the end of eve | Extension | Kernels | Sources | GPU arch | Build gate | |---|---|---|---|---| | `fastvideo_kernel._C.fastvideo_kernel_ops` | TurboDiffusion INT8 GEMM, quant, RMSNorm, LayerNorm | `csrc/turbodiffusion/` | every arch in `TORCH_CUDA_ARCH_LIST` | always built | -| same extension, optional part | ThunderKittens sliding-tile attention (`sta_fwd`) and VSA block-sparse (`block_sparse_fwd/bwd`) | `csrc/attention/*_h100.cu` | Hopper `sm_90a` only | `FASTVIDEO_KERNEL_BUILD_TK` (AUTO = ON iff `9.0a` is in the arch list; always OFF on aarch64 hosts — TK headers don't compile there) | +| same extension, optional part | ThunderKittens sliding-tile attention (`sta_fwd`) and VSA block-sparse (`block_sparse_fwd/bwd`) | `csrc/attention/*_h100.cu` | Hopper `sm_90a` only | `FASTVIDEO_KERNEL_BUILD_TK` (AUTO enables `9.0a` on supported hosts; ON requires TK and fails on aarch64 compiled-extension builds, where its headers do not compile) | | same extension, optional part | MiniMax-H3 block-sparse VSA forward (64- and 128-token blocks) | `csrc/attention/block_sparse*.cu` | Data-center Blackwell `sm_100a`/`sm_103a` | ON iff `10.0a` or `10.3a` is in `TORCH_CUDA_ARCH_LIST` | | same extension, optional part | fused NVLink Ulysses all-to-all | `csrc/comm/ulysses_all_to_all.cu` | CUDA | `FASTVIDEO_KERNEL_BUILD_ULYSSES_A2A` (AUTO = ON with NCCL 2.29+ device headers and library; always OFF on ROCm) | | `fp4attn_cuda`, `fp4quant_cuda` | FP4 attention + quantization ("attn_qat_infer", modified SageAttention3) | `attn_qat_infer/` | consumer/workstation Blackwell `sm_120a` (CUDA ≥ 12.8) + `sm_121a` (DGX Spark GB10; CUDA ≥ 13) | `FASTVIDEO_KERNEL_BUILD_ATTN_QAT_INFER` (AUTO = ON iff `12.0a` or `12.1a` is in the arch list) | @@ -46,7 +46,7 @@ Notes: ## 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. +This detects the GPU architecture and enables ThunderKittens kernels for Hopper (H100/sm_90a) on supported hosts. On aarch64, TK is unavailable and the package uses Triton fallbacks. For a compiled-extension build, set `CMAKE_ARGS="-DFASTVIDEO_KERNEL_BUILD_TK=ON"` to require TK; the build fails on aarch64 instead of silently omitting it. Before installation, set CUDA toolchain paths: diff --git a/fastvideo-kernel/build.sh b/fastvideo-kernel/build.sh index 569417297..71745e10c 100755 --- a/fastvideo-kernel/build.sh +++ b/fastvideo-kernel/build.sh @@ -210,11 +210,9 @@ if [ "${GPU_BACKEND}" = "CUDA" ] && ! (( PYTHON_ONLY )); then fi if ! has_cmake_arg "FASTVIDEO_KERNEL_BUILD_TK"; then - if [ "${cc_major}" = "9" ] && [ "${cc_minor}" = "0" ]; then - CMAKE_ARGS="${CMAKE_ARGS:-} -DFASTVIDEO_KERNEL_BUILD_TK=ON" - else - CMAKE_ARGS="${CMAKE_ARGS:-} -DFASTVIDEO_KERNEL_BUILD_TK=OFF" - fi + # Let CMake detect both the target GPU arch and host support. Explicit ON + # remains a request to fail if TK cannot be built on the host. + CMAKE_ARGS="${CMAKE_ARGS:-} -DFASTVIDEO_KERNEL_BUILD_TK=AUTO" fi fi diff --git a/fastvideo-kernel/tests/test_build_arch.py b/fastvideo-kernel/tests/test_build_arch.py index 83aa945dc..45a200709 100644 --- a/fastvideo-kernel/tests/test_build_arch.py +++ b/fastvideo-kernel/tests/test_build_arch.py @@ -14,7 +14,9 @@ def _write_executable(path: Path, contents: str) -> None: path.chmod(0o755) -def _run_build(tmp_path: Path, capability: str, cuda_release: str) -> tuple[subprocess.CompletedProcess[str], Path]: +def _run_build( + tmp_path: Path, capability: str, cuda_release: str, cmake_args: str | None = None +) -> tuple[subprocess.CompletedProcess[str], Path]: kernel_dir = tmp_path / "fastvideo-kernel" kernel_dir.mkdir() shutil.copy2(BUILD_SCRIPT, kernel_dir / "build.sh") @@ -46,6 +48,8 @@ def _run_build(tmp_path: Path, capability: str, cuda_release: str) -> tuple[subp "FV_TEST_RESULT": str(result_dir), "PATH": f"{fake_bin}{os.pathsep}{env['PATH']}", }) + if cmake_args is not None: + env["CMAKE_ARGS"] = cmake_args result = subprocess.run( ["bash", "build.sh"], cwd=kernel_dir, @@ -82,3 +86,20 @@ def test_build_script_rejects_sm121_before_cmake_on_cuda_12(tmp_path: Path) -> N assert result.returncode != 0 assert "sm_121a requires CUDA Toolkit 13.0+" in result.stderr + + +@pytest.mark.parametrize("capability", ["9.0", "10.3"]) +def test_build_script_defers_tk_detection_to_cmake(tmp_path: Path, capability: str) -> None: + result, result_dir = _run_build(tmp_path, capability, "13.0") + + assert result.returncode == 0, result.stderr + assert "-DFASTVIDEO_KERNEL_BUILD_TK=AUTO" in (result_dir / "cmake_args").read_text() + + +@pytest.mark.parametrize("option", ["ON", "OFF"]) +def test_build_script_preserves_explicit_tk_option(tmp_path: Path, option: str) -> None: + result, result_dir = _run_build(tmp_path, "9.0", "13.0", f"-DFASTVIDEO_KERNEL_BUILD_TK={option}") + + assert result.returncode == 0, result.stderr + assert (result_dir / "cmake_args").read_text().count("-DFASTVIDEO_KERNEL_BUILD_TK=") == 1 + assert f"-DFASTVIDEO_KERNEL_BUILD_TK={option}" in (result_dir / "cmake_args").read_text()