diff --git a/cmake/onnxruntime_cuda_source_filters.cmake b/cmake/onnxruntime_cuda_source_filters.cmake index e65b58a2f3f23..2963a3e697469 100644 --- a/cmake/onnxruntime_cuda_source_filters.cmake +++ b/cmake/onnxruntime_cuda_source_filters.cmake @@ -3,7 +3,7 @@ # Shared filtering logic for CUDA contrib ops .cu source lists. # Both the main CUDA provider and the plugin EP build use identical filtering -# rules for flash attention (quick build) and MoE GEMM FP4/FP8 kernels. +# rules for flash attention, XQA, and MoE GEMM FP4/FP8 kernels. # # Usage: # onnxruntime_filter_cuda_cu_sources() @@ -143,6 +143,25 @@ function(onnxruntime_extract_flash_attention_sources CU_SRC_LIST) set("${_FA_FLASH_SOURCES}" "${_flash_srcs}" PARENT_SCOPE) endfunction() +# Extract XQA CUDA source files into a separate list for SM80+ compilation. +function(onnxruntime_extract_xqa_sources CU_SRC_LIST) + cmake_parse_arguments(PARSE_ARGV 1 _XQA "" "XQA_SOURCES" "") + + set(_list "${${CU_SRC_LIST}}") + set(_xqa_srcs) + foreach(_src IN LISTS _list) + if(_src MATCHES "/bert/xqa/.*\\.cu$") + list(APPEND _xqa_srcs "${_src}") + endif() + endforeach() + if(_xqa_srcs) + list(REMOVE_ITEM _list ${_xqa_srcs}) + endif() + + set("${CU_SRC_LIST}" "${_list}" PARENT_SCOPE) + set("${_XQA_XQA_SOURCES}" "${_xqa_srcs}" PARENT_SCOPE) +endfunction() + # Extract LLM CUDA source files into separate lists for per-architecture compilation. # The LLM directory (contrib_ops/cuda/llm/) contains kernels with minimum SM75 support # (fpA_intB_gemv/gemm enforce arch >= 75). SM90-specific launchers (fpA_intB_gemm diff --git a/cmake/onnxruntime_providers_cuda.cmake b/cmake/onnxruntime_providers_cuda.cmake index 52eb24579cf8c..53a89b2ca5a0c 100644 --- a/cmake/onnxruntime_providers_cuda.cmake +++ b/cmake/onnxruntime_providers_cuda.cmake @@ -83,6 +83,9 @@ onnxruntime_extract_flash_attention_sources(onnxruntime_cuda_contrib_ops_cu_srcs FLASH_SOURCES onnxruntime_cuda_flash_attention_srcs ) + onnxruntime_extract_xqa_sources(onnxruntime_cuda_contrib_ops_cu_srcs + XQA_SOURCES onnxruntime_cuda_xqa_srcs + ) onnxruntime_extract_llm_sources(onnxruntime_cuda_contrib_ops_cu_srcs LLM_SOURCES onnxruntime_cuda_llm_srcs LLM_SM90_SOURCES onnxruntime_cuda_llm_sm90_srcs @@ -524,6 +527,26 @@ endif() endif() + # XQA kernels require SM80+. Compiling them with mixed SM75/SM80+ architectures + # causes CUDA 13.3 to emit host references to kernels omitted from the SM75 pass. + if(onnxruntime_cuda_xqa_srcs) + onnxruntime_filter_cuda_archs(_ort_xqa_cuda_architectures MIN_SM 80) + if(_ort_xqa_cuda_architectures) + onnxruntime_add_cuda_object_library( + NAME onnxruntime_providers_cuda_xqa + PARENT onnxruntime_providers_cuda + CUDA_ARCHITECTURES "${_ort_xqa_cuda_architectures}" + NVCC_THREADS "${onnxruntime_NVCC_THREADS}" + SOURCES ${onnxruntime_cuda_xqa_srcs}) + else() + if(TARGET onnxruntime_providers_cuda_obj) + target_sources(onnxruntime_providers_cuda_obj PRIVATE ${onnxruntime_cuda_xqa_srcs}) + else() + target_sources(onnxruntime_providers_cuda PRIVATE ${onnxruntime_cuda_xqa_srcs}) + endif() + endif() + endif() + if(NOT onnxruntime_DISABLE_CONTRIB_OPS) # SM90 TMA warp-specialized files use SM90-specific collective operations. # Compile at exactly 90a-real: SM120+ GPUs run SM90 native code via forward compat. diff --git a/cmake/onnxruntime_providers_cuda_plugin.cmake b/cmake/onnxruntime_providers_cuda_plugin.cmake index 9e940b0103bb8..99f1ced14ebd1 100644 --- a/cmake/onnxruntime_providers_cuda_plugin.cmake +++ b/cmake/onnxruntime_providers_cuda_plugin.cmake @@ -120,6 +120,9 @@ onnxruntime_extract_sm_specific_cuda_sources(CUDA_PLUGIN_EP_CU_SRCS onnxruntime_extract_flash_attention_sources(CUDA_PLUGIN_EP_CU_SRCS FLASH_SOURCES _cuda_plugin_flash_attention_srcs ) +onnxruntime_extract_xqa_sources(CUDA_PLUGIN_EP_CU_SRCS + XQA_SOURCES _cuda_plugin_xqa_srcs +) onnxruntime_extract_llm_sources(CUDA_PLUGIN_EP_CU_SRCS LLM_SOURCES _cuda_plugin_llm_srcs LLM_SM90_SOURCES _cuda_plugin_llm_sm90_srcs @@ -355,6 +358,21 @@ if(_cuda_plugin_flash_attention_srcs) endif() endif() +if(_cuda_plugin_xqa_srcs) + onnxruntime_filter_cuda_archs(_plugin_xqa_cuda_architectures MIN_SM 80) + if(_plugin_xqa_cuda_architectures) + onnxruntime_add_cuda_plugin_object_library( + NAME onnxruntime_providers_cuda_plugin_xqa + PARENT onnxruntime_providers_cuda_plugin + CUDA_ARCHITECTURES "${_plugin_xqa_cuda_architectures}" + NVCC_THREADS "${onnxruntime_plugin_nvcc_threads}" + COMPILE_OPTIONS ${_cuda_plugin_shared_compile_options} + SOURCES ${_cuda_plugin_xqa_srcs}) + else() + target_sources(onnxruntime_providers_cuda_plugin PRIVATE ${_cuda_plugin_xqa_srcs}) + endif() +endif() + if(NOT onnxruntime_DISABLE_CONTRIB_OPS) # SM90 TMA warp-specialized files use SM90-specific collective operations. # Also includes fpA_intB SM90 launchers (guarded by #ifndef EXCLUDE_SM_90). diff --git a/cmake/onnxruntime_unittests.cmake b/cmake/onnxruntime_unittests.cmake index 1974b5d601f46..40669a6681aaf 100644 --- a/cmake/onnxruntime_unittests.cmake +++ b/cmake/onnxruntime_unittests.cmake @@ -1036,6 +1036,9 @@ if (onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS AND NOT onnxruntime_BUILD_CUDA_EP_ if(TARGET onnxruntime_providers_cuda_flash_attention) target_link_libraries(onnxruntime_providers_cuda_ut PRIVATE onnxruntime_providers_cuda_flash_attention) endif() + if(TARGET onnxruntime_providers_cuda_xqa) + target_link_libraries(onnxruntime_providers_cuda_ut PRIVATE onnxruntime_providers_cuda_xqa) + endif() if(TARGET onnxruntime_providers_cuda_llm) target_link_libraries(onnxruntime_providers_cuda_ut PRIVATE onnxruntime_providers_cuda_llm) endif()