Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 26 additions & 2 deletions fastvideo-kernel/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -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}")
Expand All @@ -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
Expand All @@ -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()
Expand Down
5 changes: 5 additions & 0 deletions fastvideo-kernel/csrc/common_extension.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
@@ -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``
Expand All @@ -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
Expand Down
15 changes: 15 additions & 0 deletions tests/test_fasth3_packaging.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Loading