| 1 | // SPDX-License-Identifier: Apache-2.0 |
| 2 | // SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | // |
| 4 | // AMX_FP8 attention implementation for Intel Diamond Rapids (DMR). |
| 5 | // |
| 6 | // Architecture: |
| 7 | // QK phase: native FP8×FP8 MMA via _tile_dpfp8ps / _tile_dpbf8ps |
| 8 | // Q is quantized from BF16 to the KV cache FP8 type. |
| 9 | // K cache is stored in the native AMX-FP8 layout (1 byte/element). |
| 10 | // PV phase: native FP8 MMA with 64-byte VNNI4 tiles. |
| 11 | // Softmax writes P once in compact FP8 form for reuse across all |
| 12 | // output-dimension groups. |
| 13 | // |
| 14 | // Scale bookkeeping for QK: |
| 15 | // true_score = Q_fp32 · K_fp32 |
| 16 | // = (Q_fp8 * q_scale) · (K_fp8 * k_scale) |
| 17 | // = dot(Q_fp8, K_fp8) * q_scale * k_scale |
| 18 | // execute_attention() applies: scale *= q_dynamic_scale * k_scale |
| 19 | // where q_dynamic_scale is computed per-tile in copy_q_heads_tile(). |
| 20 | // |
| 21 | // Scale bookkeeping for PV: |
| 22 | // raw_C = dot(P_fp8, V_fp8) |
| 23 | // ≈ dot(P/p_scale, V_real/v_scale) |
| 24 | // = dot(P, V_real) / (p_scale * v_scale) |
| 25 | // In final_output: multiply by output_v_scale = p_scale * v_scale. |
| 26 | // |
| 27 | #ifndef CPU_ATTN_AMX_FP8_HPP |
| 28 | #define CPU_ATTN_AMX_FP8_HPP |
| 29 | |
| 30 | #ifdef CPU_CAPABILITY_AMXFP8 |
| 31 | |
| 32 | #include "cpu_attn_amx.hpp" // pulls in TileGemm224/TileGemm122 BF16 specialisations |
| 33 | #include "cpu_attn_fp8.hpp" |
| 34 | #include "cpu_attn_impl.hpp" |
| 35 | #include "cpu_types.hpp" |
| 36 | |
| 37 | #include <amxfp8intrin.h> // requires GCC 14+ / Clang 17+, -mamx-fp8 |
| 38 | |
| 39 | namespace cpu_attention { |
| 40 | |
| 41 | // --------------------------------------------------------------------------- |
| 42 | // Native-FP8 K cache layout reshape helper. |
| 43 | // |
| 44 | // Layout per 16-token group (token_num_per_group = AMX_TILE_ROW_NUM = 16): |
| 45 | // k-slice s (s = 0 .. head_dim/64 - 1): |
| 46 | // row k (k = 0..15): four consecutive K values for each of 16 tokens |
| 47 | // (VNNI4 B-operand layout) |
| 48 | // One k-slice = 16 rows × 64 bytes = 1024 bytes = 1 AMX FP8 tile. |
| 49 | // Full group = (head_dim / 64) slices × 1024 bytes. |
| 50 | // |
| 51 | // For head_dim=128: group = 2 × 1024 = 2048 bytes. |
| 52 | // Consecutive 16-token groups are contiguous within each cache block. |
| 53 | // --------------------------------------------------------------------------- |
| 54 | template <typename scalar_t, uint8_t (*quant_fn)(float, float)> |
| 55 | inline void reshape_and_cache_k_amx_fp8_impl( |
| 56 | const scalar_t* key_ptr, uint8_t* key_cache_ptr, const int64_t* slot_ptr, |
| 57 | int64_t token_num, int64_t head_num, int64_t head_dim, int64_t block_size, |
| 58 | int64_t k_stride0, int64_t k_stride1, int64_t kc_stride0, |
| 59 | int64_t kc_stride1, float k_inv) { |
| 60 | constexpr int64_t token_num_per_group = AMX_TILE_ROW_NUM; // 16 |
| 61 | constexpr int64_t fp8_per_row = AMX_TILE_ROW_BYTES; // 64 FP8/row |
| 62 | const int64_t slice_num = head_dim / fp8_per_row; |
| 63 | // Bytes per 16-token group = slice_num * (16 * 64) = head_dim * 16 |
| 64 | const int64_t group_bytes = token_num_per_group * head_dim; |
| 65 | |
| 66 | #pragma omp parallel for collapse(2) schedule(static) |
| 67 | for (int64_t tok = 0; tok < token_num; ++tok) { |
| 68 | for (int64_t h = 0; h < head_num; ++h) { |
| 69 | const int64_t slot = slot_ptr[tok]; |
| 70 | if (slot < 0) continue; |
| 71 | const int64_t block_idx = slot / block_size; |
| 72 | const int64_t block_offset = slot % block_size; |
| 73 | const int64_t group_idx = block_offset / token_num_per_group; |
| 74 | const int64_t group_offset = block_offset % token_num_per_group; |
| 75 | |
| 76 | const scalar_t* ksrc = key_ptr + tok * k_stride0 + h * k_stride1; |
| 77 | uint8_t* group_base = key_cache_ptr + block_idx * kc_stride0 + |
| 78 | h * kc_stride1 + group_idx * group_bytes; |
| 79 | |
| 80 | // Each k-slice is one 16x64-byte VNNI4 B tile. Row k stores dimensions |
| 81 | // [4*k, 4*k+4) for every token; each token occupies four bytes. |
| 82 | for (int64_t s = 0; s < slice_num; ++s) { |
| 83 | uint8_t* tile_base = group_base + s * token_num_per_group * fp8_per_row; |
| 84 | for (int64_t k = 0; k < fp8_per_row / 4; ++k) { |
| 85 | uint8_t* packed_values = |
| 86 | tile_base + k * fp8_per_row + group_offset * 4; |
| 87 | for (int64_t lane = 0; lane < 4; ++lane) { |
| 88 | packed_values[lane] = quant_fn( |
| 89 | static_cast<float>(ksrc[s * fp8_per_row + k * 4 + lane]), |
| 90 | k_inv); |
| 91 | } |
| 92 | } |
| 93 | } |
| 94 | } |
| 95 | } |
| 96 | } |
| 97 | |
| 98 | template <typename fp8_t> |
| 99 | __attribute__((target("avx10.2"), noinline)) void quantize_probability_row( |
| 100 | const float* __restrict__ src, uint8_t* __restrict__ dst, int32_t cols) { |
| 101 | constexpr float fp8_max = |
| 102 | std::is_same_v<fp8_t, c10::Float8_e4m3fn> ? 448.0f : 57344.0f; |
| 103 | const __m512 scale = _mm512_set1_ps(fp8_max); |
| 104 | for (int32_t col = 0; col < cols; col += 16) { |
| 105 | const __m512 scaled = _mm512_mul_ps(_mm512_loadu_ps(src + col), scale); |
| 106 | const __m256h fp16 = (__m256h)_mm512_cvtps_ph( |
| 107 | scaled, _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC); |
| 108 | if constexpr (std::is_same_v<fp8_t, c10::Float8_e4m3fn>) { |
| 109 | _mm_storeu_si128(reinterpret_cast<__m128i*>(dst + col), |
| 110 | _mm256_cvtph_hf8(fp16)); |
| 111 | } else { |
| 112 | _mm_storeu_si128(reinterpret_cast<__m128i*>(dst + col), |
| 113 | _mm256_cvtph_bf8(fp16)); |
| 114 | } |
| 115 | } |
| 116 | } |
| 117 | |
| 118 | template <typename fp8_t> |
| 119 | __attribute__((target("avx10.2"), noinline)) void quantize_q_tile( |
| 120 | const c10::BFloat16* __restrict__ src, uint8_t* __restrict__ dst, |
| 121 | int32_t q_num, int32_t q_heads_per_kv, int64_t q_num_stride, |
| 122 | int64_t q_head_stride, int64_t head_dim, float inv_scale) { |
| 123 | const __m512 scale = _mm512_set1_ps(inv_scale); |
| 124 | constexpr float fp8_max = |
| 125 | std::is_same_v<fp8_t, c10::Float8_e4m3fn> ? 448.0f : 57344.0f; |
| 126 | const __m512 max_value = _mm512_set1_ps(fp8_max); |
| 127 | const __m512 min_value = _mm512_set1_ps(-fp8_max); |
| 128 | const int64_t head_size_block_num = head_dim / AMX_TILE_ROW_BYTES; |
| 129 | int32_t packed_row = 0; |
| 130 | |
| 131 | for (int32_t q = 0; q < q_num; ++q) { |
| 132 | const c10::BFloat16* head = src + q * q_num_stride; |
| 133 | for (int32_t h = 0; h < q_heads_per_kv; |
| 134 | ++h, head += q_head_stride, ++packed_row) { |
| 135 | uint8_t* packed_head = |
| 136 | dst + |
| 137 | (packed_row / AMX_TILE_ROW_NUM) * head_size_block_num * |
| 138 | AMX_TILE_BYTES + |
| 139 | (packed_row % AMX_TILE_ROW_NUM) * AMX_TILE_ROW_BYTES; |
| 140 | for (int64_t e = 0; e < head_dim; e += 16) { |
| 141 | const __m256i bf16 = |
| 142 | _mm256_loadu_si256(reinterpret_cast<const __m256i*>(head + e)); |
| 143 | __m512 values = _mm512_castsi512_ps( |
| 144 | _mm512_slli_epi32(_mm512_cvtepu16_epi32(bf16), 16)); |
| 145 | values = _mm512_mul_ps(values, scale); |
| 146 | values = _mm512_min_ps(values, max_value); |
| 147 | values = _mm512_max_ps(values, min_value); |
| 148 | const __m256h fp16 = (__m256h)_mm512_cvtps_ph( |
| 149 | values, _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC); |
| 150 | const int64_t block = e / AMX_TILE_ROW_BYTES; |
| 151 | const int64_t offset = e % AMX_TILE_ROW_BYTES; |
| 152 | __m128i fp8; |
| 153 | if constexpr (std::is_same_v<fp8_t, c10::Float8_e4m3fn>) { |
| 154 | fp8 = _mm256_cvtph_hf8(fp16); |
| 155 | } else { |
| 156 | fp8 = _mm256_cvtph_bf8(fp16); |
| 157 | } |
| 158 | _mm_storeu_si128(reinterpret_cast<__m128i*>( |
| 159 | packed_head + block * AMX_TILE_BYTES + offset), |
| 160 | fp8); |
| 161 | } |
| 162 | } |
| 163 | } |
| 164 | } |
| 165 | |
| 166 | // --------------------------------------------------------------------------- |
| 167 | // --------------------------------------------------------------------------- |
| 168 | // TileGemm224<uint8_t, kv_cache_t>: 2-2-4 pattern for AMX_FP8. |
| 169 | // |
| 170 | // QK phase: native FP8 × FP8 MMA (_tile_dphf8ps / _tile_dpbf8ps). |
| 171 | // PV phase: native wide FP8 MMA (colsb=64 tiles). |
| 172 | // P is quantized once per KV tile and reused by every output-dimension group. |
| 173 | // V (FP8 E4M3) loaded directly from cache (no dequant). |
| 174 | // Result correction is folded into the final output scale. |
| 175 | // --------------------------------------------------------------------------- |
| 176 | template <typename kv_cache_t> |
| 177 | class TileGemm224<uint8_t, kv_cache_t> { |
| 178 | static_assert(std::is_same_v<kv_cache_t, c10::Float8_e4m3fn> || |
| 179 | std::is_same_v<kv_cache_t, c10::Float8_e5m2>, |
| 180 | "kv_cache_t must be Float8_e4m3fn or Float8_e5m2"); |
| 181 | |
| 182 | // FP8 tile size in uint8_t elements (= AMX_TILE_BYTES, one per byte). |
| 183 | static constexpr int64_t fp8_tile_elems = AMX_TILE_BYTES; // 1024 |
| 184 | |
| 185 | public: |
| 186 | static constexpr bool prequantize_probabilities = true; |
| 187 | |
| 188 | FORCE_INLINE static void quantize_probability_row(const float* src, |
| 189 | uint8_t* dst, |
| 190 | int32_t cols) { |
| 191 | cpu_attention::quantize_probability_row<kv_cache_t>(src, dst, cols); |
| 192 | } |
| 193 | |
| 194 | // ------------------------------------------------------------------ |
| 195 | // QK: k_times = head_dim / 64 (each FP8 tile covers 64 k-elements). |
| 196 | // PV: native wide FP8 MMA. |
| 197 | // ------------------------------------------------------------------ |
| 198 | template <AttentionGemmPhase phase, int32_t k_size> |
| 199 | FORCE_INLINE static void gemm(const int32_t m_size, void* __restrict__ a_tile, |
| 200 | void* __restrict__ b_tile, |
| 201 | float* __restrict__ c_tile, const int64_t lda, |
| 202 | const int64_t ldb, const int64_t ldc, |
| 203 | const int32_t block_size, |
| 204 | const int32_t dynamic_k_size, |
| 205 | const bool accum_c) { |
| 206 | if constexpr (phase == AttentionGemmPhase::PV) { |
| 207 | // -------- Native wide FP8 PV (FP8 × FP8 → FP32) -------- |
| 208 | // |
| 209 | // A fixed FP8-dependent P scale lets the tile accumulate all k-steps |
| 210 | // before get_output_v_scale() applies the correction. |
| 211 | |
| 212 | const int32_t m_0 = |
| 213 | std::min(m_size, static_cast<int32_t>(AMX_TILE_ROW_NUM)); |
| 214 | const int32_t m_1 = m_size - m_0; |
| 215 | |
| 216 | constexpr int32_t v_stride = AMX_TILE_ROW_BYTES; |
| 217 | const int32_t k_times = |
| 218 | dynamic_k_size / (AMX_TILE_ROW_NUM * 4 / sizeof(uint8_t)); |
| 219 | constexpr int32_t a_advance = AMX_TILE_ROW_BYTES; |
| 220 | constexpr int32_t v_advance = AMX_TILE_BYTES; |
| 221 | const int32_t c_stride = ldc * sizeof(float); |
| 222 | |
| 223 | for (int32_t output_group = 0; output_group < 2; ++output_group) { |
| 224 | uint8_t* a0 = reinterpret_cast<uint8_t*>(a_tile); |
| 225 | uint8_t* a1 = a0 + lda * AMX_TILE_ROW_NUM; |
| 226 | uint8_t* v2 = reinterpret_cast<uint8_t*>(b_tile) + |
| 227 | output_group * 2 * block_size * AMX_TILE_ROW_BYTES / 4; |
| 228 | uint8_t* v3 = v2 + block_size * AMX_TILE_ROW_BYTES / 4; |
| 229 | float* c4 = c_tile + output_group * 2 * AMX_TILE_ROW_NUM; |
| 230 | float* c5 = c4 + AMX_TILE_ROW_NUM; |
| 231 | float* c6 = c4 + AMX_TILE_ROW_NUM * ldc; |
| 232 | float* c7 = c6 + AMX_TILE_ROW_NUM; |
| 233 | |
| 234 | if (accum_c) { |
| 235 | _tile_loadd(4, c4, c_stride); |
| 236 | _tile_loadd(5, c5, c_stride); |
| 237 | if (m_1 > 0) { |
| 238 | _tile_loadd(6, c6, c_stride); |
| 239 | _tile_loadd(7, c7, c_stride); |
| 240 | } |
| 241 | } else { |
| 242 | _tile_zero(4); |
| 243 | _tile_zero(5); |
| 244 | if (m_1 > 0) { |
| 245 | _tile_zero(6); |
| 246 | _tile_zero(7); |
| 247 | } |
| 248 | } |
| 249 | |
| 250 | for (int32_t k = 0; k < k_times; ++k) { |
| 251 | _tile_loadd(0, a0, lda); |
| 252 | _tile_stream_loadd(2, v2, v_stride); |
| 253 | _tile_stream_loadd(3, v3, v_stride); |
| 254 | if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e4m3fn>) { |
| 255 | _tile_dphf8ps(4, 0, 2); |
| 256 | _tile_dphf8ps(5, 0, 3); |
| 257 | } else { |
| 258 | _tile_dpbf8ps(4, 0, 2); |
| 259 | _tile_dpbf8ps(5, 0, 3); |
| 260 | } |
| 261 | if (m_1 > 0) { |
| 262 | _tile_loadd(1, a1, lda); |
| 263 | if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e4m3fn>) { |
| 264 | _tile_dphf8ps(6, 1, 2); |
| 265 | _tile_dphf8ps(7, 1, 3); |
| 266 | } else { |
| 267 | _tile_dpbf8ps(6, 1, 2); |
| 268 | _tile_dpbf8ps(7, 1, 3); |
| 269 | } |
| 270 | } |
| 271 | a0 += a_advance; |
| 272 | a1 += a_advance; |
| 273 | v2 += v_advance; |
| 274 | v3 += v_advance; |
| 275 | } |
| 276 | |
| 277 | _tile_stored(4, c4, c_stride); |
| 278 | _tile_stored(5, c5, c_stride); |
| 279 | if (m_1 > 0) { |
| 280 | _tile_stored(6, c6, c_stride); |
| 281 | _tile_stored(7, c7, c_stride); |
| 282 | } |
| 283 | } |
| 284 | return; |
| 285 | } |
| 286 | |
| 287 | // -------- QK phase: native FP8 MMA -------- |
| 288 | // k_times: each FP8 AMX tile covers (AMX_TILE_ROW_NUM * 4) FP8 = 64 |
| 289 | // inner-product elements. |
| 290 | const int32_t k_times = |
| 291 | dynamic_k_size / (AMX_TILE_ROW_NUM * 4 / sizeof(uint8_t)); |
| 292 | |
| 293 | const int32_t b_tile_stride = AMX_TILE_ROW_BYTES; // 64 bytes |
| 294 | |
| 295 | const int32_t m_0 = |
| 296 | std::min(m_size, static_cast<int32_t>(AMX_TILE_ROW_NUM)); |
| 297 | const int32_t m_1 = m_size - m_0; |
| 298 | const int32_t c_tile_stride = ldc * sizeof(float); |
| 299 | constexpr int32_t output_group_num = 2; |
| 300 | for (int32_t output_group = 0; output_group < output_group_num; |
| 301 | ++output_group) { |
| 302 | uint8_t* __restrict__ a_tile_0 = static_cast<uint8_t*>(a_tile); |
| 303 | uint8_t* __restrict__ a_tile_1 = a_tile_0 + k_times * fp8_tile_elems; |
| 304 | uint8_t* __restrict__ b_tile_2 = |
| 305 | static_cast<uint8_t*>(b_tile) + |
| 306 | output_group * 2 * k_size * AMX_TILE_ROW_NUM; |
| 307 | uint8_t* __restrict__ b_tile_3 = b_tile_2 + (k_size * AMX_TILE_ROW_NUM); |
| 308 | float* __restrict__ c_tile_4 = c_tile + output_group * 32; |
| 309 | float* __restrict__ c_tile_5 = c_tile_4 + AMX_TILE_ROW_NUM; |
| 310 | float* __restrict__ c_tile_6 = c_tile_4 + AMX_TILE_ROW_NUM * ldc; |
| 311 | float* __restrict__ c_tile_7 = c_tile_6 + AMX_TILE_ROW_NUM; |
| 312 | |
| 313 | if (accum_c) { |
| 314 | _tile_loadd(4, c_tile_4, c_tile_stride); |
| 315 | _tile_loadd(5, c_tile_5, c_tile_stride); |
| 316 | _tile_loadd(6, c_tile_6, c_tile_stride); |
| 317 | _tile_loadd(7, c_tile_7, c_tile_stride); |
| 318 | } else { |
| 319 | _tile_zero(4); |
| 320 | _tile_zero(5); |
| 321 | _tile_zero(6); |
| 322 | _tile_zero(7); |
| 323 | } |
| 324 | |
| 325 | for (int32_t k = 0; k < k_times; ++k) { |
| 326 | _tile_loadd(0, a_tile_0, AMX_TILE_ROW_BYTES); |
| 327 | _tile_stream_loadd(2, b_tile_2, b_tile_stride); |
| 328 | if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e4m3fn>) { |
| 329 | _tile_dphf8ps(4, 0, 2); |
| 330 | } else { |
| 331 | _tile_dpbf8ps(4, 0, 2); |
| 332 | } |
| 333 | _tile_stream_loadd(3, b_tile_3, b_tile_stride); |
| 334 | if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e4m3fn>) { |
| 335 | _tile_dphf8ps(5, 0, 3); |
| 336 | } else { |
| 337 | _tile_dpbf8ps(5, 0, 3); |
| 338 | } |
| 339 | _tile_loadd(1, a_tile_1, AMX_TILE_ROW_BYTES); |
| 340 | if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e4m3fn>) { |
| 341 | _tile_dphf8ps(6, 1, 2); |
| 342 | _tile_dphf8ps(7, 1, 3); |
| 343 | } else { |
| 344 | _tile_dpbf8ps(6, 1, 2); |
| 345 | _tile_dpbf8ps(7, 1, 3); |
| 346 | } |
| 347 | |
| 348 | // Q is packed as [row group][K slice]. |
| 349 | a_tile_0 += fp8_tile_elems; |
| 350 | a_tile_1 += fp8_tile_elems; |
| 351 | // Advance K cache (native layout): 1 FP8 tile per k-step. |
| 352 | b_tile_2 += fp8_tile_elems; |
| 353 | b_tile_3 += fp8_tile_elems; |
| 354 | } |
| 355 | |
| 356 | _tile_stored(4, c_tile_4, c_tile_stride); |
| 357 | _tile_stored(5, c_tile_5, c_tile_stride); |
| 358 | _tile_stored(6, c_tile_6, c_tile_stride); |
| 359 | _tile_stored(7, c_tile_7, c_tile_stride); |
| 360 | } |
| 361 | } |
| 362 | |
| 363 | FORCE_INLINE static void init_tile_config(int32_t m, __tilecfg& config) { |
| 364 | // QK and PV share one 64-byte tile configuration. |
| 365 | const int32_t m_0 = std::min(m, static_cast<int32_t>(AMX_TILE_ROW_NUM)); |
| 366 | const int32_t m_1 = m - m_0; |
| 367 | // Use all 64 bytes per row for every tile (same as BF16 AMX). |
| 368 | // This works for both FP8 (64 FP8/row) and BF16 (32 BF16/row = 64B). |
| 369 | for (int i = 0; i < 8; ++i) config.colsb[i] = AMX_TILE_ROW_BYTES; |
| 370 | config.rows[0] = m_0; |
| 371 | config.rows[1] = m_1; |
| 372 | config.rows[2] = AMX_TILE_ROW_NUM; |
| 373 | config.rows[3] = AMX_TILE_ROW_NUM; |
| 374 | config.rows[4] = m_0; |
| 375 | config.rows[5] = m_0; |
| 376 | config.rows[6] = m_1; |
| 377 | config.rows[7] = m_1; |
| 378 | _tile_loadconfig(&config); |
| 379 | } |
| 380 | }; |
| 381 | |
| 382 | // --------------------------------------------------------------------------- |
| 383 | // TileGemm122<uint8_t, kv_cache_t>: 1-2-2 pattern for AMX_FP8. |
| 384 | // --------------------------------------------------------------------------- |
| 385 | template <typename kv_cache_t> |
| 386 | class TileGemm122<uint8_t, kv_cache_t> { |
| 387 | static_assert(std::is_same_v<kv_cache_t, c10::Float8_e4m3fn> || |
| 388 | std::is_same_v<kv_cache_t, c10::Float8_e5m2>, |
| 389 | "kv_cache_t must be Float8_e4m3fn or Float8_e5m2"); |
| 390 | |
| 391 | static constexpr int64_t fp8_tile_elems = AMX_TILE_BYTES; |
| 392 | |
| 393 | public: |
| 394 | static constexpr bool prequantize_probabilities = true; |
| 395 | |
| 396 | FORCE_INLINE static void quantize_probability_row(const float* src, |
| 397 | uint8_t* dst, |
| 398 | int32_t cols) { |
| 399 | cpu_attention::quantize_probability_row<kv_cache_t>(src, dst, cols); |
| 400 | } |
| 401 | |
| 402 | template <AttentionGemmPhase phase, int32_t k_size> |
| 403 | FORCE_INLINE static void gemm(const int32_t m_size, void* __restrict__ a_tile, |
| 404 | void* __restrict__ b_tile, |
| 405 | float* __restrict__ c_tile, const int64_t lda, |
| 406 | const int64_t ldb, const int64_t ldc, |
| 407 | const int32_t block_size, |
| 408 | const int32_t dynamic_k_size, |
| 409 | const bool accum_c) { |
| 410 | if constexpr (phase == AttentionGemmPhase::PV) { |
| 411 | // Native wide FP8 PV (1-2-2 pattern) |
| 412 | // A fixed FP8-dependent P scale enables accumulation across all k-steps. |
| 413 | constexpr int32_t v_stride = AMX_TILE_ROW_BYTES; |
| 414 | constexpr int32_t a_advance = AMX_TILE_ROW_BYTES; |
| 415 | constexpr int32_t v_advance = AMX_TILE_BYTES; |
| 416 | const int32_t k_times = |
| 417 | dynamic_k_size / (AMX_TILE_ROW_NUM * 4 / sizeof(uint8_t)); |
| 418 | |
| 419 | uint8_t* a = reinterpret_cast<uint8_t*>(a_tile); |
| 420 | uint8_t* v2 = reinterpret_cast<uint8_t*>(b_tile); |
| 421 | uint8_t* v3 = v2 + (block_size * AMX_TILE_ROW_BYTES / 4); |
| 422 | uint8_t* v4 = v3 + (block_size * AMX_TILE_ROW_BYTES / 4); |
| 423 | uint8_t* v5 = v4 + (block_size * AMX_TILE_ROW_BYTES / 4); |
| 424 | float* c4 = c_tile + 2 * AMX_TILE_ROW_NUM; |
| 425 | float* c5 = c_tile + 3 * AMX_TILE_ROW_NUM; |
| 426 | float* c6 = c_tile; |
| 427 | float* c7 = c6 + AMX_TILE_ROW_BYTES / sizeof(float); |
| 428 | |
| 429 | const int32_t c_stride = ldc * sizeof(float); |
| 430 | if (accum_c) { |
| 431 | _tile_loadd(4, c4, c_stride); |
| 432 | _tile_loadd(5, c5, c_stride); |
| 433 | _tile_loadd(6, c6, c_stride); |
| 434 | _tile_loadd(7, c7, c_stride); |
| 435 | } else { |
| 436 | _tile_zero(4); |
| 437 | _tile_zero(5); |
| 438 | _tile_zero(6); |
| 439 | _tile_zero(7); |
| 440 | } |
| 441 | |
| 442 | for (int32_t k = 0; k < k_times; ++k) { |
| 443 | _tile_loadd(0, a, lda); |
| 444 | _tile_stream_loadd(2, v2, v_stride); |
| 445 | _tile_stream_loadd(3, v3, v_stride); |
| 446 | if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e4m3fn>) { |
| 447 | _tile_dphf8ps(6, 0, 2); |
| 448 | _tile_dphf8ps(7, 0, 3); |
| 449 | } else { |
| 450 | _tile_dpbf8ps(6, 0, 2); |
| 451 | _tile_dpbf8ps(7, 0, 3); |
| 452 | } |
| 453 | _tile_stream_loadd(2, v4, v_stride); |
| 454 | _tile_stream_loadd(3, v5, v_stride); |
| 455 | if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e4m3fn>) { |
| 456 | _tile_dphf8ps(4, 0, 2); |
| 457 | _tile_dphf8ps(5, 0, 3); |
| 458 | } else { |
| 459 | _tile_dpbf8ps(4, 0, 2); |
| 460 | _tile_dpbf8ps(5, 0, 3); |
| 461 | } |
| 462 | a += a_advance; |
| 463 | v2 += v_advance; |
| 464 | v3 += v_advance; |
| 465 | v4 += v_advance; |
| 466 | v5 += v_advance; |
| 467 | } |
| 468 | |
| 469 | _tile_stored(4, c4, c_stride); |
| 470 | _tile_stored(5, c5, c_stride); |
| 471 | _tile_stored(6, c6, c_stride); |
| 472 | _tile_stored(7, c7, c_stride); |
| 473 | return; |
| 474 | } |
| 475 | |
| 476 | // -------- QK phase: native FP8 MMA (1-2-2 pattern) -------- |
| 477 | const int32_t k_times = |
| 478 | dynamic_k_size / (AMX_TILE_ROW_NUM * 4 / sizeof(uint8_t)); |
| 479 | const int32_t k_group_times = k_times / 2; |
| 480 | const bool has_tail = (k_times % 2 == 1); |
| 481 | |
| 482 | const int32_t b_stride = AMX_TILE_ROW_BYTES; |
| 483 | const int32_t c_stride = ldc * sizeof(float); |
| 484 | constexpr int32_t output_group_num = 2; |
| 485 | for (int32_t output_group = 0; output_group < output_group_num; |
| 486 | ++output_group) { |
| 487 | uint8_t* __restrict__ a_tile_0 = static_cast<uint8_t*>(a_tile); |
| 488 | uint8_t* __restrict__ a_tile_1 = a_tile_0 + fp8_tile_elems; |
| 489 | uint8_t* __restrict__ b_tile_2 = |
| 490 | static_cast<uint8_t*>(b_tile) + |
| 491 | output_group * 2 * k_size * AMX_TILE_ROW_NUM; |
| 492 | uint8_t* __restrict__ b_tile_3 = b_tile_2 + (k_size * AMX_TILE_ROW_NUM); |
| 493 | uint8_t* __restrict__ b_tile_4 = b_tile_2 + fp8_tile_elems; |
| 494 | uint8_t* __restrict__ b_tile_5 = b_tile_3 + fp8_tile_elems; |
| 495 | float* __restrict__ c_tile_6 = c_tile + output_group * 32; |
| 496 | float* __restrict__ c_tile_7 = c_tile_6 + AMX_TILE_ROW_NUM; |
| 497 | |
| 498 | if (accum_c) { |
| 499 | _tile_loadd(6, c_tile_6, c_stride); |
| 500 | _tile_loadd(7, c_tile_7, c_stride); |
| 501 | } else { |
| 502 | _tile_zero(6); |
| 503 | _tile_zero(7); |
| 504 | } |
| 505 | |
| 506 | for (int32_t k = 0; k < k_group_times; ++k) { |
| 507 | _tile_loadd(0, a_tile_0, AMX_TILE_ROW_BYTES); |
| 508 | _tile_stream_loadd(2, b_tile_2, b_stride); |
| 509 | if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e4m3fn>) { |
| 510 | _tile_dphf8ps(6, 0, 2); |
| 511 | } else { |
| 512 | _tile_dpbf8ps(6, 0, 2); |
| 513 | } |
| 514 | _tile_stream_loadd(3, b_tile_3, b_stride); |
| 515 | if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e4m3fn>) { |
| 516 | _tile_dphf8ps(7, 0, 3); |
| 517 | } else { |
| 518 | _tile_dpbf8ps(7, 0, 3); |
| 519 | } |
| 520 | _tile_loadd(1, a_tile_1, AMX_TILE_ROW_BYTES); |
| 521 | _tile_stream_loadd(2, b_tile_4, b_stride); |
| 522 | if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e4m3fn>) { |
| 523 | _tile_dphf8ps(6, 1, 2); |
| 524 | } else { |
| 525 | _tile_dpbf8ps(6, 1, 2); |
| 526 | } |
| 527 | _tile_stream_loadd(3, b_tile_5, b_stride); |
| 528 | if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e4m3fn>) { |
| 529 | _tile_dphf8ps(7, 1, 3); |
| 530 | } else { |
| 531 | _tile_dpbf8ps(7, 1, 3); |
| 532 | } |
| 533 | |
| 534 | a_tile_0 += 2 * fp8_tile_elems; |
| 535 | a_tile_1 += 2 * fp8_tile_elems; |
| 536 | b_tile_2 += 2 * fp8_tile_elems; |
| 537 | b_tile_3 += 2 * fp8_tile_elems; |
| 538 | b_tile_4 += 2 * fp8_tile_elems; |
| 539 | b_tile_5 += 2 * fp8_tile_elems; |
| 540 | } |
| 541 | |
| 542 | if (has_tail) { |
| 543 | _tile_loadd(0, a_tile_0, AMX_TILE_ROW_BYTES); |
| 544 | _tile_stream_loadd(2, b_tile_2, b_stride); |
| 545 | if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e4m3fn>) { |
| 546 | _tile_dphf8ps(6, 0, 2); |
| 547 | } else { |
| 548 | _tile_dpbf8ps(6, 0, 2); |
| 549 | } |
| 550 | _tile_stream_loadd(3, b_tile_3, b_stride); |
| 551 | if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e4m3fn>) { |
| 552 | _tile_dphf8ps(7, 0, 3); |
| 553 | } else { |
| 554 | _tile_dpbf8ps(7, 0, 3); |
| 555 | } |
| 556 | } |
| 557 | |
| 558 | _tile_stored(6, c_tile_6, c_stride); |
| 559 | _tile_stored(7, c_tile_7, c_stride); |
| 560 | } |
| 561 | } |
| 562 | |
| 563 | FORCE_INLINE static void init_tile_config(int32_t m, __tilecfg& config) { |
| 564 | for (int i = 0; i < 8; ++i) config.colsb[i] = AMX_TILE_ROW_BYTES; |
| 565 | config.rows[0] = m; |
| 566 | config.rows[1] = m; |
| 567 | config.rows[2] = AMX_TILE_ROW_NUM; |
| 568 | config.rows[3] = AMX_TILE_ROW_NUM; |
| 569 | config.rows[4] = m; |
| 570 | config.rows[5] = m; |
| 571 | config.rows[6] = m; |
| 572 | config.rows[7] = m; |
| 573 | _tile_loadconfig(&config); |
| 574 | } |
| 575 | }; |
| 576 | |
| 577 | // --------------------------------------------------------------------------- |
| 578 | // AttentionImpl<ISA::AMX_FP8, scalar_t, head_dim, kv_cache_scalar_t> |
| 579 | // |
| 580 | // scalar_t = BF16 (input query/output type) |
| 581 | // kv_cache_scalar_t = Float8_e4m3fn or Float8_e5m2 |
| 582 | // --------------------------------------------------------------------------- |
| 583 | template <typename scalar_t, int64_t head_dim, typename kv_cache_scalar_t> |
| 584 | requires(std::is_same_v<scalar_t, c10::BFloat16> && |
| 585 | (std::is_same_v<kv_cache_scalar_t, c10::Float8_e4m3fn> || |
| 586 | std::is_same_v<kv_cache_scalar_t, c10::Float8_e5m2>) && |
| 587 | (head_dim % 64 == 0)) |
| 588 | class AttentionImpl<ISA::AMX_FP8, scalar_t, head_dim, kv_cache_scalar_t> { |
| 589 | public: |
| 590 | // Q buffer holds FP8-quantized query. |
| 591 | using query_t = scalar_t; // BF16 input |
| 592 | using q_buffer_t = uint8_t; // FP8 quantized, stored in Q buffer |
| 593 | using kv_cache_t = kv_cache_scalar_t; |
| 594 | using logits_buffer_t = float; |
| 595 | using partial_output_buffer_t = float; |
| 596 | using prob_buffer_t = scalar_t; // BF16 softmax probs (PV phase) |
| 597 | |
| 598 | constexpr static int64_t BlockSizeAlignment = 64; |
| 599 | constexpr static int64_t HeadDimAlignment = 64; |
| 600 | constexpr static int64_t MaxQHeadNumPerIteration = 32; |
| 601 | constexpr static int64_t HeadDim = head_dim; |
| 602 | constexpr static ISA ISAType = ISA::AMX_FP8; |
| 603 | constexpr static bool scale_on_logits = true; |
| 604 | |
| 605 | float k_scale = 1.0f; |
| 606 | float v_scale = 1.0f; |
| 607 | // Dynamically computed per copy_q_heads_tile() call. |
| 608 | float q_dynamic_scale = 1.0f; |
| 609 | |
| 610 | public: |
| 611 | AttentionImpl() : current_q_head_num_(0) { |
| 612 | vec_op::unroll_loop<int, 8>([&](int i) { amx_tile_config_.colsb[i] = 64; }); |
| 613 | } |
| 614 | |
| 615 | ~AttentionImpl() { _tile_release(); } |
| 616 | |
| 617 | void init_from_input(const AttentionInput* input) { |
| 618 | k_scale = input->k_scale_fp8; |
| 619 | v_scale = input->v_scale_fp8; |
| 620 | // q_scale_fp8 from input is used as a hint; dynamic scale is recomputed |
| 621 | // per-tile in copy_q_heads_tile(). |
| 622 | } |
| 623 | |
| 624 | // Output scale for the PV result. |
| 625 | float get_output_v_scale() const noexcept { |
| 626 | constexpr float fp8_max = |
| 627 | std::is_same_v<kv_cache_t, c10::Float8_e4m3fn> ? 448.0f : 57344.0f; |
| 628 | return v_scale / fp8_max; |
| 629 | } |
| 630 | |
| 631 | template <template <typename tile_gemm_t> typename attention> |
| 632 | FORCE_INLINE void execute_attention(DEFINE_CPU_ATTENTION_PARAMS) { |
| 633 | // Apply QK scale correction: score_fp32 = dot(Q_fp8, K_fp8) * q_dyn * k |
| 634 | scale *= q_dynamic_scale * k_scale; |
| 635 | |
| 636 | if (q_head_num > AMX_TILE_ROW_NUM) { |
| 637 | if (q_head_num != current_q_head_num_) { |
| 638 | current_q_head_num_ = q_head_num; |
| 639 | TileGemm224<q_buffer_t, kv_cache_t>::init_tile_config(q_head_num, |
| 640 | amx_tile_config_); |
| 641 | } |
| 642 | attention<TileGemm224<q_buffer_t, kv_cache_t>> attention_iteration; |
| 643 | attention_iteration(CPU_ATTENTION_PARAMS); |
| 644 | } else { |
| 645 | if (q_head_num != current_q_head_num_) { |
| 646 | current_q_head_num_ = q_head_num; |
| 647 | TileGemm122<q_buffer_t, kv_cache_t>::init_tile_config(q_head_num, |
| 648 | amx_tile_config_); |
| 649 | } |
| 650 | attention<TileGemm122<q_buffer_t, kv_cache_t>> attention_iteration; |
| 651 | attention_iteration(CPU_ATTENTION_PARAMS); |
| 652 | } |
| 653 | } |
| 654 | |
| 655 | // K cache stride: same element-count formula as AMX_BF16. |
| 656 | // (BlockSizeAlignment * head_dim) uint8_t elements = same byte count as |
| 657 | // before since kv_cache_t is 1-byte FP8. |
| 658 | constexpr static int64_t k_cache_token_group_stride( |
| 659 | const int32_t /*block_size*/) { |
| 660 | return BlockSizeAlignment * head_dim; |
| 661 | } |
| 662 | |
| 663 | // Both native FP8 formats use VNNI4 packing. |
| 664 | constexpr static int64_t v_cache_token_group_stride( |
| 665 | const int32_t /*block_size*/) { |
| 666 | return BlockSizeAlignment * (AMX_TILE_ROW_BYTES / 4); |
| 667 | } |
| 668 | |
| 669 | constexpr static int64_t v_cache_head_group_stride(const int32_t block_size) { |
| 670 | return block_size * HeadDimAlignment; |
| 671 | } |
| 672 | |
| 673 | // Copy Q heads tile: quantize BF16 to the KV cache FP8 type and pack it. |
| 674 | // |
| 675 | // Q buffer layout (FP8): |
| 676 | // Same tile-prepacking as AMX_BF16 copy_q_heads_tile but with uint8_t |
| 677 | // (FP8) elements. Each FP8 tile row = 64 bytes = 64 FP8 elements. |
| 678 | // head_size_block_num = head_dim / 64 (vs head_dim/32 for BF16). |
| 679 | // |
| 680 | // Dynamic scale: max(abs(Q)) / fp8_max is computed over all elements, |
| 681 | // stored in q_dynamic_scale for use in execute_attention(). |
| 682 | void copy_q_heads_tile( |
| 683 | scalar_t* __restrict__ src, // [q_num, q_heads_per_kv, head_size] BF16 |
| 684 | q_buffer_t* __restrict__ q_buffer, const int32_t q_num, |
| 685 | const int32_t q_heads_per_kv, const int64_t q_num_stride, |
| 686 | const int64_t q_head_stride, const float /*scale*/) { |
| 687 | // ---- Step 1: compute dynamic per-tile Q scale ---- |
| 688 | constexpr float fp8_max = |
| 689 | std::is_same_v<kv_cache_t, c10::Float8_e4m3fn> ? 448.0f : 57344.0f; |
| 690 | float max_abs = 0.0f; |
| 691 | { |
| 692 | scalar_t* src_iter = src; |
| 693 | __m512 max_abs_vec = _mm512_setzero_ps(); |
| 694 | const __m512i abs_mask = _mm512_set1_epi32(0x7FFFFFFF); |
| 695 | for (int32_t q = 0; q < q_num; ++q, src_iter += q_num_stride) { |
| 696 | scalar_t* head_iter = src_iter; |
| 697 | for (int32_t h = 0; h < q_heads_per_kv; |
| 698 | ++h, head_iter += q_head_stride) { |
| 699 | for (int64_t e = 0; e < head_dim; e += 16) { |
| 700 | const __m256i bf16 = _mm256_loadu_si256( |
| 701 | reinterpret_cast<const __m256i*>(head_iter + e)); |
| 702 | const __m512i fp32_bits = |
| 703 | _mm512_slli_epi32(_mm512_cvtepu16_epi32(bf16), 16); |
| 704 | const __m512 abs_values = |
| 705 | _mm512_castsi512_ps(_mm512_and_si512(fp32_bits, abs_mask)); |
| 706 | max_abs_vec = _mm512_max_ps(max_abs_vec, abs_values); |
| 707 | } |
| 708 | } |
| 709 | } |
| 710 | max_abs = _mm512_reduce_max_ps(max_abs_vec); |
| 711 | } |
| 712 | if (max_abs == 0.0f) max_abs = 1.0f; // avoid div-by-zero |
| 713 | const float inv_scale = fp8_max / max_abs; |
| 714 | // Store true scale for score correction in execute_attention(). |
| 715 | q_dynamic_scale = max_abs / fp8_max; |
| 716 | |
| 717 | // ---- Step 2: quantize BF16 Q to FP8 and pack into q_buffer ---- |
| 718 | // Layout: same tile-prepacking as AMX_BF16, but with FP8 (1 byte/elem). |
| 719 | // head_size_block_num = head_dim * sizeof(uint8_t) / AMX_TILE_ROW_BYTES |
| 720 | // = head_dim / 64 |
| 721 | constexpr int64_t bytes_per_head = head_dim * sizeof(uint8_t); |
| 722 | static_assert(bytes_per_head % AMX_TILE_ROW_BYTES == 0, |
| 723 | "head_dim must be divisible by 64 for AMX_FP8"); |
| 724 | quantize_q_tile<kv_cache_t>(src, reinterpret_cast<uint8_t*>(q_buffer), |
| 725 | q_num, q_heads_per_kv, q_num_stride, |
| 726 | q_head_stride, head_dim, inv_scale); |
| 727 | } |
| 728 | |
| 729 | // reshape KV to AMX-FP8 friendly layout. |
| 730 | // K: native FP8 layout (1 byte/element, 16-token groups × k-slices). |
| 731 | // V: native VNNI4 packing. |
| 732 | static void reshape_and_cache( |
| 733 | const scalar_t* __restrict__ key, const scalar_t* __restrict__ value, |
| 734 | kv_cache_t* __restrict__ key_cache, kv_cache_t* __restrict__ value_cache, |
| 735 | const int64_t* __restrict__ slot_mapping, const int64_t token_num, |
| 736 | const int64_t key_token_num_stride, const int64_t value_token_num_stride, |
| 737 | const int64_t head_num, const int64_t key_head_num_stride, |
| 738 | const int64_t value_head_num_stride, const int64_t num_blocks, |
| 739 | const int64_t num_blocks_stride, const int64_t cache_head_num_stride, |
| 740 | const int64_t block_size, const int64_t /*block_size_stride*/, |
| 741 | const float k_inv = 0.0f, const float v_inv = 0.0f) { |
| 742 | constexpr auto qfn = select_fp8_quant_fn<kv_cache_t>(); |
| 743 | |
| 744 | // K: native AMX-FP8 layout |
| 745 | reshape_and_cache_k_amx_fp8_impl<scalar_t, qfn>( |
| 746 | key, reinterpret_cast<uint8_t*>(key_cache), slot_mapping, token_num, |
| 747 | head_num, head_dim, block_size, key_token_num_stride, |
| 748 | key_head_num_stride, num_blocks_stride, cache_head_num_stride, k_inv); |
| 749 | |
| 750 | { |
| 751 | constexpr int64_t token_num_per_sub_group = 4; |
| 752 | constexpr int64_t head_elems_per_group = 16; // AMX_TILE_ROW_BYTES/4 |
| 753 | const int64_t group_num = head_dim / head_elems_per_group; |
| 754 | const int64_t group_size = block_size * head_elems_per_group; |
| 755 | uint8_t* vc = reinterpret_cast<uint8_t*>(value_cache); |
| 756 | |
| 757 | #pragma omp parallel for collapse(2) schedule(static) |
| 758 | for (int64_t tok = 0; tok < token_num; ++tok) { |
| 759 | for (int64_t h = 0; h < head_num; ++h) { |
| 760 | const int64_t slot = slot_mapping[tok]; |
| 761 | if (slot < 0) continue; |
| 762 | const int64_t block_idx = slot / block_size; |
| 763 | const int64_t block_offset = slot % block_size; |
| 764 | const int64_t sub_group_idx = block_offset / token_num_per_sub_group; |
| 765 | const int64_t sub_group_offset = |
| 766 | block_offset % token_num_per_sub_group; |
| 767 | |
| 768 | const scalar_t* vsrc = |
| 769 | value + tok * value_token_num_stride + h * value_head_num_stride; |
| 770 | uint8_t* vdst = |
| 771 | vc + block_idx * num_blocks_stride + h * cache_head_num_stride + |
| 772 | sub_group_idx * token_num_per_sub_group * head_elems_per_group + |
| 773 | sub_group_offset; |
| 774 | |
| 775 | for (int64_t i = 0; i < group_num; ++i) { |
| 776 | for (int64_t j = 0; j < head_elems_per_group; ++j) |
| 777 | vdst[j * token_num_per_sub_group] = |
| 778 | qfn(static_cast<float>(vsrc[j]), v_inv); |
| 779 | vsrc += head_elems_per_group; |
| 780 | vdst += group_size; |
| 781 | } |
| 782 | } |
| 783 | } |
| 784 | } |
| 785 | } |
| 786 | |
| 787 | private: |
| 788 | alignas(64) __tilecfg amx_tile_config_; |
| 789 | int32_t current_q_head_num_; |
| 790 | }; |
| 791 | |
| 792 | } // namespace cpu_attention |
| 793 | |
| 794 | #endif // CPU_CAPABILITY_AMXFP8 |
| 795 | #endif // CPU_ATTN_AMX_FP8_HPP |