cmake_minimum_required(VERSION 3.26 FATAL_ERROR)
project(fastvideo-kernel LANGUAGES CXX)

# Capture any caller-provided -DCMAKE_CUDA_ARCHITECTURES *before* enable_language(CUDA)
# auto-populates it with CMake's built-in default (an old arch, e.g. sm_75 on CUDA 13).
# torch's cmake actually ignores CMAKE_CUDA_ARCHITECTURES (it drives arch selection via
# TORCH_CUDA_ARCH_LIST), so we only use this captured value to honor an explicit pin by
# translating it into TORCH_CUDA_ARCH_LIST below.
set(_FASTVIDEO_USER_CUDA_ARCH "${CMAKE_CUDA_ARCHITECTURES}")

# Prefer environment variable (used by CI or uv pip install git+repo_addr) if CMake var is not explicitly set.
if(NOT DEFINED GPU_BACKEND AND DEFINED ENV{GPU_BACKEND})
    set(GPU_BACKEND "$ENV{GPU_BACKEND}")
endif()
if(NOT GPU_BACKEND)
    set(GPU_BACKEND "CUDA")
endif()

if(GPU_BACKEND STREQUAL "ROCM")
    enable_language(HIP)
else()
    enable_language(CUDA)
    # Ensure CUDA toolkit targets (CUDA::cudart, CUDA::cuda_driver, etc.) are available.
    find_package(CUDAToolkit REQUIRED)
endif()

# Import common utils if needed, but we keep it simple for now

# Find Python and Torch
find_package(Python COMPONENTS Interpreter Development.Module REQUIRED)

