diff --git a/.buildkite/scripts/hardware_ci/run-cpu-test.sh b/.buildkite/scripts/hardware_ci/run-cpu-test.sh index 43a83161f2..a27ddd78ff 100644 --- a/.buildkite/scripts/hardware_ci/run-cpu-test.sh +++ b/.buildkite/scripts/hardware_ci/run-cpu-test.sh @@ -57,7 +57,7 @@ prune_if_disk_pressure # omits `--pull` so it already prefers a cached image over the network. A # single attempt only -- the retry loop below covers a miss here too. echo "--- :docker: Pre-fetching base image" -docker pull ubuntu:22.04 || true +docker pull ubuntu:25.04 || true # building the docker image echo "--- :docker: Building Docker image" diff --git a/benchmarks/kernels/cpu/benchmark_cpu_attn.py b/benchmarks/kernels/cpu/benchmark_cpu_attn.py index cdfcb1d404..e6538c8850 100644 --- a/benchmarks/kernels/cpu/benchmark_cpu_attn.py +++ b/benchmarks/kernels/cpu/benchmark_cpu_attn.py @@ -13,13 +13,26 @@ from vllm._custom_ops import ( cpu_attn_reshape_and_cache, ) from vllm.utils.argparse_utils import FlexibleArgumentParser -from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, set_random_seed +from vllm.utils.torch_utils import ( + STR_DTYPE_TO_TORCH_DTYPE, + is_quantized_kv_cache, + set_random_seed, +) from vllm.v1.attention.backends.cpu_attn import CPUAttentionBackend, _get_attn_isa +# Enable AMX tile data registers so isolated runs don't rely on other operations +# to trigger oneDNN's _init_amx() first (which can cause hangs otherwise). +if torch.cpu._is_amx_tile_supported(): + torch.cpu._init_amx() + +KV_CACHE_DTYPE_CHOICES = ["auto", "fp8_e4m3", "fp8_e5m2"] + def get_attn_isa( block_size: int | None = None, dtype: torch.dtype | None = None, + head_size: int | None = None, + kv_cache_dtype: str = "auto", ): # Delegate to _get_attn_isa so the fallback path applies the same arch # gating (e.g. RISC-V RVV is only chosen when the build's hardcoded @@ -28,6 +41,8 @@ def get_attn_isa( return _get_attn_isa( dtype if dtype is not None else torch.bfloat16, block_size if block_size else 32, + head_size=head_size, + kv_cache_dtype=kv_cache_dtype if kv_cache_dtype != "auto" else None, ) @@ -53,6 +68,7 @@ def main( use_sink: bool = False, enable_kv_split: bool = False, isa: str | None = None, + kv_cache_dtype: str = "auto", seed: int = 0, iters: int = 20, ) -> None: @@ -68,8 +84,18 @@ def main( scale = head_size**-0.5 token_num = sum(query_lens) + # Resolve kv_cache_dtype: "auto" means same as compute dtype (no quantization) + effective_kv_cache_dtype = kv_cache_dtype if kv_cache_dtype != "auto" else None + is_fp8_kv = ( + is_quantized_kv_cache(effective_kv_cache_dtype) + if effective_kv_cache_dtype + else False + ) + # FP8 KV cache is stored as uint8 (byte-level view of fp8 values) + kv_cache_torch_dtype = torch.uint8 if is_fp8_kv else dtype + if isa is None: - isa = get_attn_isa(block_size, dtype) + isa = get_attn_isa(block_size, dtype, kv_cache_dtype) s_aux = ( 15 * torch.rand((num_query_heads,), dtype=torch.bfloat16) if use_sink else None @@ -98,9 +124,9 @@ def main( ) key_cache, value_cache = key_value.unbind(0) - # KV cache for CPU attention + # KV cache for CPU attention; FP8 variants store quantized values as uint8 packed_key_cache = torch.empty( - num_blocks, num_kv_heads, block_size, head_size, dtype=dtype + num_blocks, num_kv_heads, block_size, head_size, dtype=kv_cache_torch_dtype ) packed_value_cache = torch.empty_like(packed_key_cache) @@ -122,6 +148,7 @@ def main( value_cache=packed_value_cache, slot_mapping=slot_mapping, isa=isa, + kv_cache_dtype=kv_cache_dtype, ) metadata = cpu_attn_get_scheduler_metadata( @@ -159,6 +186,7 @@ def main( softcap=0, scheduler_metadata=metadata, s_aux=s_aux, + kv_cache_dtype=kv_cache_dtype, ) end_time = time.perf_counter_ns() times.append((end_time - start_time) / 1e6) @@ -233,7 +261,19 @@ if __name__ == "__main__": ) parser.add_argument("--use-sink", action="store_true") parser.add_argument( - "--isa", type=str, choices=["vec", "neon", "amx", "vec16", "rvv"], default=None + "--isa", + type=str, + choices=["vec", "neon", "amx", "amx_fp8", "vec16", "rvv"], + default=None, + ) + parser.add_argument( + "--kv-cache-dtype", + type=str, + choices=KV_CACHE_DTYPE_CHOICES, + default="auto", + help="KV cache dtype: 'auto' uses compute dtype (bfloat16), " + "'fp8_e4m3'/'fp8_e5m2' enables FP8 quantized KV cache " + "(requires AMX-FP8 capable hardware for amx_fp8 ISA).", ) parser.add_argument("--seed", type=int, default=0) parser.add_argument("--iters", type=int, default=20) @@ -264,7 +304,13 @@ if __name__ == "__main__": enable_kv_split=args.enable_kv_split, isa=args.isa if args.isa is not None - else get_attn_isa(args.block_size, STR_DTYPE_TO_TORCH_DTYPE[args.dtype]), + else get_attn_isa( + args.block_size, + STR_DTYPE_TO_TORCH_DTYPE[args.dtype], + args.head_size, + args.kv_cache_dtype, + ), + kv_cache_dtype=args.kv_cache_dtype, seed=args.seed, iters=args.iters, ) diff --git a/cmake/cpu_extension.cmake b/cmake/cpu_extension.cmake index 92c3daae64..35f56f242d 100644 --- a/cmake/cpu_extension.cmake +++ b/cmake/cpu_extension.cmake @@ -1,4 +1,5 @@ include(FetchContent) +include(CheckCXXCompilerFlag) set(CMAKE_CXX_STANDARD_REQUIRED ON) set(CMAKE_CXX_STANDARD 20) @@ -134,6 +135,12 @@ if (CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|amd64" OR ENABLE_X86_ISA) CMAKE_CXX_COMPILER_VERSION VERSION_GREATER_EQUAL 12.3)) message(FATAL_ERROR "X86 backend requires gcc/g++ >= 12.3") endif() + # Work around a GCC 15 optimizer crash in oneDNN (tree-ssa-pre pass). + # Keep this scoped to GCC 15+ so older toolchains are unaffected. + if (CMAKE_CXX_COMPILER_ID STREQUAL "GNU" AND + CMAKE_CXX_COMPILER_VERSION VERSION_GREATER_EQUAL 15) + list(APPEND CXX_COMPILE_FLAGS "-fno-tree-pre") + endif() list(APPEND CXX_COMPILE_FLAGS "-mf16c") list(APPEND CXX_COMPILE_FLAGS_AVX512 ${CXX_COMPILE_FLAGS}) list(APPEND CXX_COMPILE_FLAGS_AVX2 ${CXX_COMPILE_FLAGS}) @@ -148,6 +155,24 @@ if (CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|amd64" OR ENABLE_X86_ISA) "-mamx-tile" "-mavx512bf16" "-mavx512vnni") + # Some packaged GCC 14 builds still do not accept -mamx-fp8, so gate on + # actual flag support rather than compiler version alone. + # Option to enable/disable AMX-FP8 support (default: ON) + option(ENABLE_AMX_FP8 "Enable AMX-FP8 support" ON) + + if (ENABLE_AMX_FP8) + check_cxx_compiler_flag("-mamx-fp8" COMPILER_SUPPORTS_AMX_FP8_FLAG) + if (COMPILER_SUPPORTS_AMX_FP8_FLAG) + list(APPEND CXX_COMPILE_FLAGS_AVX512_AMX "-mamx-fp8") + set(AMX_FP8_SUPPORTED TRUE) + else() + set(AMX_FP8_SUPPORTED FALSE) + message(STATUS "AMX-FP8 disabled: compiler does not support -mamx-fp8 (compiler: ${CMAKE_CXX_COMPILER_ID} ${CMAKE_CXX_COMPILER_VERSION})") + endif() + else() + set(AMX_FP8_SUPPORTED FALSE) + message(STATUS "AMX-FP8 disabled by ENABLE_AMX_FP8=OFF") + endif() list(APPEND CXX_COMPILE_FLAGS_AVX2 "-mavx2") elseif (POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND) @@ -581,6 +606,10 @@ if (ENABLE_X86_ISA) # For AMX kernels target_compile_definitions(_C PRIVATE "-DCPU_CAPABILITY_AMXBF16") + if (AMX_FP8_SUPPORTED) + target_compile_definitions(_C PRIVATE "-DCPU_CAPABILITY_AMXFP8") + message(STATUS "AMX-FP8 (Diamond Rapids) enabled") + endif() # AVX512F define_extension_target( diff --git a/csrc/cpu/cpu_attn.cpp b/csrc/cpu/cpu_attn.cpp index 52c806f4fa..b31a6faa01 100644 --- a/csrc/cpu/cpu_attn.cpp +++ b/csrc/cpu/cpu_attn.cpp @@ -1,5 +1,17 @@ #include "cpu_attn_dispatch_generated.h" +#if defined(__x86_64__) || defined(_M_X64) + #include +// Runtime check for AMX-FP8 (Diamond Rapids): CPUID leaf 7, subleaf 1, EAX[21]. +static bool runtime_has_amx_fp8() { + unsigned int eax = 0, ebx = 0, ecx = 0, edx = 0; + if (__get_cpuid_count(7, 1, &eax, &ebx, &ecx, &edx)) { + return (eax >> 21) & 1u; + } + return false; +} +#endif // x86_64 + // Maps kv_cache_dtype string to Fp8KVCacheDataType enum. // "auto" -> kAuto(0); "fp8"/"fp8_e4m3" -> kFp8E4M3; "fp8_e5m2" -> kFp8E5M2. static inline cpu_attention::Fp8KVCacheDataType parse_fp8_kv_dtype( @@ -18,6 +30,16 @@ bool cpu_attn_has_isa(const std::string& isa) { return true; #else return false; +#endif + } + if (isa == "amx_fp8") { +#ifdef CPU_CAPABILITY_AMXFP8 + // Guard with runtime CPUID check: the binary may be compiled with + // -mamx-fp8 (CPU_CAPABILITY_AMXFP8 defined) but run on a CPU that + // has AMX-BF16 without AMX-FP8 (e.g. Sapphire Rapids, Emerald Rapids). + return runtime_has_amx_fp8(); +#else + return false; #endif } return false; @@ -35,6 +57,8 @@ torch::Tensor get_scheduler_metadata( cpu_attention::ISA isa; if (isa_hint == "amx") { isa = cpu_attention::ISA::AMX; + } else if (isa_hint == "amx_fp8") { + isa = cpu_attention::ISA::AMX_FP8; } else if (isa_hint == "vec") { isa = cpu_attention::ISA::VEC; } else if (isa_hint == "vec16") { @@ -130,6 +154,8 @@ void cpu_attn_reshape_and_cache( cpu_attention::ISA isa_tag = [&]() { if (isa == "amx") { return cpu_attention::ISA::AMX; + } else if (isa == "amx_fp8") { + return cpu_attention::ISA::AMX_FP8; } else if (isa == "vec") { return cpu_attention::ISA::VEC; } else if (isa == "vec16") { @@ -149,8 +175,9 @@ void cpu_attn_reshape_and_cache( if (is_fp8) { TORCH_CHECK(isa_tag == cpu_attention::ISA::AMX || + isa_tag == cpu_attention::ISA::AMX_FP8 || isa_tag == cpu_attention::ISA::VEC, - "FP8 KV cache is only supported on x86 (AMX/VEC) ISA"); + "FP8 KV cache is only supported on x86 (AMX_FP8/AMX/VEC) ISA"); } VLLM_DISPATCH_FLOATING_TYPES( @@ -237,8 +264,9 @@ void cpu_attention_with_kv_cache( input.k_scale_fp8 = static_cast(k_scale); input.v_scale_fp8 = static_cast(v_scale); TORCH_CHECK(input.metadata->isa == cpu_attention::ISA::AMX || + input.metadata->isa == cpu_attention::ISA::AMX_FP8 || input.metadata->isa == cpu_attention::ISA::VEC, - "FP8 KV cache is only supported on x86 (AMX/VEC) ISA"); + "FP8 KV cache is only supported on x86 (AMX_FP8/AMX/VEC) ISA"); } VLLM_DISPATCH_FLOATING_TYPES( diff --git a/csrc/cpu/cpu_attn_amx_fp8.hpp b/csrc/cpu/cpu_attn_amx_fp8.hpp new file mode 100644 index 0000000000..df606c330a --- /dev/null +++ b/csrc/cpu/cpu_attn_amx_fp8.hpp @@ -0,0 +1,795 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// AMX_FP8 attention implementation for Intel Diamond Rapids (DMR). +// +// Architecture: +// QK phase: native FP8×FP8 MMA via _tile_dpfp8ps / _tile_dpbf8ps +// Q is quantized from BF16 to the KV cache FP8 type. +// K cache is stored in the native AMX-FP8 layout (1 byte/element). +// PV phase: native FP8 MMA with 64-byte VNNI4 tiles. +// Softmax writes P once in compact FP8 form for reuse across all +// output-dimension groups. +// +// Scale bookkeeping for QK: +// true_score = Q_fp32 · K_fp32 +// = (Q_fp8 * q_scale) · (K_fp8 * k_scale) +// = dot(Q_fp8, K_fp8) * q_scale * k_scale +// execute_attention() applies: scale *= q_dynamic_scale * k_scale +// where q_dynamic_scale is computed per-tile in copy_q_heads_tile(). +// +// Scale bookkeeping for PV: +// raw_C = dot(P_fp8, V_fp8) +// ≈ dot(P/p_scale, V_real/v_scale) +// = dot(P, V_real) / (p_scale * v_scale) +// In final_output: multiply by output_v_scale = p_scale * v_scale. +// +#ifndef CPU_ATTN_AMX_FP8_HPP +#define CPU_ATTN_AMX_FP8_HPP + +#ifdef CPU_CAPABILITY_AMXFP8 + + #include "cpu_attn_amx.hpp" // pulls in TileGemm224/TileGemm122 BF16 specialisations + #include "cpu_attn_fp8.hpp" + #include "cpu_attn_impl.hpp" + #include "cpu_types.hpp" + + #include // requires GCC 14+ / Clang 17+, -mamx-fp8 + +namespace cpu_attention { + +// --------------------------------------------------------------------------- +// Native-FP8 K cache layout reshape helper. +// +// Layout per 16-token group (token_num_per_group = AMX_TILE_ROW_NUM = 16): +// k-slice s (s = 0 .. head_dim/64 - 1): +// row k (k = 0..15): four consecutive K values for each of 16 tokens +// (VNNI4 B-operand layout) +// One k-slice = 16 rows × 64 bytes = 1024 bytes = 1 AMX FP8 tile. +// Full group = (head_dim / 64) slices × 1024 bytes. +// +// For head_dim=128: group = 2 × 1024 = 2048 bytes. +// Consecutive 16-token groups are contiguous within each cache block. +// --------------------------------------------------------------------------- +template +inline void reshape_and_cache_k_amx_fp8_impl( + const scalar_t* key_ptr, uint8_t* key_cache_ptr, const int64_t* slot_ptr, + int64_t token_num, int64_t head_num, int64_t head_dim, int64_t block_size, + int64_t k_stride0, int64_t k_stride1, int64_t kc_stride0, + int64_t kc_stride1, float k_inv) { + constexpr int64_t token_num_per_group = AMX_TILE_ROW_NUM; // 16 + constexpr int64_t fp8_per_row = AMX_TILE_ROW_BYTES; // 64 FP8/row + const int64_t slice_num = head_dim / fp8_per_row; + // Bytes per 16-token group = slice_num * (16 * 64) = head_dim * 16 + const int64_t group_bytes = token_num_per_group * head_dim; + + #pragma omp parallel for collapse(2) schedule(static) + for (int64_t tok = 0; tok < token_num; ++tok) { + for (int64_t h = 0; h < head_num; ++h) { + const int64_t slot = slot_ptr[tok]; + if (slot < 0) continue; + const int64_t block_idx = slot / block_size; + const int64_t block_offset = slot % block_size; + const int64_t group_idx = block_offset / token_num_per_group; + const int64_t group_offset = block_offset % token_num_per_group; + + const scalar_t* ksrc = key_ptr + tok * k_stride0 + h * k_stride1; + uint8_t* group_base = key_cache_ptr + block_idx * kc_stride0 + + h * kc_stride1 + group_idx * group_bytes; + + // Each k-slice is one 16x64-byte VNNI4 B tile. Row k stores dimensions + // [4*k, 4*k+4) for every token; each token occupies four bytes. + for (int64_t s = 0; s < slice_num; ++s) { + uint8_t* tile_base = group_base + s * token_num_per_group * fp8_per_row; + for (int64_t k = 0; k < fp8_per_row / 4; ++k) { + uint8_t* packed_values = + tile_base + k * fp8_per_row + group_offset * 4; + for (int64_t lane = 0; lane < 4; ++lane) { + packed_values[lane] = quant_fn( + static_cast(ksrc[s * fp8_per_row + k * 4 + lane]), + k_inv); + } + } + } + } + } +} + +template +__attribute__((target("avx10.2"), noinline)) void quantize_probability_row( + const float* __restrict__ src, uint8_t* __restrict__ dst, int32_t cols) { + constexpr float fp8_max = + std::is_same_v ? 448.0f : 57344.0f; + const __m512 scale = _mm512_set1_ps(fp8_max); + for (int32_t col = 0; col < cols; col += 16) { + const __m512 scaled = _mm512_mul_ps(_mm512_loadu_ps(src + col), scale); + const __m256h fp16 = (__m256h)_mm512_cvtps_ph( + scaled, _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC); + if constexpr (std::is_same_v) { + _mm_storeu_si128(reinterpret_cast<__m128i*>(dst + col), + _mm256_cvtph_hf8(fp16)); + } else { + _mm_storeu_si128(reinterpret_cast<__m128i*>(dst + col), + _mm256_cvtph_bf8(fp16)); + } + } +} + +template +__attribute__((target("avx10.2"), noinline)) void quantize_q_tile( + const c10::BFloat16* __restrict__ src, uint8_t* __restrict__ dst, + int32_t q_num, int32_t q_heads_per_kv, int64_t q_num_stride, + int64_t q_head_stride, int64_t head_dim, float inv_scale) { + const __m512 scale = _mm512_set1_ps(inv_scale); + constexpr float fp8_max = + std::is_same_v ? 448.0f : 57344.0f; + const __m512 max_value = _mm512_set1_ps(fp8_max); + const __m512 min_value = _mm512_set1_ps(-fp8_max); + const int64_t head_size_block_num = head_dim / AMX_TILE_ROW_BYTES; + int32_t packed_row = 0; + + for (int32_t q = 0; q < q_num; ++q) { + const c10::BFloat16* head = src + q * q_num_stride; + for (int32_t h = 0; h < q_heads_per_kv; + ++h, head += q_head_stride, ++packed_row) { + uint8_t* packed_head = + dst + + (packed_row / AMX_TILE_ROW_NUM) * head_size_block_num * + AMX_TILE_BYTES + + (packed_row % AMX_TILE_ROW_NUM) * AMX_TILE_ROW_BYTES; + for (int64_t e = 0; e < head_dim; e += 16) { + const __m256i bf16 = + _mm256_loadu_si256(reinterpret_cast(head + e)); + __m512 values = _mm512_castsi512_ps( + _mm512_slli_epi32(_mm512_cvtepu16_epi32(bf16), 16)); + values = _mm512_mul_ps(values, scale); + values = _mm512_min_ps(values, max_value); + values = _mm512_max_ps(values, min_value); + const __m256h fp16 = (__m256h)_mm512_cvtps_ph( + values, _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC); + const int64_t block = e / AMX_TILE_ROW_BYTES; + const int64_t offset = e % AMX_TILE_ROW_BYTES; + __m128i fp8; + if constexpr (std::is_same_v) { + fp8 = _mm256_cvtph_hf8(fp16); + } else { + fp8 = _mm256_cvtph_bf8(fp16); + } + _mm_storeu_si128(reinterpret_cast<__m128i*>( + packed_head + block * AMX_TILE_BYTES + offset), + fp8); + } + } + } +} + +// --------------------------------------------------------------------------- +// --------------------------------------------------------------------------- +// TileGemm224: 2-2-4 pattern for AMX_FP8. +// +// QK phase: native FP8 × FP8 MMA (_tile_dphf8ps / _tile_dpbf8ps). +// PV phase: native wide FP8 MMA (colsb=64 tiles). +// P is quantized once per KV tile and reused by every output-dimension group. +// V (FP8 E4M3) loaded directly from cache (no dequant). +// Result correction is folded into the final output scale. +// --------------------------------------------------------------------------- +template +class TileGemm224 { + static_assert(std::is_same_v || + std::is_same_v, + "kv_cache_t must be Float8_e4m3fn or Float8_e5m2"); + + // FP8 tile size in uint8_t elements (= AMX_TILE_BYTES, one per byte). + static constexpr int64_t fp8_tile_elems = AMX_TILE_BYTES; // 1024 + + public: + static constexpr bool prequantize_probabilities = true; + + FORCE_INLINE static void quantize_probability_row(const float* src, + uint8_t* dst, + int32_t cols) { + cpu_attention::quantize_probability_row(src, dst, cols); + } + + // ------------------------------------------------------------------ + // QK: k_times = head_dim / 64 (each FP8 tile covers 64 k-elements). + // PV: native wide FP8 MMA. + // ------------------------------------------------------------------ + template + FORCE_INLINE static void gemm(const int32_t m_size, void* __restrict__ a_tile, + void* __restrict__ b_tile, + float* __restrict__ c_tile, const int64_t lda, + const int64_t ldb, const int64_t ldc, + const int32_t block_size, + const int32_t dynamic_k_size, + const bool accum_c) { + if constexpr (phase == AttentionGemmPhase::PV) { + // -------- Native wide FP8 PV (FP8 × FP8 → FP32) -------- + // + // A fixed FP8-dependent P scale lets the tile accumulate all k-steps + // before get_output_v_scale() applies the correction. + + const int32_t m_0 = + std::min(m_size, static_cast(AMX_TILE_ROW_NUM)); + const int32_t m_1 = m_size - m_0; + + constexpr int32_t v_stride = AMX_TILE_ROW_BYTES; + const int32_t k_times = + dynamic_k_size / (AMX_TILE_ROW_NUM * 4 / sizeof(uint8_t)); + constexpr int32_t a_advance = AMX_TILE_ROW_BYTES; + constexpr int32_t v_advance = AMX_TILE_BYTES; + const int32_t c_stride = ldc * sizeof(float); + + for (int32_t output_group = 0; output_group < 2; ++output_group) { + uint8_t* a0 = reinterpret_cast(a_tile); + uint8_t* a1 = a0 + lda * AMX_TILE_ROW_NUM; + uint8_t* v2 = reinterpret_cast(b_tile) + + output_group * 2 * block_size * AMX_TILE_ROW_BYTES / 4; + uint8_t* v3 = v2 + block_size * AMX_TILE_ROW_BYTES / 4; + float* c4 = c_tile + output_group * 2 * AMX_TILE_ROW_NUM; + float* c5 = c4 + AMX_TILE_ROW_NUM; + float* c6 = c4 + AMX_TILE_ROW_NUM * ldc; + float* c7 = c6 + AMX_TILE_ROW_NUM; + + if (accum_c) { + _tile_loadd(4, c4, c_stride); + _tile_loadd(5, c5, c_stride); + if (m_1 > 0) { + _tile_loadd(6, c6, c_stride); + _tile_loadd(7, c7, c_stride); + } + } else { + _tile_zero(4); + _tile_zero(5); + if (m_1 > 0) { + _tile_zero(6); + _tile_zero(7); + } + } + + for (int32_t k = 0; k < k_times; ++k) { + _tile_loadd(0, a0, lda); + _tile_stream_loadd(2, v2, v_stride); + _tile_stream_loadd(3, v3, v_stride); + if constexpr (std::is_same_v) { + _tile_dphf8ps(4, 0, 2); + _tile_dphf8ps(5, 0, 3); + } else { + _tile_dpbf8ps(4, 0, 2); + _tile_dpbf8ps(5, 0, 3); + } + if (m_1 > 0) { + _tile_loadd(1, a1, lda); + if constexpr (std::is_same_v) { + _tile_dphf8ps(6, 1, 2); + _tile_dphf8ps(7, 1, 3); + } else { + _tile_dpbf8ps(6, 1, 2); + _tile_dpbf8ps(7, 1, 3); + } + } + a0 += a_advance; + a1 += a_advance; + v2 += v_advance; + v3 += v_advance; + } + + _tile_stored(4, c4, c_stride); + _tile_stored(5, c5, c_stride); + if (m_1 > 0) { + _tile_stored(6, c6, c_stride); + _tile_stored(7, c7, c_stride); + } + } + return; + } + + // -------- QK phase: native FP8 MMA -------- + // k_times: each FP8 AMX tile covers (AMX_TILE_ROW_NUM * 4) FP8 = 64 + // inner-product elements. + const int32_t k_times = + dynamic_k_size / (AMX_TILE_ROW_NUM * 4 / sizeof(uint8_t)); + + const int32_t b_tile_stride = AMX_TILE_ROW_BYTES; // 64 bytes + + const int32_t m_0 = + std::min(m_size, static_cast(AMX_TILE_ROW_NUM)); + const int32_t m_1 = m_size - m_0; + const int32_t c_tile_stride = ldc * sizeof(float); + constexpr int32_t output_group_num = 2; + for (int32_t output_group = 0; output_group < output_group_num; + ++output_group) { + uint8_t* __restrict__ a_tile_0 = static_cast(a_tile); + uint8_t* __restrict__ a_tile_1 = a_tile_0 + k_times * fp8_tile_elems; + uint8_t* __restrict__ b_tile_2 = + static_cast(b_tile) + + output_group * 2 * k_size * AMX_TILE_ROW_NUM; + uint8_t* __restrict__ b_tile_3 = b_tile_2 + (k_size * AMX_TILE_ROW_NUM); + float* __restrict__ c_tile_4 = c_tile + output_group * 32; + float* __restrict__ c_tile_5 = c_tile_4 + AMX_TILE_ROW_NUM; + float* __restrict__ c_tile_6 = c_tile_4 + AMX_TILE_ROW_NUM * ldc; + float* __restrict__ c_tile_7 = c_tile_6 + AMX_TILE_ROW_NUM; + + if (accum_c) { + _tile_loadd(4, c_tile_4, c_tile_stride); + _tile_loadd(5, c_tile_5, c_tile_stride); + _tile_loadd(6, c_tile_6, c_tile_stride); + _tile_loadd(7, c_tile_7, c_tile_stride); + } else { + _tile_zero(4); + _tile_zero(5); + _tile_zero(6); + _tile_zero(7); + } + + for (int32_t k = 0; k < k_times; ++k) { + _tile_loadd(0, a_tile_0, AMX_TILE_ROW_BYTES); + _tile_stream_loadd(2, b_tile_2, b_tile_stride); + if constexpr (std::is_same_v) { + _tile_dphf8ps(4, 0, 2); + } else { + _tile_dpbf8ps(4, 0, 2); + } + _tile_stream_loadd(3, b_tile_3, b_tile_stride); + if constexpr (std::is_same_v) { + _tile_dphf8ps(5, 0, 3); + } else { + _tile_dpbf8ps(5, 0, 3); + } + _tile_loadd(1, a_tile_1, AMX_TILE_ROW_BYTES); + if constexpr (std::is_same_v) { + _tile_dphf8ps(6, 1, 2); + _tile_dphf8ps(7, 1, 3); + } else { + _tile_dpbf8ps(6, 1, 2); + _tile_dpbf8ps(7, 1, 3); + } + + // Q is packed as [row group][K slice]. + a_tile_0 += fp8_tile_elems; + a_tile_1 += fp8_tile_elems; + // Advance K cache (native layout): 1 FP8 tile per k-step. + b_tile_2 += fp8_tile_elems; + b_tile_3 += fp8_tile_elems; + } + + _tile_stored(4, c_tile_4, c_tile_stride); + _tile_stored(5, c_tile_5, c_tile_stride); + _tile_stored(6, c_tile_6, c_tile_stride); + _tile_stored(7, c_tile_7, c_tile_stride); + } + } + + FORCE_INLINE static void init_tile_config(int32_t m, __tilecfg& config) { + // QK and PV share one 64-byte tile configuration. + const int32_t m_0 = std::min(m, static_cast(AMX_TILE_ROW_NUM)); + const int32_t m_1 = m - m_0; + // Use all 64 bytes per row for every tile (same as BF16 AMX). + // This works for both FP8 (64 FP8/row) and BF16 (32 BF16/row = 64B). + for (int i = 0; i < 8; ++i) config.colsb[i] = AMX_TILE_ROW_BYTES; + config.rows[0] = m_0; + config.rows[1] = m_1; + config.rows[2] = AMX_TILE_ROW_NUM; + config.rows[3] = AMX_TILE_ROW_NUM; + config.rows[4] = m_0; + config.rows[5] = m_0; + config.rows[6] = m_1; + config.rows[7] = m_1; + _tile_loadconfig(&config); + } +}; + +// --------------------------------------------------------------------------- +// TileGemm122: 1-2-2 pattern for AMX_FP8. +// --------------------------------------------------------------------------- +template +class TileGemm122 { + static_assert(std::is_same_v || + std::is_same_v, + "kv_cache_t must be Float8_e4m3fn or Float8_e5m2"); + + static constexpr int64_t fp8_tile_elems = AMX_TILE_BYTES; + + public: + static constexpr bool prequantize_probabilities = true; + + FORCE_INLINE static void quantize_probability_row(const float* src, + uint8_t* dst, + int32_t cols) { + cpu_attention::quantize_probability_row(src, dst, cols); + } + + template + FORCE_INLINE static void gemm(const int32_t m_size, void* __restrict__ a_tile, + void* __restrict__ b_tile, + float* __restrict__ c_tile, const int64_t lda, + const int64_t ldb, const int64_t ldc, + const int32_t block_size, + const int32_t dynamic_k_size, + const bool accum_c) { + if constexpr (phase == AttentionGemmPhase::PV) { + // Native wide FP8 PV (1-2-2 pattern) + // A fixed FP8-dependent P scale enables accumulation across all k-steps. + constexpr int32_t v_stride = AMX_TILE_ROW_BYTES; + constexpr int32_t a_advance = AMX_TILE_ROW_BYTES; + constexpr int32_t v_advance = AMX_TILE_BYTES; + const int32_t k_times = + dynamic_k_size / (AMX_TILE_ROW_NUM * 4 / sizeof(uint8_t)); + + uint8_t* a = reinterpret_cast(a_tile); + uint8_t* v2 = reinterpret_cast(b_tile); + uint8_t* v3 = v2 + (block_size * AMX_TILE_ROW_BYTES / 4); + uint8_t* v4 = v3 + (block_size * AMX_TILE_ROW_BYTES / 4); + uint8_t* v5 = v4 + (block_size * AMX_TILE_ROW_BYTES / 4); + float* c4 = c_tile + 2 * AMX_TILE_ROW_NUM; + float* c5 = c_tile + 3 * AMX_TILE_ROW_NUM; + float* c6 = c_tile; + float* c7 = c6 + AMX_TILE_ROW_BYTES / sizeof(float); + + const int32_t c_stride = ldc * sizeof(float); + if (accum_c) { + _tile_loadd(4, c4, c_stride); + _tile_loadd(5, c5, c_stride); + _tile_loadd(6, c6, c_stride); + _tile_loadd(7, c7, c_stride); + } else { + _tile_zero(4); + _tile_zero(5); + _tile_zero(6); + _tile_zero(7); + } + + for (int32_t k = 0; k < k_times; ++k) { + _tile_loadd(0, a, lda); + _tile_stream_loadd(2, v2, v_stride); + _tile_stream_loadd(3, v3, v_stride); + if constexpr (std::is_same_v) { + _tile_dphf8ps(6, 0, 2); + _tile_dphf8ps(7, 0, 3); + } else { + _tile_dpbf8ps(6, 0, 2); + _tile_dpbf8ps(7, 0, 3); + } + _tile_stream_loadd(2, v4, v_stride); + _tile_stream_loadd(3, v5, v_stride); + if constexpr (std::is_same_v) { + _tile_dphf8ps(4, 0, 2); + _tile_dphf8ps(5, 0, 3); + } else { + _tile_dpbf8ps(4, 0, 2); + _tile_dpbf8ps(5, 0, 3); + } + a += a_advance; + v2 += v_advance; + v3 += v_advance; + v4 += v_advance; + v5 += v_advance; + } + + _tile_stored(4, c4, c_stride); + _tile_stored(5, c5, c_stride); + _tile_stored(6, c6, c_stride); + _tile_stored(7, c7, c_stride); + return; + } + + // -------- QK phase: native FP8 MMA (1-2-2 pattern) -------- + const int32_t k_times = + dynamic_k_size / (AMX_TILE_ROW_NUM * 4 / sizeof(uint8_t)); + const int32_t k_group_times = k_times / 2; + const bool has_tail = (k_times % 2 == 1); + + const int32_t b_stride = AMX_TILE_ROW_BYTES; + const int32_t c_stride = ldc * sizeof(float); + constexpr int32_t output_group_num = 2; + for (int32_t output_group = 0; output_group < output_group_num; + ++output_group) { + uint8_t* __restrict__ a_tile_0 = static_cast(a_tile); + uint8_t* __restrict__ a_tile_1 = a_tile_0 + fp8_tile_elems; + uint8_t* __restrict__ b_tile_2 = + static_cast(b_tile) + + output_group * 2 * k_size * AMX_TILE_ROW_NUM; + uint8_t* __restrict__ b_tile_3 = b_tile_2 + (k_size * AMX_TILE_ROW_NUM); + uint8_t* __restrict__ b_tile_4 = b_tile_2 + fp8_tile_elems; + uint8_t* __restrict__ b_tile_5 = b_tile_3 + fp8_tile_elems; + float* __restrict__ c_tile_6 = c_tile + output_group * 32; + float* __restrict__ c_tile_7 = c_tile_6 + AMX_TILE_ROW_NUM; + + if (accum_c) { + _tile_loadd(6, c_tile_6, c_stride); + _tile_loadd(7, c_tile_7, c_stride); + } else { + _tile_zero(6); + _tile_zero(7); + } + + for (int32_t k = 0; k < k_group_times; ++k) { + _tile_loadd(0, a_tile_0, AMX_TILE_ROW_BYTES); + _tile_stream_loadd(2, b_tile_2, b_stride); + if constexpr (std::is_same_v) { + _tile_dphf8ps(6, 0, 2); + } else { + _tile_dpbf8ps(6, 0, 2); + } + _tile_stream_loadd(3, b_tile_3, b_stride); + if constexpr (std::is_same_v) { + _tile_dphf8ps(7, 0, 3); + } else { + _tile_dpbf8ps(7, 0, 3); + } + _tile_loadd(1, a_tile_1, AMX_TILE_ROW_BYTES); + _tile_stream_loadd(2, b_tile_4, b_stride); + if constexpr (std::is_same_v) { + _tile_dphf8ps(6, 1, 2); + } else { + _tile_dpbf8ps(6, 1, 2); + } + _tile_stream_loadd(3, b_tile_5, b_stride); + if constexpr (std::is_same_v) { + _tile_dphf8ps(7, 1, 3); + } else { + _tile_dpbf8ps(7, 1, 3); + } + + a_tile_0 += 2 * fp8_tile_elems; + a_tile_1 += 2 * fp8_tile_elems; + b_tile_2 += 2 * fp8_tile_elems; + b_tile_3 += 2 * fp8_tile_elems; + b_tile_4 += 2 * fp8_tile_elems; + b_tile_5 += 2 * fp8_tile_elems; + } + + if (has_tail) { + _tile_loadd(0, a_tile_0, AMX_TILE_ROW_BYTES); + _tile_stream_loadd(2, b_tile_2, b_stride); + if constexpr (std::is_same_v) { + _tile_dphf8ps(6, 0, 2); + } else { + _tile_dpbf8ps(6, 0, 2); + } + _tile_stream_loadd(3, b_tile_3, b_stride); + if constexpr (std::is_same_v) { + _tile_dphf8ps(7, 0, 3); + } else { + _tile_dpbf8ps(7, 0, 3); + } + } + + _tile_stored(6, c_tile_6, c_stride); + _tile_stored(7, c_tile_7, c_stride); + } + } + + FORCE_INLINE static void init_tile_config(int32_t m, __tilecfg& config) { + for (int i = 0; i < 8; ++i) config.colsb[i] = AMX_TILE_ROW_BYTES; + config.rows[0] = m; + config.rows[1] = m; + config.rows[2] = AMX_TILE_ROW_NUM; + config.rows[3] = AMX_TILE_ROW_NUM; + config.rows[4] = m; + config.rows[5] = m; + config.rows[6] = m; + config.rows[7] = m; + _tile_loadconfig(&config); + } +}; + +// --------------------------------------------------------------------------- +// AttentionImpl +// +// scalar_t = BF16 (input query/output type) +// kv_cache_scalar_t = Float8_e4m3fn or Float8_e5m2 +// --------------------------------------------------------------------------- +template + requires(std::is_same_v && + (std::is_same_v || + std::is_same_v) && + (head_dim % 64 == 0)) +class AttentionImpl { + public: + // Q buffer holds FP8-quantized query. + using query_t = scalar_t; // BF16 input + using q_buffer_t = uint8_t; // FP8 quantized, stored in Q buffer + using kv_cache_t = kv_cache_scalar_t; + using logits_buffer_t = float; + using partial_output_buffer_t = float; + using prob_buffer_t = scalar_t; // BF16 softmax probs (PV phase) + + constexpr static int64_t BlockSizeAlignment = 64; + constexpr static int64_t HeadDimAlignment = 64; + constexpr static int64_t MaxQHeadNumPerIteration = 32; + constexpr static int64_t HeadDim = head_dim; + constexpr static ISA ISAType = ISA::AMX_FP8; + constexpr static bool scale_on_logits = true; + + float k_scale = 1.0f; + float v_scale = 1.0f; + // Dynamically computed per copy_q_heads_tile() call. + float q_dynamic_scale = 1.0f; + + public: + AttentionImpl() : current_q_head_num_(0) { + vec_op::unroll_loop([&](int i) { amx_tile_config_.colsb[i] = 64; }); + } + + ~AttentionImpl() { _tile_release(); } + + void init_from_input(const AttentionInput* input) { + k_scale = input->k_scale_fp8; + v_scale = input->v_scale_fp8; + // q_scale_fp8 from input is used as a hint; dynamic scale is recomputed + // per-tile in copy_q_heads_tile(). + } + + // Output scale for the PV result. + float get_output_v_scale() const noexcept { + constexpr float fp8_max = + std::is_same_v ? 448.0f : 57344.0f; + return v_scale / fp8_max; + } + + template