Repository: ggerganov/llama.cpp
File: ggml/src/ggml-cuda/common.cuh (CUDA)

Code before:
        }
    }
    return a;

#else
    NO_DEVICE_CODE;
    return a;
#endif // FP16_AVAILABLE
}

enum warp_reduce_method {
    WARP_REDUCE_MAX,
    WARP_REDUCE_SUM,
};

template <warp_reduce_method reduce_method, const unsigned int block_size_template = 0> static __device__ float two_stage_warp_reduce(float val, float * shared_vals) {

    float (*reduce_fun)(float);
    switch (reduce_method) {
        case WARP_REDUCE_MAX:
            reduce_fun = warp_reduce_max;
            break;
        case WARP_REDUCE_SUM:
            reduce_fun = warp_reduce_sum;
            break;
    }

    val = reduce_fun(val);
    const unsigned int block_size = block_size_template == 0 ? blockDim.x : block_size_template;
    if (block_size > WARP_SIZE) {
        assert((block_size <= 1024) && (block_size % WARP_SIZE) == 0);
        const int        warp_id = threadIdx.x / WARP_SIZE;
        const int        lane_id = threadIdx.x % WARP_SIZE;
        if (lane_id == 0) {
            shared_vals[warp_id] = val;
        }
        __syncthreads();
        switch (reduce_method) {
            case WARP_REDUCE_MAX:
                val = -INFINITY;
                break;
            case WARP_REDUCE_SUM:
                val = 0.0f;
                break;
        }
        if (lane_id < (static_cast<int>(block_size) / WARP_SIZE)) {
            val = shared_vals[lane_id];
        }
        return reduce_fun(val);
    } else {
        return val;

Review comment (suggestion):
We can make this a trait

```
struct reduction_traits<MAX> {
   using reduce_fn_t = warp_reduce_max; 
   static constexpr float sentinel = -INFINITY;
}
```

Code after:
        }
    }
    return a;

#else
    NO_DEVICE_CODE;
    return a;
#endif // FP16_AVAILABLE
}

enum class block_reduce_method {
    MAX,
    SUM,
};

template<block_reduce_method method_t, typename T>
struct block_reduce_policy;

template <typename T, typename... Ts>
inline constexpr bool is_any = (std::is_same_v<T, Ts> || ...);

template<typename...>
inline constexpr bool ggml_cuda_dependent_false_v = false;

template <typename T> struct block_reduce_policy<block_reduce_method::SUM, T> {
    static __device__ T reduce(T val) {
        if constexpr(is_any<T, float, float2, half2, int>) {
            return warp_reduce_sum(val);
        } else {
            static_assert(ggml_cuda_dependent_false_v<T>, "Unsupported type for block reduce sum");
        }
    }

    static __device__ T sentinel() {
        if constexpr (std::is_same_v<T, float>) {
            return 0.0f;
        } else if constexpr (std::is_same_v<T, float2>) {
            return make_float2(0.0f, 0.0f);
        } else if constexpr (std::is_same_v<T, half2>) {
            return make_half2(0.0f, 0.0f);
        } else if constexpr (std::is_same_v<T, int>) {
            return 0;
        } else {
            static_assert(ggml_cuda_dependent_false_v<T>, "Unsupported type for block reduce sum");
        }
    }
};

template <typename T> struct block_reduce_policy<block_reduce_method::MAX, T> {
    static __device__ T reduce(T val) {
        if constexpr (is_any<T, float, half2>) {