# ---------------------------------------------------------------------------
# Resolve the target CUDA architecture, BEFORE find_package(Torch) below.
#
# torch's cmake (Caffe2 public/cuda.cmake) takes over arch selection: it emits
# the real -gencode flags from TORCH_CUDA_ARCH_LIST and forces
# CMAKE_CUDA_ARCHITECTURES to OFF. So the *effective* arch is whatever
# TORCH_CUDA_ARCH_LIST is when find_package(Torch) runs. build.sh exports it;
# standards-based builds (pip / uv pip install, sdist) don't, and torch then
# auto-detects an arch that does not match the GPU -- the kernels build but fail
# at runtime ("no kernel image is available for execution on the device").
# Resolve it here when absent (mirrors build.sh): honor a pinned
# CMAKE_CUDA_ARCHITECTURES if given, else probe the visible GPU with torch.
# ---------------------------------------------------------------------------
if(NOT GPU_BACKEND STREQUAL "ROCM")
    if(DEFINED ENV{TORCH_CUDA_ARCH_LIST})
        message(STATUS "CUDA arch: TORCH_CUDA_ARCH_LIST=$ENV{TORCH_CUDA_ARCH_LIST} (from environment)")
    elseif(TORCH_CUDA_ARCH_LIST)
        set(ENV{TORCH_CUDA_ARCH_LIST} "${TORCH_CUDA_ARCH_LIST}")
        message(STATUS "CUDA arch: TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST} (from cmake)")
    else()
        set(_FV_ARCH_LIST "")
        if(_FASTVIDEO_USER_CUDA_ARCH)
            # Caller pinned -DCMAKE_CUDA_ARCHITECTURES (which torch ignores); translate it
            # to the TORCH_CUDA_ARCH_LIST spelling: "121" -> "12.1", "90a" -> "9.0a".
            # Only numeric spellings translate; keywords like "native"/"all" would
            # otherwise be mangled into nonsense ("nativ.e").
            foreach(_fv_arch IN LISTS _FASTVIDEO_USER_CUDA_ARCH)
                if(NOT _fv_arch MATCHES "^[0-9]+[af]?$")
                    message(FATAL_ERROR
                        "fastvideo-kernel: CMAKE_CUDA_ARCHITECTURES='${_fv_arch}' is not "
                        "supported. Use a numeric arch (e.g. 90a, 121), set "
                        "TORCH_CUDA_ARCH_LIST directly (e.g. 9.0a), or unset both to "
                        "auto-detect from the visible GPU.")
                endif()
                string(REGEX MATCH "[af]$" _fv_suffix "${_fv_arch}")
                string(REGEX REPLACE "[af]$" "" _fv_num "${_fv_arch}")
                string(REGEX REPLACE "(.)$" ".\\1" _fv_num "${_fv_num}")    # dot before the last digit
                list(APPEND _FV_ARCH_LIST "${_fv_num}${_fv_suffix}")
            endforeach()
            message(STATUS "CUDA arch: TORCH_CUDA_ARCH_LIST=${_FV_ARCH_LIST} (from -DCMAKE_CUDA_ARCHITECTURES=${_FASTVIDEO_USER_CUDA_ARCH})")
        else()
            # Best-effort probe of the visible GPU (mirrors build.sh detect_with_torch).
            execute_process(
                COMMAND "${Python_EXECUTABLE}" -c "import torch; assert torch.cuda.is_available(); mj, mn = torch.cuda.get_device_capability(0); print(f'{mj}.{mn}a' if (mj, mn) in ((9, 0), (12, 0)) else f'{mj}.{mn}')"
                OUTPUT_VARIABLE _FV_ARCH_LIST
                OUTPUT_STRIP_TRAILING_WHITESPACE
                RESULT_VARIABLE _fv_detect_rc
                ERROR_QUIET
            )
            if(_fv_detect_rc EQUAL 0 AND _FV_ARCH_LIST)
                message(STATUS "CUDA arch: TORCH_CUDA_ARCH_LIST=${_FV_ARCH_LIST} (detected via torch, live GPU)")
            else()
                set(_FV_ARCH_LIST "")
            endif()
        endif()

        if(_FV_ARCH_LIST)
            set(TORCH_CUDA_ARCH_LIST "${_FV_ARCH_LIST}")
            set(ENV{TORCH_CUDA_ARCH_LIST} "${_FV_ARCH_LIST}")
        else()
            message(FATAL_ERROR
                "fastvideo-kernel: could not determine the target CUDA architecture.\n"
                "Refusing to let torch auto-detect an arch that may not run on this GPU. "
                "Fix with one of:\n"
                "  - set TORCH_CUDA_ARCH_LIST (e.g. 12.1, or 9.0a for Hopper), or\n"
                "  - pass -DCMAKE_CUDA_ARCHITECTURES=<arch> (e.g. 121), or\n"
                "  - build where the target GPU is visible to torch.\n"
                "Note: 'pip/uv pip install' builds under build isolation, which hides the "
                "GPU; set TORCH_CUDA_ARCH_LIST or add --no-build-isolation. "
                "fastvideo-kernel/build.sh sets all of this for you.")
        endif()
    endif()
endif()

# Robustly find Torch include paths using Python
execute_process(
    COMMAND "${Python_EXECUTABLE}" -c "import torch; from torch.utils.cpp_extension import include_paths; print(';'.join(include_paths()))"
    OUTPUT_VARIABLE TORCH_INCLUDE_PATHS
    OUTPUT_STRIP_TRAILING_WHITESPACE
)
list(APPEND TORCH_INCLUDE_DIRS ${TORCH_INCLUDE_PATHS})

# Find Torch package (still useful for libraries)
find_package(Torch REQUIRED)

# Include directories
include_directories(
    ${CMAKE_SOURCE_DIR}/include
    ${CMAKE_SOURCE_DIR}/include/cutlass/include
    ${CMAKE_SOURCE_DIR}/include/tk/include
    ${CMAKE_SOURCE_DIR}/include/tk/prototype
    ${CMAKE_SOURCE_DIR}/csrc
    ${CMAKE_SOURCE_DIR}/csrc/turbodiffusion
    ${TORCH_INCLUDE_DIRS}
)

