nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

clean

Browse Source
This commit is contained in:
RandNMR73
2026-04-09 00:29:59 +00:00
parent 027c324d27
commit d96de5f433
7 changed files with 152 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo-kernel/CMakeLists.txt
+123
View File
@@ -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
$<$<COMPILE_LANGUAGE:CXX>:-O3 -std=c++17>
$<$<COMPILE_LANGUAGE:CUDA>:${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
$<$<COMPILE_LANGUAGE:CXX>:-O3 -std=c++17>
$<$<COMPILE_LANGUAGE:CUDA>:${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()
fastvideo-kernel/MANIFEST.in
+1
View File
@@ -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
fastvideo-kernel/README.md
+5
View File
@@ -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.
fastvideo-kernel/pyproject.toml
+1 -1
View File
@@ -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"]
fastvideo/attention/backends/modified_sage_attn3.py
+10 -8
View File
@@ -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()
fastvideo/attention/backends/modified_sageattn_3_setup.py
+5 -2
View File
@@ -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=[
fastvideo/platforms/cuda.py
+7 -3
View File
@@ -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"