cmake_minimum_required(VERSION 3.26 FATAL_ERROR) project(sgl-kernel LANGUAGES CXX CUDA) # utils include(${CMAKE_CURRENT_LIST_DIR}/cmake/utils.cmake) include(FetchContent) # CMake cmake_policy(SET CMP0169 OLD) cmake_policy(SET CMP0177 NEW) set(CMAKE_COLOR_DIAGNOSTICS ON) set(CMAKE_VERBOSE_MAKEFILE ON CACHE BOOL "ON") set(CMAKE_POSITION_INDEPENDENT_CODE ON) set(CMAKE_SHARED_LIBRARY_PREFIX "") # GitHub Artifactory set(GITHUB_ARTIFACTORY "github.com" CACHE STRING "GitHub mirror URL") # Python find_package(Python COMPONENTS Interpreter Development.Module ${SKBUILD_SABI_COMPONENT} REQUIRED) # CXX set(CMAKE_CXX_STANDARD 17) set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O3") # CUDA enable_language(CUDA) find_package(CUDAToolkit REQUIRED) set_property(GLOBAL PROPERTY CUDA_SEPARABLE_COMPILATION ON) message(STATUS "Detected CUDA_VERSION=${CUDA_VERSION}") if ("${CUDA_VERSION}" VERSION_GREATER_EQUAL "13.0") message("CUDA_VERSION ${CUDA_VERSION} >= 13.0") elseif ("${CUDA_VERSION}" VERSION_GREATER_EQUAL "12.8") message("CUDA_VERSION ${CUDA_VERSION} >= 12.8") elseif ("${CUDA_VERSION}" VERSION_GREATER_EQUAL "12.4") message("CUDA_VERSION ${CUDA_VERSION} >= 12.4") elseif ("${CUDA_VERSION}" VERSION_GREATER_EQUAL "12.1") message("CUDA_VERSION ${CUDA_VERSION} >= 12.1") elseif ("${CUDA_VERSION}" VERSION_GREATER_EQUAL "11.8") message("CUDA_VERSION ${CUDA_VERSION} >= 11.8") endif() # Torch find_package(Torch REQUIRED) clear_cuda_arches(CMAKE_FLAG) # Third Party repos # cutlass FetchContent_Declare( repo-cutlass URL https://${GITHUB_ARTIFACTORY}/NVIDIA/cutlass/archive/57e3cfb47a2d9e0d46eb6335c3dc411498efa198.tar.gz URL_HASH SHA256=09237099a70f80bff1dc8bb80c843a674bb4fdcb46e43cc6993e711c5ca89bb5 ) FetchContent_Populate(repo-cutlass) # fmt FetchContent_Declare( repo-fmt URL https://${GITHUB_ARTIFACTORY}/fmtlib/fmt/archive/553ec11ec06fbe0beebfbb45f9dc3c9eabd83d28.tar.gz URL_HASH SHA256=c314292789d28c3c3b420e75a7b2d1706f685f7fb63289128d46aeaea2c6be71 ) FetchContent_Populate(repo-fmt) # Triton kernel FetchContent_Declare( repo-triton URL https://${GITHUB_ARTIFACTORY}/triton-lang/triton/archive/v3.6.0.tar.gz URL_HASH SHA256=be270ed11ca5a8fbd9d7941c5bbe9a23a9f6e2ffd372c8398346928bee464774 ) FetchContent_Populate(repo-triton) # flashinfer FetchContent_Declare( repo-flashinfer URL https://${GITHUB_ARTIFACTORY}/flashinfer-ai/flashinfer/archive/bc29697ba20b7e6bdb728ded98f04788e16ee021.tar.gz URL_HASH SHA256=931dfd118f4b6de8c7d98702153c7c03840139170af21a07607693bd9749744d ) FetchContent_Populate(repo-flashinfer) # flash-attention FetchContent_Declare( repo-flash-attention URL https://${GITHUB_ARTIFACTORY}/sgl-project/sgl-attn/archive/f89bc2306632d1ec5f97b014dded4254f5b4a907.tar.gz URL_HASH SHA256=418b5681584dc3efff496a1cab5ffd58d2728d89dcfe0ea16e6985d6ef35c68c ) FetchContent_Populate(repo-flash-attention) # ccache option option(ENABLE_CCACHE "Whether to use ccache" ON) find_program(CCACHE_FOUND ccache) if(CCACHE_FOUND AND ENABLE_CCACHE AND DEFINED ENV{CCACHE_DIR}) message(STATUS "Building with CCACHE enabled") set_property(GLOBAL PROPERTY RULE_LAUNCH_COMPILE "ccache") set_property(GLOBAL PROPERTY RULE_LAUNCH_LINK "ccache") endif() # Configure gencode below SM90 if(CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64") set(DEFAULT_ENABLE_BELOW_SM90 OFF) message(STATUS "For aarch64, disable gencode below SM90 by default") else() set(DEFAULT_ENABLE_BELOW_SM90 ON) endif() option(ENABLE_BELOW_SM90 "Enable gencode below SM90" ${DEFAULT_ENABLE_BELOW_SM90}) set(DEFAULT_SGL_KERNEL_ENABLE_FA3 OFF) if (CMAKE_SYSTEM_PROCESSOR MATCHES "^(aarch64|arm64)$") message(STATUS "For aarch64, disable FA3 by default") endif() if ("${CUDA_VERSION}" VERSION_GREATER_EQUAL "12.4" AND NOT CMAKE_SYSTEM_PROCESSOR MATCHES "^(aarch64|arm64)$") set(DEFAULT_SGL_KERNEL_ENABLE_FA3 ON) endif() include_directories( ${PROJECT_SOURCE_DIR}/include ${PROJECT_SOURCE_DIR}/csrc ) set(SGL_KERNEL_CUDA_FLAGS "-DNDEBUG" "-DOPERATOR_NAMESPACE=sgl-kernel" "-O3" "-Xcompiler" "-fPIC" "-gencode=arch=compute_90,code=sm_90" "-std=c++17" "-DFLASHINFER_ENABLE_F16" "-DCUTE_USE_PACKED_TUPLE=1" "-DCUTLASS_ENABLE_TENSOR_CORE_MMA=1" "-DCUTLASS_VERSIONS_GENERATED" "-DCUTLASS_TEST_LEVEL=0" "-DCUTLASS_TEST_ENABLE_CACHED_RESULTS=1" "-DCUTLASS_DEBUG_TRACE_LEVEL=0" "--expt-relaxed-constexpr" "--expt-extended-lambda" # The following flag leads to the CMAKE_BUILD_PARALLEL_LEVEL breaking, # it triggers OOM with low memory host. Extract the threads number to # option named SGL_KERNEL_COMPILE_THREADS, default value 32. # "--threads=32" # Supress warnings "-Xcompiler=-Wno-clang-format-violations" "-Xcompiler=-Wno-conversion" "-Xcompiler=-Wno-deprecated-declarations" "-Xcompiler=-Wno-terminate" "-Xcompiler=-Wfatal-errors" "-Xcompiler=-ftemplate-backtrace-limit=1" "-Xcudafe=--diag_suppress=177" # variable was declared but never referenced "-Xcudafe=--diag_suppress=2361" # invalid narrowing conversion from "char" to "signed char" # uncomment to debug # "--ptxas-options=-v" # "--ptxas-options=--verbose,--register-usage-level=10,--warn-on-local-memory-usage" ) set(SGL_KERNEL_COMPILE_THREADS 32 CACHE STRING "Set compilation threads, default 32") # When SGL_KERNEL_COMPILE_THREADS value is less than 1, set it to 1 if (NOT SGL_KERNEL_COMPILE_THREADS MATCHES "^[0-9]+$") message(FATAL_ERROR "SGL_KERNEL_COMPILE_THREADS must be an integer, but was set to '${SGL_KERNEL_COMPILE_THREADS}'.") elseif (SGL_KERNEL_COMPILE_THREADS LESS 1) message(STATUS "SGL_KERNEL_COMPILE_THREADS was set to a value less than 1. Using 1 instead.") set(SGL_KERNEL_COMPILE_THREADS 1) endif() list(APPEND SGL_KERNEL_CUDA_FLAGS "--threads=${SGL_KERNEL_COMPILE_THREADS}" ) option(SGL_KERNEL_ENABLE_BF16 "Enable BF16" ON) option(SGL_KERNEL_ENABLE_FP8 "Enable FP8" ON) option(SGL_KERNEL_ENABLE_FP4 "Enable FP4" OFF) option(SGL_KERNEL_ENABLE_FA3 "Enable FA3" ${DEFAULT_SGL_KERNEL_ENABLE_FA3}) option(SGL_KERNEL_ENABLE_FA3_SPARSE_MASK "Enable FA3 sparse mask kernels" OFF) option(SGL_KERNEL_ENABLE_SM90A "Enable SM90A" OFF) option(SGL_KERNEL_ENABLE_SM100A "Enable SM100A" OFF) if (SGL_KERNEL_ENABLE_BF16) list(APPEND SGL_KERNEL_CUDA_FLAGS "-DFLASHINFER_ENABLE_BF16" ) endif() if (SGL_KERNEL_ENABLE_FP8) list(APPEND SGL_KERNEL_CUDA_FLAGS "-DFLASHINFER_ENABLE_FP8" "-DFLASHINFER_ENABLE_FP8_E4M3" "-DFLASHINFER_ENABLE_FP8_E5M2" ) endif() if (ENABLE_BELOW_SM90) list(APPEND SGL_KERNEL_CUDA_FLAGS "-gencode=arch=compute_80,code=sm_80" "-gencode=arch=compute_89,code=sm_89" ) if (CMAKE_SYSTEM_PROCESSOR STREQUAL "aarch64") list(APPEND SGL_KERNEL_CUDA_FLAGS "-gencode=arch=compute_87,code=sm_87" ) endif() endif() if ("${CUDA_VERSION}" VERSION_GREATER_EQUAL "12.8" OR SGL_KERNEL_ENABLE_SM100A) list(APPEND SGL_KERNEL_CUDA_FLAGS "-gencode=arch=compute_100a,code=sm_100a" "-gencode=arch=compute_120a,code=sm_120a" ) # refer sm_121, sm_110 and sm_101 description https://github.com/pytorch/pytorch/pull/156176 if ("${CUDA_VERSION}" VERSION_GREATER_EQUAL "13.0") list(APPEND SGL_KERNEL_CUDA_FLAGS "-gencode=arch=compute_103a,code=sm_103a" "--compress-mode=size" ) if (CMAKE_SYSTEM_PROCESSOR STREQUAL "aarch64") list(APPEND SGL_KERNEL_CUDA_FLAGS "-gencode=arch=compute_110a,code=sm_110a" "-gencode=arch=compute_121a,code=sm_121a" ) endif() else() if (CMAKE_SYSTEM_PROCESSOR STREQUAL "aarch64") list(APPEND SGL_KERNEL_CUDA_FLAGS "-gencode=arch=compute_101a,code=sm_101a" ) endif() endif() endif() if ("${CUDA_VERSION}" VERSION_GREATER_EQUAL "12.4" AND SGL_KERNEL_ENABLE_FA3) list(APPEND SGL_KERNEL_CUDA_FLAGS "-gencode=arch=compute_90a,code=sm_90a" ) endif() if ("${CUDA_VERSION}" VERSION_GREATER_EQUAL "12.8" OR SGL_KERNEL_ENABLE_FP4) list(APPEND SGL_KERNEL_CUDA_FLAGS "-DENABLE_NVFP4=1" ) endif() # All source files # NOTE: Please sort the filenames alphabetically set(SOURCES "csrc/allreduce/custom_all_reduce.cu" "csrc/attention/cutlass_mla_kernel.cu" "csrc/attention/merge_attn_states.cu" "csrc/attention/vertical_slash_index.cu" "csrc/common_extension.cc" "csrc/elementwise/activation.cu" "csrc/elementwise/concat_mla.cu" "csrc/elementwise/copy.cu" "csrc/elementwise/dsv4_norm_rope.cu" "csrc/elementwise/fused_add_rms_norm_kernel.cu" "csrc/elementwise/pos_enc.cu" "csrc/elementwise/topk.cu" "csrc/expert_specialization/es_fp8_blockwise.cu" "csrc/expert_specialization/es_sm100_mxfp8_blockscaled.cu" "csrc/expert_specialization/es_sm100_mxfp8_blockscaled_group_quant.cu" "csrc/gemm/awq_kernel.cu" "csrc/gemm/bmm_fp8.cu" "csrc/gemm/dsv3_fused_a_gemm.cu" "csrc/gemm/fp8_blockwise_gemm_kernel.cu" "csrc/gemm/fp8_gemm_kernel.cu" "csrc/gemm/int8_gemm_kernel.cu" "csrc/gemm/per_token_group_quant_8bit.cu" "csrc/gemm/per_token_group_quant_8bit_v2.cu" "csrc/gemm/per_token_quant_fp8.cu" "csrc/gemm/qserve_w4a8_per_chn_gemm.cu" "csrc/gemm/qserve_w4a8_per_group_gemm.cu" "csrc/gemm/gptq/gptq_kernel.cu" "csrc/grammar/apply_token_bitmask_inplace_cuda.cu" "csrc/infllm_v2/max_pooling.cu" "csrc/kvcacheio/transfer.cu" "csrc/mamba/causal_conv1d.cu" "csrc/memory/weak_ref_tensor.cpp" "csrc/moe/cutlass_moe/w4a8/scaled_mm_entry.cu" "csrc/moe/cutlass_moe/w4a8/w4a8_moe_data.cu" "csrc/moe/cutlass_moe/w4a8/w4a8_grouped_mm_c3x.cu" "csrc/moe/moe_align_kernel.cu" "csrc/moe/fused_qknorm_rope_kernel.cu" "csrc/moe/moe_sum.cu" "csrc/moe/moe_sum_reduce.cu" "csrc/moe/moe_topk_softmax_kernels.cu" "csrc/moe/moe_topk_sigmoid_kernels.cu" "csrc/moe/fp8_blockwise_moe_kernel.cu" "csrc/moe/prepare_moe_input.cu" "csrc/quantization/gguf/gguf_kernel.cu" "csrc/speculative/eagle_utils.cu" "csrc/speculative/ngram_utils.cu" "csrc/speculative/packbit.cu" "csrc/speculative/speculative_sampling.cu" "${repo-flashinfer_SOURCE_DIR}/csrc/norm.cu" "${repo-flashinfer_SOURCE_DIR}/csrc/renorm.cu" "${repo-flash-attention_SOURCE_DIR}/csrc/flash_attn/src/flash_fwd_sparse_hdim128_bf16_causal_sm80.cu" "${repo-flash-attention_SOURCE_DIR}/csrc/flash_attn/src/flash_fwd_sparse_hdim128_bf16_sm80.cu" "${repo-flash-attention_SOURCE_DIR}/csrc/flash_attn/src/flash_fwd_sparse_hdim128_fp16_causal_sm80.cu" "${repo-flash-attention_SOURCE_DIR}/csrc/flash_attn/src/flash_fwd_sparse_hdim128_fp16_sm80.cu" "${repo-flash-attention_SOURCE_DIR}/csrc/flash_attn/flash_sparse_api.cpp" ) set(INCLUDES ${repo-cutlass_SOURCE_DIR}/include ${repo-cutlass_SOURCE_DIR}/tools/util/include ${repo-flashinfer_SOURCE_DIR}/include ${repo-flashinfer_SOURCE_DIR}/csrc ${repo-cutlass_SOURCE_DIR}/examples/77_blackwell_fmha ${repo-cutlass_SOURCE_DIR}/examples/common ${repo-flash-attention_SOURCE_DIR}/csrc/flash_attn/src ) # =========================== Common SM90 Build ============================= # # Build SM90 library with fast math optimization (same namespace, different directory) Python_add_library(common_ops_sm90_build MODULE USE_SABI ${SKBUILD_SABI_VERSION} WITH_SOABI ${SOURCES}) target_compile_options(common_ops_sm90_build PRIVATE $<$:${SGL_KERNEL_CUDA_FLAGS} -use_fast_math> ) target_include_directories(common_ops_sm90_build PRIVATE ${INCLUDES}) # Set output name and separate build directory to avoid conflicts set_target_properties(common_ops_sm90_build PROPERTIES OUTPUT_NAME "common_ops" LIBRARY_OUTPUT_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/sm90" ) # =========================== Common SM100+ Build ============================= # # Build SM100+ library with precise math (same namespace, different directory) Python_add_library(common_ops_sm100_build MODULE USE_SABI ${SKBUILD_SABI_VERSION} WITH_SOABI ${SOURCES}) target_compile_options(common_ops_sm100_build PRIVATE $<$:${SGL_KERNEL_CUDA_FLAGS}> ) target_include_directories(common_ops_sm100_build PRIVATE ${INCLUDES}) # Set output name and separate build directory to avoid conflicts set_target_properties(common_ops_sm100_build PROPERTIES OUTPUT_NAME "common_ops" LIBRARY_OUTPUT_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/sm100" ) find_package(Python3 COMPONENTS Interpreter REQUIRED) execute_process( COMMAND ${Python3_EXECUTABLE} -c "import torch; print(int(torch._C._GLIBCXX_USE_CXX11_ABI))" OUTPUT_VARIABLE TORCH_CXX11_ABI OUTPUT_STRIP_TRAILING_WHITESPACE ) if(TORCH_CXX11_ABI STREQUAL "0") message(STATUS "Using old C++ ABI (-D_GLIBCXX_USE_CXX11_ABI=0)") set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -D_GLIBCXX_USE_CXX11_ABI=0") set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -D_GLIBCXX_USE_CXX11_ABI=0") else() message(STATUS "Using new C++11 ABI (-D_GLIBCXX_USE_CXX11_ABI=1)") set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -D_GLIBCXX_USE_CXX11_ABI=1") set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -D_GLIBCXX_USE_CXX11_ABI=1") endif() target_link_libraries(common_ops_sm90_build PRIVATE ${TORCH_LIBRARIES} c10 cuda cublas cublasLt) target_link_libraries(common_ops_sm100_build PRIVATE ${TORCH_LIBRARIES} c10 cuda cublas cublasLt) # sparse flash attention target_compile_definitions(common_ops_sm90_build PRIVATE FLASHATTENTION_DISABLE_BACKWARD FLASHATTENTION_DISABLE_DROPOUT FLASHATTENTION_DISABLE_UNEVEN_K ) target_compile_definitions(common_ops_sm100_build PRIVATE FLASHATTENTION_DISABLE_BACKWARD FLASHATTENTION_DISABLE_DROPOUT FLASHATTENTION_DISABLE_UNEVEN_K ) # Install to different subdirectories # CMake will find the built libraries in their respective LIBRARY_OUTPUT_DIRECTORY locations # and install them to the specified destinations install(TARGETS common_ops_sm90_build LIBRARY DESTINATION sgl_kernel/sm90) install(TARGETS common_ops_sm100_build LIBRARY DESTINATION sgl_kernel/sm100) # ============================ Optional Install: FA3 ============================= # # set flash-attention sources file # Now FA3 support sm80/sm86/sm90 if (SGL_KERNEL_ENABLE_FA3) set(SGL_FLASH_KERNEL_CUDA_FLAGS "-DNDEBUG" "-DOPERATOR_NAMESPACE=sgl-kernel" "-O3" "-Xcompiler" "-fPIC" "-gencode=arch=compute_90a,code=sm_90a" "-std=c++17" "-DCUTE_USE_PACKED_TUPLE=1" "-DCUTLASS_ENABLE_TENSOR_CORE_MMA=1" "-DCUTLASS_VERSIONS_GENERATED" "-DCUTLASS_TEST_LEVEL=0" "-DCUTLASS_TEST_ENABLE_CACHED_RESULTS=1" "-DCUTLASS_DEBUG_TRACE_LEVEL=0" "-DCUTLASS_ENABLE_GDC_FOR_SM90" # For PDL "-DCUTE_SM90_EXTENDED_MMA_SHAPES_ENABLED" # Necessary for the WGMMA shapes that we use "--expt-relaxed-constexpr" "--expt-extended-lambda" "--use_fast_math" "-Xcompiler=-Wconversion" "-Xcompiler=-fno-strict-aliasing" ) if (ENABLE_BELOW_SM90) list(APPEND SGL_FLASH_KERNEL_CUDA_FLAGS "-gencode=arch=compute_80,code=sm_80" "-gencode=arch=compute_86,code=sm_86" ) # SM8X Logic file(GLOB FA3_SM8X_GEN_SRCS "${repo-flash-attention_SOURCE_DIR}/hopper/instantiations/flash_fwd_hdim*_sm80.cu") endif() file(GLOB FA3_BF16_GEN_SRCS "${repo-flash-attention_SOURCE_DIR}/hopper/instantiations/flash_fwd_hdim[0-9]*_bf16*_sm90.cu") file(GLOB FA3_BF16_GEN_SRCS_ "${repo-flash-attention_SOURCE_DIR}/hopper/instantiations/flash_fwd_hdimdiff_bf16*_sm90.cu") list(APPEND FA3_BF16_GEN_SRCS ${FA3_BF16_GEN_SRCS_}) # FP16 source files - use individual hdim files instead of hdimall to avoid ptxas crash file(GLOB FA3_FP16_GEN_SRCS "${repo-flash-attention_SOURCE_DIR}/hopper/instantiations/flash_fwd_hdim[0-9]*_fp16*_sm90.cu") file(GLOB FA3_FP16_GEN_SRCS_ "${repo-flash-attention_SOURCE_DIR}/hopper/instantiations/flash_fwd_hdimdiff_fp16*_sm90.cu") list(APPEND FA3_FP16_GEN_SRCS ${FA3_FP16_GEN_SRCS_}) # FP8 source files file(GLOB FA3_FP8_GEN_SRCS "${repo-flash-attention_SOURCE_DIR}/hopper/instantiations/flash_fwd_hdim[0-9]*_e4m3*_sm90.cu") file(GLOB FA3_FP8_GEN_SRCS_ "${repo-flash-attention_SOURCE_DIR}/hopper/instantiations/flash_fwd_hdimdiff_e4m3*_sm90.cu") list(APPEND FA3_FP8_GEN_SRCS ${FA3_FP8_GEN_SRCS_}) set(FA3_GEN_SRCS ${FA3_BF16_GEN_SRCS} ${FA3_FP16_GEN_SRCS} ${FA3_FP8_GEN_SRCS} ${FA3_SM8X_GEN_SRCS}) set(FLASH_SOURCES "csrc/flash_extension.cc" "${repo-flash-attention_SOURCE_DIR}/hopper/flash_prepare_scheduler.cu" "${repo-flash-attention_SOURCE_DIR}/hopper/flash_api.cpp" "${repo-flash-attention_SOURCE_DIR}/hopper/flash_fwd_combine.cu" "${FA3_GEN_SRCS}" ) Python_add_library(flash_ops MODULE USE_SABI ${SKBUILD_SABI_VERSION} WITH_SOABI ${FLASH_SOURCES}) target_compile_options(flash_ops PRIVATE $<$:${SGL_FLASH_KERNEL_CUDA_FLAGS}>) target_include_directories(flash_ops PRIVATE ${repo-cutlass_SOURCE_DIR}/include ${repo-cutlass_SOURCE_DIR}/tools/util/include ${repo-flash-attention_SOURCE_DIR}/hopper ) target_link_libraries(flash_ops PRIVATE ${TORCH_LIBRARIES} c10 cuda) install(TARGETS flash_ops LIBRARY DESTINATION "sgl_kernel") set(FLASH_OPS_COMPILE_DEFS FLASHATTENTION_DISABLE_BACKWARD FLASHATTENTION_DISABLE_DROPOUT FLASHATTENTION_DISABLE_UNEVEN_K FLASHATTENTION_VARLEN_ONLY ) if(NOT ENABLE_BELOW_SM90) list(APPEND FLASH_OPS_COMPILE_DEFS FLASHATTENTION_DISABLE_SM8x) endif() if(NOT SGL_KERNEL_ENABLE_FA3_SPARSE_MASK) list(APPEND FLASH_OPS_COMPILE_DEFS FLASHATTENTION_DISABLE_SPARSE_MASK) endif() target_compile_definitions(flash_ops PRIVATE ${FLASH_OPS_COMPILE_DEFS}) endif() # ===================== InfLLM-V2 FlashAttention backend ===================== # # Standalone pybind extension `infllm_ops`, vendored from # 3rdparty/infllmv2_cuda_impl. Kept as its own module so its `flash::` symbols # stay isolated from sgl-kernel's own flash attention. Mirrors the original # setup.py: only hdim 64/128 bf16 forward instantiations are compiled (the # vendored static_switch.h forces bf16 and dispatches headdim to {64, 128} # only). Backward kernels are intentionally omitted because SGLang only uses # these ops for inference. set(INFLLM_FLASH_CUDA_FLAGS "-DNDEBUG" "-O3" "-std=c++17" "-Xcompiler" "-fPIC" "-U__CUDA_NO_HALF_OPERATORS__" "-U__CUDA_NO_HALF_CONVERSIONS__" "-U__CUDA_NO_HALF2_OPERATORS__" "-U__CUDA_NO_BFLOAT16_CONVERSIONS__" "--expt-relaxed-constexpr" "--expt-extended-lambda" "--use_fast_math" "-DFLASHATTENTION_DISABLE_DROPOUT" "-DFLASHATTENTION_DISABLE_ALIBI" "-DFLASHATTENTION_DISABLE_SOFTCAP" "-DFLASHATTENTION_DISABLE_UNEVEN_K" "-DFLASHATTENTION_DISABLE_LOCAL" "--threads=${SGL_KERNEL_COMPILE_THREADS}" ) # Arch gencodes: match the original setup.py auto-detection # (80 always; 90 for CUDA>=11.8; 120 for CUDA>=12.8). if (ENABLE_BELOW_SM90) list(APPEND INFLLM_FLASH_CUDA_FLAGS "-gencode=arch=compute_80,code=sm_80") endif() list(APPEND INFLLM_FLASH_CUDA_FLAGS "-gencode=arch=compute_90,code=sm_90") if ("${CUDA_VERSION}" VERSION_GREATER_EQUAL "12.8" OR SGL_KERNEL_ENABLE_SM100A) list(APPEND INFLLM_FLASH_CUDA_FLAGS "-gencode=arch=compute_120a,code=sm_120a") endif() set(INFLLM_FLASH_SOURCES "csrc/infllm_v2/flash_extension.cc" "csrc/infllm_v2/flash_attn/flash_api.cpp" "csrc/infllm_v2/flash_attn/src/flash_fwd_split_hdim64_bf16_sm80.cu" "csrc/infllm_v2/flash_attn/src/flash_fwd_split_hdim128_bf16_sm80.cu" "csrc/infllm_v2/flash_attn/src/flash_fwd_split_hdim64_bf16_causal_sm80.cu" "csrc/infllm_v2/flash_attn/src/flash_fwd_split_hdim128_bf16_causal_sm80.cu" ) Python_add_library(infllm_ops MODULE WITH_SOABI ${INFLLM_FLASH_SOURCES}) target_compile_options(infllm_ops PRIVATE $<$:${INFLLM_FLASH_CUDA_FLAGS}>) target_include_directories(infllm_ops PRIVATE ${repo-cutlass_SOURCE_DIR}/include ${repo-cutlass_SOURCE_DIR}/tools/util/include ${CMAKE_CURRENT_LIST_DIR}/csrc/infllm_v2/flash_attn ${CMAKE_CURRENT_LIST_DIR}/csrc/infllm_v2/flash_attn/src ) # The pybind module binds functions taking at::Generator, which pulls in # THPGeneratorClass from libtorch_python (not part of TORCH_LIBRARIES). find_library(TORCH_PYTHON_LIBRARY torch_python PATHS "${TORCH_INSTALL_PREFIX}/lib" REQUIRED) target_link_libraries(infllm_ops PRIVATE ${TORCH_LIBRARIES} ${TORCH_PYTHON_LIBRARY} c10 cuda) install(TARGETS infllm_ops LIBRARY DESTINATION "sgl_kernel") # Build spatial_ops as a separate, optional extension for green contexts set(SPATIAL_SOURCES "csrc/spatial/greenctx_stream.cu" "csrc/spatial_extension.cc" ) Python_add_library(spatial_ops MODULE USE_SABI ${SKBUILD_SABI_VERSION} WITH_SOABI ${SPATIAL_SOURCES}) target_compile_options(spatial_ops PRIVATE $<$:${SGL_KERNEL_CUDA_FLAGS}>) target_link_libraries(spatial_ops PRIVATE ${TORCH_LIBRARIES} c10 cuda) install(TARGETS spatial_ops LIBRARY DESTINATION sgl_kernel) # ============================ Extra Install: FLashMLA ============================= # include(${CMAKE_CURRENT_LIST_DIR}/cmake/flashmla.cmake) # ============================ Extra Install: triton kernels ============================= # install(DIRECTORY "${repo-triton_SOURCE_DIR}/python/triton_kernels/triton_kernels/" DESTINATION "triton_kernels" PATTERN ".git*" EXCLUDE PATTERN "__pycache__" EXCLUDE)