[bugfix] Honor explicit ThunderKittens build requests on aarch64 (#2062)

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()

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:
 

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
 

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()