From 8acd4d7d7e6f436601ef3ae51678f130fdc04d25 Mon Sep 17 00:00:00 2001
From: Yi Zhong <207368749+vincentzed@users.noreply.github.com>
Date: Tue, 27 Jan 2026 11:12:47 -0500
Subject: [PATCH] Make flashMLA work on: Cu13, B300 (#17600)

Signed-off-by: vincentzed <207368749+vincentzed@users.noreply.github.com>
---
 sgl-kernel/cmake/flashmla.cmake | 53 +++++++++++++++++++++++++++++++++
 1 file changed, 53 insertions(+)

diff --git a/sgl-kernel/cmake/flashmla.cmake b/sgl-kernel/cmake/flashmla.cmake
index c17266af2..b1546b151 100644
--- a/sgl-kernel/cmake/flashmla.cmake
+++ b/sgl-kernel/cmake/flashmla.cmake
@@ -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