quant_utils.cuh 26.2 KB
Newer Older
1
#pragma once
2
#ifndef USE_ROCM
3
#include <hip/hip_fp8.h>
4
#endif
5
6
7
8
9

#include <hip/hip_fp16.h>
#include <hip/hip_bf16.h>
#include <hip/hip_bfloat16.h>

10
#include "../../../../attention/attention_dtypes.h"
11

12
namespace vllm {
13
14
15
#ifdef USE_ROCM

namespace fp8 {
zhuwenwen's avatar
zhuwenwen committed
16
  // #ifdef ENABLE_FP8
17

zhuwenwen's avatar
zhuwenwen committed
18
19
20
21
22
23
24
25
26
// KV-CACHE int8
static inline __device__ float fp8_to_float(uint8_t input) {
  const uint32_t w = (uint32_t)input << 24;
  const uint32_t sign = w & UINT32_C(0x80000000);
  const uint32_t nonsign = w & UINT32_C(0x7FFFFFFF);
  uint32_t renorm_shift = __clz(nonsign);
  renorm_shift = renorm_shift > 4 ? renorm_shift - 4 : 0;
  uint32_t result = sign | ((nonsign << renorm_shift >> 4) + ((0x78 - renorm_shift) << 23));
  return c10::detail::fp32_from_bits(result);
27
28
}

zhuwenwen's avatar
zhuwenwen committed
29
// float -> fp8
30
static inline __device__ uint8_t float_to_fp8_e4m3(float f) {
zhuwenwen's avatar
zhuwenwen committed
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
  constexpr uint32_t fp8_max = UINT32_C(1087) << 20;
  constexpr uint32_t denorm_mask = UINT32_C(141) << 23;
  uint32_t f_bits = c10::detail::fp32_to_bits(f);
  uint8_t result = 0u;
  const uint32_t sign = f_bits & UINT32_C(0x80000000);
  f_bits ^= sign;
  if (f_bits >= fp8_max) {
    result = 0x7f;
  } else {
    if (f_bits < (UINT32_C(121) << 23)) {
      f_bits =
        c10::detail::fp32_to_bits(c10::detail::fp32_from_bits(f_bits) + c10::detail::fp32_from_bits(denorm_mask));
      result = static_cast<uint8_t>(f_bits - denorm_mask);
    } else {
      uint8_t mant_odd = (f_bits >> 20) & 1;
      f_bits += ((uint32_t)(7 - 127) << 23) + 0x7FFFF;
      f_bits += mant_odd;
      result = static_cast<uint8_t>(f_bits >> 20);
    }
  }
51

zhuwenwen's avatar
zhuwenwen committed
52
53
  result |= static_cast<uint8_t>(sign >> 24);
  return result;
54
55
}

56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
static inline __device__ uint8_t float_to_fp8_e5m2(float f) {
  constexpr uint32_t fp32_inf = UINT32_C(255) << 23;
  constexpr uint32_t fp8_max = UINT32_C(143) << 23;
  constexpr uint32_t denorm_mask = UINT32_C(134) << 23;
  uint32_t f_bits = c10::detail::fp32_to_bits(f);
  uint8_t result = 0u;
  const uint32_t sign = f_bits & UINT32_C(0x80000000);
  f_bits ^= sign;
  if (f_bits >= fp8_max) {
    result = f_bits > fp32_inf ? UINT8_C(0x7F) : UINT8_C(0x7C);
  } else {
    if (f_bits < (UINT32_C(113) << 23)) {
      f_bits = c10::detail::fp32_to_bits(c10::detail::fp32_from_bits(f_bits)
               + c10::detail::fp32_from_bits(denorm_mask));
      result = static_cast<uint8_t>(f_bits - denorm_mask);
    } else {
      uint32_t mant_odd = (f_bits >> 21) & 1;
      f_bits += ((uint32_t)(15 - 127) << 23) + 0xFFFFF;
      f_bits += mant_odd;
      result = static_cast<uint8_t>(f_bits >> 21);
    }
  }
  result |= static_cast<uint8_t>(sign >> 24);
  return result;
}

zhuwenwen's avatar
zhuwenwen committed
82
83
84
85
// template <typename Tout, typename Tin>
// __inline__ __device__ Tout vec_conversion(const Tin& x) {
//   return x;
// }
86
87

template <typename Tout, typename Tin>
88
__inline__ __device__ Tout scaled_vec_conversion(const Tin& x,
89
                                                 const float scale, Fp8KVCacheDataType kv_type) {
90
  return x;
91
92
}

zhuwenwen's avatar
zhuwenwen committed
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
//     #if HIP_FP8_TYPE_OCP
// using fp8_type = __hip_fp8_e4m3;
// using fp8x2_type = __hip_fp8x2_e4m3;
//     #else
// using fp8_type = __hip_fp8_e4m3_fnuz;
// using fp8x2_type = __hip_fp8x2_e4m3_fnuz;
//     #endif

// // fp8 -> half
// template <>
// __inline__ __device__ uint16_t
// vec_conversion<uint16_t, uint8_t>(const uint8_t& a) {
//   return __hip_cvt_fp8_to_halfraw(a, fp8_type::__default_interpret).x;
// }

// // fp8x2 -> half2
// template <>
// __inline__ __device__ uint32_t
// vec_conversion<uint32_t, uint16_t>(const uint16_t& a) {
//   union {
//     __half2_raw h2r;
//     uint32_t ui32;
//   } tmp;
//   tmp.h2r = __hip_cvt_fp8x2_to_halfraw2(a, fp8_type::__default_interpret);
//   return tmp.ui32;
// }

// // fp8x4 -> half2x2
// template <>
// __inline__ __device__ uint2 vec_conversion<uint2, uint32_t>(const uint32_t& a) {
//   union {
//     uint2 u32x2;
//     uint32_t u32[2];
//   } tmp;
//   tmp.u32[0] = vec_conversion<uint32_t, uint16_t>((uint16_t)a);
//   tmp.u32[1] = vec_conversion<uint32_t, uint16_t>((uint16_t)(a >> 16U));
//   return tmp.u32x2;
// }

// // fp8x8 -> half2x4
// template <>
// __inline__ __device__ uint4 vec_conversion<uint4, uint2>(const uint2& a) {
//   union {
//     uint4 u64x2;
//     uint2 u64[2];
//   } tmp;
//   tmp.u64[0] = vec_conversion<uint2, uint32_t>(a.x);
//   tmp.u64[1] = vec_conversion<uint2, uint32_t>(a.y);
//   return tmp.u64x2;
// }

// using __nv_bfloat16 = __hip_bfloat16;

// // fp8 -> __nv_bfloat16
// template <>
// __inline__ __device__ __nv_bfloat16
// vec_conversion<__nv_bfloat16, uint8_t>(const uint8_t& a) {
//   fp8_type f8;
//   f8.__x = a;
//   return __float2bfloat16(static_cast<float>(f8));
// }

// using __nv_bfloat162 = __hip_bfloat162;

// // fp8x2 -> __nv_bfloat162
// template <>
// __inline__ __device__ __nv_bfloat162
// vec_conversion<__nv_bfloat162, uint16_t>(const uint16_t& a) {
//   __nv_bfloat162 res;
//   res.x = vec_conversion<__nv_bfloat16, uint8_t>((uint8_t)a);
//   res.y = vec_conversion<__nv_bfloat16, uint8_t>((uint8_t)(a >> 8U));
//   return res;
// }

// // fp8x4 -> bf16_4_t
// template <>
// __inline__ __device__ bf16_4_t
// vec_conversion<bf16_4_t, uint32_t>(const uint32_t& a) {
//   bf16_4_t res;
//   res.x = vec_conversion<__nv_bfloat162, uint16_t>((uint16_t)a);
//   res.y = vec_conversion<__nv_bfloat162, uint16_t>((uint16_t)(a >> 16U));
//   return res;
// }

// // fp8x8 -> bf16_8_t
// template <>
// __inline__ __device__ bf16_8_t vec_conversion<bf16_8_t, uint2>(const uint2& a) {
//   bf16_4_t tmp1, tmp2;
//   tmp1 = vec_conversion<bf16_4_t, uint32_t>(a.x);
//   tmp2 = vec_conversion<bf16_4_t, uint32_t>(a.y);
//   bf16_8_t res;
//   res.x = tmp1.x;
//   res.y = tmp1.y;
//   res.z = tmp2.x;
//   res.w = tmp2.y;
//   return res;
// }

// // fp8 -> float
// template <>
// __inline__ __device__ float vec_conversion<float, uint8_t>(const uint8_t& a) {
//   fp8_type f8;
//   f8.__x = a;
//   return static_cast<float>(f8);
// }

// // fp8x2 -> float2
// template <>
// __inline__ __device__ float2
// vec_conversion<float2, uint16_t>(const uint16_t& a) {
//   fp8x2_type f8x2;
//   f8x2.__x = a;
//   return static_cast<float2>(f8x2);
// }

// // fp8x4 -> float4
// template <>
// __inline__ __device__ Float4_
// vec_conversion<Float4_, uint32_t>(const uint32_t& a) {
//   Float4_ res;
//   res.x = vec_conversion<float2, uint16_t>((uint16_t)a);
//   res.y = vec_conversion<float2, uint16_t>((uint16_t)(a >> 16U));
//   return res;
// }

// // fp8x4 -> float4
// template <>
// __inline__ __device__ float4
// vec_conversion<float4, uint32_t>(const uint32_t& a) {
//   Float4_ tmp = vec_conversion<Float4_, uint32_t>(a);
//   float4 res = make_float4(tmp.x.x, tmp.x.y, tmp.y.x, tmp.y.y);
//   return res;
// }

// // fp8x8 -> float8
// template <>
// __inline__ __device__ Float8_ vec_conversion<Float8_, uint2>(const uint2& a) {
//   Float4_ tmp1, tmp2;
//   tmp1 = vec_conversion<Float4_, uint32_t>(a.x);
//   tmp2 = vec_conversion<Float4_, uint32_t>(a.y);
//   Float8_ res;
//   res.x = tmp1.x;
//   res.y = tmp1.y;
//   res.z = tmp2.x;
//   res.w = tmp2.y;
//   return res;
// }

// // half -> fp8
// template <>
// __inline__ __device__ uint8_t
// vec_conversion<uint8_t, uint16_t>(const uint16_t& a) {
//   __half_raw tmp;
//   tmp.x = a;
//   return __hip_cvt_halfraw_to_fp8(tmp, fp8_type::__default_saturation,
//                                   fp8_type::__default_interpret);
// }

// template <>
// __inline__ __device__ uint16_t
// vec_conversion<uint16_t, uint32_t>(const uint32_t& a) {
//   union {
//     uint32_t ui32;
//     __half2_raw h2r;
//   } tmp;
//   tmp.ui32 = a;
//   return __hip_cvt_halfraw2_to_fp8x2(tmp.h2r, fp8_type::__default_saturation,
//                                      fp8_type::__default_interpret);
// }

// // bf16 -> fp8
// template <>
// __inline__ __device__ uint8_t
// vec_conversion<uint8_t, __nv_bfloat16>(const __nv_bfloat16& a) {
//   return __hip_cvt_float_to_fp8(__bfloat162float(a),
//                                 fp8_type::__default_saturation,
//                                 fp8_type::__default_interpret);
// }

// // float -> fp8
// template <>
// __inline__ __device__ uint8_t vec_conversion<uint8_t, float>(const float& a) {
//   return __hip_cvt_float_to_fp8(a, fp8_type::__default_saturation,
//                                 fp8_type::__default_interpret);
// }

// // float2 -> half2
// template <>
// __inline__ __device__ uint32_t
// vec_conversion<uint32_t, float2>(const float2& a) {
//   union {
//     half2 float16;
//     uint32_t uint32;
//   };

//   float16 = __float22half2_rn(a);
//   return uint32;
// }

// // Float4 -> half2x2
// template <>
// __inline__ __device__ uint2 vec_conversion<uint2, Float4_>(const Float4_& a) {
//   uint2 b;
//   float2 val;
//   val.x = a.x.x;
//   val.y = a.x.y;
//   b.x = vec_conversion<uint32_t, float2>(val);

//   val.x = a.y.x;
//   val.y = a.y.y;
//   b.y = vec_conversion<uint32_t, float2>(val);
//   return b;
// }

// // Float4 -> float4
// template <>
// __inline__ __device__ float4 vec_conversion<float4, Float4_>(const Float4_& a) {
//   float4 b;
//   b.x = a.x.x;
//   b.y = a.x.y;
//   b.z = a.y.x;
//   b.w = a.y.y;
//   return b;
// }

// // Float8 -> half2x4
// template <>
// __inline__ __device__ uint4 vec_conversion<uint4, Float8_>(const Float8_& a) {
//   uint4 b;
//   b.x = vec_conversion<uint32_t, float2>(a.x);
//   b.y = vec_conversion<uint32_t, float2>(a.y);
//   b.z = vec_conversion<uint32_t, float2>(a.z);
//   b.w = vec_conversion<uint32_t, float2>(a.w);
//   return b;
// }

// // float2 -> bfloat162
// template <>
// __inline__ __device__ __nv_bfloat162
// vec_conversion<__nv_bfloat162, float2>(const float2& a) {
//   __nv_bfloat162 b = __float22bfloat162_rn(a);
//   return b;
// }

// // Float4 -> bfloat162x2
// template <>
// __inline__ __device__ bf16_4_t
// vec_conversion<bf16_4_t, Float4_>(const Float4_& a) {
//   bf16_4_t b;
//   b.x = __float22bfloat162_rn(a.x);
//   b.y = __float22bfloat162_rn(a.y);
//   return b;
// }

// // Float8 -> bfloat162x4
// template <>
// __inline__ __device__ bf16_8_t
// vec_conversion<bf16_8_t, Float8_>(const Float8_& a) {
//   bf16_8_t b;
//   b.x = __float22bfloat162_rn(a.x);
//   b.y = __float22bfloat162_rn(a.y);
//   b.z = __float22bfloat162_rn(a.z);
//   b.w = __float22bfloat162_rn(a.w);
//   return b;
// }
358

359
360
/* Scaled and vectorized conversions, for data exchange between high and low
   precision domains
361

362
363
364
   Convention of the scale in API, e.g: FP8_data = Quantization(
   High_Precision_data / scale ) s.t. Quantize(HP / scale) => FP8 Dequant(FP8) *
   scale =>  HP
365
366
367
368
369
370
371

 */

using __nv_bfloat16 = __hip_bfloat16;

// fp8 -> __nv_bfloat16
template <>
372
__inline__ __device__ __nv_bfloat16
373
374
375
376
scaled_vec_conversion<__nv_bfloat16, uint8_t>(const uint8_t& a, float scale, Fp8KVCacheDataType kv_type) {
  if (kv_type == vllm::Fp8KVCacheDataType::kFp8E5M2) {
    assert(false);
  }
zhuwenwen's avatar
zhuwenwen committed
377
378
379
380
381

  return __float2bfloat16(fp8_to_float(a) * scale);
  // fp8_type f8;
  // f8.__x = a;
  // return __float2bfloat16(static_cast<float>(f8) * scale);
382
383
384
385
}

// fp8x2 -> __nv_bfloat162
template <>
386
387
__inline__ __device__ __nv_bfloat162
scaled_vec_conversion<__nv_bfloat162, uint16_t>(const uint16_t& a,
388
                                                float scale, Fp8KVCacheDataType kv_type) {
389
  __nv_bfloat162 res;
390
  res.x = scaled_vec_conversion<__nv_bfloat16, uint8_t>((uint8_t)a, scale, kv_type);
391
  res.y =
392
      scaled_vec_conversion<__nv_bfloat16, uint8_t>((uint8_t)(a >> 8U), scale, kv_type);
393
  return res;
394
395
396
397
}

// fp8x4 -> bf16_4_t
template <>
398
__inline__ __device__ bf16_4_t
399
scaled_vec_conversion<bf16_4_t, uint32_t>(const uint32_t& a, float scale, Fp8KVCacheDataType kv_type) {
400
  bf16_4_t res;
401
  res.x = scaled_vec_conversion<__nv_bfloat162, uint16_t>((uint16_t)a, scale, kv_type);
402
  res.y = scaled_vec_conversion<__nv_bfloat162, uint16_t>((uint16_t)(a >> 16U),
403
                                                          scale, kv_type);
404
  return res;
405
406
407
408
}

// fp8x8 -> bf16_8_t
template <>
409
__inline__ __device__ bf16_8_t
410
scaled_vec_conversion<bf16_8_t, uint2>(const uint2& a, float scale, Fp8KVCacheDataType kv_type) {
411
  bf16_4_t tmp1, tmp2;
412
413
  tmp1 = scaled_vec_conversion<bf16_4_t, uint32_t>(a.x, scale, kv_type);
  tmp2 = scaled_vec_conversion<bf16_4_t, uint32_t>(a.y, scale, kv_type);
414
415
416
417
418
419
  bf16_8_t res;
  res.x = tmp1.x;
  res.y = tmp1.y;
  res.z = tmp2.x;
  res.w = tmp2.y;
  return res;
420
421
422
423
}

// fp8 -> float
template <>
424
__inline__ __device__ float scaled_vec_conversion<float, uint8_t>(
425
426
427
428
    const uint8_t& a, float scale, Fp8KVCacheDataType kv_type) {
    if (kv_type == vllm::Fp8KVCacheDataType::kFp8E5M2) {
      assert(false);
    }
zhuwenwen's avatar
zhuwenwen committed
429
430
431
432
    return fp8_to_float(a) * scale;
  // fp8_type f8;
  // f8.__x = a;
  // return static_cast<float>(f8) * scale;
433
434
435
436
}

// fp8x2 -> float2
template <>
437
__inline__ __device__ float2
438
scaled_vec_conversion<float2, uint16_t>(const uint16_t& a, float scale, Fp8KVCacheDataType kv_type) {
zhuwenwen's avatar
zhuwenwen committed
439
    float2 f2r;
440
441
    f2r.x = scaled_vec_conversion<float, uint8_t>((uint8_t)a, scale, kv_type);
    f2r.y = scaled_vec_conversion<float, uint8_t>((uint8_t)(a >> 8U), scale, kv_type);
zhuwenwen's avatar
zhuwenwen committed
442
443
444
445
446
    return f2r;
    // [[maybe_unused]] 
  // fp8x2_type f8x2;
  // f8x2.__x = a;
  // return static_cast<float2>(f8x2) * scale;
447
448
449
450
}

// fp8x4 -> float4
template <>
451
__inline__ __device__ Float4_
452
scaled_vec_conversion<Float4_, uint32_t>(const uint32_t& a, const float scale, Fp8KVCacheDataType kv_type) {
453
  Float4_ res;
454
455
  res.x = scaled_vec_conversion<float2, uint16_t>((uint16_t)a, scale, kv_type);
  res.y = scaled_vec_conversion<float2, uint16_t>((uint16_t)(a >> 16U), scale, kv_type);
456
  return res;
457
458
}

459
460
461
// fp8x4 -> float4
template <>
__inline__ __device__ float4
462
463
scaled_vec_conversion<float4, uint32_t>(const uint32_t& a, float scale, Fp8KVCacheDataType kv_type) {
  Float4_ res = scaled_vec_conversion<Float4_, uint32_t>(a, scale, kv_type);
464
465
466
  return {res.x.x, res.x.y, res.y.x, res.y.y};
}

467
468
// fp8x8 -> float8
template <>
469
__inline__ __device__ Float8_
470
scaled_vec_conversion<Float8_, uint2>(const uint2& a, float scale, Fp8KVCacheDataType kv_type) {
471
  Float4_ tmp1, tmp2;
472
473
  tmp1 = scaled_vec_conversion<Float4_, uint32_t>(a.x, scale, kv_type);
  tmp2 = scaled_vec_conversion<Float4_, uint32_t>(a.y, scale, kv_type);
474
475
476
477
478
479
  Float8_ res;
  res.x = tmp1.x;
  res.y = tmp1.y;
  res.z = tmp2.x;
  res.w = tmp2.y;
  return res;
480
481
}

482
483
484
// fp8 -> half
template <>
__inline__ __device__ uint16_t
485
486
487
488
scaled_vec_conversion<uint16_t, uint8_t>(const uint8_t& a, float scale, Fp8KVCacheDataType kv_type) {
  if (kv_type == vllm::Fp8KVCacheDataType::kFp8E5M2) {
    assert(false);
  }
zhuwenwen's avatar
zhuwenwen committed
489
490
491
492
493
  float res = fp8_to_float(a) * scale;
  return float_to_half(res);
  // __half_raw res;
  // res.data = scaled_vec_conversion<float, uint8_t>(a, scale);
  // return res.x;
494
495
496
497
498
}

// fp8x2 -> half2
template <>
__inline__ __device__ uint32_t
499
scaled_vec_conversion<uint32_t, uint16_t>(const uint16_t& a, float scale, Fp8KVCacheDataType kv_type) {
500
  union {
zhuwenwen's avatar
zhuwenwen committed
501
502
503
    uint16_t u16[2];
    uint32_t u32;
  } res;
504
505
  res.u16[0] = scaled_vec_conversion<uint16_t, uint8_t>((uint8_t)a, scale, kv_type);
  res.u16[1] = scaled_vec_conversion<uint16_t, uint8_t>((uint8_t)(a >> 8U), scale, kv_type);
zhuwenwen's avatar
zhuwenwen committed
506
507
508
509
510
511
512
513
514
515
516
  return res.u32;
  // [[maybe_unused]] __half2_raw h2r =
  //     __hip_cvt_fp8x2_to_halfraw2(a, fp8_type::__default_interpret);
  // union {
  //   __half2_raw h2r;
  //   uint32_t ui32;
  // } tmp;
  // tmp.h2r = __hip_cvt_fp8x2_to_halfraw2(a, fp8_type::__default_interpret);
  // tmp.h2r.x.data *= scale;
  // tmp.h2r.y.data *= scale;
  // return tmp.ui32;
517
518
519
520
521
}

// fp8x4 -> half2x2
template <>
__inline__ __device__ uint2
522
scaled_vec_conversion<uint2, uint32_t>(const uint32_t& a, float scale, Fp8KVCacheDataType kv_type) {
523
524
525
526
  union {
    uint2 u32x2;
    uint32_t u32[2];
  } tmp;
527
528
  tmp.u32[0] = scaled_vec_conversion<uint32_t, uint16_t>((uint16_t)a, scale, kv_type);
  tmp.u32[1] = scaled_vec_conversion<uint32_t, uint16_t>((uint16_t)(a >> 16U), scale, kv_type);
529
530
  return tmp.u32x2;
}
531

532
533
534
// fp8x8 -> half2x4
template <>
__inline__ __device__ uint4 scaled_vec_conversion<uint4, uint2>(const uint2& a,
535
                                                                float scale, Fp8KVCacheDataType kv_type) {
536
537
538
539
  union {
    uint4 u64x2;
    uint2 u64[2];
  } tmp;
540
541
  tmp.u64[0] = scaled_vec_conversion<uint2, uint32_t>(a.x, scale, kv_type);
  tmp.u64[1] = scaled_vec_conversion<uint2, uint32_t>(a.y, scale, kv_type);
542
543
  return tmp.u64x2;
}
544
545
546

// half -> fp8
template <>
547
__inline__ __device__ uint8_t
548
scaled_vec_conversion<uint8_t, uint16_t>(const uint16_t& a, float scale, Fp8KVCacheDataType kv_type) {
zhuwenwen's avatar
zhuwenwen committed
549
  float res_f = half_to_float(a) / scale;
550
551
552
553
554
  if (kv_type == vllm::Fp8KVCacheDataType::kFp8E4M3) {
    return float_to_fp8_e4m3(res_f);
  } else {
    return float_to_fp8_e5m2(res_f);
  }
zhuwenwen's avatar
zhuwenwen committed
555
556
557
558
559
  // __half_raw tmp;
  // tmp.x = a;
  // tmp.data /= scale;
  // return __hip_cvt_halfraw_to_fp8(tmp, fp8_type::__default_saturation,
  //                                 fp8_type::__default_interpret);
560
}
561

562
563
564
// halfx2 -> fp8x2
template <>
__inline__ __device__ uint16_t
565
scaled_vec_conversion<uint16_t, uint32_t>(const uint32_t& a, float scale, Fp8KVCacheDataType kv_type) {
566
  union {
zhuwenwen's avatar
zhuwenwen committed
567
568
    uint8_t ui8[2];
    uint16_t ui16;
569
  } tmp;
zhuwenwen's avatar
zhuwenwen committed
570
571
572
573
574
  union {
    uint32_t ui32;
    half2 h2r;
  } tmp_a;
  tmp_a.ui32 = a;
575
576
  tmp.ui8[0] = scaled_vec_conversion<uint8_t, uint16_t>(tmp_a.h2r.data[0], scale, kv_type);
  tmp.ui8[1] = scaled_vec_conversion<uint8_t, uint16_t>(tmp_a.h2r.data[1], scale, kv_type);
zhuwenwen's avatar
zhuwenwen committed
577
578
579
580
581
582
583
584
585
586
  return tmp.ui16;
  // union {
  //   uint32_t ui32;
  //   __half2_raw h2r;
  // } tmp;
  // tmp.ui32 = a;
  // tmp.h2r.x.data /= scale;
  // tmp.h2r.y.data /= scale;
  // return __hip_cvt_halfraw2_to_fp8x2(tmp.h2r, fp8_type::__default_saturation,
  //                                    fp8_type::__default_interpret);
587
588
589
590
591
}

// half2x2 -> fp8x4
template <>
__inline__ __device__ uint32_t
592
scaled_vec_conversion<uint32_t, uint2>(const uint2& a, float scale, Fp8KVCacheDataType kv_type) {
593
594
595
596
  union {
    uint16_t ui16[2];
    uint32_t ui32;
  } tmp;
597
598
  tmp.ui16[0] = scaled_vec_conversion<uint16_t, uint32_t>(a.x, scale, kv_type);
  tmp.ui16[1] = scaled_vec_conversion<uint16_t, uint32_t>(a.y, scale, kv_type);
599
600
601
602
603
604
  return tmp.ui32;
}

// half2x4 -> fp8x8
template <>
__inline__ __device__ uint2 scaled_vec_conversion<uint2, uint4>(const uint4& a,
605
                                                                float scale, Fp8KVCacheDataType kv_type) {
606
607
608
609
610
611
  union {
    uint2 ui2[2];
    uint4 ui4;
  } tmp;
  tmp.ui4 = a;
  uint2 res;
612
613
  res.x = scaled_vec_conversion<uint32_t, uint2>(tmp.ui2[0], scale, kv_type);
  res.y = scaled_vec_conversion<uint32_t, uint2>(tmp.ui2[1], scale, kv_type);
614
  return res;
615
616
617
618
}

// bf16 -> fp8
template <>
619
__inline__ __device__ uint8_t scaled_vec_conversion<uint8_t, __nv_bfloat16>(
620
    const __nv_bfloat16& a, float scale, Fp8KVCacheDataType kv_type) {
zhuwenwen's avatar
zhuwenwen committed
621
      float res_f = (static_cast<float>(a)) / scale;
622
623
624
625
626
      if (kv_type == vllm::Fp8KVCacheDataType::kFp8E4M3) {
        return float_to_fp8_e4m3(res_f);
      } else {
        return float_to_fp8_e5m2(res_f);
      }
zhuwenwen's avatar
zhuwenwen committed
627
628
629
  // return __hip_cvt_float_to_fp8(__bfloat162float(a) / scale,
  //                               fp8_type::__default_saturation,
  //                               fp8_type::__default_interpret);
630
631
632
633
634
}

// bf16x2 -> fp8x2
template <>
__inline__ __device__ uint16_t scaled_vec_conversion<uint16_t, __nv_bfloat162>(
635
    const __nv_bfloat162& a, float scale, Fp8KVCacheDataType kv_type) {
636
637
638
639
  union {
    uint8_t ui8[2];
    uint16_t ui16;
  } tmp;
640
641
  tmp.ui8[0] = scaled_vec_conversion<uint8_t, __nv_bfloat16>(a.x, scale, kv_type);
  tmp.ui8[1] = scaled_vec_conversion<uint8_t, __nv_bfloat16>(a.y, scale, kv_type);
642
643
644
645
646
647
  return tmp.ui16;
}

// bf16x4 -> fp8x4
template <>
__inline__ __device__ uint32_t
648
scaled_vec_conversion<uint32_t, bf16_4_t>(const bf16_4_t& a, float scale, Fp8KVCacheDataType kv_type) {
649
650
651
652
  union {
    uint16_t ui16[2];
    uint32_t ui32;
  } tmp;
653
654
  tmp.ui16[0] = scaled_vec_conversion<uint16_t, __nv_bfloat162>(a.x, scale, kv_type);
  tmp.ui16[1] = scaled_vec_conversion<uint16_t, __nv_bfloat162>(a.y, scale, kv_type);
655
656
657
658
659
660
  return tmp.ui32;
}

// bf16x8 -> fp8x8
template <>
__inline__ __device__ uint2
661
scaled_vec_conversion<uint2, bf16_8_t>(const bf16_8_t& a, float scale, Fp8KVCacheDataType kv_type) {
662
  uint2 res;
663
664
  res.x = scaled_vec_conversion<uint32_t, bf16_4_t>({a.x, a.y}, scale, kv_type);
  res.y = scaled_vec_conversion<uint32_t, bf16_4_t>({a.z, a.w}, scale, kv_type);
665
  return res;
666
667
668
669
}

// float -> fp8
template <>
670
__inline__ __device__ uint8_t
671
672
673
674
675
676
scaled_vec_conversion<uint8_t, float>(const float& a, float scale, Fp8KVCacheDataType kv_type) {
  if (kv_type == vllm::Fp8KVCacheDataType::kFp8E4M3) {
    return float_to_fp8_e4m3(a / scale);
  } else {
    return float_to_fp8_e5m2(a / scale);
  }
zhuwenwen's avatar
zhuwenwen committed
677
678
  // return __hip_cvt_float_to_fp8(a / scale, fp8_type::__default_saturation,
  //                               fp8_type::__default_interpret);
679
680
}

681
// floatx2 -> fp8x2
682
template <>
683
__inline__ __device__ uint16_t
684
scaled_vec_conversion<uint16_t, float2>(const float2& a, float scale, Fp8KVCacheDataType kv_type) {
zhuwenwen's avatar
zhuwenwen committed
685
686
687
688
  union {
    uint8_t ui8[2];
    uint16_t ui16;
  } tmp;
689
690
  tmp.ui8[0] = scaled_vec_conversion<uint8_t, float>(a.x, scale, kv_type);
  tmp.ui8[1] = scaled_vec_conversion<uint8_t, float>(a.y, scale, kv_type);
zhuwenwen's avatar
zhuwenwen committed
691
692
693
  return tmp.ui16;
  // return __hip_cvt_float2_to_fp8x2(a / scale, fp8_type::__default_saturation,
  //                                  fp8_type::__default_interpret);
694
695
696
697
698
}

// floatx4 -> fp8x4
template <>
__inline__ __device__ uint32_t
699
scaled_vec_conversion<uint32_t, float4>(const float4& a, float scale, Fp8KVCacheDataType kv_type) {
700
701
702
703
  union {
    uint16_t ui16[2];
    uint32_t ui32;
  } tmp;
704
705
  tmp.ui16[0] = scaled_vec_conversion<uint16_t, float2>({a.x, a.y}, scale, kv_type);
  tmp.ui16[1] = scaled_vec_conversion<uint16_t, float2>({a.z, a.w}, scale, kv_type);
706
  return tmp.ui32;
707
}
zhuwenwen's avatar
zhuwenwen committed
708
  // #endif  // ENABLE_FP8
709

zhuwenwen's avatar
zhuwenwen committed
710
711
712
713
714
715
716
717
718
719
// template <typename Tout, typename Tin, Fp8KVCacheDataType kv_dt>
// __inline__ __device__ Tout convert(const Tin& x) {
//   #ifdef ENABLE_FP8
//   if constexpr (kv_dt == Fp8KVCacheDataType::kFp8E4M3) {
//     return vec_conversion<Tout, Tin>(x);
//   }
//   #endif
//   assert(false);
//   return {};  // Squash missing return statement warning
// }
720
721

template <typename Tout, typename Tin, Fp8KVCacheDataType kv_dt>
722
__inline__ __device__ Tout scaled_convert(const Tin& x, const float scale) {
zhuwenwen's avatar
zhuwenwen committed
723
  // #ifdef ENABLE_FP8
724
725
726
  if constexpr (kv_dt == Fp8KVCacheDataType::kFp8E4M3 || kv_dt == Fp8KVCacheDataType::kFp8E5M2) {
    return scaled_vec_conversion<Tout, Tin>(x, scale, kv_dt);
  }
zhuwenwen's avatar
zhuwenwen committed
727
  // #endif
728
  assert(false);
729
  return {};  // Squash missing return statement warning
730
731
}

732
733
734
735
736
737
  // The following macro is used to dispatch the conversion function based on
  // the data type of the key and value cache. The FN is a macro that calls a
  // function with template<typename scalar_t, typename cache_t,
  // Fp8KVCacheDataType kv_dt>.
  #define DISPATCH_BY_KV_CACHE_DTYPE(SRC_DTYPE, KV_DTYPE, FN)                  \
    if (KV_DTYPE == "auto") {                                                  \
738
      if (SRC_DTYPE == at::ScalarType::Float) {                                \
739
        FN(float, float, vllm::Fp8KVCacheDataType::kAuto);                     \
740
      } else if (SRC_DTYPE == at::ScalarType::Half) {                          \
741
        FN(uint16_t, uint16_t, vllm::Fp8KVCacheDataType::kAuto);               \
742
      } else if (SRC_DTYPE == at::ScalarType::BFloat16) {                      \
743
        FN(__nv_bfloat16, __nv_bfloat16, vllm::Fp8KVCacheDataType::kAuto);     \
744
745
746
      } else {                                                                 \
        TORCH_CHECK(false, "Unsupported input type of kv cache: ", SRC_DTYPE); \
      }                                                                        \
xiabo's avatar
xiabo committed
747
748
749
750
751
752
753
754
755
756
    } else if (KV_DTYPE == "int8") {                                           \
      if (SRC_DTYPE == at::ScalarType::Float) {                                \
        FN(float, uint8_t, vllm::Fp8KVCacheDataType::kInt8);                   \
      } else if (SRC_DTYPE == at::ScalarType::Half) {                          \
        FN(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kInt8);                \
      } else if (SRC_DTYPE == at::ScalarType::BFloat16) {                      \
        FN(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kInt8);           \
      } else {                                                                 \
        TORCH_CHECK(false,"Unsupported input type of kv cache: ", SRC_DTYPE);  \
      }                                                                        \
757
    } else {                                                                   \
zhuwenwen's avatar
zhuwenwen committed
758
      if (KV_DTYPE == "fp8" || KV_DTYPE == "fp8_e4m3" || KV_DTYPE == "fp8_ds_mla") {                       \
759
760
761
762
763
764
765
766
767
768
        if (SRC_DTYPE == at::ScalarType::Float) {                              \
          FN(float, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3);              \
        } else if (SRC_DTYPE == at::ScalarType::Half) {                        \
          FN(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3);           \
        } else if (SRC_DTYPE == at::ScalarType::BFloat16) {                    \
          FN(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3);      \
        } else {                                                               \
          TORCH_CHECK(false,                                                   \
                      "Unsupported input type of kv cache: ", SRC_DTYPE);      \
        }                                                                      \
769
770
771
772
773
774
775
776
777
778
779
780
      }                                                                        \
      else if (KV_DTYPE == "fp8_e5m2") {                                       \
        if (SRC_DTYPE == at::ScalarType::Float) {                              \
          FN(float, uint8_t, vllm::Fp8KVCacheDataType::kFp8E5M2);              \
        } else if (SRC_DTYPE == at::ScalarType::Half) {                        \
          FN(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kFp8E5M2);           \
        } else if (SRC_DTYPE == at::ScalarType::BFloat16) {                    \
          FN(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E5M2);      \
        } else {                                                               \
          TORCH_CHECK(false,                                                   \
                      "Unsupported input type of kv cache: ", SRC_DTYPE);      \
        }                                                                      \
781
782
783
784
      } else {                                                                 \
        TORCH_CHECK(false, "Unsupported data type of kv cache: ", KV_DTYPE);   \
      }                                                                        \
    }
785

786
787
}  // namespace fp8
#endif  // USE_ROCM
zhuwenwen's avatar
zhuwenwen committed
788
}  // namespace vllm