| 1 | // SPDX-License-Identifier: Apache-2.0 |
| 2 | // SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | #pragma once |
| 4 | #include <algorithm> |
| 5 | #include <cmath> |
| 6 | #include <cstdint> |
| 7 | #include <limits> |
| 8 | #include <type_traits> |
| 9 | |
| 10 | #include "cpu/utils.hpp" |
| 11 | |
| 12 | typedef uint32_t __attribute__((__may_alias__)) u32_alias_t; |
| 13 | typedef uint16_t __attribute__((__may_alias__)) u16_alias_t; |
| 14 | typedef float __attribute__((__may_alias__)) f32_alias_t; |
| 15 | |
| 16 | // Reference scalar dequant — used to verify vectorized AMX dequant. |
| 17 | inline float fp8e4m3_to_float_scalar(uint8_t b, float scale) noexcept { |
| 18 | // NaN encoding in E4M3 |
| 19 | if ((b & 0x7F) == 0x7F) return std::numeric_limits<float>::quiet_NaN(); |
| 20 | uint32_t b_u32 = static_cast<uint32_t>(b); |
| 21 | uint32_t sign = (b_u32 & 0x80) << 24; |
| 22 | uint32_t payload = (b_u32 & 0x7F) << 20; |
| 23 | uint32_t bits = sign | payload; |
| 24 | float b_f32_unscaled = *reinterpret_cast<const f32_alias_t*>(&bits); |
| 25 | float b_f32_scaled = b_f32_unscaled * scale * 0x1p120f; |
| 26 | return b_f32_scaled; |
| 27 | } |
| 28 | |
| 29 | inline uint8_t float_to_fp8e4m3_scalar(float v, float inv_scale) noexcept { |
| 30 | v *= inv_scale; |
| 31 | constexpr float fp8_max = 448.0f; |
| 32 | v = std::max(-fp8_max, std::min(fp8_max, v)); |
| 33 | if (v == 0.0f) return 0; |
| 34 | |
| 35 | // Inverse mapping of fp8e4m3_to_float_scalar: shift the effective exponent |
| 36 | // bias from fp32 (127) back to fp8 e4m3 (7), then pack sign|payload. |
| 37 | float v_f32_unscaled = v * 0x1p-120f; |
| 38 | uint32_t bits = *reinterpret_cast<const u32_alias_t*>(&v_f32_unscaled); |
| 39 | uint8_t sign = static_cast<uint8_t>((bits >> 24) & 0x80); |
| 40 | uint8_t payload = static_cast<uint8_t>((bits >> 20) & 0x7F); |
| 41 | if (payload == 0) return sign; |
| 42 | payload = std::min<uint8_t>(payload, 0x7E); // keep 0x7F as NaN encoding |
| 43 | return static_cast<uint8_t>(sign | payload); |
| 44 | } |
| 45 | |
| 46 | // --------------------------------------------------------------------------- |
| 47 | // AMX reshape impl — parameterised on the quantisation function. |
| 48 | // Writes key/value into uint8 FP8 KV cache using the AMX tile-friendly layout. |
| 49 | // K: halfword-packed (2 FP8 per uint16, token_num_per_group=16). |
| 50 | // V: sub-group packing (token_num_per_sub_group=2, head_elems_per_group=16). |
| 51 | // block_size must be divisible by 32. |
| 52 | // --------------------------------------------------------------------------- |
| 53 | template <typename scalar_t, uint8_t (*quant_fn)(float, float)> |
| 54 | inline void reshape_and_cache_fp8_amx_impl( |
| 55 | const scalar_t* key_ptr, const scalar_t* value_ptr, uint8_t* key_cache_ptr, |
| 56 | uint8_t* value_cache_ptr, const int64_t* slot_ptr, int64_t token_num, |
| 57 | int64_t head_num, int64_t head_dim, int64_t block_size, int64_t k_stride0, |
| 58 | int64_t k_stride1, int64_t v_stride0, int64_t v_stride1, int64_t kc_stride0, |
| 59 | int64_t kc_stride1, int64_t vc_stride0, int64_t vc_stride1, float k_inv, |
| 60 | float v_inv) { |
| 61 | constexpr int64_t token_num_per_group = 16; // AMX_TILE_ROW_NUM |
| 62 | const int64_t halfword_num = head_dim / 2; // 2 FP8 per uint16 |
| 63 | const int64_t halfword_num_per_group = token_num_per_group * halfword_num; |
| 64 | constexpr int64_t head_elems_per_group = 16; |
| 65 | constexpr int64_t token_num_per_sub_group = 2; // = 4 / sizeof(BF16) |
| 66 | const int64_t group_num = head_dim / head_elems_per_group; |
| 67 | const int64_t group_size = block_size * head_elems_per_group; |
| 68 | |
| 69 | #pragma omp parallel for collapse(2) schedule(static) |
| 70 | for (int64_t tok = 0; tok < token_num; ++tok) { |
| 71 | for (int64_t h = 0; h < head_num; ++h) { |
| 72 | const int64_t slot = slot_ptr[tok]; |
| 73 | if (slot < 0) continue; |
| 74 | const int64_t block_idx = slot / block_size; |
| 75 | const int64_t block_offset = slot % block_size; |
| 76 | |
| 77 | // Key: halfword-packed, 2 FP8 per uint16 |
| 78 | { |
| 79 | const scalar_t* ksrc = key_ptr + tok * k_stride0 + h * k_stride1; |
| 80 | const int64_t group_idx = block_offset / token_num_per_group; |
| 81 | const int64_t group_offset = block_offset % token_num_per_group; |
| 82 | uint16_t* kdst = |
| 83 | reinterpret_cast<uint16_t*>(key_cache_ptr + block_idx * kc_stride0 + |
| 84 | h * kc_stride1) + |
| 85 | group_idx * halfword_num_per_group + group_offset; |
| 86 | for (int64_t j = 0; j < halfword_num; ++j) { |
| 87 | uint8_t fp8_0 = quant_fn(static_cast<float>(ksrc[j * 2]), k_inv); |
| 88 | uint8_t fp8_1 = quant_fn(static_cast<float>(ksrc[j * 2 + 1]), k_inv); |
| 89 | uint8_t bytes[2] = {fp8_0, fp8_1}; |
| 90 | uint16_t hw = *reinterpret_cast<const u16_alias_t*>(bytes); |
| 91 | kdst[j * token_num_per_group] = hw; |
| 92 | } |
| 93 | } |
| 94 | |
| 95 | // Value: sub-group packing (token_num_per_sub_group = 2) |
| 96 | { |
| 97 | const scalar_t* vsrc = value_ptr + tok * v_stride0 + h * v_stride1; |
| 98 | const int64_t sub_group_idx = block_offset / token_num_per_sub_group; |
| 99 | const int64_t sub_group_offset = block_offset % token_num_per_sub_group; |
| 100 | uint8_t* vdst = |
| 101 | value_cache_ptr + block_idx * vc_stride0 + h * vc_stride1 + |
| 102 | sub_group_idx * token_num_per_sub_group * head_elems_per_group + |
| 103 | sub_group_offset; |
| 104 | for (int64_t i = 0; i < group_num; ++i) { |
| 105 | for (int64_t j = 0; j < head_elems_per_group; ++j) |
| 106 | vdst[j * token_num_per_sub_group] = |
| 107 | quant_fn(static_cast<float>(vsrc[j]), v_inv); |
| 108 | vsrc += head_elems_per_group; |
| 109 | vdst += group_size; |
| 110 | } |
| 111 | } |
| 112 | } |
| 113 | } |
| 114 | } |
| 115 | |
| 116 | // --------------------------------------------------------------------------- |
| 117 | // FP8 E5M2 scalar helpers |
| 118 | // --------------------------------------------------------------------------- |
| 119 | inline uint8_t float_to_fp8e5m2_scalar(float v, float inv_scale) noexcept { |
| 120 | v *= inv_scale; |
| 121 | constexpr float fp8_e5m2_max = 57344.0f; |
| 122 | v = std::max(-fp8_e5m2_max, std::min(fp8_e5m2_max, v)); |
| 123 | if (v == 0.0f) return 0; |
| 124 | uint32_t bits = *reinterpret_cast<const u32_alias_t*>(&v); |
| 125 | const uint8_t sign = static_cast<uint8_t>((bits >> 24) & 0x80); |
| 126 | const int32_t exp_fp32 = static_cast<int32_t>((bits >> 23) & 0xFF) - 127; |
| 127 | const uint8_t mant2 = static_cast<uint8_t>((bits >> 21) & 0x03); |
| 128 | if (exp_fp32 < -14) { // subnormal in E5M2 |
| 129 | const int shift = -14 - exp_fp32; |
| 130 | if (shift + 21 >= 32) |
| 131 | return sign; // underflow: too small for E5M2 subnormal |
| 132 | const uint32_t m = (0x800000u | (bits & 0x7FFFFFu)) >> (shift + 21); |
| 133 | return sign | static_cast<uint8_t>(std::min<uint32_t>(m, 3u)); |
| 134 | } |
| 135 | const uint8_t exp5 = static_cast<uint8_t>(exp_fp32 + 15); |
| 136 | return sign | (exp5 << 2) | mant2; |
| 137 | } |
| 138 | |
| 139 | // --------------------------------------------------------------------------- |
| 140 | // Select the FP8 quant function at compile time based on kv_cache_t. |
| 141 | // --------------------------------------------------------------------------- |
| 142 | template <typename kv_cache_t> |
| 143 | constexpr auto select_fp8_quant_fn() { |
| 144 | if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e5m2>) |
| 145 | return float_to_fp8e5m2_scalar; |
| 146 | else |
| 147 | return float_to_fp8e4m3_scalar; |
| 148 | } |
| 149 | |
| 150 | // --------------------------------------------------------------------------- |
| 151 | // VEC reshape impl — parameterised on the quantisation function. |
| 152 | // Writes key (column-major) and value (row-major) into uint8 FP8 KV cache. |
| 153 | // The pragma omp must live outside VLLM_DISPATCH_FLOATING_TYPES because |
| 154 | // #pragma cannot appear inside variadic macro arguments. |
| 155 | // --------------------------------------------------------------------------- |
| 156 | template <typename scalar_t, uint8_t (*quant_fn)(float, float)> |
| 157 | inline void reshape_and_cache_fp8_vec_impl( |
| 158 | const scalar_t* key_ptr, const scalar_t* value_ptr, uint8_t* key_cache_ptr, |
| 159 | uint8_t* value_cache_ptr, const int64_t* slot_ptr, int64_t token_num, |
| 160 | int64_t head_num, int64_t head_dim, int64_t block_size, int64_t k_stride0, |
| 161 | int64_t k_stride1, int64_t v_stride0, int64_t v_stride1, int64_t kc_stride0, |
| 162 | int64_t kc_stride1, int64_t vc_stride0, int64_t vc_stride1, float k_inv, |
| 163 | float v_inv) { |
| 164 | #pragma omp parallel for collapse(2) schedule(static) |
| 165 | for (int64_t tok = 0; tok < token_num; ++tok) { |
| 166 | for (int64_t h = 0; h < head_num; ++h) { |
| 167 | const int64_t slot = slot_ptr[tok]; |
| 168 | if (slot < 0) continue; |
| 169 | const int64_t block_idx = slot / block_size; |
| 170 | const int64_t block_offset = slot % block_size; |
| 171 | |
| 172 | // Key layout: column-major within block |
| 173 | const scalar_t* ksrc = key_ptr + tok * k_stride0 + h * k_stride1; |
| 174 | uint8_t* kdst = key_cache_ptr + block_idx * kc_stride0 + h * kc_stride1 + |
| 175 | block_offset; |
| 176 | for (int64_t i = 0; i < head_dim; ++i) |
| 177 | kdst[i * block_size] = quant_fn(static_cast<float>(ksrc[i]), k_inv); |
| 178 | |
| 179 | // Value layout: row-major within block (contiguous head_dim bytes) |
| 180 | const scalar_t* vsrc = value_ptr + tok * v_stride0 + h * v_stride1; |
| 181 | uint8_t* vdst = value_cache_ptr + block_idx * vc_stride0 + |
| 182 | h * vc_stride1 + block_offset * head_dim; |
| 183 | for (int64_t i = 0; i < head_dim; ++i) |
| 184 | vdst[i] = quant_fn(static_cast<float>(vsrc[i]), v_inv); |
| 185 | } |
| 186 | } |
| 187 | } |