| 1 | #ifndef CPU_ATTN_AMX_HPP |
| 2 | #define CPU_ATTN_AMX_HPP |
| 3 | |
| 4 | #include "cpu_attn_fp8.hpp" |
| 5 | #include "cpu_attn_impl.hpp" |
| 6 | |
| 7 | namespace cpu_attention { |
| 8 | namespace { |
| 9 | // AMX specific |
| 10 | constexpr static int64_t AMX_TILE_ROW_BYTES = 64; |
| 11 | constexpr static int64_t AMX_TILE_ROW_NUM = 16; |
| 12 | constexpr static int64_t AMX_TILE_BYTES = AMX_TILE_ROW_BYTES * AMX_TILE_ROW_NUM; |
| 13 | |
| 14 | typedef struct __tile_config { |
| 15 | uint8_t palette_id = 1; |
| 16 | uint8_t start_row = 0; |
| 17 | uint8_t reserved_0[14] = {0}; |
| 18 | uint16_t colsb[16] = {0}; |
| 19 | uint8_t rows[16] = {0}; |
| 20 | } __tilecfg; |
| 21 | |
| 22 | // 2-2-4 pattern, for 16 < m <= 32 |
| 23 | // TILE 0, 1: load A matrix, row num should be 16, m - 16 |
| 24 | // TILE 2, 3: load B matrix, row num should be 16 |
| 25 | // TILE 4, 5, 6, 7: store results C matrix, row num should be 16, 16, |
| 26 | // m - 16, m - 16 |
| 27 | // q_buffer_t: A (Q/P) tile type; kv_cache_t: B (K/V cache) tile type. |
| 28 | template <typename q_buffer_t, typename kv_cache_t> |
| 29 | class TileGemm224 { |
| 30 | public: |
| 31 | template <AttentionGemmPhase phase, int32_t k_size> |
| 32 | FORCE_INLINE static void gemm(const int32_t m_size, void* __restrict__ a_tile, |
| 33 | void* __restrict__ b_tile, |
| 34 | float* __restrict__ c_tile, const int64_t lda, |
| 35 | const int64_t ldb, const int64_t ldc, |
| 36 | const int32_t block_size, |
| 37 | const int32_t dynamic_k_size, |
| 38 | const bool accum_c) { |
| 39 | TORCH_CHECK(false, "Unsupported kv cache type for TileGemm224"); |
| 40 | } |
| 41 | |
| 42 | FORCE_INLINE static void init_tile_config(int32_t m, __tilecfg& config) { |
| 43 | TORCH_CHECK(false, "Unsupported kv cache type for TileGemm224"); |
| 44 | } |
| 45 | }; |
| 46 | |
| 47 | // Dequantize one FP8 tile (AMX_TILE_ROW_NUM rows x 32 cols) to BF16. |
| 48 | template <typename kv_cache_t> |
| 49 | FORCE_INLINE void deq_tile_amx(const uint8_t* src, c10::BFloat16* dst) { |
| 50 | for (int r = 0; r < AMX_TILE_ROW_NUM; ++r) { |
| 51 | if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e4m3fn>) { |
| 52 | vec_op::BF16Vec32(src + r * 32, vec_op::fp8_bf16_e4m3_tag{}) |
| 53 | .save(dst + r * 32); |
| 54 | } else { |
| 55 | vec_op::BF16Vec32(src + r * 32, vec_op::fp8_bf16_e5m2_tag{}) |
| 56 | .save(dst + r * 32); |
| 57 | } |
| 58 | } |
| 59 | } |
| 60 | |
| 61 | // For FP8: dequant src into scratch and return scratch. |
| 62 | // For BF16: return src directly (scratch is unused; the compiler elides it). |
| 63 | template <typename kv_cache_t> |
| 64 | FORCE_INLINE const c10::BFloat16* prepare_b_tile(const kv_cache_t* src, |
| 65 | c10::BFloat16* scratch) { |
| 66 | if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e4m3fn> || |
| 67 | std::is_same_v<kv_cache_t, c10::Float8_e5m2>) { |
| 68 | deq_tile_amx<kv_cache_t>(reinterpret_cast<const uint8_t*>(src), scratch); |
| 69 | return scratch; |
| 70 | } else { |
| 71 | return reinterpret_cast<const c10::BFloat16*>(src); |
| 72 | } |
| 73 | } |
| 74 | |
| 75 | // Handles both BF16 and FP8 KV cache (2-2-4 pattern). |
| 76 | template <typename kv_cache_t> |
| 77 | class TileGemm224<c10::BFloat16, kv_cache_t> { |
| 78 | static_assert(std::is_same_v<kv_cache_t, c10::BFloat16> || |
| 79 | std::is_same_v<kv_cache_t, c10::Float8_e4m3fn> || |
| 80 | std::is_same_v<kv_cache_t, c10::Float8_e5m2>, |
| 81 | "kv_cache_t must be BFloat16, Float8_e4m3fn, or Float8_e5m2"); |
| 82 | |
| 83 | static constexpr bool fp8_kv = |
| 84 | std::is_same_v<kv_cache_t, c10::Float8_e4m3fn> || |
| 85 | std::is_same_v<kv_cache_t, c10::Float8_e5m2>; |
| 86 | |
| 87 | static constexpr int64_t tile_elems = AMX_TILE_BYTES / sizeof(c10::BFloat16); |
| 88 | // BF16 path: scratch_elems=1 so the scratch array is eliminated by the |
| 89 | // compiler. |
| 90 | static constexpr int64_t scratch_elems = fp8_kv ? tile_elems : 1; |
| 91 | |
| 92 | public: |
| 93 | template <AttentionGemmPhase phase, int32_t k_size> |
| 94 | FORCE_INLINE static void gemm(const int32_t m_size, |
| 95 | c10::BFloat16* __restrict__ a_tile, |
| 96 | kv_cache_t* __restrict__ b_tile, |
| 97 | float* __restrict__ c_tile, const int64_t lda, |
| 98 | const int64_t ldb, const int64_t ldc, |
| 99 | const int32_t block_size, |
| 100 | const int32_t dynamic_k_size, |
| 101 | const bool accum_c) { |
| 102 | const int32_t k_times = |
| 103 | dynamic_k_size / (AMX_TILE_ROW_NUM * 4 / sizeof(c10::BFloat16)); |
| 104 | |
| 105 | c10::BFloat16* __restrict__ a_tile_0 = a_tile; |
| 106 | c10::BFloat16* __restrict__ a_tile_1 = a_tile + lda * AMX_TILE_ROW_NUM; |
| 107 | const int64_t a_tile_stride = [&]() { |
| 108 | if constexpr (phase == AttentionGemmPhase::QK) { |
| 109 | // q_buffer is prepacked |
| 110 | return AMX_TILE_ROW_BYTES; |
| 111 | } else if constexpr (phase == AttentionGemmPhase::PV) { |
| 112 | // logits_buffer is row-major |
| 113 | return lda * sizeof(c10::BFloat16); |
| 114 | } else { |
| 115 | TORCH_CHECK(false, "Unreachable"); |
| 116 | } |
| 117 | }(); |
| 118 | |
| 119 | kv_cache_t* __restrict__ b_tile_2 = b_tile; |
| 120 | kv_cache_t* __restrict__ b_tile_3 = [&]() { |
| 121 | if constexpr (phase == AttentionGemmPhase::QK) { |
| 122 | // k_cache is prepacked |
| 123 | return b_tile + (k_size * AMX_TILE_ROW_BYTES / 4); |
| 124 | } else if constexpr (phase == AttentionGemmPhase::PV) { |
| 125 | // v_cache is prepacked |
| 126 | return b_tile + (block_size * AMX_TILE_ROW_BYTES / 4); |
| 127 | } else { |
| 128 | TORCH_CHECK(false, "Unreachable"); |
| 129 | } |
| 130 | }(); |
| 131 | // k_cache, v_cache are prepacked |
| 132 | const int32_t b_tile_stride = AMX_TILE_ROW_BYTES; |
| 133 | |
| 134 | // logits_buffer, output_buffer are not prepacked |
| 135 | float* __restrict__ c_tile_4 = c_tile; |
| 136 | float* __restrict__ c_tile_5 = |
| 137 | c_tile_4 + AMX_TILE_ROW_BYTES / sizeof(float); |
| 138 | float* __restrict__ c_tile_6 = c_tile + AMX_TILE_ROW_NUM * ldc; |
| 139 | float* __restrict__ c_tile_7 = |
| 140 | c_tile_6 + AMX_TILE_ROW_BYTES / sizeof(float); |
| 141 | const int32_t c_tile_stride = ldc * sizeof(float); |
| 142 | |
| 143 | if (accum_c) { |
| 144 | _tile_loadd(4, c_tile_4, c_tile_stride); |
| 145 | _tile_loadd(5, c_tile_5, c_tile_stride); |
| 146 | _tile_loadd(6, c_tile_6, c_tile_stride); |
| 147 | _tile_loadd(7, c_tile_7, c_tile_stride); |
| 148 | } else { |
| 149 | _tile_zero(4); |
| 150 | _tile_zero(5); |
| 151 | _tile_zero(6); |
| 152 | _tile_zero(7); |
| 153 | } |
| 154 | |
| 155 | alignas(64) c10::BFloat16 scratch_2[scratch_elems]; |
| 156 | alignas(64) c10::BFloat16 scratch_3[scratch_elems]; |
| 157 | for (int32_t k = 0; k < k_times; ++k) { |
| 158 | const c10::BFloat16* load_2 = prepare_b_tile(b_tile_2, scratch_2); |
| 159 | const c10::BFloat16* load_3 = prepare_b_tile(b_tile_3, scratch_3); |
| 160 | |
| 161 | _tile_loadd(0, a_tile_0, a_tile_stride); |
| 162 | _tile_stream_loadd(2, const_cast<c10::BFloat16*>(load_2), b_tile_stride); |
| 163 | _tile_dpbf16ps(4, 0, 2); |
| 164 | _tile_stream_loadd(3, const_cast<c10::BFloat16*>(load_3), b_tile_stride); |
| 165 | _tile_dpbf16ps(5, 0, 3); |
| 166 | _tile_loadd(1, a_tile_1, a_tile_stride); |
| 167 | _tile_dpbf16ps(6, 1, 2); |
| 168 | _tile_dpbf16ps(7, 1, 3); |
| 169 | |
| 170 | // update ptrs |
| 171 | if constexpr (phase == AttentionGemmPhase::QK) { |
| 172 | // Q buffer is prepacked |
| 173 | a_tile_0 += AMX_TILE_BYTES / sizeof(c10::BFloat16); |
| 174 | a_tile_1 += AMX_TILE_BYTES / sizeof(c10::BFloat16); |
| 175 | } else if constexpr (phase == AttentionGemmPhase::PV) { |
| 176 | // P buffer is not prepacked |
| 177 | a_tile_0 += AMX_TILE_ROW_BYTES / sizeof(c10::BFloat16); |
| 178 | a_tile_1 += AMX_TILE_ROW_BYTES / sizeof(c10::BFloat16); |
| 179 | } else { |
| 180 | TORCH_CHECK(false, "Unreachable"); |
| 181 | } |
| 182 | b_tile_2 += AMX_TILE_BYTES / sizeof(c10::BFloat16); |
| 183 | b_tile_3 += AMX_TILE_BYTES / sizeof(c10::BFloat16); |
| 184 | } |
| 185 | |
| 186 | _tile_stored(4, c_tile_4, c_tile_stride); |
| 187 | _tile_stored(5, c_tile_5, c_tile_stride); |
| 188 | _tile_stored(6, c_tile_6, c_tile_stride); |
| 189 | _tile_stored(7, c_tile_7, c_tile_stride); |
| 190 | } |
| 191 | |
| 192 | FORCE_INLINE static void init_tile_config(int32_t m, __tilecfg& config) { |
| 193 | const int32_t m_0 = AMX_TILE_ROW_NUM; |
| 194 | const int32_t m_1 = m - AMX_TILE_ROW_NUM; |
| 195 | config.rows[0] = m_0; |
| 196 | config.rows[1] = m_1; |
| 197 | config.rows[2] = AMX_TILE_ROW_NUM; |
| 198 | config.rows[3] = AMX_TILE_ROW_NUM; |
| 199 | config.rows[4] = m_0; |
| 200 | config.rows[5] = m_0; |
| 201 | config.rows[6] = m_1; |
| 202 | config.rows[7] = m_1; |
| 203 | _tile_loadconfig(&config); |
| 204 | } |
| 205 | }; |
| 206 | |
| 207 | // 1-2-2 pattern, for 0 < m <= 16 |
| 208 | // TILE 0, (1): load A matrix, use extra 1 tile for prefetch, row num should |
| 209 | // be m, m |
| 210 | // TILE 2, 3, (4, 5): load B matrix, use extra 2 tiles for prefetch, row num |
| 211 | // should be 16 |
| 212 | // TILE 6, 7: store results C matrix, row num should be m |
| 213 | // q_buffer_t: A (Q/P) tile type; kv_cache_t: B (K/V cache) tile type. |
| 214 | template <typename q_buffer_t, typename kv_cache_t> |
| 215 | class TileGemm122 { |
| 216 | public: |
| 217 | template <AttentionGemmPhase phase, int32_t k_size> |
| 218 | FORCE_INLINE static void gemm(const int32_t m_size, void* __restrict__ a_tile, |
| 219 | void* __restrict__ b_tile, |
| 220 | float* __restrict__ c_tile, const int64_t lda, |
| 221 | const int64_t ldb, const int64_t ldc, |
| 222 | const int32_t block_size, |
| 223 | const int32_t dynamic_k_size, |
| 224 | const bool accum_c) { |
| 225 | TORCH_CHECK(false, "Unsupported kv cache type for TileGemm122"); |
| 226 | } |
| 227 | |
| 228 | FORCE_INLINE static void init_tile_config(int32_t m, __tilecfg& config) { |
| 229 | TORCH_CHECK(false, "Unsupported kv cache type for TileGemm122"); |
| 230 | } |
| 231 | }; |
| 232 | |
| 233 | // Handles both BF16 and FP8 KV cache (1-2-2 pattern). |
| 234 | template <typename kv_cache_t> |
| 235 | class TileGemm122<c10::BFloat16, kv_cache_t> { |
| 236 | static_assert(std::is_same_v<kv_cache_t, c10::BFloat16> || |
| 237 | std::is_same_v<kv_cache_t, c10::Float8_e4m3fn> || |
| 238 | std::is_same_v<kv_cache_t, c10::Float8_e5m2>, |
| 239 | "kv_cache_t must be BFloat16, Float8_e4m3fn, or Float8_e5m2"); |
| 240 | |
| 241 | static constexpr bool fp8_kv = |
| 242 | std::is_same_v<kv_cache_t, c10::Float8_e4m3fn> || |
| 243 | std::is_same_v<kv_cache_t, c10::Float8_e5m2>; |
| 244 | |
| 245 | static constexpr int64_t tile_elems = AMX_TILE_BYTES / sizeof(c10::BFloat16); |
| 246 | static constexpr int64_t scratch_elems = fp8_kv ? tile_elems : 1; |
| 247 | |
| 248 | public: |
| 249 | template <AttentionGemmPhase phase, int32_t k_size> |
| 250 | FORCE_INLINE static void gemm(const int32_t m_size, |
| 251 | c10::BFloat16* __restrict__ a_tile, |
| 252 | kv_cache_t* __restrict__ b_tile, |
| 253 | float* __restrict__ c_tile, const int64_t lda, |
| 254 | const int64_t ldb, const int64_t ldc, |
| 255 | const int32_t block_size, |
| 256 | const int32_t dynamic_k_size, |
| 257 | const bool accum_c) { |
| 258 | c10::BFloat16* __restrict__ a_tile_0 = a_tile; |
| 259 | c10::BFloat16* __restrict__ a_tile_1 = [&]() { |
| 260 | if constexpr (phase == AttentionGemmPhase::QK) { |
| 261 | // q_buffer is prepacked |
| 262 | return a_tile + AMX_TILE_BYTES / sizeof(c10::BFloat16); |
| 263 | } else if constexpr (phase == AttentionGemmPhase::PV) { |
| 264 | // logits_buffer is row-major |
| 265 | return a_tile + AMX_TILE_ROW_BYTES / sizeof(c10::BFloat16); |
| 266 | } else { |
| 267 | TORCH_CHECK(false, "Unreachable"); |
| 268 | } |
| 269 | }(); |
| 270 | const int64_t a_tile_stride = [&]() { |
| 271 | if constexpr (phase == AttentionGemmPhase::QK) { |
| 272 | // q_buffer is prepacked |
| 273 | return AMX_TILE_ROW_BYTES; |
| 274 | } else if constexpr (phase == AttentionGemmPhase::PV) { |
| 275 | // logits_buffer is row-major |
| 276 | return lda * sizeof(c10::BFloat16); |
| 277 | } else { |
| 278 | TORCH_CHECK(false, "Unreachable"); |
| 279 | } |
| 280 | }(); |
| 281 | |
| 282 | kv_cache_t* __restrict__ b_tile_2 = b_tile; |
| 283 | kv_cache_t* __restrict__ b_tile_3 = [&]() { |
| 284 | if constexpr (phase == AttentionGemmPhase::QK) { |
| 285 | return b_tile + (k_size * AMX_TILE_ROW_BYTES / 4); |
| 286 | } else if constexpr (phase == AttentionGemmPhase::PV) { |
| 287 | return b_tile + (block_size * AMX_TILE_ROW_BYTES / 4); |
| 288 | } else { |
| 289 | TORCH_CHECK(false, "Unreachable"); |
| 290 | } |
| 291 | }(); |
| 292 | kv_cache_t* __restrict__ b_tile_4 = |
| 293 | b_tile_2 + AMX_TILE_BYTES / sizeof(c10::BFloat16); |
| 294 | kv_cache_t* __restrict__ b_tile_5 = |
| 295 | b_tile_3 + AMX_TILE_BYTES / sizeof(c10::BFloat16); |
| 296 | int64_t b_stride = AMX_TILE_ROW_BYTES; |
| 297 | |
| 298 | float* __restrict__ c_tile_6 = c_tile; |
| 299 | float* __restrict__ c_tile_7 = c_tile + AMX_TILE_ROW_BYTES / sizeof(float); |
| 300 | int64_t c_stride = ldc * sizeof(float); |
| 301 | |
| 302 | const int32_t k_times = |
| 303 | dynamic_k_size / (AMX_TILE_ROW_NUM * 4 / sizeof(c10::BFloat16)); |
| 304 | const int32_t k_group_times = k_times / 2; |
| 305 | const bool has_tail = (k_times % 2 == 1); |
| 306 | |
| 307 | if (accum_c) { |
| 308 | _tile_loadd(6, c_tile_6, c_stride); |
| 309 | _tile_loadd(7, c_tile_7, c_stride); |
| 310 | } else { |
| 311 | _tile_zero(6); |
| 312 | _tile_zero(7); |
| 313 | } |
| 314 | |
| 315 | alignas(64) c10::BFloat16 scratch_2[scratch_elems]; |
| 316 | alignas(64) c10::BFloat16 scratch_3[scratch_elems]; |
| 317 | alignas(64) c10::BFloat16 scratch_4[scratch_elems]; |
| 318 | alignas(64) c10::BFloat16 scratch_5[scratch_elems]; |
| 319 | for (int32_t k = 0; k < k_group_times; ++k) { |
| 320 | const c10::BFloat16* load_2 = prepare_b_tile(b_tile_2, scratch_2); |
| 321 | const c10::BFloat16* load_3 = prepare_b_tile(b_tile_3, scratch_3); |
| 322 | const c10::BFloat16* load_4 = prepare_b_tile(b_tile_4, scratch_4); |
| 323 | const c10::BFloat16* load_5 = prepare_b_tile(b_tile_5, scratch_5); |
| 324 | |
| 325 | _tile_loadd(0, a_tile_0, a_tile_stride); |
| 326 | _tile_stream_loadd(2, const_cast<c10::BFloat16*>(load_2), b_stride); |
| 327 | _tile_dpbf16ps(6, 0, 2); |
| 328 | _tile_stream_loadd(3, const_cast<c10::BFloat16*>(load_3), b_stride); |
| 329 | _tile_dpbf16ps(7, 0, 3); |
| 330 | _tile_loadd(1, a_tile_1, a_tile_stride); |
| 331 | _tile_stream_loadd(4, const_cast<c10::BFloat16*>(load_4), b_stride); |
| 332 | _tile_dpbf16ps(6, 1, 4); |
| 333 | _tile_stream_loadd(5, const_cast<c10::BFloat16*>(load_5), b_stride); |
| 334 | _tile_dpbf16ps(7, 1, 5); |
| 335 | |
| 336 | // update ptrs |
| 337 | if constexpr (phase == AttentionGemmPhase::QK) { |
| 338 | // Q buffer is prepacked |
| 339 | a_tile_0 += 2 * AMX_TILE_BYTES / sizeof(c10::BFloat16); |
| 340 | a_tile_1 += 2 * AMX_TILE_BYTES / sizeof(c10::BFloat16); |
| 341 | } else if constexpr (phase == AttentionGemmPhase::PV) { |
| 342 | // P buffer is not prepacked |
| 343 | a_tile_0 += 2 * AMX_TILE_ROW_BYTES / sizeof(c10::BFloat16); |
| 344 | a_tile_1 += 2 * AMX_TILE_ROW_BYTES / sizeof(c10::BFloat16); |
| 345 | } |
| 346 | b_tile_2 += 2 * AMX_TILE_BYTES / sizeof(c10::BFloat16); |
| 347 | b_tile_3 += 2 * AMX_TILE_BYTES / sizeof(c10::BFloat16); |
| 348 | b_tile_4 += 2 * AMX_TILE_BYTES / sizeof(c10::BFloat16); |
| 349 | b_tile_5 += 2 * AMX_TILE_BYTES / sizeof(c10::BFloat16); |
| 350 | } |
| 351 | |
| 352 | if (has_tail) { |
| 353 | const c10::BFloat16* load_2 = prepare_b_tile(b_tile_2, scratch_2); |
| 354 | const c10::BFloat16* load_3 = prepare_b_tile(b_tile_3, scratch_3); |
| 355 | |
| 356 | _tile_loadd(0, a_tile_0, a_tile_stride); |
| 357 | _tile_stream_loadd(2, const_cast<c10::BFloat16*>(load_2), b_stride); |
| 358 | _tile_dpbf16ps(6, 0, 2); |
| 359 | _tile_stream_loadd(3, const_cast<c10::BFloat16*>(load_3), b_stride); |
| 360 | _tile_dpbf16ps(7, 0, 3); |
| 361 | } |
| 362 | |
| 363 | _tile_stored(6, c_tile_6, c_stride); |
| 364 | _tile_stored(7, c_tile_7, c_stride); |
| 365 | } |
| 366 | |
| 367 | FORCE_INLINE static void init_tile_config(int32_t m, __tilecfg& config) { |
| 368 | config.rows[0] = m; |
| 369 | config.rows[1] = m; |
| 370 | config.rows[2] = AMX_TILE_ROW_NUM; |
| 371 | config.rows[3] = AMX_TILE_ROW_NUM; |
| 372 | config.rows[4] = AMX_TILE_ROW_NUM; |
| 373 | config.rows[5] = AMX_TILE_ROW_NUM; |
| 374 | config.rows[6] = m; |
| 375 | config.rows[7] = m; |
| 376 | _tile_loadconfig(&config); |
| 377 | } |
| 378 | }; |
| 379 | |
| 380 | } // namespace |
| 381 | |
| 382 | template <typename scalar_t, int64_t head_dim, typename kv_cache_scalar_t> |
| 383 | class AttentionImpl<ISA::AMX, scalar_t, head_dim, kv_cache_scalar_t> { |
| 384 | static constexpr bool fp8_kv = |
| 385 | std::is_same_v<kv_cache_scalar_t, c10::Float8_e4m3fn> || |
| 386 | std::is_same_v<kv_cache_scalar_t, c10::Float8_e5m2>; |
| 387 | |
| 388 | public: |
| 389 | using query_t = scalar_t; |
| 390 | using q_buffer_t = scalar_t; |
| 391 | using kv_cache_t = kv_cache_scalar_t; |
| 392 | using logits_buffer_t = float; |
| 393 | using partial_output_buffer_t = float; |
| 394 | using prob_buffer_t = scalar_t; |
| 395 | |
| 396 | constexpr static int64_t BlockSizeAlignment = |
| 397 | 32; // AMX_TILE_ROW_NUM = 16 tokens/tile; 32 = 2 tiles |
| 398 | constexpr static int64_t HeadDimAlignment = |
| 399 | 2 * (AMX_TILE_ROW_BYTES / 4); // headdim num unit of PV phase |
| 400 | constexpr static int64_t MaxQHeadNumPerIteration = 32; |
| 401 | constexpr static int64_t HeadDim = head_dim; |
| 402 | constexpr static ISA ISAType = ISA::AMX; |
| 403 | constexpr static bool scale_on_logits = true; |
| 404 | |
| 405 | float k_scale = 1.0f; |
| 406 | float v_scale = 1.0f; |
| 407 | |
| 408 | public: |
| 409 | AttentionImpl() : current_q_head_num_(0) { |
| 410 | // Use all columns in AMX tiles |
| 411 | vec_op::unroll_loop<int, 8>([&](int i) { amx_tile_config_.colsb[i] = 64; }); |
| 412 | } |
| 413 | |
| 414 | ~AttentionImpl() { _tile_release(); } |
| 415 | |
| 416 | void init_from_input(const AttentionInput* input) { |
| 417 | if constexpr (fp8_kv) { |
| 418 | k_scale = input->k_scale_fp8; |
| 419 | v_scale = input->v_scale_fp8; |
| 420 | } |
| 421 | } |
| 422 | |
| 423 | float get_output_v_scale() const noexcept { |
| 424 | if constexpr (fp8_kv) { |
| 425 | // AMX dequant places FP8 payload into a BF16 field (exponent bias 127). |
| 426 | // Correction = 2^(127 - FP8_bias): E4M3 bias=7 → 2^120, E5M2 bias=15 → |
| 427 | // 2^112. |
| 428 | constexpr float bias = |
| 429 | std::is_same_v<kv_cache_t, c10::Float8_e5m2> ? 0x1p112f : 0x1p120f; |
| 430 | return v_scale * bias; |
| 431 | } |
| 432 | return 1.0f; |
| 433 | } |
| 434 | |
| 435 | template <template <typename tile_gemm_t> typename attention> |
| 436 | FORCE_INLINE void execute_attention(DEFINE_CPU_ATTENTION_PARAMS) { |
| 437 | if constexpr (fp8_kv) { |
| 438 | // Same bias correction as get_output_v_scale: AMX FP8→BF16 dequant |
| 439 | // shifts the exponent bias from FP8 to BF16 (127), so we multiply by |
| 440 | // 2^(127-FP8_bias) to recover the true value. E4M3: 2^120, E5M2: 2^112. |
| 441 | const float bias = |
| 442 | std::is_same_v<kv_cache_t, c10::Float8_e5m2> ? 0x1p112f : 0x1p120f; |
| 443 | scale *= k_scale * bias; |
| 444 | } |
| 445 | if (q_head_num > AMX_TILE_ROW_NUM) { |
| 446 | if (q_head_num != current_q_head_num_) { |
| 447 | current_q_head_num_ = q_head_num; |
| 448 | TileGemm224<q_buffer_t, kv_cache_t>::init_tile_config(q_head_num, |
| 449 | amx_tile_config_); |
| 450 | } |
| 451 | attention<TileGemm224<q_buffer_t, kv_cache_t>> attention_iteration; |
| 452 | attention_iteration(CPU_ATTENTION_PARAMS); |
| 453 | } else { |
| 454 | if (q_head_num != current_q_head_num_) { |
| 455 | current_q_head_num_ = q_head_num; |
| 456 | TileGemm122<q_buffer_t, kv_cache_t>::init_tile_config(q_head_num, |
| 457 | amx_tile_config_); |
| 458 | } |
| 459 | attention<TileGemm122<q_buffer_t, kv_cache_t>> attention_iteration; |
| 460 | attention_iteration(CPU_ATTENTION_PARAMS); |
| 461 | } |
| 462 | } |
| 463 | |
| 464 | // k_cache_token_group_stride: stride of K cache when move to next |
| 465 | // BlockSizeAlignment tokens in a block |
| 466 | constexpr static int64_t k_cache_token_group_stride( |
| 467 | const int32_t block_size) { |
| 468 | return BlockSizeAlignment * head_dim; |
| 469 | } |
| 470 | |
| 471 | // v_cache_token_group_stride: stride of V cache when move to next |
| 472 | // BlockSizeAlignment tokens in a block |
| 473 | constexpr static int64_t v_cache_token_group_stride( |
| 474 | const int32_t block_size) { |
| 475 | return BlockSizeAlignment * (AMX_TILE_ROW_BYTES / 4); |
| 476 | } |
| 477 | |
| 478 | // v_cache_head_group_stride: stride of V cache when move to next |
| 479 | // HeadDimAlignment head dims in a block |
| 480 | constexpr static int64_t v_cache_head_group_stride(const int32_t block_size) { |
| 481 | return block_size * HeadDimAlignment; |
| 482 | } |
| 483 | |
| 484 | static void copy_q_heads_tile( |
| 485 | scalar_t* __restrict__ src, // [q_num, q_heads_per_kv, head_size] |
| 486 | scalar_t* __restrict__ q_buffer, const int32_t q_num, |
| 487 | const int32_t q_heads_per_kv, const int64_t q_num_stride, |
| 488 | const int64_t q_head_stride, const float scale) { |
| 489 | constexpr int64_t bytes_per_head = head_dim * sizeof(scalar_t); |
| 490 | static_assert(bytes_per_head % AMX_TILE_ROW_BYTES == 0); |
| 491 | constexpr int64_t head_size_block_num = bytes_per_head / AMX_TILE_ROW_BYTES; |
| 492 | constexpr int64_t head_elem_num_pre_block = |
| 493 | AMX_TILE_ROW_BYTES / sizeof(scalar_t); |
| 494 | |
| 495 | int32_t idx = 0; |
| 496 | int8_t* __restrict__ q_buffer_iter = reinterpret_cast<int8_t*>(q_buffer); |
| 497 | for (int32_t q_num_idx = 0; q_num_idx < q_num; |
| 498 | ++q_num_idx, src += q_num_stride) { |
| 499 | scalar_t* __restrict__ src_iter = src; |
| 500 | for (int32_t q_head_idx = 0; q_head_idx < q_heads_per_kv; |
| 501 | ++q_head_idx, src_iter += q_head_stride) { |
| 502 | vec_op::unroll_loop<int32_t, head_size_block_num>( |
| 503 | [&](int32_t head_size_block_idx) { |
| 504 | // Use INT8Vec64 for 64 bytes block |
| 505 | vec_op::INT8Vec64 vec(src_iter + head_size_block_idx * |
| 506 | head_elem_num_pre_block); |
| 507 | vec.save(q_buffer_iter + head_size_block_idx * AMX_TILE_BYTES); |
| 508 | }); |
| 509 | |
| 510 | ++idx; |
| 511 | q_buffer_iter += AMX_TILE_ROW_BYTES; |
| 512 | if ((idx & (AMX_TILE_ROW_NUM - 1)) == 0) { |
| 513 | // head is in another amx tile |
| 514 | q_buffer_iter -= AMX_TILE_ROW_NUM * AMX_TILE_ROW_BYTES; |
| 515 | q_buffer_iter += head_size_block_num * AMX_TILE_BYTES; |
| 516 | } |
| 517 | } |
| 518 | } |
| 519 | } |
| 520 | |
| 521 | // reshape KV to AMX friendly layout |
| 522 | static void reshape_and_cache( |
| 523 | const scalar_t* __restrict__ key, const scalar_t* __restrict__ value, |
| 524 | kv_cache_t* __restrict__ key_cache, kv_cache_t* __restrict__ value_cache, |
| 525 | const int64_t* __restrict__ slot_mapping, const int64_t token_num, |
| 526 | const int64_t key_token_num_stride, const int64_t value_token_num_stride, |
| 527 | const int64_t head_num, const int64_t key_head_num_stride, |
| 528 | const int64_t value_head_num_stride, const int64_t num_blocks, |
| 529 | const int64_t num_blocks_stride, const int64_t cache_head_num_stride, |
| 530 | const int64_t block_size, const int64_t block_size_stride, |
| 531 | const float k_inv = 0.0f, const float v_inv = 0.0f) { |
| 532 | if constexpr (fp8_kv) { |
| 533 | constexpr auto qfn = select_fp8_quant_fn<kv_cache_t>(); |
| 534 | reshape_and_cache_fp8_amx_impl<scalar_t, qfn>( |
| 535 | key, value, reinterpret_cast<uint8_t*>(key_cache), |
| 536 | reinterpret_cast<uint8_t*>(value_cache), slot_mapping, token_num, |
| 537 | head_num, head_dim, block_size, key_token_num_stride, |
| 538 | key_head_num_stride, value_token_num_stride, value_head_num_stride, |
| 539 | num_blocks_stride, cache_head_num_stride, num_blocks_stride, |
| 540 | cache_head_num_stride, k_inv, v_inv); |
| 541 | return; |
| 542 | } |
| 543 | |
| 544 | // For AMX 2D tiles, size of each line is 64 bytes |
| 545 | constexpr int64_t amx_tile_row_size = AMX_TILE_ROW_BYTES; |
| 546 | // For AMX B matrix, N always is 16 |
| 547 | constexpr int64_t amx_b_tile_n_size = AMX_TILE_ROW_BYTES / 4; |
| 548 | constexpr int64_t amx_b_tile_k_size = amx_tile_row_size / sizeof(scalar_t); |
| 549 | // For now suppose block_size is divisible by amx_tile_column_num |
| 550 | TORCH_CHECK_EQ(block_size % amx_b_tile_k_size, 0); |
| 551 | |
| 552 | scalar_t* __restrict__ kc = reinterpret_cast<scalar_t*>(key_cache); |
| 553 | scalar_t* __restrict__ vc = reinterpret_cast<scalar_t*>(value_cache); |
| 554 | |
| 555 | #pragma omp parallel for collapse(2) |
| 556 | for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) { |
| 557 | for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) { |
| 558 | const int64_t pos = slot_mapping[token_idx]; |
| 559 | if (pos < 0) { |
| 560 | // skip |
| 561 | continue; |
| 562 | } |
| 563 | |
| 564 | const int64_t block_idx = pos / block_size; |
| 565 | const int64_t block_offset = pos % block_size; |
| 566 | { |
| 567 | // Write Key |
| 568 | // Head elements should be packed as quand-words and stored in token |
| 569 | // groups with (quadword_stride/4) tokens |
| 570 | constexpr int64_t token_num_per_group = amx_tile_row_size / 4; |
| 571 | static_assert(head_dim % (4 / sizeof(scalar_t)) == 0); |
| 572 | constexpr int64_t quadword_num = head_dim / (4 / sizeof(scalar_t)); |
| 573 | const int32_t* key_start_quadword_ptr = |
| 574 | reinterpret_cast<const int32_t*>( |
| 575 | key + token_idx * key_token_num_stride + |
| 576 | head_idx * key_head_num_stride); |
| 577 | const int64_t group_idx = block_offset / token_num_per_group; |
| 578 | const int64_t group_offset = block_offset % token_num_per_group; |
| 579 | constexpr int64_t quadword_num_per_group = |
| 580 | token_num_per_group * quadword_num; |
| 581 | int32_t* key_cache_start_ptr = |
| 582 | reinterpret_cast<int32_t*>(kc + block_idx * num_blocks_stride + |
| 583 | head_idx * cache_head_num_stride) + |
| 584 | group_idx * quadword_num_per_group + group_offset; |
| 585 | |
| 586 | #pragma GCC unroll 8 |
| 587 | for (int64_t i = 0, j = 0; j < quadword_num; |
| 588 | i += token_num_per_group, ++j) { |
| 589 | key_cache_start_ptr[i] = key_start_quadword_ptr[j]; |
| 590 | } |
| 591 | } |
| 592 | { |
| 593 | // Write Value |
| 594 | // Different from Key, block_size dimension is packed rather than |
| 595 | // head_size dimension block_size dimension is packed as quand-words; |
| 596 | constexpr int64_t token_num_per_sub_group = 4 / sizeof(scalar_t); |
| 597 | const int64_t token_num_per_group = block_size; |
| 598 | constexpr int64_t head_elems_per_group = amx_b_tile_n_size; |
| 599 | const int64_t group_size = token_num_per_group * head_elems_per_group; |
| 600 | // For now suppose head_dim is divisible by amx_b_tile_n_size |
| 601 | static_assert(head_dim % head_elems_per_group == 0); |
| 602 | constexpr int64_t group_num = head_dim / head_elems_per_group; |
| 603 | const int64_t sub_group_idx = block_offset / token_num_per_sub_group; |
| 604 | const int64_t sub_group_offset = |
| 605 | block_offset % token_num_per_sub_group; |
| 606 | |
| 607 | const scalar_t* value_start_ptr = value + |
| 608 | token_idx * value_token_num_stride + |
| 609 | head_idx * value_head_num_stride; |
| 610 | scalar_t* value_cache_start_ptr = |
| 611 | vc + block_idx * num_blocks_stride + |
| 612 | head_idx * cache_head_num_stride + |
| 613 | sub_group_idx * token_num_per_sub_group * amx_b_tile_n_size + |
| 614 | sub_group_offset; |
| 615 | |
| 616 | for (int64_t i = 0; i < group_num; ++i) { |
| 617 | #pragma GCC unroll head_elems_per_group |
| 618 | for (int64_t j = 0, k = 0; j < head_elems_per_group; |
| 619 | ++j, k += token_num_per_sub_group) { |
| 620 | value_cache_start_ptr[k] = value_start_ptr[j]; |
| 621 | } |
| 622 | value_start_ptr += head_elems_per_group; |
| 623 | value_cache_start_ptr += group_size; |
| 624 | } |
| 625 | } |
| 626 | } |
| 627 | } |
| 628 | } |
| 629 | |
| 630 | private: |
| 631 | alignas(64) __tilecfg amx_tile_config_; |
| 632 | int32_t current_q_head_num_; |
| 633 | }; |
| 634 | } // namespace cpu_attention |
| 635 | |
| 636 | #endif |