chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Make flashMLA work on: Cu13, B300 (#17600)

Browse Source 
Signed-off-by: vincentzed <207368749+vincentzed@users.noreply.github.com>
This commit is contained in:
Yi Zhong
2026-01-27 11:12:47 -05:00
committed by GitHub
parent 547e2d037e
commit 8acd4d7d7e
1 changed files with 53 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

53
sgl-kernel/cmake/flashmla.cmake
View File

@@ -30,6 +30,59 @@ if(${CUDA_VERSION} VERSION_GREATER 12.8)
"-gencode=arch=compute_100a,code=sm_100a"
)
endif()
if(${CUDA_VERSION} VERSION_GREATER_EQUAL "13.0")
# Patch FlashMLA sources for SM103a support.
# These patches are only needed (and only valid) with CUDA 13+.
# Patch flashmla_utils.h: widen IS_SM100 to cover the full SM100 family
set(FLASHMLA_UTILS_FILE "${repo-flashmla_SOURCE_DIR}/csrc/flashmla_utils.h")
file(READ "${FLASHMLA_UTILS_FILE}" FLASHMLA_UTILS_CONTENT)
string(REPLACE
"#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ == 1000)
#define IS_SM100 1"
"#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) && (__CUDA_ARCH__ < 1100)
#define IS_SM100 1"
FLASHMLA_UTILS_CONTENT "${FLASHMLA_UTILS_CONTENT}")
file(WRITE "${FLASHMLA_UTILS_FILE}" "${FLASHMLA_UTILS_CONTENT}")
message(STATUS "Patched flashmla_utils.h for SM103a support")
# Patch cutlass/arch/config.h: add SM103 architecture defines.
# The new block is inserted right before the existing "// SM101 and SM101a"
# anchor in the upstream header.
set(CUTLASS_CONFIG_FILE "${repo-flashmla_SOURCE_DIR}/csrc/cutlass/include/cutlass/arch/config.h")
file(READ "${CUTLASS_CONFIG_FILE}" CUTLASS_CONFIG_CONTENT)
string(FIND "${CUTLASS_CONFIG_CONTENT}" "SM103" SM103_FOUND)
if(SM103_FOUND EQUAL -1)
string(REPLACE
"// SM101 and SM101a"
"// SM103 and SM103a
#if !CUTLASS_CLANG_CUDA && (__CUDACC_VER_MAJOR__ >= 13)
#define CUTLASS_ARCH_MMA_SM103_SUPPORTED 1
#if (!defined(CUTLASS_ARCH_MMA_SM103_ENABLED) && defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 1030)
#define CUTLASS_ARCH_MMA_SM103_ENABLED 1
#if !defined(CUTLASS_ARCH_MMA_SM100A_ENABLED)
#define CUTLASS_ARCH_MMA_SM100A_ENABLED 1
#endif
#if !defined(CUTLASS_ARCH_MMA_SM100F_ENABLED)
#define CUTLASS_ARCH_MMA_SM100F_ENABLED 1
#endif
#endif
#endif
/////////////////////////////////////////////////////////////////////////////////////////////////
// SM101 and SM101a"
CUTLASS_CONFIG_CONTENT "${CUTLASS_CONFIG_CONTENT}")
file(WRITE "${CUTLASS_CONFIG_FILE}" "${CUTLASS_CONFIG_CONTENT}")
message(STATUS "Patched cutlass/arch/config.h for SM103a support")
else()
message(STATUS "cutlass/arch/config.h already patched for SM103a")
endif()
list(APPEND FLASHMLA_CUDA_FLAGS
"-gencode=arch=compute_103a,code=sm_103a"
)
endif()
set(FlashMLA_SOURCES