bugfix: Fix compilation with FP16_QK_REDUCTION enabled. (#962)

As described in #806 and #936, setting the cmake build flag
`FLASHINFER_GEN_USE_FP16_QK_REDUCTIONS` to "true" causes a build failure
due to `cuda_fp16.h` not supporting `constexpr` cast from `__half` type
to `float`. Note that the issue is not just a CMake/C++ configuration
issue the issue will be triggered even in the flashinfer JIT code
compilation path as reported in #915.

The PR fixes #806 and #936 by adding a modified version of the FP16
header from the [FP16 library](https://github.com/Maratyszcza/FP16) that
supports `constexpr` versions of the conversion functions. To make the
conversion functions `constexpr`, I am using `std::bit_cast` that is the
reason for bumping the required standard to 20.

With these changes I am able to build the C++ API with both
`FLASHINFER_GEN_USE_FP16_QK_REDUCTIONS` ON and OFF.

Fixes #936 
Fixes #806

Author: diptorupd

CMakeLists.txt

@@ -3,8 +3,8 @@ project(flashinfer CUDA CXX)
 
 include(cmake/utils/Utils.cmake)
 
-set(CMAKE_CXX_STANDARD 17)
-set(CMAKE_CUDA_STANDARD 17)
+set(CMAKE_CXX_STANDARD 20)
+set(CMAKE_CUDA_STANDARD 20)
 
 if(EXISTS ${CMAKE_BINARY_DIR}/config.cmake)
   include(${CMAKE_BINARY_DIR}/config.cmake)
@@ -63,7 +63,7 @@ flashinfer_option(FLASHINFER_GEN_HEAD_DIMS "Head dims to enable" 64 128 256)
 flashinfer_option(FLASHINFER_GEN_POS_ENCODING_MODES "Pos encodings to enable" 0
                   1 2)
 flashinfer_option(FLASHINFER_GEN_USE_FP16_QK_REDUCTIONS
-                  "QK reductions to enable" "false" "true")
+                  "QK reductions to enable" OFF)
 flashinfer_option(FLASHINFER_GEN_MASK_MODES "Mask modes to enable" 0 1 2)
 
 if(DEFINED FLASHINFER_CUDA_ARCHITECTURES)
@@ -125,25 +125,77 @@ set(POS_ENCODING_MODES ${FLASHINFER_GEN_POS_ENCODING_MODES})
 set(USE_FP16_QK_REDUCTIONS ${FLASHINFER_GEN_USE_FP16_QK_REDUCTIONS})
 set(MASK_MODES ${FLASHINFER_GEN_MASK_MODES})
 
+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 options
 message(STATUS "FLASHINFER_HEAD_DIMS=${HEAD_DIMS}")
 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}")
 
+# 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 ${PROJECT_SOURCE_DIR}/src/generated)
 file(MAKE_DIRECTORY ${PROJECT_SOURCE_DIR}/src/generated)
 
+if(FLASHINFER_GEN_USE_FP16_QK_REDUCTIONS)
+  # ----------------------------- Dependencies -------------------------------#
+  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)
+
 set(AOT_GENERATE_COMMAND
     ${Python3_EXECUTABLE} -m aot_build_utils.generate --path
-    ${PROJECT_SOURCE_DIR}/src/generated --head_dims ${HEAD_DIMS}
-    --pos_encoding_modes ${POS_ENCODING_MODES} --use_fp16_qk_reductions
-    ${USE_FP16_QK_REDUCTIONS} --mask_modes ${MASK_MODES} --enable_f16
-    ${FLASHINFER_ENABLE_F16} --enable_bf16 ${FLASHINFER_ENABLE_BF16}
-    --enable_fp8_e4m3 ${FLASHINFER_ENABLE_FP8_E4M3} --enable_fp8_e5m2
+    ${GENERATED_SOURCE_DIR} --head_dims ${HEAD_DIMS} --pos_encoding_modes
+    ${POS_ENCODING_MODES} --use_fp16_qk_reductions ${USE_FP16_QK_REDUCTIONS}
+    --mask_modes ${MASK_MODES} --enable_f16 ${FLASHINFER_ENABLE_F16}
+    --enable_bf16 ${FLASHINFER_ENABLE_BF16} --enable_fp8_e4m3
+    ${FLASHINFER_ENABLE_FP8_E4M3} --enable_fp8_e5m2
     ${FLASHINFER_ENABLE_FP8_E5M2})
 
+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 ${HEAD_DIMS_SM90} --pos_encoding_modes
+    ${POS_ENCODING_MODES} --use_fp16_qk_reductions ${USE_FP16_QK_REDUCTIONS}
+    --mask_modes ${MASK_MODES})
+
 execute_process(COMMAND ${AOT_GENERATE_COMMAND}
                 WORKING_DIRECTORY ${PROJECT_SOURCE_DIR})
