| 1 | #pragma once |
| 2 | |
| 3 | #include "attention_generic.cuh" |
| 4 | #include "torch_utils.h" |
| 5 | |
| 6 | #include <stdint.h> |
| 7 | #ifdef ENABLE_FP8 |
| 8 | #ifndef USE_ROCM |
| 9 | #include <cuda_fp8.h> |
| 10 | #endif // USE_ROCM |
| 11 | #endif // ENABLE_FP8 |
| 12 | |
| 13 | namespace vllm { |
| 14 | |
| 15 | enum class Fp8KVCacheDataType { |
| 16 | kAuto = 0, |
| 17 | kFp8E4M3 = 1, |
| 18 | kFp8E5M2 = 2, |
| 19 | }; |
| 20 | |
| 21 | inline Fp8KVCacheDataType get_fp8_kv_cache_data_type( |
| 22 | const std::string& dtype_str) { |
| 23 | // dtype_str refers to CacheDType at vllm.config.cache.CacheDType |
| 24 | if (dtype_str == "auto" || dtype_str == "float16" || |
| 25 | dtype_str == "bfloat16") { |
| 26 | // unquantized kv cache |
| 27 | return Fp8KVCacheDataType::kAuto; |
| 28 | } else if (dtype_str == "fp8" || dtype_str == "fp8_ds_mla" || |
| 29 | dtype_str == "fp8_e4m3") { |
| 30 | return Fp8KVCacheDataType::kFp8E4M3; |
| 31 | } else if (dtype_str == "fp8_e5m2") { |
| 32 | return Fp8KVCacheDataType::kFp8E5M2; |
| 33 | } |
| 34 | TORCH_UTILS_CHECK(false, "Unsupported fp8 kv cache data type: ", dtype_str); |
| 35 | } |
| 36 | |
| 37 | // fp8 vector types for quantization of kv cache |
| 38 | template <> |
| 39 | struct Vec<uint8_t, 1> { |
| 40 | using Type = uint8_t; |
| 41 | }; |
| 42 | |
| 43 | template <> |
| 44 | struct Vec<uint8_t, 2> { |
| 45 | using Type = uint16_t; |
| 46 | }; |
| 47 | |
| 48 | template <> |
| 49 | struct Vec<uint8_t, 4> { |
| 50 | using Type = uint32_t; |
| 51 | }; |
| 52 | |
| 53 | template <> |
| 54 | struct Vec<uint8_t, 8> { |
| 55 | using Type = uint2; |
| 56 | }; |
| 57 | |
| 58 | } // namespace vllm |