# ---------------------------
# ThunderKittens (TK) toggles
# ---------------------------
# AUTO: enable TK only when we can confidently target Hopper (sm_90a).
# ON:   force-enable TK kernels (intended for release wheels/images; does NOT require a GPU).
# OFF:  never build TK kernels.
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_ATTN_QAT_INFER_DEFAULT "AUTO")
if(DEFINED FASTVIDEO_KERNEL_BUILD_MODIFIED_SAGE3 AND NOT DEFINED CACHE{FASTVIDEO_KERNEL_BUILD_ATTN_QAT_INFER})
    set(_FASTVIDEO_KERNEL_BUILD_ATTN_QAT_INFER_DEFAULT "${FASTVIDEO_KERNEL_BUILD_MODIFIED_SAGE3}")
endif()

set(FASTVIDEO_KERNEL_BUILD_ATTN_QAT_INFER "${_FASTVIDEO_KERNEL_BUILD_ATTN_QAT_INFER_DEFAULT}" CACHE STRING
    "Build attn_qat_infer Blackwell inference kernels: AUTO/ON/OFF")
set_property(CACHE FASTVIDEO_KERNEL_BUILD_ATTN_QAT_INFER PROPERTY STRINGS AUTO ON OFF)

if(DEFINED FASTVIDEO_KERNEL_BUILD_MODIFIED_SAGE3)
    message(DEPRECATION
        "FASTVIDEO_KERNEL_BUILD_MODIFIED_SAGE3 is deprecated. "
        "Use FASTVIDEO_KERNEL_BUILD_ATTN_QAT_INFER instead.")
endif()

# 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}")
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_ATTN_QAT_INFER: ${FASTVIDEO_KERNEL_BUILD_ATTN_QAT_INFER}")

set(ENABLE_TK_KERNELS OFF)
if(FASTVIDEO_KERNEL_BUILD_TK STREQUAL "ON")
    set(ENABLE_TK_KERNELS ON)
elseif(FASTVIDEO_KERNEL_BUILD_TK STREQUAL "OFF")
    set(ENABLE_TK_KERNELS OFF)
else()
    # AUTO: detect Hopper if possible.
    if(TORCH_CUDA_ARCH_LIST)
        # Accept common spellings: 9.0a, 90a, sm_90a.
        string(REGEX MATCH "(^|[; ,])((9\\.0a)|(90a)|(sm_90a))([; ,]|$)" _HAS_90A "${TORCH_CUDA_ARCH_LIST}")
        if(_HAS_90A)
            set(ENABLE_TK_KERNELS ON)
        endif()
    else()
        # Best-effort local detection (works when a CUDA device is visible).
        execute_process(
            COMMAND "${Python_EXECUTABLE}" -c "import torch; import sys; \nprint('1' if (torch.cuda.is_available() and torch.version.cuda and torch.cuda.get_device_capability()[0] >= 9) else '0')"
            OUTPUT_VARIABLE _LOCAL_HAS_HOPPER
            OUTPUT_STRIP_TRAILING_WHITESPACE
            ERROR_QUIET
        )
        if(_LOCAL_HAS_HOPPER STREQUAL "1")
            set(ENABLE_TK_KERNELS ON)
        endif()
    endif()
endif()

# ThunderKittens headers don't compile on aarch64 hosts: plain char is unsigned
# 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)$")
    message(STATUS "ThunderKittens kernels: forced OFF on ${CMAKE_SYSTEM_PROCESSOR} (tk headers are not aarch64-clean)")
    set(ENABLE_TK_KERNELS OFF)
endif()

if(ENABLE_TK_KERNELS)
    message(STATUS "ThunderKittens kernels: ENABLED")
else()
    message(STATUS "ThunderKittens kernels: DISABLED (will use Triton fallbacks at runtime)")
endif()

set(ENABLE_ATTN_QAT_INFER OFF)
if(GPU_BACKEND STREQUAL "ROCM")
    message(STATUS "attn_qat_infer kernels: DISABLED (ROCm build)")
