A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 22:15 UTC
[ROCm][Quantization][Kernel] Using HIP FP8 header (#12593)
Gregory Shtrasberg committed
Feb 25, 2025 at 03:39 UTC
aabeb2688fba861a39547a2a33649e6330caaafd
6 files changed
+267
-634
CMakeLists.txt
+19
@@ -174,6 +174,25 @@ include(FetchContent)
174
file(MAKE_DIRECTORY ${FETCHCONTENT_BASE_DIR}) # Ensure the directory exists
175
message(STATUS "FetchContent base directory: ${FETCHCONTENT_BASE_DIR}")
176
177
+#
178
+# Set rocm version dev int.
179
+#
180
+if(VLLM_GPU_LANG STREQUAL "HIP")
181
+ #
182
+ # Overriding the default -O set up by cmake, adding ggdb3 for the most verbose devug info
183
+ #
184
+ set(CMAKE_${VLLM_GPU_LANG}_FLAGS_DEBUG "${CMAKE_${VLLM_GPU_LANG}_FLAGS_DEBUG} -O0 -ggdb3")
185
+ set(CMAKE_CXX_FLAGS_DEBUG "${CMAKE_CXX_FLAGS_DEBUG} -O0 -ggdb3")
186
+
187
+
188
+ #
189
+ # Certain HIP functions are marked as [[nodiscard]], yet vllm ignores the result which generates
190
+ # a lot of warnings that always mask real issues. Suppressing until this is properly addressed.
191
+ #
192
+ set(CMAKE_${VLLM_GPU_LANG}_FLAGS "${CMAKE_${VLLM_GPU_LANG}_FLAGS} -Wno-unused-result")
193
+ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-unused-result")
194
+endif()
195
+
196
#
197
# Define other extension targets
198
#
csrc/quantization/fp8/amd/hip_float8.h
deleted
-137
@@ -1,137 +0,0 @@
1
-#pragma once
2
-
3
-#ifdef __HIPCC__
4
- #include <hip/hip_runtime.h>
5
-#else
6
- #include <type_traits>
7
- #include <stdint.h>
8
- #include <math.h>
9
- #include <iostream>
10
-#endif
11
-
12
-#include "hip_float8_impl.h"
13
-
14
-struct alignas(1) hip_fp8 {
15
- struct from_bits_t {};
16
- HIP_FP8_HOST_DEVICE static constexpr from_bits_t from_bits() {
17
- return from_bits_t();
18
- }
19
- uint8_t data;
20
-
21
- hip_fp8() = default;
22
- HIP_FP8_HOST_DEVICE constexpr hip_fp8(const hip_fp8&) = default;
23
- HIP_FP8_HOST_DEVICE constexpr hip_fp8(uint8_t v) = delete;
24
- explicit HIP_FP8_HOST_DEVICE constexpr hip_fp8(uint8_t v, from_bits_t)
25
- : data(v) {}
26
-
27
-#ifdef __HIP__MI300__
28
- // NOTE: ON-DEVICE... always optimal bias
29
- explicit HIP_FP8_DEVICE hip_fp8(float v)
30
- : data(hip_fp8_impl::to_fp8_from_fp32(v)) {}
31
-
32
- explicit HIP_FP8_DEVICE hip_fp8(_Float16 v)
33
- : hip_fp8(static_cast<float>(v)) {}
34
-
35
- // Host only implementation using s/w simulation
36
- explicit HIP_FP8_HOST
37
-#else // __HIP__MI300__
38
- // both Host and DEVICE for non-MI300 using s/w simulation
39
- explicit HIP_FP8_HOST_DEVICE
40
-#endif // __HIP__MI300__
41
- hip_fp8(float v) {
42
- data = hip_fp8_impl::to_float8<4, 3, float, true /*negative_zero_nan*/,
43
- true /*clip*/>(v);
44
- }
45
-
46
- explicit HIP_FP8_HOST_DEVICE hip_fp8(double v)
47
- : hip_fp8(static_cast<float>(v)) {}
48
-
49
-#ifdef __HIP__MI300__
50
- // upcast using device specific intrinsic
51
- explicit inline HIP_FP8_DEVICE operator float() const {
52
- float fval;
53
- uint32_t i32val = static_cast<uint32_t>(data);
54
-
55
- // upcast
56
- asm volatile("v_cvt_f32_fp8 %0, %1 src0_sel:BYTE_0"
57
- : "=v"(fval)
58
- : "v"(i32val));
59
-
60
- return fval;
61
- }
62
-
63
- explicit inline HIP_FP8_HOST operator float() const
64
-#else // __HIP__MI300__
65
- explicit inline HIP_FP8_HOST_DEVICE operator float() const
66
-#endif // __HIP__MI300__
67
- {
68
- return hip_fp8_impl::from_float8<4, 3, float, true /*negative_zero_nan*/>(
69
- data);
70
- }
71
-};
72
-
73
-namespace std {
74
-inline hip_fp8 sin(hip_fp8 a) { return hip_fp8(sinf(float(a))); }
75
-inline hip_fp8 cos(hip_fp8 a) { return hip_fp8(cosf(float(a))); }
76
-HIP_FP8_HOST_DEVICE constexpr hip_fp8 real(const hip_fp8& a) { return a; }
77
-} // namespace std
78
-
79
-// Special operator overloading
80
-inline std::ostream& operator<<(std::ostream& os, const hip_fp8& f8) {
81
- return os << float(f8);
82
-}
83
-
84
-// all + operator overloading with mixed types
85
-// mixed types, always converts to f32, does computation in f32, and returns
86
-// float
87
-inline HIP_FP8_HOST_DEVICE float operator+(const float fa, hip_fp8 b) {
88
- return (fa + float(b));
89
-}
90
-
91
-inline HIP_FP8_HOST_DEVICE float operator+(hip_fp8 a, const float fb) {
92
- return (float(a) + fb);
93
-}
94
-
95
-inline HIP_FP8_HOST_DEVICE hip_fp8 operator+(hip_fp8 a, hip_fp8 b) {
96
- return hip_fp8(float(a) + float(b));
97
-}
98
-
99
-inline HIP_FP8_HOST_DEVICE hip_fp8& operator+=(hip_fp8& a, hip_fp8 b) {
100
- return a = hip_fp8(float(a) + float(b));
101
-}
102
-
103
-// overloading multiplication, always returns float,
104
-inline HIP_FP8_HOST_DEVICE float operator*(hip_fp8 a, hip_fp8 b) {
105
- return float(a) * float(b);
106
-}
107
-
108
-inline HIP_FP8_HOST_DEVICE float operator*(float a, hip_fp8 b) {
109
- return (a * float(b));
110
-}
111
-
112
-inline HIP_FP8_HOST_DEVICE float operator*(hip_fp8 a, float b) {
113
- return (float(a) * b);
114
-}
115
-
116
-inline HIP_FP8_HOST_DEVICE float operator*(int32_t a, hip_fp8 b) {
117
- return ((float)a * float(b));
118
-}
119
-
120
-inline HIP_FP8_HOST_DEVICE float operator*(double a, hip_fp8 b) {
121
- return ((float)a * float(b));
122
-}
123
-
124
-// overloading for compare
125
-inline HIP_FP8_HOST_DEVICE bool operator==(hip_fp8 a, hip_fp8 b) {
126
- return (a.data == b.data);
127
-}
128
-inline HIP_FP8_HOST_DEVICE bool operator!=(hip_fp8 a, hip_fp8 b) {
129
- return (a.data != b.data);
130
-}
131
-
132
-inline HIP_FP8_HOST_DEVICE bool operator>=(hip_fp8 a, hip_fp8 b) {
133
- return static_cast<float>(a) >= static_cast<float>(b);
134
-}
135
-inline HIP_FP8_HOST_DEVICE bool operator>(hip_fp8 a, hip_fp8 b) {
136
- return static_cast<float>(a) > static_cast<float>(b);
137
-}
csrc/quantization/fp8/amd/hip_float8_impl.h
deleted
-315
@@ -1,315 +0,0 @@
1
-#pragma once
2
-
3
-#if defined(__HIPCC__) && defined(__gfx942__)
4
- #define __HIP__MI300__
5
-#endif
6
-
7
-#ifdef __HIPCC__
8
- #define HIP_FP8_HOST_DEVICE __host__ __device__
9
- #define HIP_FP8_HOST __host__
10
- #define HIP_FP8_DEVICE __device__
11
-#else
12
- #define HIP_FP8_HOST_DEVICE
13
- #define HIP_FP8_HOST
14
- #define HIP_FP8_DEVICE
15
-#endif
16
-
17
-namespace hip_fp8_impl {
18
-
19
-#ifdef __HIP__MI300__
20
-HIP_FP8_DEVICE uint8_t to_fp8_from_fp32(float v) {
21
- uint8_t i8data;
22
- union {
23
- float fval;
24
- uint32_t i32val;
25
- uint8_t i8val[4]; // NOTE: not endian independent
26
- } val;
27
-
28
- uint32_t ival = 0;
29
- val.fval = v;
30
-
31
- if ((val.i32val & 0x7F800000) !=
32
- 0x7F800000) { /// propagate NAN/INF, no clipping
33
- val.fval = __builtin_amdgcn_fmed3f(val.fval, 240.0, -240.0);
34
- }
35
-
36
- ival = __builtin_amdgcn_cvt_pk_fp8_f32(val.fval, val.fval, ival,
37
- false); // false -> WORD0
38
- val.i32val = ival;
39
- i8data = val.i8val[0];
40
-
41
- return i8data;
42
-}
43
-#endif // __HIP__MI300__
44
-
45
-HIP_FP8_HOST inline int clz(uint32_t x) { return __builtin_clz(x); }
46
-#if defined(__HIPCC__) || defined(__CUDA_ARCH__)
47
-HIP_FP8_DEVICE inline int clz(uint32_t x) { return __clz(x); }
48
-#endif
49
-
50
-template <int we, int wm, typename T, bool negative_zero_nan, bool clip>
51
-HIP_FP8_HOST_DEVICE uint8_t to_float8(T _x, bool stoch = false,
52
- uint32_t rng = 0) {
53
-#ifdef __HIPCC__
54
- constexpr bool is_half = std::is_same<T, _Float16>::value;
55
-#else
56
- constexpr bool is_half = false;
57
-#endif
58
- constexpr bool is_float = std::is_same<T, float>::value;
59
- static_assert(wm + we == 7, "wm+we==7");
60
- static_assert(is_half || is_float, "Only half and float can be cast to f8");
61
-
62
- const int mfmt = (sizeof(T) == 4) ? 23 : 10;
63
- uint32_t x;
64
- if (sizeof(T) == 4) {
65
- x = reinterpret_cast<uint32_t&>(_x);
66
- } else {
67
- x = reinterpret_cast<uint16_t&>(_x);
68
- }
69
-
70
- uint32_t head, mantissa;
71
- int exponent, bias;
72
- uint32_t sign;
73
-
74
- if (sizeof(T) == 4) {
75
- head = x & 0xFF800000;
76
- mantissa = x & 0x7FFFFF;
77
- exponent = (head >> 23) & 0xFF;
78
- sign = head >> 31;
79
- bias = 127;
80
- } else {
81
- head = x & 0xFC00;
82
- mantissa = x & 0x3FF;
83
- exponent = (head >> 10) & 0x1F;
84
- sign = head >> 15;
85
- bias = 15;
86
- }
87
-
88
- uint32_t signed_inf = (sign << 7) + (((1 << we) - 1) << wm);
89
-
90
- // Deal with inf and NaNs
91
- if (negative_zero_nan) {
92
- if (sizeof(T) == 4) {
93
- if ((x & 0x7F800000) == 0x7F800000) {
94
- return 0x80;
95
- }
96
- } else {
97
- // if(__hisinf(x) || __hisnan(x))
98
- if ((x & 0x7C00) == 0x7C00) {
99
- return 0x80;
100
- }
101
- }
102
- } else {
103
- if (sizeof(T) == 4) {
104
- if ((x & 0x7F800000) == 0x7F800000) {
105
- return signed_inf + (mantissa != 0 ? 1 : 0);
106
- }
107
- } else {
108
- if ((x & 0x7C00) == 0x7C00) {
109
- return signed_inf + (mantissa != 0 ? 1 : 0);
110
- }
111
- }
112
- }
113
- if (x == 0) {
114
- return 0;
115
- }
116
-
117
- // First need to check if it is normal or denorm as there is a difference of
118
- // implicit 1 Then need to adjust the exponent to align with the F8 exponent,
119
- // in the meanwhile, shift The mantissa. Then for stochastic rounding, add rng
120
- // to mantissa and truncate. And for RNE, no need to add rng. Then probably
121
- // need to check whether there is carry and adjust exponent and mantissa again
122
-
123
- // For IEEE bias mode, the bias is 2^(k-1) -1 where k is the width of exponent
124
- // bits
125
- const int f8_bias = (1 << (we - 1)) - 1 + (negative_zero_nan ? 1 : 0);
126
- const int f8_denormal_act_exponent =
127
- 1 - f8_bias; // actual exponent of f8 denormal
128
- // act_exponent is the actual exponent of fp32/fp16 (after subtracting bias)
129
- // f8_exponent is the converted f8 exponent with bias encoding
130
- // exponent_diff is the diff between fp32/fp16 exponent and f8 exponent,
131
- // the difference needs to be adjusted and mantissa shifted
132
- int act_exponent, f8_exponent, exponent_diff;
133
-
134
- if (exponent == 0) { // fp32/fp16 is in denormal.
135
- /* fp32 denormal is below 2^-127 so it is usually not a concern here, we
136
-mostly concern fp16 here. In this case, f8 is usually in denormal. But there
137
-could be exceptions. fp16 denormal has exponent bias 15 while bf8 with NANOO has
138
-exponent bias 16. It means that there are some numbers in fp16 denormal but they
139
-are bf8 (NANOO) normals - smallest bf8 (NANOO) normal is 2^-15. fp16 numbers
140
-where exponent==0 (actual exponent -14) and highest bit of mantissa is 1 are bf8
141
-(NANOO) normal. In this case, the fp16 mantissa should be shift left by 1 */
142
- act_exponent = exponent - bias + 1;
143
- exponent_diff =
144
- f8_denormal_act_exponent -
145
- act_exponent; // actual exponent is exponent-bias+1 as it is denormal
146
- } else { // fp32/fp16 is normal with implicit 1
147
- act_exponent = exponent - bias;
148
- if (act_exponent <= f8_denormal_act_exponent) {
149
- /* This is the case where fp32/fp16 is normal but it is in f8 denormal
150
-range. For example fp8 nanoo mode, denormal exponent is -7, but if the
151
-fp32/fp16 actual exponent is -7, it is actually larger due to the implicit 1,
152
-Therefore it needs to be adjust to -6 and mantissa shift right by 1.
153
-So for fp32/fp16, exponent -8 is the cut point to convert to fp8 nanoo */
154
- exponent_diff = f8_denormal_act_exponent - act_exponent;
155
- } else { // both fp32/fp16 and f8 are in normal range
156
- exponent_diff = 0; // exponent_diff=0 does not mean there is no
157
- // difference for this case, act_exponent could be
158
- // larger. Just that it does not need shift mantissa
159
- }
160
- mantissa += (1 << mfmt); // Add the implicit 1 into mantissa
161
- }
162
-
163
- bool midpoint = (mantissa & ((1 << (mfmt - wm + exponent_diff)) - 1)) ==
164
- static_cast<uint32_t>(1 << (mfmt - wm + exponent_diff - 1));
165
- /* This part is a bit tricky. The judgment of whether it is a tie needs to be
166
- done before we shift right as shift right could rip off some residual part
167
- and make something not midpoint look like midpoint. For example, the fp16
168
- number 0x1002 (0 00100 0000000010), it is larger than midpoint, but after
169
- shift right by 4 bits, it would look like midpoint.
170
-*/
171
-
172
- if (exponent_diff > 0) {
173
- mantissa >>= exponent_diff;
174
- } else if (exponent_diff == -1) {
175
- mantissa <<= -exponent_diff;
176
- }
177
- bool implicit_one = mantissa & (1 << mfmt);
178
- // if there is no implicit 1, it means the f8 is denormal and need to adjust
179
- // to denorm exponent
180
- f8_exponent = (act_exponent + exponent_diff) /*actual f8 exponent*/ +
181
- f8_bias - (implicit_one ? 0 : 1);
182
-
183
- // Now we have the exponent and mantissa adjusted
184
- uint32_t drop_mask = (1 << (mfmt - wm)) - 1;
185
- bool odd = mantissa & (1 << (mfmt - wm)); // if the least significant bit
186
- // that is not truncated is 1
187
- mantissa +=
188
- (stoch ? rng : (midpoint ? (odd ? mantissa : mantissa - 1) : mantissa)) &
189
- drop_mask;
190
-
191
- // Now we deal with overflow
192
- if (f8_exponent == 0) {
193
- if ((1 << mfmt) & mantissa) {
194
- f8_exponent = 1; // denormal overflow to become normal, promote exponent
195
- }
196
- } else {
197
- if ((1 << (mfmt + 1)) & mantissa) {
198
- mantissa >>= 1;
199
- f8_exponent++;
200
- }
201
- }
202
-
203
- mantissa >>= (mfmt - wm);
204
-
205
- // above range: quantize to maximum possible float of the same sign
206
- const int max_exp = (1 << we) - (negative_zero_nan ? 1 : 2);
207
- if (f8_exponent > max_exp) {
208
- if (clip) {
209
- mantissa = (1 << wm) - 1;
210
- f8_exponent = max_exp;
211
- } else {
212
- return signed_inf;
213
- }
214
- }
215
-
216
- if (f8_exponent == 0 && mantissa == 0) {
217
- return negative_zero_nan ? 0 : (sign << 7);
218
- }
219
- mantissa &= (1 << wm) - 1;
220
- return (sign << 7) | (f8_exponent << wm) | mantissa;
221
-}
222
-
223
-template <int we, int wm, typename T = float, bool negative_zero_nan = true>
224
-inline HIP_FP8_HOST_DEVICE T from_float8(uint8_t x) {
225
-#ifdef __HIPCC__
226
- constexpr bool is_half = std::is_same<T, _Float16>::value;
227
-#else
228
- constexpr bool is_half = false;
229
-#endif
230
- constexpr bool is_float = std::is_same<T, float>::value;
231
- static_assert(is_half || is_float, "only half and float are supported");
232
-
233
- constexpr int weo = is_half ? 5 : 8;
234
- constexpr int wmo = is_half ? 10 : (is_float ? 23 : 7);
235
-
236
- T fInf, fNegInf, fNaN, fNeg0;
237
-
238
-#ifdef __HIPCC__
239
- if (is_half) {
240
- const uint16_t ihInf = 0x7C00;
241
- const uint16_t ihNegInf = 0xFC00;
242
- const uint16_t ihNaN = 0x7C01;
243
- const uint16_t ihNeg0 = 0x8000;
244
- fInf = reinterpret_cast<const _Float16&>(ihInf);
245
- fNegInf = reinterpret_cast<const _Float16&>(ihNegInf);
246
- fNaN = reinterpret_cast<const _Float16&>(ihNaN);
247
- fNeg0 = reinterpret_cast<const _Float16&>(ihNeg0);
248
- } else
249
-#endif
250
- if (is_float) {
251
- const uint32_t ifInf = 0x7F800000;
252
- const uint32_t ifNegInf = 0xFF800000;
253
- const uint32_t ifNaN = 0x7F800001;
254
- const uint32_t ifNeg0 = 0x80000000;
255
- fInf = reinterpret_cast<const float&>(ifInf);
256
- fNegInf = reinterpret_cast<const float&>(ifNegInf);
257
- fNaN = reinterpret_cast<const float&>(ifNaN);
258
- fNeg0 = reinterpret_cast<const float&>(ifNeg0);
259
- }
260
-
261
- if (x == 0) {
262
- return 0;
263
- }
264
-
265
- uint32_t sign = x >> 7;
266
- uint32_t mantissa = x & ((1 << wm) - 1);
267
- int exponent = (x & 0x7F) >> wm;
268
- if (negative_zero_nan) {
269
- if (x == 0x80) {
270
- return fNaN;
271
- }
272
- } else {
273
- if (x == 0x80) {
274
- return fNeg0;
275
- }
276
- if (exponent == ((1 << we) - 1)) {
277
- return (mantissa == 0) ? (sign ? fNegInf : fInf) : fNaN;
278
- }
279
- }
280
- typename std::conditional<sizeof(T) == 2, uint16_t, uint32_t>::type retval;
281
- if (we == 5 && is_half && !negative_zero_nan) {
282
- retval = x << 8;
283
- return reinterpret_cast<const T&>(retval);
284
- }
285
-
286
- const int exp_low_cutoff =
287
- (1 << (weo - 1)) - (1 << (we - 1)) + 1 - (negative_zero_nan ? 1 : 0);
288
-
289
- // subnormal input
290
- if (exponent == 0) {
291
- // guaranteed mantissa!=0 since cases 0x0 and 0x80 are handled above
292
- int sh = 1 + clz(mantissa) - (32 - wm);
293
- mantissa <<= sh;
294
- exponent += 1 - sh;
295
- mantissa &= ((1 << wm) - 1);
296
- }
297
- exponent += exp_low_cutoff - 1;
298
- mantissa <<= wmo - wm;
299
-
300
- // subnormal output (occurs when T=half, we=5, negative_zero_nan=true)
301
- if (exponent <= 0) {
302
- mantissa |= 1 << wmo;
303
- mantissa >>= 1 - exponent;
304
- exponent = 0;
305
- }
306
-
307
- if (sizeof(T) == 2) {
308
- retval = (sign << 15) | (exponent << 10) | mantissa;
309
- } else {
310
- retval = (sign << 31) | (exponent << 23) | mantissa;
311
- }
312
- return reinterpret_cast<const T&>(retval);
313
-}
314
-
315
-} // namespace hip_fp8_impl
csrc/quantization/fp8/amd/quant_utils.cuh
+230
-168
@@ -1,13 +1,11 @@
1
#pragma once
2
-#include "hip_float8.h"
2
+#include <hip/hip_fp8.h>
3
4
#include <hip/hip_fp16.h>
5
#include <hip/hip_bf16.h>
6
#include <hip/hip_bfloat16.h>
7
8
-#include "../../../attention/dtype_fp8.cuh"
9
-#include "../../../attention/dtype_float32.cuh"
10
-#include "../../../attention/dtype_bfloat16.cuh"
8
+#include "../../../attention/attention_dtypes.h"
9
10
namespace vllm {
11
#ifdef USE_ROCM
@@ -26,40 +24,31 @@ __inline__ __device__ Tout scaled_vec_conversion(const Tin& x,
24
return x;
25
}
26
27
+ #if HIP_FP8_TYPE_FNUZ
28
+using fp8_type = __hip_fp8_e4m3_fnuz;
29
+using fp8x2_type = __hip_fp8x2_e4m3_fnuz;
30
+ #elif HIP_FP8_TYPE_OCP
31
+using fp8_type = __hip_fp8_e4m3;
32
+using fp8x2_type = __hip_fp8x2_e4m3;
33
+ #endif
34
+
35
// fp8 -> half
36
template <>
37
__inline__ __device__ uint16_t
38
vec_conversion<uint16_t, uint8_t>(const uint8_t& a) {
33
- hip_fp8 f8{a, hip_fp8::from_bits()};
34
- __half_raw res;
35
- res.data = static_cast<float>(f8);
36
- return res.x;
39
+ return __hip_cvt_fp8_to_halfraw(a, fp8_type::__default_interpret).x;
40
}
41
42
// fp8x2 -> half2
43
template <>
44
__inline__ __device__ uint32_t
45
vec_conversion<uint32_t, uint16_t>(const uint16_t& a) {
43
- #if defined(__HIP__MI300__) && \
44
- defined(__HIP_FP8_EXPERIMENTAL_BULK_CONVERT__)
45
- const auto& f2 = __builtin_amdgcn_cvt_pk_f32_fp8(a, 0);
46
union {
47
__half2_raw h2r;
48
uint32_t ui32;
49
} tmp;
50
- tmp.h2r.x.data = f2[0];
51
- tmp.h2r.y.data = f2[1];
50
+ tmp.h2r = __hip_cvt_fp8x2_to_halfraw2(a, fp8_type::__default_interpret);
51
return tmp.ui32;
53
- #else
54
- union {
55
- uint16_t u16[2];
56
- uint32_t u32;
57
- } tmp;
58
-
59
- tmp.u16[0] = vec_conversion<uint16_t, uint8_t>(static_cast<uint8_t>(a));
60
- tmp.u16[1] = vec_conversion<uint16_t, uint8_t>(static_cast<uint8_t>(a >> 8U));
61
- return tmp.u32;
62
- #endif
52
}
53
54
// fp8x4 -> half2x2
@@ -92,9 +81,9 @@ using __nv_bfloat16 = __hip_bfloat16;
81
template <>
82
__inline__ __device__ __nv_bfloat16
83
vec_conversion<__nv_bfloat16, uint8_t>(const uint8_t& a) {
95
- hip_fp8 f8{a, hip_fp8::from_bits()};
96
- float f{f8};
97
- return __float2bfloat16(f);
84
+ fp8_type f8;
85
+ f8.__x = a;
86
+ return __float2bfloat16(static_cast<float>(f8));
87
}
88
89
using __nv_bfloat162 = __hip_bfloat162;
@@ -136,27 +125,18 @@ __inline__ __device__ bf16_8_t vec_conversion<bf16_8_t, uint2>(const uint2& a) {
125
// fp8 -> float
126
template <>
127
__inline__ __device__ float vec_conversion<float, uint8_t>(const uint8_t& a) {
139
- hip_fp8 fp8{a, hip_fp8::from_bits()};
140
- return static_cast<float>(fp8);
128
+ fp8_type f8;
129
+ f8.__x = a;
130
+ return static_cast<float>(f8);
131
}
132
133
// fp8x2 -> float2
134
template <>
135
__inline__ __device__ float2
136
vec_conversion<float2, uint16_t>(const uint16_t& a) {
147
- #if defined(__HIP__MI300__) && \
148
- defined(__HIP_FP8_EXPERIMENTAL_BULK_CONVERT__)
149
- float2 res;
150
- const auto& f2 = __builtin_amdgcn_cvt_pk_f32_fp8(a, 0);
151
- res.x = f2[0];
152
- res.y = f2[1];
153
- return res;
154
- #else
155
- float2 res;
156
- res.x = vec_conversion<float, uint8_t>(static_cast<uint8_t>(a));
157
- res.y = vec_conversion<float, uint8_t>(static_cast<uint8_t>(a >> 8U));
158
- return res;
159
- #endif
137
+ fp8x2_type f8x2;
138
+ f8x2.__x = a;
139
+ return static_cast<float2>(f8x2);
140
}
141
142
// fp8x4 -> float4
@@ -169,6 +149,15 @@ vec_conversion<Float4_, uint32_t>(const uint32_t& a) {
149
return res;
150
}
151
152
+// fp8x4 -> float4
153
+template <>
154
+__inline__ __device__ float4
155
+vec_conversion<float4, uint32_t>(const uint32_t& a) {
156
+ Float4_ tmp = vec_conversion<Float4_, uint32_t>(a);
157
+ float4 res = make_float4(tmp.x.x, tmp.x.y, tmp.y.x, tmp.y.y);
158
+ return res;
159
+}
160
+
161
// fp8x8 -> float8
162
template <>
163
__inline__ __device__ Float8_ vec_conversion<Float8_, uint2>(const uint2& a) {
@@ -189,33 +178,36 @@ __inline__ __device__ uint8_t
178
vec_conversion<uint8_t, uint16_t>(const uint16_t& a) {
179
__half_raw tmp;
180
tmp.x = a;
181
+ return __hip_cvt_halfraw_to_fp8(tmp, fp8_type::__default_saturation,
182
+ fp8_type::__default_interpret);
183
+}
184
193
- hip_fp8 f8{static_cast<float>(tmp.data)};
194
- return f8.data;
185
+template <>
186
+__inline__ __device__ uint16_t
187
+vec_conversion<uint16_t, uint32_t>(const uint32_t& a) {
188
+ union {
189
+ uint32_t ui32;
190
+ __half2_raw h2r;
191
+ } tmp;
192
+ tmp.ui32 = a;
193
+ return __hip_cvt_halfraw2_to_fp8x2(tmp.h2r, fp8_type::__default_saturation,
194
+ fp8_type::__default_interpret);
195
}
196
197
// bf16 -> fp8
198
template <>
199
__inline__ __device__ uint8_t
200
vec_conversion<uint8_t, __nv_bfloat16>(const __nv_bfloat16& a) {
201
- hip_fp8 res{__bfloat162float(a)};
202
- return res.data;
201
+ return __hip_cvt_float_to_fp8(__bfloat162float(a),
202
+ fp8_type::__default_saturation,
203
+ fp8_type::__default_interpret);
204
}
205
206
// float -> fp8
207
template <>
208
__inline__ __device__ uint8_t vec_conversion<uint8_t, float>(const float& a) {
208
- hip_fp8 f8(a);
209
- return f8.data;
210
-}
211
-
212
-// fp8x4 -> float4
213
-template <>
214
-__inline__ __device__ float4
215
-vec_conversion<float4, uint32_t>(const uint32_t& a) {
216
- Float4_ tmp = vec_conversion<Float4_, uint32_t>(a);
217
- float4 res = make_float4(tmp.x.x, tmp.x.y, tmp.y.x, tmp.y.y);
218
- return res;
209
+ return __hip_cvt_float_to_fp8(a, fp8_type::__default_saturation,
210
+ fp8_type::__default_interpret);
211
}
212
213
// float2 -> half2
@@ -307,90 +299,22 @@ vec_conversion<bf16_8_t, Float8_>(const Float8_& a) {
299
300
*/
301
310
-// fp8 -> half
311
-template <>
312
-__inline__ __device__ uint16_t
313
-scaled_vec_conversion<uint16_t, uint8_t>(const uint8_t& a, const float scale) {
314
- hip_fp8 f8{a, hip_fp8::from_bits()};
315
- __half_raw res;
316
- res.data = static_cast<float>(f8) * scale;
317
- return res.x;
318
-}
319
-
320
-// fp8x2 -> half2
321
-template <>
322
-__inline__ __device__ uint32_t scaled_vec_conversion<uint32_t, uint16_t>(
323
- const uint16_t& a, const float scale) {
324
- #if defined(__HIP__MI300__) && \
325
- defined(__HIP_FP8_EXPERIMENTAL_BULK_CONVERT__)
326
- const auto& f2 = __builtin_amdgcn_cvt_pk_f32_fp8(a, 0);
327
- union {
328
- __half2_raw h2r;
329
- uint32_t ui32;
330
- } tmp;
331
- tmp.h2r.x.data = f2[0] * scale;
332
- tmp.h2r.y.data = f2[1] * scale;
333
- return tmp.ui32;
334
- #else
335
- union {
336
- uint16_t u16[2];
337
- uint32_t u32;
338
- } tmp;
339
-
340
- tmp.u16[0] =
341
- scaled_vec_conversion<uint16_t, uint8_t>(static_cast<uint8_t>(a), scale);
342
- tmp.u16[1] = scaled_vec_conversion<uint16_t, uint8_t>(
343
- static_cast<uint8_t>(a >> 8U), scale);
344
- return tmp.u32;
345
- #endif
346
-}
347
-
348
-// fp8x4 -> half2x2
349
-template <>
350
-__inline__ __device__ uint2
351
-scaled_vec_conversion<uint2, uint32_t>(const uint32_t& a, const float scale) {
352
- union {
353
- uint2 u32x2;
354
- uint32_t u32[2];
355
- } tmp;
356
- tmp.u32[0] = scaled_vec_conversion<uint32_t, uint16_t>((uint16_t)a, scale);
357
- tmp.u32[1] =
358
- scaled_vec_conversion<uint32_t, uint16_t>((uint16_t)(a >> 16U), scale);
359
- return tmp.u32x2;
360
-}
361
-
362
-// fp8x8 -> half2x4
363
-template <>
364
-__inline__ __device__ uint4
365
-scaled_vec_conversion<uint4, uint2>(const uint2& a, const float scale) {
366
- union {
367
- uint4 u64x2;
368
- uint2 u64[2];
369
- } tmp;
370
- tmp.u64[0] = scaled_vec_conversion<uint2, uint32_t>(a.x, scale);
371
- tmp.u64[1] = scaled_vec_conversion<uint2, uint32_t>(a.y, scale);
372
- return tmp.u64x2;
373
-}
374
-
302
using __nv_bfloat16 = __hip_bfloat16;
303
304
// fp8 -> __nv_bfloat16
305
template <>
306
__inline__ __device__ __nv_bfloat16
380
-scaled_vec_conversion<__nv_bfloat16, uint8_t>(const uint8_t& a,
381
- const float scale) {
382
- hip_fp8 f8{a, hip_fp8::from_bits()};
383
- float f{f8};
384
- return __float2bfloat16(f * scale);
307
+scaled_vec_conversion<__nv_bfloat16, uint8_t>(const uint8_t& a, float scale) {
308
+ fp8_type f8;
309
+ f8.__x = a;
310
+ return __float2bfloat16(static_cast<float>(f8) * scale);
311
}
312
387
-using __nv_bfloat162 = __hip_bfloat162;
388
-
313
// fp8x2 -> __nv_bfloat162
314
template <>
315
__inline__ __device__ __nv_bfloat162
316
scaled_vec_conversion<__nv_bfloat162, uint16_t>(const uint16_t& a,
393
- const float scale) {
317
+ float scale) {
318
__nv_bfloat162 res;
319
res.x = scaled_vec_conversion<__nv_bfloat16, uint8_t>((uint8_t)a, scale);
320
res.y =
@@ -400,8 +324,8 @@ scaled_vec_conversion<__nv_bfloat162, uint16_t>(const uint16_t& a,
324
325
// fp8x4 -> bf16_4_t
326
template <>
403
-__inline__ __device__ bf16_4_t scaled_vec_conversion<bf16_4_t, uint32_t>(
404
- const uint32_t& a, const float scale) {
327
+__inline__ __device__ bf16_4_t
328
+scaled_vec_conversion<bf16_4_t, uint32_t>(const uint32_t& a, float scale) {
329
bf16_4_t res;
330
res.x = scaled_vec_conversion<__nv_bfloat162, uint16_t>((uint16_t)a, scale);
331
res.y = scaled_vec_conversion<__nv_bfloat162, uint16_t>((uint16_t)(a >> 16U),
@@ -412,7 +336,7 @@ __inline__ __device__ bf16_4_t scaled_vec_conversion<bf16_4_t, uint32_t>(
336
// fp8x8 -> bf16_8_t
337
template <>
338
__inline__ __device__ bf16_8_t
415
-scaled_vec_conversion<bf16_8_t, uint2>(const uint2& a, const float scale) {
339
+scaled_vec_conversion<bf16_8_t, uint2>(const uint2& a, float scale) {
340
bf16_4_t tmp1, tmp2;
341
tmp1 = scaled_vec_conversion<bf16_4_t, uint32_t>(a.x, scale);
342
tmp2 = scaled_vec_conversion<bf16_4_t, uint32_t>(a.y, scale);
@@ -427,29 +351,19 @@ scaled_vec_conversion<bf16_8_t, uint2>(const uint2& a, const float scale) {
351
// fp8 -> float
352
template <>
353
__inline__ __device__ float scaled_vec_conversion<float, uint8_t>(
430
- const uint8_t& a, const float scale) {
431
- hip_fp8 fp8{a, hip_fp8::from_bits()};
432
- return static_cast<float>(fp8) * scale;
354
+ const uint8_t& a, float scale) {
355
+ fp8_type f8;
356
+ f8.__x = a;
357
+ return static_cast<float>(f8) * scale;
358
}
359
360
// fp8x2 -> float2
361
template <>
362
__inline__ __device__ float2
438
-scaled_vec_conversion<float2, uint16_t>(const uint16_t& a, const float scale) {
439
- #if defined(__HIP__MI300__) && \
440
- defined(__HIP_FP8_EXPERIMENTAL_BULK_CONVERT__)
441
- float2 res;
442
- const auto& f2 = __builtin_amdgcn_cvt_pk_f32_fp8(a, 0);
443
- res.x = f2[0] * scale;
444
- res.y = f2[1] * scale;
445
- return res;
446
- #else
447
- float2 res;
448
- res.x = scaled_vec_conversion<float, uint8_t>(static_cast<uint8_t>(a), scale);
449
- res.y = scaled_vec_conversion<float, uint8_t>(static_cast<uint8_t>(a >> 8U),
450
- scale);
451
- return res;
452
- #endif
363
+scaled_vec_conversion<float2, uint16_t>(const uint16_t& a, float scale) {
364
+ fp8x2_type f8x2;
365
+ f8x2.__x = a;
366
+ return static_cast<float2>(f8x2) * scale;
367
}
368
369
// fp8x4 -> float4
@@ -462,10 +376,18 @@ scaled_vec_conversion<Float4_, uint32_t>(const uint32_t& a, const float scale) {
376
return res;
377
}
378
379
+// fp8x4 -> float4
380
+template <>
381
+__inline__ __device__ float4
382
+scaled_vec_conversion<float4, uint32_t>(const uint32_t& a, float scale) {
383
+ Float4_ res = scaled_vec_conversion<Float4_, uint32_t>(a, scale);
384
+ return {res.x.x, res.x.y, res.y.x, res.y.y};
385
+}
386
+
387
// fp8x8 -> float8
388
template <>
389
__inline__ __device__ Float8_
468
-scaled_vec_conversion<Float8_, uint2>(const uint2& a, const float scale) {
390
+scaled_vec_conversion<Float8_, uint2>(const uint2& a, float scale) {
391
Float4_ tmp1, tmp2;
392
tmp1 = scaled_vec_conversion<Float4_, uint32_t>(a.x, scale);
393
tmp2 = scaled_vec_conversion<Float4_, uint32_t>(a.y, scale);
@@ -477,44 +399,184 @@ scaled_vec_conversion<Float8_, uint2>(const uint2& a, const float scale) {
399
return res;
400
}
401
480
-/* Quantize(HP / scale) => FP8 */
402
+// fp8 -> half
403
+template <>
404
+__inline__ __device__ uint16_t
405
+scaled_vec_conversion<uint16_t, uint8_t>(const uint8_t& a, float scale) {
406
+ __half_raw res;
407
+ res.data = scaled_vec_conversion<float, uint8_t>(a, scale);
408
+ return res.x;
409
+}
410
+
411
+// fp8x2 -> half2
412
+template <>
413
+__inline__ __device__ uint32_t
414
+scaled_vec_conversion<uint32_t, uint16_t>(const uint16_t& a, float scale) {
415
+ __half2_raw h2r =
416
+ __hip_cvt_fp8x2_to_halfraw2(a, fp8_type::__default_interpret);
417
+ union {
418
+ __half2_raw h2r;
419
+ uint32_t ui32;
420
+ } tmp;
421
+ tmp.h2r = __hip_cvt_fp8x2_to_halfraw2(a, fp8_type::__default_interpret);
422
+ tmp.h2r.x.data *= scale;
423
+ tmp.h2r.y.data *= scale;
424
+ return tmp.ui32;
425
+}
426
+
427
+// fp8x4 -> half2x2
428
+template <>
429
+__inline__ __device__ uint2
430
+scaled_vec_conversion<uint2, uint32_t>(const uint32_t& a, float scale) {
431
+ union {
432
+ uint2 u32x2;
433
+ uint32_t u32[2];
434
+ } tmp;
435
+ tmp.u32[0] = scaled_vec_conversion<uint32_t, uint16_t>((uint16_t)a, scale);
436
+ tmp.u32[1] =
437
+ scaled_vec_conversion<uint32_t, uint16_t>((uint16_t)(a >> 16U), scale);
438
+ return tmp.u32x2;
439
+}
440
482
-// TODO(Hai): vectorized to add
441
+// fp8x8 -> half2x4
442
+template <>
443
+__inline__ __device__ uint4 scaled_vec_conversion<uint4, uint2>(const uint2& a,
444
+ float scale) {
445
+ union {
446
+ uint4 u64x2;
447
+ uint2 u64[2];
448
+ } tmp;
449
+ tmp.u64[0] = scaled_vec_conversion<uint2, uint32_t>(a.x, scale);
450
+ tmp.u64[1] = scaled_vec_conversion<uint2, uint32_t>(a.y, scale);
451
+ return tmp.u64x2;
452
+}
453
454
// half -> fp8
455
template <>
456
__inline__ __device__ uint8_t
487
-scaled_vec_conversion<uint8_t, uint16_t>(const uint16_t& a, const float scale) {
457
+scaled_vec_conversion<uint8_t, uint16_t>(const uint16_t& a, float scale) {
458
__half_raw tmp;
459
tmp.x = a;
460
+ tmp.data /= scale;
461
+ return __hip_cvt_halfraw_to_fp8(tmp, fp8_type::__default_saturation,
462
+ fp8_type::__default_interpret);
463
+}
464
491
- hip_fp8 f8{static_cast<float>(tmp.data) / scale};
492
- return f8.data;
465
+// halfx2 -> fp8x2
466
+template <>
467
+__inline__ __device__ uint16_t
468
+scaled_vec_conversion<uint16_t, uint32_t>(const uint32_t& a, float scale) {
469
+ union {
470
+ uint32_t ui32;
471
+ __half2_raw h2r;
472
+ } tmp;
473
+ tmp.ui32 = a;
474
+ tmp.h2r.x.data /= scale;
475
+ tmp.h2r.y.data /= scale;
476
+ return __hip_cvt_halfraw2_to_fp8x2(tmp.h2r, fp8_type::__default_saturation,
477
+ fp8_type::__default_interpret);
478
+}
479
+
480
+// half2x2 -> fp8x4
481
+template <>
482
+__inline__ __device__ uint32_t
483
+scaled_vec_conversion<uint32_t, uint2>(const uint2& a, float scale) {
484
+ union {
485
+ uint16_t ui16[2];
486
+ uint32_t ui32;
487
+ } tmp;
488
+ tmp.ui16[0] = scaled_vec_conversion<uint16_t, uint32_t>(a.x, scale);
489
+ tmp.ui16[1] = scaled_vec_conversion<uint16_t, uint32_t>(a.y, scale);
490
+ return tmp.ui32;
491
+}
492
+
493
+// half2x4 -> fp8x8
494
+template <>
495
+__inline__ __device__ uint2 scaled_vec_conversion<uint2, uint4>(const uint4& a,
496
+ float scale) {
497
+ union {
498
+ uint2 ui2[2];
499
+ uint4 ui4;
500
+ } tmp;
501
+ tmp.ui4 = a;
502
+ uint2 res;
503
+ res.x = scaled_vec_conversion<uint32_t, uint2>(tmp.ui2[0], scale);
504
+ res.y = scaled_vec_conversion<uint32_t, uint2>(tmp.ui2[1], scale);
505
+ return res;
506
}
507
508
// bf16 -> fp8
509
template <>
510
__inline__ __device__ uint8_t scaled_vec_conversion<uint8_t, __nv_bfloat16>(
498
- const __nv_bfloat16& a, const float scale) {
499
- hip_fp8 res{__bfloat162float(a) / scale};
500
- return res.data;
511
+ const __nv_bfloat16& a, float scale) {
512
+ return __hip_cvt_float_to_fp8(__bfloat162float(a) / scale,
513
+ fp8_type::__default_saturation,
514
+ fp8_type::__default_interpret);
515
+}
516
+
517
+// bf16x2 -> fp8x2
518
+template <>
519
+__inline__ __device__ uint16_t scaled_vec_conversion<uint16_t, __nv_bfloat162>(
520
+ const __nv_bfloat162& a, float scale) {
521
+ union {
522
+ uint8_t ui8[2];
523
+ uint16_t ui16;
524
+ } tmp;
525
+ tmp.ui8[0] = scaled_vec_conversion<uint8_t, __nv_bfloat16>(a.x, scale);
526
+ tmp.ui8[1] = scaled_vec_conversion<uint8_t, __nv_bfloat16>(a.y, scale);
527
+ return tmp.ui16;
528
+}
529
+
530
+// bf16x4 -> fp8x4
531
+template <>
532
+__inline__ __device__ uint32_t
533
+scaled_vec_conversion<uint32_t, bf16_4_t>(const bf16_4_t& a, float scale) {
534
+ union {
535
+ uint16_t ui16[2];
536
+ uint32_t ui32;
537
+ } tmp;
538
+ tmp.ui16[0] = scaled_vec_conversion<uint16_t, __nv_bfloat162>(a.x, scale);
539
+ tmp.ui16[1] = scaled_vec_conversion<uint16_t, __nv_bfloat162>(a.y, scale);
540
+ return tmp.ui32;
541
+}
542
+
543
+// bf16x8 -> fp8x8
544
+template <>
545
+__inline__ __device__ uint2
546
+scaled_vec_conversion<uint2, bf16_8_t>(const bf16_8_t& a, float scale) {
547
+ uint2 res;
548
+ res.x = scaled_vec_conversion<uint32_t, bf16_4_t>({a.x, a.y}, scale);
549
+ res.y = scaled_vec_conversion<uint32_t, bf16_4_t>({a.z, a.w}, scale);
550
+ return res;
551
}
552
553
// float -> fp8
554
template <>
555
__inline__ __device__ uint8_t
506
-scaled_vec_conversion<uint8_t, float>(const float& a, const float scale) {
507
- hip_fp8 f8(a / scale);
508
- return f8.data;
556
+scaled_vec_conversion<uint8_t, float>(const float& a, float scale) {
557
+ return __hip_cvt_float_to_fp8(a / scale, fp8_type::__default_saturation,
558
+ fp8_type::__default_interpret);
559
}
560
511
-// fp8x4 -> float4
561
+// floatx2 -> fp8x2
562
template <>
513
-__inline__ __device__ float4
514
-scaled_vec_conversion<float4, uint32_t>(const uint32_t& a, const float scale) {
515
- Float4_ tmp = scaled_vec_conversion<Float4_, uint32_t>(a, scale);
516
- float4 res = make_float4(tmp.x.x, tmp.x.y, tmp.y.x, tmp.y.y);
517
- return res;
563
+__inline__ __device__ uint16_t
564
+scaled_vec_conversion<uint16_t, float2>(const float2& a, float scale) {
565
+ return __hip_cvt_float2_to_fp8x2(a / scale, fp8_type::__default_saturation,
566
+ fp8_type::__default_interpret);
567
+}
568
+
569
+// floatx4 -> fp8x4
570
+template <>
571
+__inline__ __device__ uint32_t
572
+scaled_vec_conversion<uint32_t, float4>(const float4& a, float scale) {
573
+ union {
574
+ uint16_t ui16[2];
575
+ uint32_t ui32;
576
+ } tmp;
577
+ tmp.ui16[0] = scaled_vec_conversion<uint16_t, float2>({a.x, a.y}, scale);
578
+ tmp.ui16[1] = scaled_vec_conversion<uint16_t, float2>({a.z, a.w}, scale);
579
+ return tmp.ui32;
580
}
581
#endif // ENABLE_FP8
582
csrc/quantization/fp8/common.cuh
+5
-3
@@ -12,7 +12,7 @@ C10_HOST_DEVICE constexpr auto FP8_E4M3_MAX =
12
std::numeric_limits<FP8_TYPE>::max();
13
#else
14
#include <c10/util/Float8_e4m3fnuz.h>
15
- #include "amd/hip_float8.h"
15
+ #include "amd/quant_utils.cuh"
16
using FP8_TYPE = c10::Float8_e4m3fnuz;
17
// Using the default max value from pytorch (240.0) will cause accuracy
18
// issue when running dynamic quantization. Here use 224.0f for rocm.
@@ -47,8 +47,10 @@ __device__ __forceinline__ FP8_TYPE scaled_fp8_conversion(float const val,
47
return static_cast<c10::Float8_e4m3fn>(r);
48
#else
49
// Use hardware cvt instruction for fp8 on rocm
50
- return c10::Float8_e4m3fnuz(hip_fp8(r).data,
51
- c10::Float8_e4m3fnuz::from_bits());
50
+ return c10::Float8_e4m3fnuz(
51
+ __hip_cvt_float_to_fp8(r, fp8::fp8_type::__default_saturation,
52
+ fp8::fp8_type::__default_interpret),
53
+ c10::Float8_e4m3fnuz::from_bits());
54
#endif
55
}
56
tests/kernels/test_cache.py
+13
-11
@@ -159,19 +159,20 @@ def test_reshape_and_cache(
159
device)
160
key_cache, value_cache = key_caches[0], value_caches[0]
161
162
+ # Using default kv_scale
163
+ k_scale = (key.amax() / 64.0).to(torch.float32)
164
+ v_scale = (value.amax() / 64.0).to(torch.float32)
165
+
166
# Clone the KV caches.
167
if kv_cache_dtype == "fp8":
168
cloned_key_cache = torch.empty_like(key_cache, dtype=torch.float16)
165
- ops.convert_fp8(cloned_key_cache, key_cache)
169
+ ops.convert_fp8(cloned_key_cache, key_cache, k_scale.item())
170
cloned_value_cache = torch.empty_like(value_cache, dtype=torch.float16)
167
- ops.convert_fp8(cloned_value_cache, value_cache)
171
+ ops.convert_fp8(cloned_value_cache, value_cache, v_scale.item())
172
else:
173
cloned_key_cache = key_cache.clone()
174
cloned_value_cache = value_cache.clone()
175
172
- # Using default kv_scale
173
- k_scale = v_scale = torch.tensor(1.0, dtype=torch.float32, device=device)
174
-
176
# Call the reshape_and_cache kernel.
177
opcheck(torch.ops._C_cache_ops.reshape_and_cache,
178
(key, value, key_cache, value_cache, slot_mapping, kv_cache_dtype,
@@ -182,9 +183,9 @@ def test_reshape_and_cache(
183
184
if kv_cache_dtype == "fp8":
185
result_key_cache = torch.empty_like(key_cache, dtype=torch.float16)
185
- ops.convert_fp8(result_key_cache, key_cache)
186
+ ops.convert_fp8(result_key_cache, key_cache, k_scale.item())
187
result_value_cache = torch.empty_like(value_cache, dtype=torch.float16)
187
- ops.convert_fp8(result_value_cache, value_cache)
188
+ ops.convert_fp8(result_value_cache, value_cache, v_scale.item())
189
190
# Run the reference implementation.
191
reshaped_key = key.reshape(num_tokens, *key_cache[0, :, :, 0, :].shape)
@@ -268,15 +269,16 @@ def test_reshape_and_cache_flash(
269
del key_caches
270
del value_caches
271
271
- k_scale = (key.amax() / 256.0).to(torch.float32)
272
- v_scale = (value.amax() / 256.0).to(torch.float32)
272
+ k_scale = (key.amax() / 64.0).to(torch.float32)
273
+ v_scale = (value.amax() / 64.0).to(torch.float32)
274
275
# Clone the KV caches.
276
if kv_cache_dtype == "fp8":
277
cloned_key_cache = torch.empty_like(key_cache, dtype=torch.float16)
277
- ops.convert_fp8(cloned_key_cache, key_cache, k_scale, kv_cache_dtype)
278
+ ops.convert_fp8(cloned_key_cache, key_cache, k_scale.item(),
279
+ kv_cache_dtype)
280
cloned_value_cache = torch.empty_like(value_cache, dtype=torch.float16)
279
- ops.convert_fp8(cloned_value_cache, value_cache, v_scale,
281
+ ops.convert_fp8(cloned_value_cache, value_cache, v_scale.item(),
282
kv_cache_dtype)
283
else:
284
cloned_key_cache = key_cache.clone()