fix(sgl-kernel): sm90 compile flashmla failed (#24130)
This commit is contained in:
@@ -17,6 +17,8 @@ set(FLASHMLA_CUDA_FLAGS
|
||||
"-Xcudafe=--diag_suppress=177" # variable was declared but never referenced
|
||||
)
|
||||
|
||||
set(FLASHMLA_ENABLE_SM100 OFF)
|
||||
|
||||
# The FlashMLA kernels only work on hopper and require CUDA 12.4 or later.
|
||||
# Only build FlashMLA kernels if we are building for something compatible with
|
||||
# sm90a
|
||||
@@ -29,6 +31,7 @@ if(${CUDA_VERSION} VERSION_GREATER 12.8)
|
||||
list(APPEND FLASHMLA_CUDA_FLAGS
|
||||
"-gencode=arch=compute_100a,code=sm_100a"
|
||||
)
|
||||
set(FLASHMLA_ENABLE_SM100 ON)
|
||||
endif()
|
||||
if(${CUDA_VERSION} VERSION_GREATER_EQUAL "13.0")
|
||||
# Patch FlashMLA sources for SM103a support.
|
||||
@@ -113,33 +116,40 @@ set(FlashMLA_SOURCES
|
||||
${repo-flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k576.cu
|
||||
${repo-flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k576_topklen.cu
|
||||
|
||||
# sm100 dense prefill/bwd.
|
||||
${repo-flashmla_SOURCE_DIR}/csrc/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu
|
||||
${repo-flashmla_SOURCE_DIR}/csrc/sm100/prefill/dense/fmha_cutlass_bwd_sm100.cu
|
||||
|
||||
# sm100 sparse prefill.
|
||||
${repo-flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head64/instantiations/phase1_k512.cu
|
||||
${repo-flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head64/instantiations/phase1_k576.cu
|
||||
${repo-flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k512.cu
|
||||
${repo-flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k576.cu
|
||||
${repo-flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_prefill_k512.cu
|
||||
|
||||
# sm100 sparse decode.
|
||||
${repo-flashmla_SOURCE_DIR}/csrc/sm100/decode/head64/instantiations/v32.cu
|
||||
${repo-flashmla_SOURCE_DIR}/csrc/sm100/decode/head64/instantiations/model1.cu
|
||||
${repo-flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512.cu
|
||||
|
||||
${repo-flashmla_SOURCE_DIR}/csrc/extension/sm90/dense_fp8/dense_fp8_python_api.cpp
|
||||
${repo-flashmla_SOURCE_DIR}/csrc/extension/sm90/dense_fp8/flash_fwd_mla_fp8_sm90.cu
|
||||
${repo-flashmla_SOURCE_DIR}/csrc/extension/sm90/dense_fp8/flash_fwd_mla_metadata.cu
|
||||
)
|
||||
|
||||
if(FLASHMLA_ENABLE_SM100)
|
||||
list(APPEND FlashMLA_SOURCES
|
||||
# sm100 dense prefill/bwd.
|
||||
${repo-flashmla_SOURCE_DIR}/csrc/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu
|
||||
${repo-flashmla_SOURCE_DIR}/csrc/sm100/prefill/dense/fmha_cutlass_bwd_sm100.cu
|
||||
|
||||
# sm100 sparse prefill.
|
||||
${repo-flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head64/instantiations/phase1_k512.cu
|
||||
${repo-flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head64/instantiations/phase1_k576.cu
|
||||
${repo-flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k512.cu
|
||||
${repo-flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k576.cu
|
||||
${repo-flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_prefill_k512.cu
|
||||
|
||||
# sm100 sparse decode.
|
||||
${repo-flashmla_SOURCE_DIR}/csrc/sm100/decode/head64/instantiations/v32.cu
|
||||
${repo-flashmla_SOURCE_DIR}/csrc/sm100/decode/head64/instantiations/model1.cu
|
||||
${repo-flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512.cu
|
||||
)
|
||||
endif()
|
||||
|
||||
Python_add_library(flashmla_ops MODULE USE_SABI ${SKBUILD_SABI_VERSION} WITH_SOABI ${FlashMLA_SOURCES})
|
||||
target_compile_options(flashmla_ops PRIVATE
|
||||
$<$<COMPILE_LANGUAGE:CXX>:-std=c++20>
|
||||
$<$<COMPILE_LANGUAGE:CUDA>:-std=c++20>
|
||||
$<$<COMPILE_LANGUAGE:CUDA>:${FLASHMLA_CUDA_FLAGS}>
|
||||
)
|
||||
if(FLASHMLA_ENABLE_SM100)
|
||||
target_compile_definitions(flashmla_ops PRIVATE FLASHMLA_ENABLE_SM100)
|
||||
endif()
|
||||
target_include_directories(flashmla_ops PRIVATE
|
||||
${repo-flashmla_SOURCE_DIR}/csrc
|
||||
${repo-flashmla_SOURCE_DIR}/csrc/kerutils/include
|
||||
|
||||
Reference in New Issue
Block a user