From d96de5f43320d3cd0b1cbb35c4d1984ebe286412 Mon Sep 17 00:00:00 2001 From: RandNMR73 Date: Thu, 9 Apr 2026 00:29:59 +0000 Subject: [PATCH] clean --- fastvideo-kernel/CMakeLists.txt | 123 ++++++++++++++++++ fastvideo-kernel/MANIFEST.in | 1 + fastvideo-kernel/README.md | 5 + fastvideo-kernel/pyproject.toml | 2 +- .../attention/backends/modified_sage_attn3.py | 18 +-- .../backends/modified_sageattn_3_setup.py | 7 +- fastvideo/platforms/cuda.py | 10 +- 7 files changed, 152 insertions(+), 14 deletions(-) diff --git a/fastvideo-kernel/CMakeLists.txt b/fastvideo-kernel/CMakeLists.txt index ca8b67c5a..676960c7e 100644 --- a/fastvideo-kernel/CMakeLists.txt +++ b/fastvideo-kernel/CMakeLists.txt @@ -50,6 +50,10 @@ include_directories( 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) +set(FASTVIDEO_KERNEL_BUILD_MODIFIED_SAGE3 "AUTO" CACHE STRING + "Build modified SageAttention3 kernels: AUTO/ON/OFF") +set_property(CACHE FASTVIDEO_KERNEL_BUILD_MODIFIED_SAGE3 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}") @@ -57,6 +61,7 @@ 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 "FASTVIDEO_KERNEL_BUILD_MODIFIED_SAGE3: ${FASTVIDEO_KERNEL_BUILD_MODIFIED_SAGE3}") set(ENABLE_TK_KERNELS OFF) if(FASTVIDEO_KERNEL_BUILD_TK STREQUAL "ON") @@ -91,6 +96,54 @@ else() message(STATUS "ThunderKittens kernels: DISABLED (will use Triton fallbacks at runtime)") endif() +set(ENABLE_MODIFIED_SAGE3 OFF) +if(GPU_BACKEND STREQUAL "ROCM") + message(STATUS "Modified SageAttention3 kernels: DISABLED (ROCm build)") +else() + set(_WANTS_MODIFIED_SAGE3 OFF) + if(FASTVIDEO_KERNEL_BUILD_MODIFIED_SAGE3 STREQUAL "ON") + set(_WANTS_MODIFIED_SAGE3 ON) + elseif(FASTVIDEO_KERNEL_BUILD_MODIFIED_SAGE3 STREQUAL "AUTO") + if(TORCH_CUDA_ARCH_LIST) + string(REGEX MATCH + "(^|[; ,])((12\\.0a)|(120a)|(sm_120a))([; ,]|$)" + _HAS_120A "${TORCH_CUDA_ARCH_LIST}") + if(_HAS_120A) + set(_WANTS_MODIFIED_SAGE3 ON) + endif() + else() + execute_process( + COMMAND "${Python_EXECUTABLE}" -c + "import torch; print('1' if (torch.cuda.is_available() and torch.version.cuda and torch.cuda.get_device_capability()[0] >= 12) else '0')" + OUTPUT_VARIABLE _LOCAL_HAS_BLACKWELL + OUTPUT_STRIP_TRAILING_WHITESPACE + ERROR_QUIET + ) + if(_LOCAL_HAS_BLACKWELL STREQUAL "1") + set(_WANTS_MODIFIED_SAGE3 ON) + endif() + endif() + endif() + + if(_WANTS_MODIFIED_SAGE3) + if(CUDAToolkit_VERSION VERSION_LESS 12.8) + message(WARNING + "Modified SageAttention3 kernels require CUDA Toolkit 12.8+. " + "Skipping because CUDAToolkit_VERSION=${CUDAToolkit_VERSION}.") + else() + set(ENABLE_MODIFIED_SAGE3 ON) + endif() + endif() + + if(ENABLE_MODIFIED_SAGE3) + message(STATUS "Modified SageAttention3 kernels: ENABLED") + else() + message(STATUS + "Modified SageAttention3 kernels: DISABLED " + "(requires CUDA 12.8+ and Blackwell sm_120a)") + endif() +endif() + # Always try to build the extension if CUDA is available, but conditionally add sources/flags set(BUILD_CXX_KERNELS ON) @@ -183,3 +236,73 @@ if(BUILD_CXX_KERNELS) install(TARGETS fastvideo_kernel_ops LIBRARY DESTINATION fastvideo_kernel/_C) endif() +if(ENABLE_MODIFIED_SAGE3) + set(MODIFIED_SAGEATTN_DIR ${CMAKE_SOURCE_DIR}/modified_sageattn) + set(MODIFIED_SAGE3_INCLUDE_DIRS + ${MODIFIED_SAGEATTN_DIR} + ${CMAKE_SOURCE_DIR}/include/cutlass/include + ${CMAKE_SOURCE_DIR}/include/cutlass/tools/util/include + ${TORCH_INCLUDE_DIRS} + ) + set(MODIFIED_SAGE3_CUDA_FLAGS + "-O3" + "-std=c++17" + "-U__CUDA_NO_HALF_OPERATORS__" + "-U__CUDA_NO_HALF_CONVERSIONS__" + "-U__CUDA_NO_BFLOAT16_OPERATORS__" + "-U__CUDA_NO_BFLOAT16_CONVERSIONS__" + "-U__CUDA_NO_BFLOAT162_OPERATORS__" + "-U__CUDA_NO_BFLOAT162_CONVERSIONS__" + "--expt-relaxed-constexpr" + "--expt-extended-lambda" + "--use_fast_math" + "--ptxas-options=--verbose,--warn-on-local-memory-usage" + "-lineinfo" + "-DCUTLASS_DEBUG_TRACE_LEVEL=0" + "-DNDEBUG" + "-DQBLKSIZE=128" + "-DKBLKSIZE=128" + "-DCTA256" + "-DDQINRMEM" + ) + + Python_add_library(fp4attn_cuda MODULE WITH_SOABI + modified_sageattn/blackwell/api.cu + ) + target_include_directories(fp4attn_cuda PRIVATE ${MODIFIED_SAGE3_INCLUDE_DIRS}) + target_compile_definitions(fp4attn_cuda PRIVATE TORCH_EXTENSION_NAME=fp4attn_cuda) + target_compile_options(fp4attn_cuda PRIVATE + $<$:-O3 -std=c++17> + $<$:${MODIFIED_SAGE3_CUDA_FLAGS}> + ) + set_target_properties(fp4attn_cuda PROPERTIES + CUDA_ARCHITECTURES "120a" + CXX_STANDARD 17 + CUDA_STANDARD 17 + ) + target_link_libraries(fp4attn_cuda PRIVATE ${TORCH_LIBRARIES} CUDA::cudart CUDA::cuda_driver) + + Python_add_library(fp4quant_cuda MODULE WITH_SOABI + modified_sageattn/quantization/fp4_quantization_4d.cu + ) + target_include_directories(fp4quant_cuda PRIVATE ${MODIFIED_SAGE3_INCLUDE_DIRS}) + target_compile_definitions(fp4quant_cuda PRIVATE TORCH_EXTENSION_NAME=fp4quant_cuda) + target_compile_options(fp4quant_cuda PRIVATE + $<$:-O3 -std=c++17> + $<$:${MODIFIED_SAGE3_CUDA_FLAGS}> + ) + set_target_properties(fp4quant_cuda PROPERTIES + CUDA_ARCHITECTURES "120a" + CXX_STANDARD 17 + CUDA_STANDARD 17 + ) + target_link_libraries(fp4quant_cuda PRIVATE ${TORCH_LIBRARIES} CUDA::cudart CUDA::cuda_driver) + + if(TORCH_PYTHON_LIBRARY_PATH) + target_link_libraries(fp4attn_cuda PRIVATE "${TORCH_PYTHON_LIBRARY_PATH}") + target_link_libraries(fp4quant_cuda PRIVATE "${TORCH_PYTHON_LIBRARY_PATH}") + endif() + + install(TARGETS fp4attn_cuda LIBRARY DESTINATION .) + install(TARGETS fp4quant_cuda LIBRARY DESTINATION .) +endif() diff --git a/fastvideo-kernel/MANIFEST.in b/fastvideo-kernel/MANIFEST.in index 695242784..630963a5d 100644 --- a/fastvideo-kernel/MANIFEST.in +++ b/fastvideo-kernel/MANIFEST.in @@ -2,5 +2,6 @@ include LICENSE include README.md include pyproject.toml recursive-include python/fastvideo_kernel *.py +recursive-include modified_sageattn *.py *.cu *.cuh *.cpp *.h recursive-include csrc *.cu *.cuh *.cpp *.h recursive-include include/tk *.cu *.cuh *.cpp *.h *.src diff --git a/fastvideo-kernel/README.md b/fastvideo-kernel/README.md index 82d3fe81d..327692b6f 100644 --- a/fastvideo-kernel/README.md +++ b/fastvideo-kernel/README.md @@ -20,6 +20,11 @@ cd fastvideo-kernel ./build.sh ``` +On supported Blackwell environments, the same install also packages +`modified_sageattn` and builds its `fp4attn_cuda` / `fp4quant_cuda` +extensions automatically. This path requires CUDA Toolkit 12.8+ and +targets `sm_120a`. + ### Rocm Build If you are in a rocm environment without the compilation toolchaine of CUDA. diff --git a/fastvideo-kernel/pyproject.toml b/fastvideo-kernel/pyproject.toml index 155f9811d..7bbd1af2a 100644 --- a/fastvideo-kernel/pyproject.toml +++ b/fastvideo-kernel/pyproject.toml @@ -32,4 +32,4 @@ dependencies = [ [tool.scikit-build] cmake.build-type = "Release" minimum-version = "build-system.requires" -wheel.packages = ["python/fastvideo_kernel"] +wheel.packages = ["python/fastvideo_kernel", "modified_sageattn"] diff --git a/fastvideo/attention/backends/modified_sage_attn3.py b/fastvideo/attention/backends/modified_sage_attn3.py index 8ed662a5a..a7464804c 100644 --- a/fastvideo/attention/backends/modified_sage_attn3.py +++ b/fastvideo/attention/backends/modified_sage_attn3.py @@ -17,11 +17,7 @@ try: # Prefer the in-repo kernel implementation during local development. from modified_sageattn import sageattn_blackwell except ImportError: - # Fall back to installed package if local version not available. - try: - from sageattn import sageattn_blackwell - except ImportError: - sageattn_blackwell = None + sageattn_blackwell = None from fastvideo.attention.backends.abstract import ( AttentionBackend, @@ -34,6 +30,10 @@ from fastvideo.logger import init_logger logger = init_logger(__name__) +def is_modified_sageattn_available() -> bool: + return sageattn_blackwell is not None + + class ModifiedSageAttention3Backend(AttentionBackend): accept_output_buffer: bool = True @@ -82,9 +82,11 @@ class ModifiedSageAttention3Impl(AttentionImpl): value: torch.Tensor, attn_metadata: AttentionMetadata, ) -> torch.Tensor: - if sageattn_blackwell is None: - raise ImportError("modified_sageattn is not available. " - "Please ensure the modified SageAttention3 kernel is installed.") + if not is_modified_sageattn_available(): + raise ImportError( + "modified_sageattn is not available. Please ensure the " + "modified SageAttention3 kernel is installed." + ) query = query.transpose(1, 2).contiguous() key = key.transpose(1, 2).contiguous() diff --git a/fastvideo/attention/backends/modified_sageattn_3_setup.py b/fastvideo/attention/backends/modified_sageattn_3_setup.py index 15d1ee5e4..95a2cf4d7 100644 --- a/fastvideo/attention/backends/modified_sageattn_3_setup.py +++ b/fastvideo/attention/backends/modified_sageattn_3_setup.py @@ -10,6 +10,8 @@ import torch from torch.utils.cpp_extension import BuildExtension, CppExtension, CUDAExtension, CUDA_HOME this_dir = os.path.dirname(os.path.abspath(__file__)) +repo_dir = Path(this_dir).resolve() +package_root = repo_dir.parents[2] / "fastvideo-kernel" PACKAGE_NAME = "modified_sageattn" @@ -68,9 +70,8 @@ if not SKIP_CUDA_BUILD: # https://github.com/pytorch/pytorch/blob/8472c24e3b5b60150096486616d98b7bea01500b/torch/utils/cpp_extension.py#L920 if FORCE_CXX11_ABI: torch._C._GLIBCXX_USE_CXX11_ABI = True - repo_dir = Path(this_dir).resolve() modified_sageattn_dir = ( - repo_dir.parents[2] / "fastvideo-kernel" / "modified_sageattn" + package_root / "modified_sageattn" ) cutlass_dir = repo_dir / "csrc" / "cutlass" (repo_dir / "csrc").mkdir(parents=True, exist_ok=True) @@ -149,6 +150,7 @@ setup( name=PACKAGE_NAME, version="1.0.0", packages=find_packages( + where=str(package_root), exclude=( "build", "csrc", @@ -158,6 +160,7 @@ setup( "benchmarks", ) ), + package_dir={"": str(package_root)}, description="FP4FlashAttention", long_description_content_type="text/markdown", classifiers=[ diff --git a/fastvideo/platforms/cuda.py b/fastvideo/platforms/cuda.py index adfac628b..12477d67e 100644 --- a/fastvideo/platforms/cuda.py +++ b/fastvideo/platforms/cuda.py @@ -142,10 +142,14 @@ class CudaPlatformBase(Platform): logger.info("Sage Attention 3 backend is not installed. Fall back to Flash Attention.") elif selected_backend == AttentionBackendEnum.MODIFIED_SAGE_ATTN_THREE: try: - from modified_sageattn import sageattn_blackwell # noqa: F401 - from fastvideo.attention.backends.modified_sage_attn3 import ( # noqa: F401 - ModifiedSageAttention3Backend) + ModifiedSageAttention3Backend, + is_modified_sageattn_available, + ) + if not is_modified_sageattn_available(): + raise ImportError( + "modified_sageattn could not be imported." + ) logger.info("Using modified Sage Attention 3 backend.") return "fastvideo.attention.backends.modified_sage_attn3.ModifiedSageAttention3Backend"