Pull in changes from flashinfer-ai#962 into flashinfer-ai#944.

Author: diptorupd

libflashinfer/CMakeLists.txt

@@ -14,11 +14,6 @@ FetchContent_Declare(
   GIT_REPOSITORY https://github.com/microsoft/mscclpp.git
   GIT_TAG 11e62024d3eb190e005b4689f8c8443d91a6c82e)
 
-set(BOOST_ENABLE_CMAKE ON)
-FetchContent_Declare(boost_math
-                     GIT_REPOSITORY https://github.com/boostorg/math.git)
-FetchContent_MakeAvailable(boost_math)
-
 # -----------------------------------------------------------------------------#
 
 find_package(Python3 REQUIRED)
@@ -81,10 +76,21 @@ if(FLASHINFER_ENABLE_BF16)
   add_definitions(-DFLASHINFER_ENABLE_BF16)
 endif(FLASHINFER_ENABLE_BF16)
 
+
 if(FLASHINFER_GEN_USE_FP16_QK_REDUCTIONS)
+  # --- Dependencies specific to FLASHINFER_GEN_USE_FP16_QK_REDUCTIONS -------#
+  include(FetchContent)
+
+  set(BOOST_ENABLE_CMAKE ON)
+  FetchContent_Declare(boost_math
+                       GIT_REPOSITORY https://github.com/boostorg/math.git)
+  FetchContent_MakeAvailable(boost_math)
+  # --------------------------------------------------------------------------#
   set(USE_FP16_QK_REDUCTIONS "true")
+  message(STATUS "USE_FP16_QK_REDUCTIONS=${USE_FP16_QK_REDUCTIONS}")
 else(FLASHINFER_GEN_USE_FP16_QK_REDUCTIONS)
   set(USE_FP16_QK_REDUCTIONS "false")
+  message(STATUS "USE_FP16_QK_REDUCTIONS=${USE_FP16_QK_REDUCTIONS}")
 endif(FLASHINFER_GEN_USE_FP16_QK_REDUCTIONS)
 
 # generate kernel inst
@@ -98,6 +104,34 @@ message(STATUS "FLASHINFER_POS_ENCODING_MODES=${POS_ENCODING_MODES}")
 message(STATUS "FLASHINFER_USE_FP16_QK_REDUCTIONS=${USE_FP16_QK_REDUCTIONS}")
 message(STATUS "FLASHINFER_MASK_MODES=${MASK_MODES}")
 
+#----------------------- SM90 head dims computation ------------------------#
+set(SM90_ALLOWED_HEAD_DIMS "64,64" "128,128" "256,256" "192,128")
+set(HEAD_DIMS_SM90 "")
+
+foreach(DIM_VAL ${HEAD_DIMS})
+  string(CONCAT TUPLE_VAL "${DIM_VAL}" "," "${DIM_VAL}")
+  list(FIND SM90_ALLOWED_HEAD_DIMS ${TUPLE_VAL} RESULT)
+  if(NOT ${RESULT} EQUAL -1)
+    list(APPEND HEAD_DIMS_SM90 ${TUPLE_VAL})
+  endif(NOT ${RESULT} EQUAL -1)
+endforeach(DIM_VAL)
+
+foreach(TUPLE_VAL ${SM90_ALLOWED_HEAD_DIMS})
+  string(REPLACE "," ";" HEAD_DIMS_LIST ${TUPLE_VAL})
+  list(GET HEAD_DIMS_LIST 0 K)
+  list(GET HEAD_DIMS_LIST 1 V)
+  if(NOT K EQUAL V)
+    list(APPEND HEAD_DIMS_SM90 ${TUPLE_VAL})
+  endif(NOT K EQUAL V)
+endforeach(TUPLE_VAL)
+
+list(REMOVE_DUPLICATES HEAD_DIMS_SM90)
+#---------------------------------------------------------------------------#
+
+# Log SM90_ALLOWED_HEAD_DIMS and HEAD_DIMS_SM90
+message(STATUS "SM90_ALLOWED_HEAD_DIMS=${SM90_ALLOWED_HEAD_DIMS}")
+message(STATUS "HEAD_DIMS_SM90=${HEAD_DIMS_SM90}")
+
 set(GENERATED_SOURCE_DIR ${CMAKE_CURRENT_BINARY_DIR}/src/generated)
 cmake_path(GET GENERATED_SOURCE_DIR PARENT_PATH GENERATED_SOURCE_DIR_ROOT)
 file(MAKE_DIRECTORY ${GENERATED_SOURCE_DIR})
@@ -114,7 +148,7 @@ set(AOT_GENERATE_COMMAND
 set(AOT_GENERATE_DISPATCH_INC_COMMAND
     ${Python3_EXECUTABLE} -m aot_build_utils.generate_dispatch_inc --path
     "${GENERATED_SOURCE_DIR}/dispatch.inc" --head_dims ${HEAD_DIMS}
-    --head_dims_sm90 "64,64" # FIXME Make this configurable.
+    --head_dims ${HEAD_DIMS_SM90}
     --pos_encoding_modes ${POS_ENCODING_MODES} --use_fp16_qk_reductions
     ${USE_FP16_QK_REDUCTIONS} --mask_modes ${MASK_MODES})
 
@@ -168,7 +202,6 @@ string(
          "--threads=1 "
          "-Xfatbin=-compress-all "
          "-use_fast_math "
-         "-DFLASHINFER_ENABLE_F16 "
          "--expt-relaxed-constexpr ")
 string(CONCAT FLASHINFER_CXXFLAGS "${WARNING_FLAGS}" "${SECURITY_FLAGS}")
 
@@ -179,11 +212,16 @@ set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} ${NVCC_FLAGS}")
 
 add_library(decode_kernels STATIC ${DECODE_KERNELS_SRCS})
 target_include_directories(decode_kernels PRIVATE ${FLASHINFER_INCLUDE_DIR})
-target_link_libraries(decode_kernels PRIVATE Boost::math)
+if(FLASHINFER_GEN_USE_FP16_QK_REDUCTIONS)
+  target_link_libraries(decode_kernels PRIVATE Boost::math)
+endif(FLASHINFER_GEN_USE_FP16_QK_REDUCTIONS)
 
 add_library(prefill_kernels STATIC ${PREFILL_KERNELS_SRCS})
 target_include_directories(prefill_kernels PRIVATE ${FLASHINFER_INCLUDE_DIR})
-target_link_libraries(prefill_kernels PRIVATE Boost::math)
+if(FLASHINFER_GEN_USE_FP16_QK_REDUCTIONS)
+  add_definitions(-DFP16_QK_REDUCTION_SUPPORTED)
+  target_link_libraries(prefill_kernels PRIVATE Boost::math)
+endif(FLASHINFER_GEN_USE_FP16_QK_REDUCTIONS)
 
 if(FLASHINFER_UNITTESTS)
   add_subdirectory(tests)