fix(cmake): derive conservative minimum CUDA architecture

.github/workflows/ci.yml

@@ -23,6 +23,11 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6.0.3
+      - name: Test CUDA architecture parsing
+        run: |
+          cmake -S server/cmake/tests -B server/cmake/tests/build
+          ctest --test-dir server/cmake/tests/build --output-on-failure --no-tests=error
+
       - uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 # v8.2.0
         with:
           version: "0.11.x"

server/CMakeLists.txt

@@ -1,8 +1,17 @@
 cmake_minimum_required(VERSION 3.21)  # 3.21 adds first-class HIP language support (project(LANGUAGES ... HIP))
+include("${CMAKE_CURRENT_LIST_DIR}/cmake/DflashCudaArchitectures.cmake")
+
 set(DFLASH27B_GPU_BACKEND "cuda" CACHE STRING "GPU backend to build: cuda or hip")
 set_property(CACHE DFLASH27B_GPU_BACKEND PROPERTY STRINGS cuda hip)
 string(TOLOWER "${DFLASH27B_GPU_BACKEND}" DFLASH27B_GPU_BACKEND)
 if(DFLASH27B_GPU_BACKEND STREQUAL "cuda")
+    # Validate an explicit user list before CUDA compiler detection so malformed
+    # values fail with the same deterministic message on every toolchain.
+    if(DEFINED CMAKE_CUDA_ARCHITECTURES)
+        dflash_parse_cuda_architectures(
+            _dflash_user_cuda_min_sm _dflash_user_cuda_numeric_archs
+            "${CMAKE_CUDA_ARCHITECTURES}")
+    endif()
     set(DFLASH27B_USER_CUDA_ARCHITECTURES "${CMAKE_CUDA_ARCHITECTURES}")
     project(dflash LANGUAGES C CXX CUDA)
 elseif(DFLASH27B_GPU_BACKEND STREQUAL "hip")
@@ -133,6 +142,8 @@ if(DFLASH27B_GPU_BACKEND STREQUAL "cuda")
             list(APPEND _dflash_archs "121")
         endif()
     endif()
+    dflash_parse_cuda_architectures(
+        _dflash_cuda_min_sm _dflash_cuda_numeric_archs "${_dflash_archs}")
     # Keep the vendored ggml-cuda build aligned with the resolved dflash arch list.
     # Otherwise ggml may inherit CMake's compiler-default arch (for example sm_52),
     # which triggers massive first-request PTX JIT on newer GPUs even though
@@ -169,8 +180,7 @@ if(DFLASH27B_GPU_BACKEND STREQUAL "cuda" AND NOT DEFINED _dflash_is_consumer_bla
     set(_dflash_is_consumer_blackwell OFF)
     # Iterate the resolved dflash_common arch list, not raw CMAKE_CUDA_ARCHITECTURES,
     # which is empty on the default path (the project supplies its own list above).
-    foreach(_arch IN LISTS _dflash_archs)
-        string(REGEX REPLACE "[^0-9]" "" _dflash_arch_num "${_arch}")
+    foreach(_dflash_arch_num IN LISTS _dflash_cuda_numeric_archs)
         if(_dflash_arch_num MATCHES "^12[0-9]$")
             set(_dflash_is_consumer_blackwell ON)
             break()
@@ -369,9 +379,6 @@ if(DFLASH27B_GPU_BACKEND STREQUAL "cuda")
         src/cuda_cross_device_copy.cpp
         src/deepseek4/deepseek4_hc_cuda.cu)
     set_target_properties(dflash_common PROPERTIES CUDA_ARCHITECTURES "${_dflash_archs}")
