[sgl-kernel] support custom fp8 flashmla kernel (#13087)

This commit is contained in:
Fan Yin
2025-11-13 12:45:21 -08:00
committed by GitHub
parent 8779100768
commit 2966367a31
5 changed files with 214 additions and 17 deletions
+6 -1
View File
@@ -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
)