[sgl-kernel] support custom fp8 flashmla kernel (#13087)
This commit is contained in:
@@ -4,7 +4,7 @@ include(FetchContent)
|
||||
FetchContent_Declare(
|
||||
repo-flashmla
|
||||
GIT_REPOSITORY https://github.com/sgl-project/FlashMLA
|
||||
GIT_TAG bc8576abc3e507425cf6498f3d3393df7733ce37
|
||||
GIT_TAG 6fde8086398bc5ed144c1e9ed0de30cd43be2cda
|
||||
GIT_SHALLOW OFF
|
||||
)
|
||||
FetchContent_Populate(repo-flashmla)
|
||||
@@ -44,6 +44,10 @@ set(FlashMLA_SOURCES
|
||||
${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
|
||||
${repo-flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd.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
|
||||
)
|
||||
|
||||
Python_add_library(flashmla_ops MODULE USE_SABI ${SKBUILD_SABI_VERSION} WITH_SOABI ${FlashMLA_SOURCES})
|
||||
@@ -51,6 +55,7 @@ target_compile_options(flashmla_ops PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:${FLASHML
|
||||
target_include_directories(flashmla_ops PRIVATE
|
||||
${repo-flashmla_SOURCE_DIR}/csrc
|
||||
${repo-flashmla_SOURCE_DIR}/csrc/sm90
|
||||
${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
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user