-    list(GET _dflash_archs 0 _dflash_cuda_min_sm)
-    # Strip any trailing 'a' suffix (e.g. "121a" -> "121")
-    string(REGEX REPLACE "[^0-9]" "" _dflash_cuda_min_sm "${_dflash_cuda_min_sm}")
     target_compile_definitions(dflash_common PRIVATE
         DFLASH27B_BACKEND_CUDA=1
         DFLASH27B_CUDA_MIN_SM=${_dflash_cuda_min_sm}
@@ -469,10 +476,11 @@ elseif(DFLASH27B_GPU_BACKEND STREQUAL "cuda")
     set(_dflash_has_sm80 OFF)
     set(_dflash_has_sm70 OFF)
     set(_dflash_has_sm60 OFF)
-    foreach(_arch IN LISTS _dflash_archs)
-        string(REGEX REPLACE "[^0-9]" "" _arch_num "${_arch}")
+    set(_dflash_sm80_archs "")
+    foreach(_arch _arch_num IN ZIP_LISTS _dflash_archs _dflash_cuda_numeric_archs)
         if(_arch_num GREATER_EQUAL 80)
             set(_dflash_has_sm80 ON)
+            list(APPEND _dflash_sm80_archs "${_arch}")
         endif()
         if(_arch_num GREATER_EQUAL 70 AND _arch_num LESS 80)
             set(_dflash_has_sm70 ON)
@@ -492,16 +500,20 @@ elseif(DFLASH27B_GPU_BACKEND STREQUAL "cuda")
     endif()
 
     if(_dflash_has_sm80)
+        # CUDA_ARCHITECTURES is a target property, not a source property.  Keep
+        # the BF16 WMMA translation unit in its own object target so a mixed
+        # fatbinary (for example 86;70) never compiles it for an unsupported SM.
+        add_library(dflash_flashprefill_sm80 OBJECT src/flashprefill_kernels.cu)
+        set_target_properties(dflash_flashprefill_sm80 PROPERTIES
+            CUDA_ARCHITECTURES "${_dflash_sm80_archs}"
+            POSITION_INDEPENDENT_CODE ON)
+        target_link_libraries(dflash_flashprefill_sm80 PRIVATE CUDA::cudart)
         target_sources(dflash_common PRIVATE
-            src/flashprefill_kernels.cu
+            $<TARGET_OBJECTS:dflash_flashprefill_sm80>
             src/pflash_ggml_adapter.cpp)
         target_compile_definitions(dflash_common PRIVATE
             DFLASH27B_HAVE_CUDA_WMMA_FLASHPREFILL=1
             DFLASH27B_HAVE_SM80_FLASHPREFILL=1)
-        # BF16 WMMA kernels require sm_80+.  Restrict compilation to those
-        # arches so nvcc never instantiates m16n16k16 BF16 WMMA on Volta/Pascal.
-        set_source_files_properties(src/flashprefill_kernels.cu PROPERTIES
-            CUDA_ARCHITECTURES "80;86;90;100;110;120;121")
     endif()
     if(_dflash_has_sm70)
         target_sources(dflash_common PRIVATE src/flashprefill_f16.cu)
@@ -537,10 +549,10 @@ if(DFLASH27B_GPU_BACKEND STREQUAL "hip" AND DFLASH27B_ENABLE_BSA AND NOT DFLASH2
     set(DFLASH27B_ENABLE_BSA OFF)
 endif()
 if(DFLASH27B_ENABLE_BSA)
-    foreach(_arch IN LISTS _dflash_archs)
-        if(_arch LESS 80)
+    foreach(_arch_num IN LISTS _dflash_cuda_numeric_archs)
+        if(_arch_num LESS 80)
             message(WARNING
-                "DFLASH27B_ENABLE_BSA=ON requested but CUDA_ARCHITECTURES contains '${_arch}' (<80); "
+                "DFLASH27B_ENABLE_BSA=ON requested but CUDA_ARCHITECTURES contains an SM below 80; "
                 "disabling BSA (the spec-prefill path will fall back to the WMMA kernel).")
             set(DFLASH27B_ENABLE_BSA OFF)
             break()

server/cmake/DflashCudaArchitectures.cmake

@@ -0,0 +1,47 @@
+include_guard(GLOBAL)
+
+# Parse a CMAKE_CUDA_ARCHITECTURES-style list into numeric SM values and find
+# its true minimum. Numeric architectures may use CUDA's architecture-specific
+# "a"/"f" suffixes and CMake's code-generation "-real"/"-virtual" suffixes.
+# Special values such as "all" and "native" cannot produce a deterministic
+# compile-time minimum and are intentionally rejected.
+function(dflash_parse_cuda_architectures out_min_sm out_numeric_arches arch_list)
+    if(NOT ARGC EQUAL 3)
+        message(FATAL_ERROR
+            "dflash_parse_cuda_architectures expects an output minimum, an "
+            "output numeric list, and one quoted CUDA architecture list")
+    endif()
+    if("${arch_list}" STREQUAL "")
+        message(FATAL_ERROR "CUDA architecture list must not be empty")
+    endif()
+
+    set(_dflash_numeric_arches "")
+    set(_dflash_min_sm "")
+    foreach(_dflash_arch IN LISTS arch_list)
+        if("${_dflash_arch}" STREQUAL "")
+            message(FATAL_ERROR
+                "CUDA architecture list contains an empty entry: '${arch_list}'")
+        endif()
+        if(NOT "${_dflash_arch}" MATCHES "^([1-9][0-9]*)(a|f)?(-(real|virtual))?$")
+            message(FATAL_ERROR
+                "Unsupported CUDA architecture '${_dflash_arch}' in '${arch_list}'. "
+                "Use a positive numeric SM with an optional a/f architecture suffix "
+                "and an optional -real/-virtual code-generation suffix.")
+        endif()
+
+        set(_dflash_arch_sm "${CMAKE_MATCH_1}")
+        list(APPEND _dflash_numeric_arches "${_dflash_arch_sm}")
+        if("${_dflash_min_sm}" STREQUAL "")
+            set(_dflash_min_sm "${_dflash_arch_sm}")
+        elseif(_dflash_arch_sm LESS _dflash_min_sm)
+            set(_dflash_min_sm "${_dflash_arch_sm}")
+        endif()
+    endforeach()
+
+    if("${_dflash_min_sm}" STREQUAL "")
+        message(FATAL_ERROR "CUDA architecture list must contain at least one entry")
+    endif()
+
+    set(${out_min_sm} "${_dflash_min_sm}" PARENT_SCOPE)
+    set(${out_numeric_arches} "${_dflash_numeric_arches}" PARENT_SCOPE)
+endfunction()

server/cmake/tests/CMakeLists.txt

@@ -0,0 +1,17 @@
+cmake_minimum_required(VERSION 3.21)
+project(dflash_cuda_architecture_tests LANGUAGES NONE)
+
+enable_testing()
+
+foreach(_case IN ITEMS permutations suffixes duplicates)
+    add_test(
+        NAME cuda_arch_${_case}
+        COMMAND "${CMAKE_COMMAND}"
+            "-DCASE=${_case}"
+            -P "${CMAKE_CURRENT_LIST_DIR}/cuda_architectures_case.cmake")
+endforeach()
+
+add_test(
+    NAME cuda_arch_errors
+    COMMAND "${CMAKE_COMMAND}"
+        -P "${CMAKE_CURRENT_LIST_DIR}/cuda_architectures_errors.cmake")

server/cmake/tests/cuda_architectures_case.cmake

@@ -0,0 +1,28 @@
+cmake_minimum_required(VERSION 3.21)
+include("${CMAKE_CURRENT_LIST_DIR}/../DflashCudaArchitectures.cmake")
+
+function(assert_arch_parse arch_list expected_min expected_numeric)
+    dflash_parse_cuda_architectures(actual_min actual_numeric "${arch_list}")
+    if(NOT "${actual_min}" STREQUAL "${expected_min}")
+        message(FATAL_ERROR
+            "minimum for '${arch_list}': expected '${expected_min}', got '${actual_min}'")
+    endif()
+    if(NOT "${actual_numeric}" STREQUAL "${expected_numeric}")
+        message(FATAL_ERROR
+            "numeric list for '${arch_list}': expected '${expected_numeric}', "
+            "got '${actual_numeric}'")
+    endif()
+endfunction()
+
+if(CASE STREQUAL "permutations")
+    assert_arch_parse("60;61;62;70;75;86;120" "60" "60;61;62;70;75;86;120")
+    assert_arch_parse("120;86;75;70;62;61;60" "60" "120;86;75;70;62;61;60")
+    assert_arch_parse("86;60;120;70;61;75;62" "60" "86;60;120;70;61;75;62")
+elseif(CASE STREQUAL "suffixes")
+    assert_arch_parse("121a-real;120f-virtual;90-virtual;100" "90" "121;120;90;100")
+    assert_arch_parse("120a;86-virtual" "86" "120;86")
+elseif(CASE STREQUAL "duplicates")
+    assert_arch_parse("120a;86;120f-real;86-virtual" "86" "120;86;120;86")
+else()
+    message(FATAL_ERROR "unknown test CASE '${CASE}'")
+endif()

server/cmake/tests/cuda_architectures_errors.cmake

@@ -0,0 +1,33 @@
+cmake_minimum_required(VERSION 3.21)
+
+function(assert_rejected label arch_list expected_error)
+    execute_process(
+        COMMAND "${CMAKE_COMMAND}"
+            "-DARCH_LIST=${arch_list}"
+            -P "${CMAKE_CURRENT_LIST_DIR}/cuda_architectures_reject.cmake"
+        RESULT_VARIABLE result
+        OUTPUT_VARIABLE stdout
+        ERROR_VARIABLE stderr)
+    if(result EQUAL 0)
+        message(FATAL_ERROR "${label}: malformed list '${arch_list}' was accepted")
+    endif()
+    set(output "${stdout}\n${stderr}")
+    if(NOT output MATCHES "${expected_error}")
+        message(FATAL_ERROR
+            "${label}: expected error /${expected_error}/, got:\n${output}")
+    endif()
+endfunction()
+
+assert_rejected("empty list" "" "must not be empty")
+assert_rejected("leading empty entry" ";86" "contains an empty entry")
+assert_rejected("middle empty entry" "86;;75" "contains an empty entry")
+assert_rejected("trailing empty entry" "86;" "contains an empty entry")
+assert_rejected("sm prefix" "sm_86" "Unsupported CUDA architecture")
+assert_rejected("unknown suffix" "86-ptx" "Unsupported CUDA architecture")
+assert_rejected("unknown architecture suffix" "120b" "Unsupported CUDA architecture")
+assert_rejected("multiple architecture suffixes" "120af" "Unsupported CUDA architecture")
+assert_rejected("suffixes out of order" "120-real-a" "Unsupported CUDA architecture")
+assert_rejected("extra suffix" "90-real-extra" "Unsupported CUDA architecture")
+assert_rejected("special native" "native" "Unsupported CUDA architecture")
+assert_rejected("special all" "all" "Unsupported CUDA architecture")
+assert_rejected("zero" "0" "Unsupported CUDA architecture")

server/cmake/tests/cuda_architectures_reject.cmake

@@ -0,0 +1,4 @@
+cmake_minimum_required(VERSION 3.21)
+include("${CMAKE_CURRENT_LIST_DIR}/../DflashCudaArchitectures.cmake")
+
+dflash_parse_cuda_architectures(actual_min actual_numeric "${ARCH_LIST}")

server/src/flashprefill_kernels.cu

@@ -101,24 +101,25 @@ __global__ void compute_mean_vector_kernel_bf16(
 }
 
 // Public launcher (called from C++).
-extern "C" void launch_compute_mean_vector_bf16(
+extern "C" int launch_compute_mean_vector_bf16(
     const void * K, void * mean_K,
     int batch, int seq_len, int n_kv_heads, int head_dim, int block_size,
     int s_K_b, int s_K_n, int s_K_h, int s_K_d,
     int s_mK_b, int s_mK_m, int s_mK_h, int s_mK_d,
     cudaStream_t stream)
 {
+    if (head_dim != 128 || block_size != 128) {
+        return -1;
+    }
     const int n_k_blocks = (seq_len + block_size - 1) / block_size;
     dim3 grid(n_k_blocks, batch * n_kv_heads, 1);
     dim3 block(head_dim, 1, 1);
-    if (head_dim == 128 && block_size == 128) {
-        compute_mean_vector_kernel_bf16<128, 128><<<grid, block, 0, stream>>>(
-            (const __nv_bfloat16 *)K, (__nv_bfloat16 *)mean_K,
-            batch, seq_len, n_kv_heads,
-            s_K_b, s_K_n, s_K_h, s_K_d,
-            s_mK_b, s_mK_m, s_mK_h, s_mK_d);
-    }
-    // Only D_HEAD=128 BLOCK=128 dispatched here. Add other combos when new heads/blocks needed.
+    compute_mean_vector_kernel_bf16<128, 128><<<grid, block, 0, stream>>>(
+        (const __nv_bfloat16 *)K, (__nv_bfloat16 *)mean_K,
+        batch, seq_len, n_kv_heads,
+        s_K_b, s_K_n, s_K_h, s_K_d,
+        s_mK_b, s_mK_m, s_mK_h, s_mK_d);
+    return 0;
 }
 
 // ---- Kernel 2: compute_block_score ----
@@ -229,7 +230,7 @@ __global__ void compute_block_score_kernel_bf16(
     }
 }
 
-extern "C" void launch_compute_block_score_bf16(
+extern "C" int launch_compute_block_score_bf16(
     const void * Q, const void * mean_K, float sm_scale,
     void * score, void * score_max,
     int batch, int n_q_heads, int n_k_heads,
@@ -240,20 +241,22 @@ extern "C" void launch_compute_block_score_bf16(
     int s_M_b, int s_M_m, int s_M_n, int s_M_h,
     cudaStream_t stream)
 {
+    if (head_dim != 128 || block_size != 128) {
+        return -1;
+    }
     const int M = (seq_len + block_size - 1) / block_size;
     dim3 grid(M, batch * n_q_heads, 1);
     dim3 block(block_size, 1, 1);
     size_t smem = block_size * sizeof(float);
-    if (head_dim == 128 && block_size == 128) {
-        compute_block_score_kernel_bf16<128, 128, 1><<<grid, block, smem, stream>>>(
-            (const __nv_bfloat16 *)Q, (const __nv_bfloat16 *)mean_K, sm_scale,
-            (float *)score, (float *)score_max,
-            batch, n_q_heads, n_k_heads, M, M,
-            s_Q_b, s_Q_n, s_Q_h, s_Q_d,
-            s_mK_b, s_mK_m, s_mK_h, s_mK_d,
-            s_S_b, s_S_m, s_S_n, s_S_h,
-            s_M_b, s_M_m, s_M_n, s_M_h);
-    }
+    compute_block_score_kernel_bf16<128, 128, 1><<<grid, block, smem, stream>>>(
+        (const __nv_bfloat16 *)Q, (const __nv_bfloat16 *)mean_K, sm_scale,
+        (float *)score, (float *)score_max,
+        batch, n_q_heads, n_k_heads, M, M,
+        s_Q_b, s_Q_n, s_Q_h, s_Q_d,
+        s_mK_b, s_mK_m, s_mK_h, s_mK_d,
+        s_S_b, s_S_m, s_S_n, s_S_h,
+        s_M_b, s_M_m, s_M_n, s_M_h);
+    return 0;
 }
 
 // ---- Kernel 4: sparse_flash_forward ----
@@ -944,7 +947,6 @@ extern "C" void launch_sparse_flash_forward_bf16(
 #endif
 }
 
-
 // ---- Kernel 3: block_select on GPU ----
 //
 // One warp per (B, M, H). Each warp scans n in [0, m] in chunks of 32, takes