diff --git a/fastvideo-kernel/CMakeLists.txt b/fastvideo-kernel/CMakeLists.txt index ac4b30aea7..be70730d62 100644 --- a/fastvideo-kernel/CMakeLists.txt +++ b/fastvideo-kernel/CMakeLists.txt @@ -397,6 +397,7 @@ if(BUILD_CXX_KERNELS) # the cmake variable stays empty, so testing that alone silently skips the kernel and # leaves a build that succeeds with the op missing. set(ENABLE_VSA_SM100_FAMILY OFF) + set(ENABLE_VSA_SM103A OFF) set(_VSA_ARCH_LIST "${TORCH_CUDA_ARCH_LIST}") if(NOT _VSA_ARCH_LIST AND DEFINED ENV{TORCH_CUDA_ARCH_LIST}) set(_VSA_ARCH_LIST "$ENV{TORCH_CUDA_ARCH_LIST}") @@ -405,12 +406,32 @@ if(BUILD_CXX_KERNELS) set(ENABLE_VSA_SM100_FAMILY ON) endif() if(ENABLE_VSA_SM100_FAMILY) - message(STATUS "fastvideo-kernel: building block_sparse_sm100a (sm_100a/sm_103a, 64- and 128-token blocks)") + set(_VSA_COMPILE_OPTIONS + "-gencode" + "arch=compute_100a,code=sm_100a" + ) + set(_VSA_NATIVE_ARCHS "sm_100a") + # CUDA 12.9 introduced the SM103 compiler target. Keep CUDA 12.8 + # GB200 builds working instead of passing nvcc an unknown compute_103a. + if(NOT CUDAToolkit_VERSION VERSION_LESS 12.9) + list(APPEND _VSA_COMPILE_OPTIONS + "-gencode" + "arch=compute_103a,code=sm_103a" + ) + list(APPEND _VSA_NATIVE_ARCHS "sm_103a") + set(ENABLE_VSA_SM103A ON) + endif() + list(APPEND _VSA_COMPILE_OPTIONS "-DVSA_BHSD=true") + string(JOIN "/" _VSA_NATIVE_ARCHS_DISPLAY ${_VSA_NATIVE_ARCHS}) + message(STATUS + "fastvideo-kernel: building block_sparse_sm100a " + "(${_VSA_NATIVE_ARCHS_DISPLAY}, 64- and 128-token blocks)" + ) list(APPEND EXTENSION_SOURCES csrc/attention/block_sparse_sm100a.cu csrc/attention/block_sparse_blk128_sm100a.cu) set_source_files_properties(csrc/attention/block_sparse_sm100a.cu csrc/attention/block_sparse_blk128_sm100a.cu PROPERTIES - COMPILE_OPTIONS "-gencode;arch=compute_100a,code=sm_100a;-gencode;arch=compute_103a,code=sm_103a;-DVSA_BHSD=true") + COMPILE_OPTIONS "${_VSA_COMPILE_OPTIONS}") endif() Python_add_library(fastvideo_kernel_ops MODULE USE_SABI ${SKBUILD_SABI_VERSION} WITH_SOABI @@ -431,6 +452,9 @@ if(BUILD_CXX_KERNELS) if(ENABLE_VSA_SM100_FAMILY) list(APPEND COMPILE_DEFS TK_COMPILE_BLOCK_SPARSE_VSA_SM100A) endif() + if(ENABLE_VSA_SM103A) + list(APPEND COMPILE_DEFS TK_COMPILE_BLOCK_SPARSE_VSA_SM103A) + endif() if(ENABLE_TK_KERNELS) list(APPEND COMPILE_DEFS TK_COMPILE_ST_ATTN TK_COMPILE_BLOCK_SPARSE) endif() diff --git a/fastvideo-kernel/csrc/common_extension.cpp b/fastvideo-kernel/csrc/common_extension.cpp index 0243b70480..0943f81228 100644 --- a/fastvideo-kernel/csrc/common_extension.cpp +++ b/fastvideo-kernel/csrc/common_extension.cpp @@ -54,6 +54,11 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("block_sparse_sm100a_blk128_fwd", torch::wrap_pybind_function(block_sparse_sm100a_blk128_fwd), "VSA block-sparse attention forward, 128-token blocks (Blackwell sm100a/sm103a)"); +#ifdef TK_COMPILE_BLOCK_SPARSE_VSA_SM103A + m.attr("_has_vsa_sm103a") = true; +#else + m.attr("_has_vsa_sm103a") = false; +#endif #endif #ifdef TK_COMPILE_ST_ATTN diff --git a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_sm100a.py b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_sm100a.py index e8347450a5..8ef5804589 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_sm100a.py +++ b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_sm100a.py @@ -1,8 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 """Data-center Blackwell CUDA block-sparse VSA forward. -The historical ``sm100a`` module and symbol names are retained for compatibility, but the -extension carries native sm_100a and sm_103a images and supports both device generations. +The historical ``sm100a`` module and symbol names are retained for compatibility. The +extension always carries a native sm_100a image and also carries sm_103a when built with +CUDA 12.9 or newer. A third backend behind the same VSA op as the Triton and CuTe-DSL paths. Forward only: it returns ``(out, lse)`` with ``lse`` in exactly the form ``triton_block_sparse_attn_forward`` @@ -29,12 +30,16 @@ 128: getattr(_C, "block_sparse_sm100a_blk128_fwd", None), } _HAS_VSA_SM100A = any(_FWD_BY_BLOCK.values()) + _HAS_VSA_SM103A = bool(getattr(_C, "_has_vsa_sm103a", False)) except ImportError: # pragma: no cover - extension not built _C = None _FWD_BY_BLOCK = {} _HAS_VSA_SM100A = False + _HAS_VSA_SM103A = False -_SUPPORTED_COMPUTE_CAPABILITIES = {(10, 0), (10, 3)} +_SUPPORTED_COMPUTE_CAPABILITIES = {(10, 0)} +if _HAS_VSA_SM103A: + _SUPPORTED_COMPUTE_CAPABILITIES.add((10, 3)) HEAD_DIM = 128 # Must match the -DVSA_BHSD the extension was compiled with (see CMakeLists). BHSD = True diff --git a/tests/test_fasth3_packaging.py b/tests/test_fasth3_packaging.py index c3d2ccdc4c..d0f8654fa8 100644 --- a/tests/test_fasth3_packaging.py +++ b/tests/test_fasth3_packaging.py @@ -39,3 +39,18 @@ def test_kernel_release_matrix_can_publish_data_center_blackwell_wheels(): assert "arch=compute_100a,code=sm_100a" in cmake assert "arch=compute_103a,code=sm_103a" in cmake assert "patchelf==0.17.2.4" in workflow + + +def test_sm103a_gencode_requires_cuda_12_9(): + cmake = (REPO_ROOT / "fastvideo-kernel" / "CMakeLists.txt").read_text(encoding="utf-8") + extension = (REPO_ROOT / "fastvideo-kernel" / "csrc" / "common_extension.cpp").read_text(encoding="utf-8") + backend = (REPO_ROOT / "fastvideo-kernel" / "python" / "fastvideo_kernel" / + "block_sparse_attn_sm100a.py").read_text(encoding="utf-8") + vsa_options = cmake.index("set(_VSA_COMPILE_OPTIONS") + sm103_guard = cmake.index("if(NOT CUDAToolkit_VERSION VERSION_LESS 12.9)", vsa_options) + sm103_gencode = cmake.index("arch=compute_103a,code=sm_103a", sm103_guard) + + assert sm103_guard < sm103_gencode < cmake.index("endif()", sm103_guard) + assert "list(APPEND COMPILE_DEFS TK_COMPILE_BLOCK_SPARSE_VSA_SM103A)" in cmake + assert 'm.attr("_has_vsa_sm103a") = true' in extension + assert 'getattr(_C, "_has_vsa_sm103a", False)' in backend