diff --git a/sgl-kernel/cmake/flashmla.cmake b/sgl-kernel/cmake/flashmla.cmake index a1b1353a8..7387f3e5d 100644 --- a/sgl-kernel/cmake/flashmla.cmake +++ b/sgl-kernel/cmake/flashmla.cmake @@ -1,8 +1,9 @@ # flash_mla +# sm90 dense decode HEAD_DIM_K=512 support (sgl-project/FlashMLA#9, merged). FetchContent_Declare( repo-flashmla - URL https://${GITHUB_ARTIFACTORY}/sgl-project/FlashMLA/archive/df022ebafb88578eab9f0300606ee765608d8b5c.tar.gz - URL_HASH SHA256=45992d7de7d051dc897aff33156a2b7515d745fc489008adc7924ab884578d52 + URL https://${GITHUB_ARTIFACTORY}/sgl-project/FlashMLA/archive/05e26647fe840b8baedae486c2d86d5ce4efeb7c.tar.gz + URL_HASH SHA256=ce369489bbfc42cdfbba9aa949de0270e64469d530748dea9f4f60b3c69dea9b ) FetchContent_Populate(repo-flashmla) @@ -156,6 +157,16 @@ target_compile_options(flashmla_ops PRIVATE if(FLASHMLA_ENABLE_SM100) target_compile_definitions(flashmla_ops PRIVATE FLASHMLA_ENABLE_SM100) endif() + +# CUDA 13 moved cuda/std/* under cccl/cuda/std/*. The vendored cutlass routes +# to when __CUDACC_VER_MAJOR__ >= 13, so the +# host C++ TU (compiled by g++, where that macro is unset for the legacy path) +# needs the cccl include root on the search path. +if(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL "13.0") + find_path(FLASHMLA_CCCL_INCLUDE NAMES cuda/std/utility + HINTS ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES} + ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}/cccl) +endif() target_include_directories(flashmla_ops PRIVATE ${repo-flashmla_SOURCE_DIR}/csrc ${repo-flashmla_SOURCE_DIR}/csrc/kerutils/include @@ -163,6 +174,7 @@ target_include_directories(flashmla_ops PRIVATE ${repo-flashmla_SOURCE_DIR}/csrc/extension/sm90/dense_fp8/ ${repo-flashmla_SOURCE_DIR}/csrc/cutlass/include ${repo-flashmla_SOURCE_DIR}/csrc/cutlass/tools/util/include + ${FLASHMLA_CCCL_INCLUDE} ) target_link_libraries(flashmla_ops PRIVATE ${TORCH_LIBRARIES} c10 cuda)