else()
    set(_WANTS_ATTN_QAT_INFER OFF)
    if(FASTVIDEO_KERNEL_BUILD_ATTN_QAT_INFER STREQUAL "ON")
        set(_WANTS_ATTN_QAT_INFER ON)
    elseif(FASTVIDEO_KERNEL_BUILD_ATTN_QAT_INFER 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_ATTN_QAT_INFER 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_ATTN_QAT_INFER ON)
            endif()
        endif()
    endif()

    if(_WANTS_ATTN_QAT_INFER)
        if(CUDAToolkit_VERSION VERSION_LESS 12.8)
            message(WARNING
                "attn_qat_infer kernels require CUDA Toolkit 12.8+. "
                "Skipping because CUDAToolkit_VERSION=${CUDAToolkit_VERSION}.")
        else()
            set(ENABLE_ATTN_QAT_INFER ON)
        endif()
    endif()

    if(ENABLE_ATTN_QAT_INFER)
        message(STATUS "attn_qat_infer kernels: ENABLED")
    else()
        message(STATUS
            "attn_qat_infer 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)

# ---------------------------------------------------------------------------
# Per-arch split for the Blackwell FP4 (attn_qat_infer) build
# ---------------------------------------------------------------------------
# The FP4 kernels are sm_120a-only (they emit `cvt.e2m1x2` etc.), while the main
# extension (Hopper-only TK + generic turbodiffusion) targets the full arch list.
# find_package(Torch) injects ONE global -gencode list into CMAKE_CUDA_FLAGS that
# forces every target onto every arch, so the FP4 sources also get the sm_90a pass
# and ptxas rejects their Blackwell instructions. Strip that global list and drive
# arch per target via CUDA_ARCHITECTURES instead (the fp4* targets pin 120a below;
# the main extension gets the full list). Only do this for the FP4 build with an
# explicit arch list, so the cu126 / local autodetect paths stay untouched.
if(ENABLE_ATTN_QAT_INFER AND TORCH_CUDA_ARCH_LIST)
    message(STATUS "[per-arch] CMAKE_CUDA_FLAGS before strip: ${CMAKE_CUDA_FLAGS}")
    string(REGEX REPLACE "-gencode[ =]+arch=[^ ]+" "" CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS}")
    string(REGEX REPLACE " +" " " CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS}")
    message(STATUS "[per-arch] CMAKE_CUDA_FLAGS after strip:  ${CMAKE_CUDA_FLAGS}")
    # Convert TORCH_CUDA_ARCH_LIST ("9.0a;12.0a") to CMake form ("90a;120a").
    set(FASTVIDEO_MAIN_CUDA_ARCHS "${TORCH_CUDA_ARCH_LIST}")
    string(REPLACE "sm_" "" FASTVIDEO_MAIN_CUDA_ARCHS "${FASTVIDEO_MAIN_CUDA_ARCHS}")
    string(REPLACE "." "" FASTVIDEO_MAIN_CUDA_ARCHS "${FASTVIDEO_MAIN_CUDA_ARCHS}")
    message(STATUS "[per-arch] main extension archs=${FASTVIDEO_MAIN_CUDA_ARCHS}, fp4* archs=120a")
endif()

# Compiler flags
set(CUDA_FLAGS
    "-DNDEBUG"
    "-O3"
    "-std=c++20"
    "--use_fast_math"
    "--expt-extended-lambda"
    "--expt-relaxed-constexpr"
    "-Xcompiler=-fno-strict-aliasing"
    "-Xcompiler=-fPIC"
    "-DTORCH_COMPILE"
    "-Xnvlink=--verbose"
    "-Xptxas=--verbose"
    "-Xptxas=--warn-on-spills"
)

# If TK is enabled, ensure we target Hopper. This is required even on GPU-less builders (CI).
if(ENABLE_TK_KERNELS)
    if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES OR CMAKE_CUDA_ARCHITECTURES STREQUAL "")
        set(CMAKE_CUDA_ARCHITECTURES "90a" CACHE STRING "CUDA architectures" FORCE)
    endif()
    list(APPEND CUDA_FLAGS "-DKITTENS_HOPPER")
    message(STATUS "CMAKE_CUDA_ARCHITECTURES: ${CMAKE_CUDA_ARCHITECTURES}")
