diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 8c01e7685..dfc1fbc7e 100644 --- a/.github/workflows/ci.yml +++ b/.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" diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 3e8defbe1..7fc929b6e 100644 --- a/server/CMakeLists.txt +++ b/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 + $ 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() diff --git a/server/cmake/DflashCudaArchitectures.cmake b/server/cmake/DflashCudaArchitectures.cmake new file mode 100644 index 000000000..332250718 --- /dev/null +++ b/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() diff --git a/server/cmake/tests/CMakeLists.txt b/server/cmake/tests/CMakeLists.txt new file mode 100644 index 000000000..8bdc0ba6b --- /dev/null +++ b/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") diff --git a/server/cmake/tests/cuda_architectures_case.cmake b/server/cmake/tests/cuda_architectures_case.cmake new file mode 100644 index 000000000..eaf7aa4af --- /dev/null +++ b/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() diff --git a/server/cmake/tests/cuda_architectures_errors.cmake b/server/cmake/tests/cuda_architectures_errors.cmake new file mode 100644 index 000000000..4430c10a5 --- /dev/null +++ b/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") diff --git a/server/cmake/tests/cuda_architectures_reject.cmake b/server/cmake/tests/cuda_architectures_reject.cmake new file mode 100644 index 000000000..a9f378952 --- /dev/null +++ b/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}") diff --git a/server/src/flashprefill_kernels.cu b/server/src/flashprefill_kernels.cu index 22dd72bd6..bbbcdff38 100644 --- a/server/src/flashprefill_kernels.cu +++ b/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><<>>( - (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><<>>( + (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><<>>( - (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><<>>( + (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