执行摘要

• 一句话：修复 FlashMLA 在 Hopper GPU 上编译失败的架构条件问题
• 推荐动作：值得精读，特别是对 sgl-kernel 的 CMake 架构和多架构条件编译感兴趣的开发者。该 PR 展示了一个清晰的 CMake + C++ 条件编译的协作模式，可作为类似问题的参考。

功能与动机

关闭 issue #24126：在 H20 GPU 上使用 CUDA 12.8 编译 sgl-kernel 时，FlashMLA 因 Blackwell SM100 源码文件被无条件编译而失败，而 SM100 gencode 仅在 CUDA_VERSION > 12.8 时添加，导致工具链不兼容。

实现拆解

1. CMake 变量引入：在 sgl-kernel/cmake/flashmla.cmake 中新增 set(FLASHMLA_ENABLE_SM100 OFF) 作为默认值。
2. 条件启用：在 CUDA 版本大于 12.8 时，将 FLASHMLA_ENABLE_SM100 设置为 ON，并添加 SM100 gencode 标志。
3. 源文件条件编译：将所有 SM100 相关的源文件（如 csrc/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu 等）包裹在 if(FLASHMLA_ENABLE_SM100) 块中，仅当开关开启时添加到编译列表。
4. C++ 注册适配：在 sgl-kernel/csrc/flashmla_extension.cc 中，为依赖 SM100 内核符号的 dense_prefill_fwd 注册添加 #ifdef FLASHMLA_ENABLE_SM100 预处理指令，避免链接错误。

关键文件：

• sgl-kernel/cmake/flashmla.cmake（模块 构建系统；类别 other；类型 core-logic）: 核心构建配置变更，引入 FLASHMLA_ENABLE_SM100 开关并条件化 SM100 源文件编译。
• sgl-kernel/csrc/flashmla_extension.cc（模块 内核注册；类别 source；类型 core-logic）: C++ 注册代码需条件编译以避免链接未定义符号。

关键符号：未识别

关键源码片段

sgl-kernel/cmake/flashmla.cmake

核心构建配置变更，引入 FLASHMLA_ENABLE_SM100 开关并条件化 SM100 源文件编译。

# 在 sgl-kernel/cmake/flashmla.cmake 中
# 默认关闭 SM100 编译，避免在 Hopper GPU 上编译 Blackwell 源码
set(FLASHMLA_ENABLE_SM100 OFF)
​
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()
​
# ... 其他源文件 ...
​
# 仅当 SM100 启用时添加相关源文件
if(FLASHMLA_ENABLE_SM100)
  list(APPEND 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
    # ... 其他 SM100 源文件 ...
  )
endif()

sgl-kernel/csrc/flashmla_extension.cc

C++ 注册代码需条件编译以避免链接未定义符号。

// sgl-kernel/csrc/flashmla_extension.cc 中的 TORCH_LIBRARY_FRAGMENT 片段
// 仅在 SM100 启用时注册 sm100 内核，避免链接错误
#ifdef FLASHMLA_ENABLE_SM100
  m.def(
      "dense_prefill_fwd(Tensor workspace_buffer, Tensor q, Tensor k, Tensor v, Tensor cumulative_seqlen_q, Tensor "
      "cumulative_seqlen_kv, Tensor o, Tensor lse, int mask_mode_code, float softmax_scale, int max_seqlen_q, int "
      "max_seqlen_kv, bool is_varlen) -> ()");
  m.impl("dense_prefill_fwd", torch::kCUDA, &FMHACutlassSM100FwdRun);
#endif

评论区精华

Review 评论中，gemini-code-assist[bot] 指出仅排除 SM100 源文件会导致链接错误，因为 flashmla_extension.cc 仍然引用了 FMHACutlassSM100FwdRun 符号。建议在 CMake 中添加编译定义并更新 C++ 源码以条件注册。该建议被采纳，最终提交中已在 C++ 文件中加入 #ifdef FLASHMLA_ENABLE_SM100。

• 链接错误风险：仅排除源文件不够，C++ 注册也需条件化 (correctness): 采纳建议：在 flashmla_extension.cc 中添加 #ifdef FLASHMLA_ENABLE_SM100 保护 dense_prefill_fwd 注册。

风险与影响

• 风险：低风险。变更仅影响构建时的架构选择逻辑，不涉及运行时内核或数学运算。主要风险是若 CMake 变量 FLASHMLA_ENABLE_SM100 在其他 CMakeLists 中被意外覆盖可能导致不一致，但当前仅在 flashmla.cmake 内部使用。无回归风险。
• 影响：影响范围限于使用 Hopper GPU（如 H20）且 CUDA 版本为 12.8 的用户，修复后能够成功编译 sgl-kernel。对 Blackwell 用户无影响。无功能或性能影响。
• 风险标记：构建系统变更

关联脉络

• PR #25326 chore: bump sgl-kernel version to 0.4.2.post2: 同样涉及 sgl-kernel 构建和 CMake 配置