endif()

if(BUILD_CXX_KERNELS)
    # Source files
    set(EXTENSION_SOURCES 
        csrc/common_extension.cpp
        csrc/turbodiffusion/gemm/gemm.cu
        csrc/turbodiffusion/norm/rmsnorm.cu
        csrc/turbodiffusion/norm/layernorm.cu
        csrc/turbodiffusion/quant/quant.cu
    )
    
    # Conditionally add TK kernels
    if(ENABLE_TK_KERNELS)
        list(APPEND EXTENSION_SOURCES
            csrc/attention/st_attn_h100.cu
            csrc/attention/block_sparse_h100.cu
        )
    endif()

    # Combined FastVideo Extension
    # Using name 'fastvideo_kernel_ops' to distinguish from the python package namespace
    # ---------------------------------------------------------------------------
    # VSA block-sparse attention forward, Blackwell (sm_100a) only.
    #
    # NOTE the "a" suffix: -arch=sm_100a is NOT enough -- it emits a plain sm_100 target and
    # ptxas rejects every tcgen05 / setmaxnreg instruction. The explicit gencode spelling
    # below is required, and matches the 10.0a entry in TORCH_CUDA_ARCH_LIST.
    #
    # Built for 64-token sparse blocks -- FastVideo's default (4,4,4) tiling, so no
    # tile-size change and no top-k granularity change is needed. VSA_BHSD
    # selects [B, H, S, D]. Both are compile-time; the Python is_supported() checks incoming
    # tensors against them so callers fall back to Triton rather than getting a wrong answer.
    # Read the ENVIRONMENT as well as the cache variable. When TORCH_CUDA_ARCH_LIST is
    # exported (build.sh, and `pip install` with it set) the branch above only prints it --
    # 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_SM100A 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}")
    endif()
    if(_VSA_ARCH_LIST MATCHES "(^|[; ,])(10\\.0a|100a|sm_100a)([; ,]|$)")
        set(ENABLE_VSA_SM100A ON)
    endif()
    if(ENABLE_VSA_SM100A)
        message(STATUS "fastvideo-kernel: building block_sparse_sm100a (Blackwell, 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;-DVSA_BHSD=true")
    endif()

    Python_add_library(fastvideo_kernel_ops MODULE USE_SABI ${SKBUILD_SABI_VERSION} WITH_SOABI
        ${EXTENSION_SOURCES}
    )

    # When the per-arch split is active (FP4 build), torch's global gencode was
    # stripped above, so set this target's arch explicitly. TK is guarded down to
    # sm_90a in source; the turbodiffusion kernels are generic, so the main
    # extension targets the full list. (No-FP4 builds keep torch's global gencode.)
    if(ENABLE_ATTN_QAT_INFER AND FASTVIDEO_MAIN_CUDA_ARCHS)
        set_target_properties(fastvideo_kernel_ops PROPERTIES
            CUDA_ARCHITECTURES "${FASTVIDEO_MAIN_CUDA_ARCHS}")
    endif()

    # Build compile definitions list
    set(COMPILE_DEFS TORCH_EXTENSION_NAME=fastvideo_kernel_ops)
    if(ENABLE_VSA_SM100A)
        list(APPEND COMPILE_DEFS TK_COMPILE_BLOCK_SPARSE_VSA_SM100A)
    endif()
    if(ENABLE_TK_KERNELS)
        list(APPEND COMPILE_DEFS TK_COMPILE_ST_ATTN TK_COMPILE_BLOCK_SPARSE)
    endif()


    target_compile_definitions(fastvideo_kernel_ops PRIVATE ${COMPILE_DEFS})

    target_compile_options(fastvideo_kernel_ops PRIVATE
        $<$<COMPILE_LANGUAGE:CUDA>:${CUDA_FLAGS}>
    )

    # Link against Torch libraries to avoid undefined symbols at import time
    # (e.g., torch::autograd vtables) when loading the extension module.
    target_link_libraries(fastvideo_kernel_ops PRIVATE ${TORCH_LIBRARIES})

    # Also link against libtorch_python to satisfy Python-binding symbols
    # (e.g., torch::PyWarningHandler) required by torch/extension.h.
    execute_process(
        COMMAND "${Python_EXECUTABLE}" -c "import torch; from pathlib import Path; p=Path(torch.__file__).parent/'lib'; m=sorted(p.glob('libtorch_python*')); print(str(m[0]) if m else '')"
        OUTPUT_VARIABLE TORCH_PYTHON_LIBRARY_PATH
        OUTPUT_STRIP_TRAILING_WHITESPACE
        ERROR_QUIET
    )
    if(TORCH_PYTHON_LIBRARY_PATH)
        message(STATUS "TORCH_PYTHON_LIBRARY_PATH: ${TORCH_PYTHON_LIBRARY_PATH}")
        target_link_libraries(fastvideo_kernel_ops PRIVATE "${TORCH_PYTHON_LIBRARY_PATH}")
    else()
        message(WARNING "Could not locate libtorch_python; fastvideo_kernel_ops may fail to import.")
    endif()

    # Link CUDA runtime + driver explicitly (fixes missing symbols like cuGetErrorString at import time)
    if(NOT GPU_BACKEND STREQUAL "ROCM")
        target_link_libraries(fastvideo_kernel_ops PRIVATE CUDA::cudart CUDA::cuda_driver)
    endif()

    # We install it to fastvideo_kernel/_C so we can load it to register the ops
    install(TARGETS fastvideo_kernel_ops LIBRARY DESTINATION fastvideo_kernel/_C)