+execute_process(COMMAND ${AOT_GENERATE_DISPATCH_INC_COMMAND}
+                WORKING_DIRECTORY ${PROJECT_SOURCE_DIR})
 
 file(GLOB_RECURSE FLASHINFER_GENERATORS
      ${PROJECT_SOURCE_DIR}/aot_build_utils/*.py)
@@ -157,21 +209,33 @@ file(GLOB_RECURSE DISPATCH_INC_FILE
 add_custom_command(
   OUTPUT ${DECODE_KERNELS_SRCS} ${PREFILL_KERNELS_SRCS} ${DISPATCH_INC_FILE}
   COMMAND ${AOT_GENERATE_COMMAND}
+  COMMAND ${AOT_GENERATE_DISPATCH_INC_COMMAND}
   DEPENDS ${FLASHINFER_GENERATORS}
   WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}
   COMMENT "Generating kernel sources"
   VERBATIM)
 add_custom_target(dispatch_inc DEPENDS ${DISPATCH_INC_FILE})
 
+string(CONCAT CXX_FLAGS "-fpic " "-fPIC ")
+
+string(CONCAT NVCC_FLAGS "-O3 " "--threads=1 " "-Xfatbin=-compress-all "
+              "-use_fast_math " "--expt-relaxed-constexpr ")
+
+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${CXX_FLAGS}")
+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_compile_options(decode_kernels PRIVATE -Xcompiler=-fPIC --fatbin-options
-                                              -compress-all)
+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_compile_options(prefill_kernels PRIVATE -Xcompiler=-fPIC
-                                               --fatbin-options -compress-all)
+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_DECODE)
   message(STATUS "Compile single decode kernel benchmarks.")

include/flashinfer/attention/prefill.cuh

@@ -15,6 +15,7 @@
  */
 #ifndef FLASHINFER_PREFILL_CUH_
 #define FLASHINFER_PREFILL_CUH_
+
 #include <cooperative_groups.h>
 #include <cuda_bf16.h>
 #include <cuda_fp16.h>
@@ -23,6 +24,9 @@
 
 #include "../cp_async.cuh"
 #include "../fastdiv.cuh"
+#ifdef FP16_QK_REDUCTION_SUPPORTED
+#include "../fp16.h"
+#endif
 #include "../frag_layout_swizzle.cuh"
 #include "../math.cuh"
 #include "../mma.cuh"
@@ -33,7 +37,6 @@
 #include "cascade.cuh"
 #include "mask.cuh"
 #include "variants.cuh"
-
 namespace flashinfer {
 
 DEFINE_HAS_MEMBER(maybe_q_rope_offset)
@@ -133,9 +136,25 @@ struct KernelTraits {
 
   using SharedStorage = SharedStorageQKVO<NUM_WARPS_KV, CTA_TILE_Q, CTA_TILE_KV, HEAD_DIM_QK,
                                           HEAD_DIM_VO, DTypeQ, DTypeKV, DTypeO>;
+#ifdef FP16_QK_REDUCTION_SUPPORTED
+  template <typename DT>
+  static constexpr DT getNegInf() {
+    if constexpr (std::is_same<DT, __half>::value) {
+      return std::bit_cast<half>(fp16_ieee_from_fp32_value(-math::inf));
+    } else {
+      return static_cast<DTypeQKAccum>(-math::inf);
+    }
+  }
 
+  static constexpr DTypeQKAccum MaskFillValue =
+      AttentionVariant::use_softmax ? getNegInf<DTypeQKAccum>() : DTypeQKAccum(0.f);
+#else
+  static_assert(!std::is_same<DTypeQKAccum, __half>::value,
+                "Set -DFP16_QK_REDUCTION_SUPPORTED and install boost_math "
+                "then recompile to support fp16 reduction");
   static constexpr DTypeQKAccum MaskFillValue =
       AttentionVariant::use_softmax ? DTypeQKAccum(-math::inf) : DTypeQKAccum(0.f);
+#endif
 };
 
 namespace {
@@ -672,6 +691,8 @@ __device__ __forceinline__ void logits_transform(
     const uint32_t kv_head_idx = blockIdx.z) {
   const uint32_t lane_idx = tid.x;
   uint32_t q[KTraits::NUM_MMA_Q][2], r[KTraits::NUM_MMA_Q][2];
+  float logits = 0., logitsTransformed = 0.;
+
 #pragma unroll
   for (uint32_t mma_q = 0; mma_q < KTraits::NUM_MMA_Q; ++mma_q) {
 #pragma unroll
@@ -691,9 +712,31 @@ __device__ __forceinline__ void logits_transform(
                                                                     2 * (lane_idx % 4) +
                                                                     8 * (reg_id / 4) + reg_id % 2;
         const uint32_t qo_head_idx = kv_head_idx * group_size + r[mma_q][(reg_id % 4) / 2];
-        s_frag[mma_q][mma_kv][reg_id] =
-            variant.LogitsTransform(params, s_frag[mma_q][mma_kv][reg_id], batch_idx, q_idx, kv_idx,
-                                    qo_head_idx, kv_head_idx);
+
+#ifdef FP16_QK_REDUCTION_SUPPORTED
+        if constexpr (std::is_same<DTypeQKAccum, __half>::value) {
+          logits = std::bit_cast<float>(fp16_ieee_to_fp32_value(s_frag[mma_q][mma_kv][reg_id]));
+        } else if constexpr (!std::is_same<DTypeQKAccum, __half>::value) {
+          logits = s_frag[mma_q][mma_kv][reg_id];
+        }
+#else
+        static_assert(!std::is_same<DTypeQKAccum, __half>::value,
+                      "Set -DFP16_QK_REDUCTION_SUPPORTED and install boost_math "
+                      "then recompile to support fp16 reduction");
+        logits = s_frag[mma_q][mma_kv][reg_id];
+#endif
+        logitsTransformed = variant.LogitsTransform(params, logits, batch_idx, q_idx, kv_idx,
+                                                    qo_head_idx, kv_head_idx);
+#ifdef FP16_QK_REDUCTION_SUPPORTED
+        if constexpr (std::is_same<DTypeQKAccum, __half>::value) {
+          s_frag[mma_q][mma_kv][reg_id] =
+              std::bit_cast<half>(fp16_ieee_from_fp32_value(logitsTransformed));
+        } else if constexpr (!std::is_same<DTypeQKAccum, __half>::value) {
+          s_frag[mma_q][mma_kv][reg_id] = logitsTransformed;
+        }
+#else
+        s_frag[mma_q][mma_kv][reg_id] = logitsTransformed;
+#endif
       }
     }
   }