Skip to content

Commit c93f647

Browse files
yzh119Antoni BaumCody Yu
authored
perf: faster fp8->fp16 dequantization for pre sm_90 arch (#439)
hardware fp8->fp16 fast conversion instruction is not available for sm_80 & sm_89, which makes #420 slow for these architectures. this pr uses marlin's fast fp8->fp16x4 conversion algorithm (copied from vllm project) to accelerate such cases. Co-authored-by: Antoni Baum <antoni@anyscale.com> Co-authored-by: Cody Yu <cody@anyscale.com>
1 parent adcf701 commit c93f647

4 files changed

Lines changed: 253 additions & 15 deletions

File tree

‎CMakeLists.txt‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,8 @@ flashinfer_option(FLASHINFER_CASCADE "Whether to compile cascade kernel tests/be
3232
flashinfer_option(FLASHINFER_SAMPLING "Whether to compile sampling kernel tests/benchmarks or not." OFF)
3333
flashinfer_option(FLASHINFER_NORM "Whether to compile normalization kernel tests/benchmarks or not." OFF)
3434
flashinfer_option(FLASHINFER_DISTRIBUTED "Whether to compile distributed kernel tests/benchmarks or not." OFF)
35+
flashinfer_option(FLASHINFER_FASTDIV_TEST "Whether to compile fastdiv kernel tests or not." OFF)
36+
flashinfer_option(FLASHINFER_FASTDEQAUNT_TEST "Whether to compile fast dequant kernel tests or not." OFF)
3537
flashinfer_option(FLASHINFER_TVM_BINDING "Whether to compile tvm binding or not." OFF)
3638
flashinfer_option(FLASHINFER_TVM_SOURCE_DIR "The path to tvm for building tvm binding." "")
3739

@@ -477,6 +479,17 @@ if(FLASHINFER_FASTDIV_TEST)
477479
target_link_libraries(test_fastdiv PRIVATE gtest gtest_main)
478480
endif(FLASHINFER_FASTDIV_TEST)
479481

482+
if(FLASHINFER_FASTDEQUANT_TEST)
483+
message(STATUS "Compile fast dequant test.")
484+
file(GLOB_RECURSE TEST_FAST_DEQUANT_SRCS ${PROJECT_SOURCE_DIR}/src/test_fast_dequant.cu)
485+
add_executable(test_fast_dequant ${TEST_FAST_DEQUANT_SRCS})
486+
target_include_directories(test_fast_dequant PRIVATE ${FLASHINFER_INCLUDE_DIR})
487+
target_include_directories(test_fast_dequant PRIVATE ${gtest_SOURCE_DIR}/include ${gtest_SOURCE_DIR})
488+
target_link_libraries(test_fast_dequant PRIVATE gtest gtest_main)
489+
endif(FLASHINFER_FASTDIV_TEST)
490+
491+
492+
480493
if (FLASHINFER_DISTRIBUTED)
481494
find_package(MPI REQUIRED)
482495

‎cmake/config.cmake‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,8 @@ set(FLASHINFER_SAMPLING ON)
1818
set(FLASHINFER_NORMALIZATION ON)
1919
# Whether to compile fastdiv tests
2020
set(FLASHINFER_FASTDIV_TEST ON)
21+
# Whether to compile fastdequant tests
22+
set(FLASHINFER_FASTDEQUANT_TEST ON)
2123
# Whether to compile distributed tests
2224
set(FLASHINFER_DISTRIBUTED ON)
2325
# The following configurations can impact the binary

‎include/flashinfer/vec_dtypes.cuh‎

Lines changed: 167 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -16,19 +16,19 @@
1616
#ifndef VEC_DTYPES_CUH_
1717
#define VEC_DTYPES_CUH_
1818

19-
#ifdef FLASHINFER_ENABLE_BF16
2019
#include <cuda_bf16.h>
21-
#endif
2220
#include <cuda_fp16.h>
23-
#ifdef FLASHINFER_ENABLE_FP8
2421
#include <cuda_fp8.h>
25-
#endif
2622
#include <cuda_runtime.h>
2723

2824
#include <type_traits>
2925

3026
namespace flashinfer {
3127

28+
#if (!defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 900))
29+
#define FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED
30+
#endif
31+
3232
#define FLASHINFER_INLINE inline __attribute__((always_inline)) __device__
3333

3434
/******************* vec_t type cast *******************/
@@ -74,11 +74,130 @@ struct vec_cast<half, float> {
7474
}
7575
};
7676

77-
#if (!defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 900))
77+
template <typename T>
78+
constexpr FLASHINFER_INLINE int get_exponent_bits() {
79+
if constexpr (std::is_same<T, __nv_fp8_e4m3>::value) {
80+
return 4;
81+
} else if constexpr (std::is_same<T, __nv_fp8_e5m2>::value) {
82+
return 5;
83+
} else if constexpr (std::is_same<T, half>::value) {
84+
return 5;
85+
} else if constexpr (std::is_same<T, nv_bfloat16>::value) {
86+
return 8;
87+
}
88+
}
89+
90+
template <typename T>
91+
constexpr FLASHINFER_INLINE int get_mantissa_bits() {
92+
if constexpr (std::is_same<T, __nv_fp8_e4m3>::value) {
93+
return 3;
94+
} else if constexpr (std::is_same<T, __nv_fp8_e5m2>::value) {
95+
return 2;
96+
} else if constexpr (std::is_same<T, half>::value) {
97+
return 11;
98+
} else if constexpr (std::is_same<T, nv_bfloat16>::value) {
99+
return 7;
100+
}
101+
}
102+
103+
/*!
104+
* \brief Fallback to software fast dequant implementation if hardware dequantization is not
105+
* available.
106+
* \note Inspired by Marlin's fast dequantization, but here we don't have to permute
107+
* weights order.
108+
* \ref
109+
* https://github.com/vllm-project/vllm/blob/6dffa4b0a6120159ef2fe44d695a46817aff65bc/csrc/quantization/fp8/fp8_marlin.cu#L120
110+
*/
111+
template <typename fp8_dtype, typename fp16_dtype>
112+
__device__ void fast_dequant_f8f16x4(uint32_t* input, uint2* output) {
113+
uint32_t q = *input;
114+
if constexpr (std::is_same<fp8_dtype, __nv_fp8_e5m2>::value &&
115+
std::is_same<fp16_dtype, half>::value) {
116+
output->x = __byte_perm(0U, q, 0x5140);
117+
output->y = __byte_perm(0U, q, 0x7362);
118+
} else {
119+
constexpr int FP8_EXPONENT = get_exponent_bits<fp8_dtype>();
120+
constexpr int FP8_MANTISSA = get_mantissa_bits<fp8_dtype>();
121+
constexpr int FP16_EXPONENT = get_exponent_bits<fp16_dtype>();
122+
123+
constexpr int RIGHT_SHIFT = FP16_EXPONENT - FP8_EXPONENT;
124+
// Calculate MASK for extracting mantissa and exponent
125+
constexpr int MASK1 = 0x80000000;
126+
constexpr int MASK2 = MASK1 >> (FP8_EXPONENT + FP8_MANTISSA);
127+
constexpr int MASK3 = MASK2 & 0x7fffffff;
128+
constexpr int MASK = MASK3 | (MASK3 >> 16);
129+
// Final MASK value: 0x7F007F00
130+
q = __byte_perm(q, q, 0x1302);
131+
132+
// Extract and shift FP8 values to FP16 format
133+
uint32_t Out1 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT);
134+
uint32_t Out2 = ((q << 8) & 0x80008000) | (((q << 8) & MASK) >> RIGHT_SHIFT);
135+
136+
constexpr int BIAS_OFFSET = (1 << (FP16_EXPONENT - 1)) - (1 << (FP8_EXPONENT - 1));
137+
// Construct and apply exponent bias
138+
if (std::is_same<fp16_dtype, half>::value) {
139+
const half2 bias_reg = __float2half2_rn(float(1 << BIAS_OFFSET));
140+
141+
// Convert to half2 and apply bias
142+
*(half2*)&(output->x) = __hmul2(*reinterpret_cast<const half2*>(&Out1), bias_reg);
143+
*(half2*)&(output->y) = __hmul2(*reinterpret_cast<const half2*>(&Out2), bias_reg);
144+
} else {
145+
constexpr uint32_t BIAS = (BIAS_OFFSET + 127) << 23;
146+
const nv_bfloat162 bias_reg = __float2bfloat162_rn(*reinterpret_cast<const float*>(&BIAS));
147+
// Convert to bfloat162 and apply bias
148+
*(nv_bfloat162*)&(output->x) =
149+
__hmul2(*reinterpret_cast<const nv_bfloat162*>(&Out1), bias_reg);
150+
*(nv_bfloat162*)&(output->y) =
151+
__hmul2(*reinterpret_cast<const nv_bfloat162*>(&Out2), bias_reg);
152+
}
153+
}
154+
}
155+
156+
template <>
157+
struct vec_cast<nv_bfloat16, __nv_fp8_e4m3> {
158+
template <size_t vec_size>
159+
FLASHINFER_INLINE static void cast(nv_bfloat16* dst, const __nv_fp8_e4m3* src) {
160+
if constexpr (vec_size == 1) {
161+
dst[0] = nv_bfloat16(src[0]);
162+
} else if constexpr (vec_size == 2) {
163+
dst[0] = nv_bfloat16(src[0]);
164+
dst[1] = nv_bfloat16(src[1]);
165+
} else {
166+
static_assert(vec_size % 4 == 0, "vec_size must be a multiple of 4");
167+
#pragma unroll
168+
for (uint32_t i = 0; i < vec_size / 4; ++i) {
169+
fast_dequant_f8f16x4<__nv_fp8_e4m3, nv_bfloat16>((uint32_t*)&src[i * 4],
170+
(uint2*)&dst[i * 4]);
171+
}
172+
}
173+
}
174+
};
175+
176+
template <>
177+
struct vec_cast<nv_bfloat16, __nv_fp8_e5m2> {
178+
template <size_t vec_size>
179+
FLASHINFER_INLINE static void cast(nv_bfloat16* dst, const __nv_fp8_e5m2* src) {
180+
if constexpr (vec_size == 1) {
181+
dst[0] = nv_bfloat16(src[0]);
182+
} else if constexpr (vec_size == 2) {
183+
dst[0] = nv_bfloat16(src[0]);
184+
dst[1] = nv_bfloat16(src[1]);
185+
} else {
186+
static_assert(vec_size % 4 == 0, "vec_size must be a multiple of 4");
187+
#pragma unroll
188+
for (uint32_t i = 0; i < vec_size / 4; ++i) {
189+
fast_dequant_f8f16x4<__nv_fp8_e5m2, nv_bfloat16>((uint32_t*)&src[i * 4],
190+
(uint2*)&dst[i * 4]);
191+
}
192+
}
193+
}
194+
};
195+
78196
template <>
79197
struct vec_cast<__nv_fp8_e4m3, half> {
80198
template <size_t vec_size>
81199
FLASHINFER_INLINE static void cast(__nv_fp8_e4m3* dst, const half* src) {
200+
#ifdef FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED
82201
if constexpr (vec_size == 1) {
83202
dst[0] = __nv_fp8_e4m3(src[0]);
84203
} else {
@@ -90,13 +209,20 @@ struct vec_cast<__nv_fp8_e4m3, half> {
90209
*(uint16_t*)&dst[i * 2] = y;
91210
}
92211
}
212+
#else
213+
#pragma unroll
214+
for (size_t i = 0; i < vec_size; ++i) {
215+
dst[i] = __nv_fp8_e4m3(src[i]);
216+
}
217+
#endif // FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED
93218
}
94219
};
95220

96221
template <>
97222
struct vec_cast<__nv_fp8_e5m2, half> {
98223
template <size_t vec_size>
99224
FLASHINFER_INLINE static void cast(__nv_fp8_e5m2* dst, const half* src) {
225+
#ifdef FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED
100226
if constexpr (vec_size == 1) {
101227
dst[0] = __nv_fp8_e5m2(src[0]);
102228
} else {
@@ -108,13 +234,20 @@ struct vec_cast<__nv_fp8_e5m2, half> {
108234
*(uint16_t*)&dst[i * 2] = y;
109235
}
110236
}
237+
#else
238+
#pragma unroll
239+
for (size_t i = 0; i < vec_size; ++i) {
240+
dst[i] = __nv_fp8_e5m2(src[i]);
241+
}
242+
#endif // FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED
111243
}
112244
};
113245

114246
template <>
115247
struct vec_cast<half, __nv_fp8_e4m3> {
116248
template <size_t vec_size>
117249
FLASHINFER_INLINE static void cast(half* dst, const __nv_fp8_e4m3* src) {
250+
#ifdef FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED
118251
if constexpr (vec_size == 1) {
119252
dst[0] = half(src[0]);
120253
} else {
@@ -126,13 +259,28 @@ struct vec_cast<half, __nv_fp8_e4m3> {
126259
*(uint32_t*)&dst[i * 2] = y;
127260
}
128261
}
262+
#else
263+
if constexpr (vec_size == 1) {
264+
dst[0] = half(src[0]);
265+
} else if constexpr (vec_size == 2) {
266+
dst[0] = half(src[0]);
267+
dst[1] = half(src[1]);
268+
} else {
269+
static_assert(vec_size % 4 == 0, "vec_size must be a multiple of 4");
270+
#pragma unroll
271+
for (uint32_t i = 0; i < vec_size / 4; ++i) {
272+
fast_dequant_f8f16x4<__nv_fp8_e4m3, half>((uint32_t*)&src[i * 4], (uint2*)&dst[i * 4]);
273+
}
274+
}
275+
#endif // FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED
129276
}
130277
};
131278

132279
template <>
133280
struct vec_cast<half, __nv_fp8_e5m2> {
134281
template <size_t vec_size>
135282
FLASHINFER_INLINE static void cast(half* dst, const __nv_fp8_e5m2* src) {
283+
#ifdef FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED
136284
if constexpr (vec_size == 1) {
137285
dst[0] = half(src[0]);
138286
} else {
@@ -144,13 +292,23 @@ struct vec_cast<half, __nv_fp8_e5m2> {
144292
*(uint32_t*)&dst[i * 2] = y;
145293
}
146294
}
295+
#else
296+
if constexpr (vec_size == 1) {
297+
dst[0] = half(src[0]);
298+
} else if constexpr (vec_size == 2) {
299+
dst[0] = half(src[0]);
300+
dst[1] = half(src[1]);
301+
} else {
302+
static_assert(vec_size % 4 == 0, "vec_size must be a multiple of 4");
303+
#pragma unroll
304+
for (uint32_t i = 0; i < vec_size / 4; ++i) {
305+
fast_dequant_f8f16x4<__nv_fp8_e5m2, half>((uint32_t*)&src[i * 4], (uint2*)&dst[i * 4]);
306+
}
307+
}
308+
#endif // FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED
147309
}
148310
};
149311

150-
#endif // !defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 900)
151-
152-
#ifdef FLASHINFER_ENABLE_BF16
153-
154312
template <>
155313
struct vec_cast<float, nv_bfloat16> {
156314
template <size_t vec_size>
@@ -180,7 +338,6 @@ struct vec_cast<nv_bfloat16, float> {
180338
}
181339
}
182340
};
183-
#endif // FLASHINFER_ENABLE_BF16
184341

185342
template <typename float_t, size_t vec_size>
186343
struct vec_t {
@@ -230,7 +387,6 @@ FLASHINFER_INLINE void cast_store_impl(tgt_float_t* dst_ptr,
230387
}
231388
}
232389

233-
#ifdef FLASHINFER_ENABLE_FP8
234390
/******************* vec_t<__nv_fp8_e4m3> *******************/
235391

236392
// __nv_fp8_e4m3 x 1
@@ -724,7 +880,6 @@ struct vec_t<__nv_fp8_e5m2, vec_size> {
724880
}
725881
}
726882
};
727-
#endif
728883

729884
/******************* vec_t<half> *******************/
730885

@@ -889,7 +1044,6 @@ struct vec_t<half, vec_size> {
8891044
}
8901045
};
8911046

892-
#ifdef FLASHINFER_ENABLE_BF16
8931047
/******************* vec_t<nv_bfloat16> *******************/
8941048

8951049
// nv_bfloat16 x 1
@@ -1071,8 +1225,6 @@ struct vec_t<nv_bfloat16, vec_size> {
10711225
}
10721226
};
10731227

1074-
#endif
1075-
10761228
/******************* vec_t<float> *******************/
10771229

10781230
// float x 1

0 commit comments

Comments
 (0)