endif()

if(ENABLE_ATTN_QAT_INFER)
    set(ATTN_QAT_INFER_DIR ${CMAKE_SOURCE_DIR}/attn_qat_infer)
    set(ATTN_QAT_INFER_INCLUDE_DIRS
        ${ATTN_QAT_INFER_DIR}
        ${CMAKE_SOURCE_DIR}/include/cutlass/include
        ${CMAKE_SOURCE_DIR}/include/cutlass/tools/util/include
        ${TORCH_INCLUDE_DIRS}
    )
    set(ATTN_QAT_INFER_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
        attn_qat_infer/blackwell/api.cu
    )
    target_include_directories(fp4attn_cuda PRIVATE ${ATTN_QAT_INFER_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>:${ATTN_QAT_INFER_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
        attn_qat_infer/quantization/fp4_quantization_4d.cu
    )
    target_include_directories(fp4quant_cuda PRIVATE ${ATTN_QAT_INFER_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>:${ATTN_QAT_INFER_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()

# One-look answer to "what is this build producing?" — kept last so it is the
# final thing configure prints. The per-kernel matrix lives in README.md.
message(STATUS "============== fastvideo-kernel build summary ==============")
message(STATUS "host / backend:       ${CMAKE_SYSTEM_PROCESSOR} / ${GPU_BACKEND}")
message(STATUS "TORCH_CUDA_ARCH_LIST: ${TORCH_CUDA_ARCH_LIST}")
message(STATUS "fastvideo_kernel_ops: ON  (turbodiffusion int8-gemm/quant/rmsnorm/layernorm, all listed archs)")
message(STATUS "  + TK sta/block_sparse (sm_90a only):          ${ENABLE_TK_KERNELS}")
message(STATUS "fp4attn/fp4quant (sm_120a only, CUDA >= 12.8):  ${ENABLE_ATTN_QAT_INFER}")
message(STATUS "Triton fallbacks ship in python/fastvideo_kernel/triton_kernels regardless.")
message(STATUS "============================================================")