[CPU] [Feat] Add native AMX-FP8 attention impl for Diamond Rapids (#49410)
Signed-off-by: zhejiangxiaomai <zhenhui.zhao@intel.com>
This commit is contained in:
zhao, zhenhui
@@ -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"
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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(
|
||||
|
||||
+30
-2
@@ -1,5 +1,17 @@
|
||||
#include "cpu_attn_dispatch_generated.h"
|
||||
|
||||
#if defined(__x86_64__) || defined(_M_X64)
|
||||
#include <cpuid.h>
|
||||
// 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<float>(k_scale);
|
||||
input.v_scale_fp8 = static_cast<float>(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(
|
||||
|
||||
@@ -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 <amxfp8intrin.h> // 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 <typename scalar_t, uint8_t (*quant_fn)(float, float)>
|
||||
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<float>(ksrc[s * fp8_per_row + k * 4 + lane]),
|
||||
k_inv);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename fp8_t>
|
||||
__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<fp8_t, c10::Float8_e4m3fn> ? 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<fp8_t, c10::Float8_e4m3fn>) {
|
||||
_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 <typename fp8_t>
|
||||
__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<fp8_t, c10::Float8_e4m3fn> ? 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<const __m256i*>(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_t, c10::Float8_e4m3fn>) {
|
||||
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<uint8_t, kv_cache_t>: 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 <typename kv_cache_t>
|
||||
class TileGemm224<uint8_t, kv_cache_t> {
|
||||
static_assert(std::is_same_v<kv_cache_t, c10::Float8_e4m3fn> ||
|
||||
std::is_same_v<kv_cache_t, c10::Float8_e5m2>,
|
||||
"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<kv_cache_t>(src, dst, cols);
|
||||
}
|
||||
|
||||
// ------------------------------------------------------------------
|
||||
// QK: k_times = head_dim / 64 (each FP8 tile covers 64 k-elements).
|
||||
// PV: native wide FP8 MMA.
|
||||
// ------------------------------------------------------------------
|
||||
template <AttentionGemmPhase phase, int32_t k_size>
|
||||
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<int32_t>(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<uint8_t*>(a_tile);
|
||||
uint8_t* a1 = a0 + lda * AMX_TILE_ROW_NUM;
|
||||
uint8_t* v2 = reinterpret_cast<uint8_t*>(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<kv_cache_t, c10::Float8_e4m3fn>) {
|
||||
_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<kv_cache_t, c10::Float8_e4m3fn>) {
|
||||
_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<int32_t>(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<uint8_t*>(a_tile);
|
||||
uint8_t* __restrict__ a_tile_1 = a_tile_0 + k_times * fp8_tile_elems;
|
||||
uint8_t* __restrict__ b_tile_2 =
|
||||
static_cast<uint8_t*>(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<kv_cache_t, c10::Float8_e4m3fn>) {
|
||||
_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<kv_cache_t, c10::Float8_e4m3fn>) {
|
||||
_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<kv_cache_t, c10::Float8_e4m3fn>) {
|
||||
_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<int32_t>(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<uint8_t, kv_cache_t>: 1-2-2 pattern for AMX_FP8.
|
||||
// ---------------------------------------------------------------------------
|
||||
template <typename kv_cache_t>
|
||||
class TileGemm122<uint8_t, kv_cache_t> {
|
||||
static_assert(std::is_same_v<kv_cache_t, c10::Float8_e4m3fn> ||
|
||||
std::is_same_v<kv_cache_t, c10::Float8_e5m2>,
|
||||
"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<kv_cache_t>(src, dst, cols);
|
||||
}
|
||||
|
||||
template <AttentionGemmPhase phase, int32_t k_size>
|
||||
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<uint8_t*>(a_tile);
|
||||
uint8_t* v2 = reinterpret_cast<uint8_t*>(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<kv_cache_t, c10::Float8_e4m3fn>) {
|
||||
_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<kv_cache_t, c10::Float8_e4m3fn>) {
|
||||
_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<uint8_t*>(a_tile);
|
||||
uint8_t* __restrict__ a_tile_1 = a_tile_0 + fp8_tile_elems;
|
||||
uint8_t* __restrict__ b_tile_2 =
|
||||
static_cast<uint8_t*>(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<kv_cache_t, c10::Float8_e4m3fn>) {
|
||||
_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<kv_cache_t, c10::Float8_e4m3fn>) {
|
||||
_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<kv_cache_t, c10::Float8_e4m3fn>) {
|
||||
_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<kv_cache_t, c10::Float8_e4m3fn>) {
|
||||
_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<kv_cache_t, c10::Float8_e4m3fn>) {
|
||||
_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<kv_cache_t, c10::Float8_e4m3fn>) {
|
||||
_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<ISA::AMX_FP8, scalar_t, head_dim, kv_cache_scalar_t>
|
||||
//
|
||||
// scalar_t = BF16 (input query/output type)
|
||||
// kv_cache_scalar_t = Float8_e4m3fn or Float8_e5m2
|
||||
// ---------------------------------------------------------------------------
|
||||
template <typename scalar_t, int64_t head_dim, typename kv_cache_scalar_t>
|
||||
requires(std::is_same_v<scalar_t, c10::BFloat16> &&
|
||||
(std::is_same_v<kv_cache_scalar_t, c10::Float8_e4m3fn> ||
|
||||
std::is_same_v<kv_cache_scalar_t, c10::Float8_e5m2>) &&
|
||||
(head_dim % 64 == 0))
|
||||
class AttentionImpl<ISA::AMX_FP8, scalar_t, head_dim, kv_cache_scalar_t> {
|
||||
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, 8>([&](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<kv_cache_t, c10::Float8_e4m3fn> ? 448.0f : 57344.0f;
|
||||
return v_scale / fp8_max;
|
||||
}
|
||||
|
||||
template <template <typename tile_gemm_t> typename attention>
|
||||
FORCE_INLINE void execute_attention(DEFINE_CPU_ATTENTION_PARAMS) {
|
||||
// Apply QK scale correction: score_fp32 = dot(Q_fp8, K_fp8) * q_dyn * k
|
||||
scale *= q_dynamic_scale * k_scale;
|
||||
|
||||
if (q_head_num > AMX_TILE_ROW_NUM) {
|
||||
if (q_head_num != current_q_head_num_) {
|
||||
current_q_head_num_ = q_head_num;
|
||||
TileGemm224<q_buffer_t, kv_cache_t>::init_tile_config(q_head_num,
|
||||
amx_tile_config_);
|
||||
}
|
||||
attention<TileGemm224<q_buffer_t, kv_cache_t>> attention_iteration;
|
||||
attention_iteration(CPU_ATTENTION_PARAMS);
|
||||
} else {
|
||||
if (q_head_num != current_q_head_num_) {
|
||||
current_q_head_num_ = q_head_num;
|
||||
TileGemm122<q_buffer_t, kv_cache_t>::init_tile_config(q_head_num,
|
||||
amx_tile_config_);
|
||||
}
|
||||
attention<TileGemm122<q_buffer_t, kv_cache_t>> attention_iteration;
|
||||
attention_iteration(CPU_ATTENTION_PARAMS);
|
||||
}
|
||||
}
|
||||
|
||||
// K cache stride: same element-count formula as AMX_BF16.
|
||||
// (BlockSizeAlignment * head_dim) uint8_t elements = same byte count as
|
||||
// before since kv_cache_t is 1-byte FP8.
|
||||
constexpr static int64_t k_cache_token_group_stride(
|
||||
const int32_t /*block_size*/) {
|
||||
return BlockSizeAlignment * head_dim;
|
||||
}
|
||||
|
||||
// Both native FP8 formats use VNNI4 packing.
|
||||
constexpr static int64_t v_cache_token_group_stride(
|
||||
const int32_t /*block_size*/) {
|
||||
return BlockSizeAlignment * (AMX_TILE_ROW_BYTES / 4);
|
||||
}
|
||||
|
||||
constexpr static int64_t v_cache_head_group_stride(const int32_t block_size) {
|
||||
return block_size * HeadDimAlignment;
|
||||
}
|
||||
|
||||
// Copy Q heads tile: quantize BF16 to the KV cache FP8 type and pack it.
|
||||
//
|
||||
// Q buffer layout (FP8):
|
||||
// Same tile-prepacking as AMX_BF16 copy_q_heads_tile but with uint8_t
|
||||
// (FP8) elements. Each FP8 tile row = 64 bytes = 64 FP8 elements.
|
||||
// head_size_block_num = head_dim / 64 (vs head_dim/32 for BF16).
|
||||
//
|
||||
// Dynamic scale: max(abs(Q)) / fp8_max is computed over all elements,
|
||||
// stored in q_dynamic_scale for use in execute_attention().
|
||||
void copy_q_heads_tile(
|
||||
scalar_t* __restrict__ src, // [q_num, q_heads_per_kv, head_size] BF16
|
||||
q_buffer_t* __restrict__ q_buffer, const int32_t q_num,
|
||||
const int32_t q_heads_per_kv, const int64_t q_num_stride,
|
||||
const int64_t q_head_stride, const float /*scale*/) {
|
||||
// ---- Step 1: compute dynamic per-tile Q scale ----
|
||||
constexpr float fp8_max =
|
||||
std::is_same_v<kv_cache_t, c10::Float8_e4m3fn> ? 448.0f : 57344.0f;
|
||||
float max_abs = 0.0f;
|
||||
{
|
||||
scalar_t* src_iter = src;
|
||||
__m512 max_abs_vec = _mm512_setzero_ps();
|
||||
const __m512i abs_mask = _mm512_set1_epi32(0x7FFFFFFF);
|
||||
for (int32_t q = 0; q < q_num; ++q, src_iter += q_num_stride) {
|
||||
scalar_t* head_iter = src_iter;
|
||||
for (int32_t h = 0; h < q_heads_per_kv;
|
||||
++h, head_iter += q_head_stride) {
|
||||
for (int64_t e = 0; e < head_dim; e += 16) {
|
||||
const __m256i bf16 = _mm256_loadu_si256(
|
||||
reinterpret_cast<const __m256i*>(head_iter + e));
|
||||
const __m512i fp32_bits =
|
||||
_mm512_slli_epi32(_mm512_cvtepu16_epi32(bf16), 16);
|
||||
const __m512 abs_values =
|
||||
_mm512_castsi512_ps(_mm512_and_si512(fp32_bits, abs_mask));
|
||||
max_abs_vec = _mm512_max_ps(max_abs_vec, abs_values);
|
||||
}
|
||||
}
|
||||
}
|
||||
max_abs = _mm512_reduce_max_ps(max_abs_vec);
|
||||
}
|
||||
if (max_abs == 0.0f) max_abs = 1.0f; // avoid div-by-zero
|
||||
const float inv_scale = fp8_max / max_abs;
|
||||
// Store true scale for score correction in execute_attention().
|
||||
q_dynamic_scale = max_abs / fp8_max;
|
||||
|
||||
// ---- Step 2: quantize BF16 Q to FP8 and pack into q_buffer ----
|
||||
// Layout: same tile-prepacking as AMX_BF16, but with FP8 (1 byte/elem).
|
||||
// head_size_block_num = head_dim * sizeof(uint8_t) / AMX_TILE_ROW_BYTES
|
||||
// = head_dim / 64
|
||||
constexpr int64_t bytes_per_head = head_dim * sizeof(uint8_t);
|
||||
static_assert(bytes_per_head % AMX_TILE_ROW_BYTES == 0,
|
||||
"head_dim must be divisible by 64 for AMX_FP8");
|
||||
quantize_q_tile<kv_cache_t>(src, reinterpret_cast<uint8_t*>(q_buffer),
|
||||
q_num, q_heads_per_kv, q_num_stride,
|
||||
q_head_stride, head_dim, inv_scale);
|
||||
}
|
||||
|
||||
// reshape KV to AMX-FP8 friendly layout.
|
||||
// K: native FP8 layout (1 byte/element, 16-token groups × k-slices).
|
||||
// V: native VNNI4 packing.
|
||||
static void reshape_and_cache(
|
||||
const scalar_t* __restrict__ key, const scalar_t* __restrict__ value,
|
||||
kv_cache_t* __restrict__ key_cache, kv_cache_t* __restrict__ value_cache,
|
||||
const int64_t* __restrict__ slot_mapping, const int64_t token_num,
|
||||
const int64_t key_token_num_stride, const int64_t value_token_num_stride,
|
||||
const int64_t head_num, const int64_t key_head_num_stride,
|
||||
const int64_t value_head_num_stride, const int64_t num_blocks,
|
||||
const int64_t num_blocks_stride, const int64_t cache_head_num_stride,
|
||||
const int64_t block_size, const int64_t /*block_size_stride*/,
|
||||
const float k_inv = 0.0f, const float v_inv = 0.0f) {
|
||||
constexpr auto qfn = select_fp8_quant_fn<kv_cache_t>();
|
||||
|
||||
// K: native AMX-FP8 layout
|
||||
reshape_and_cache_k_amx_fp8_impl<scalar_t, qfn>(
|
||||
key, reinterpret_cast<uint8_t*>(key_cache), slot_mapping, token_num,
|
||||
head_num, head_dim, block_size, key_token_num_stride,
|
||||
key_head_num_stride, num_blocks_stride, cache_head_num_stride, k_inv);
|
||||
|
||||
{
|
||||
constexpr int64_t token_num_per_sub_group = 4;
|
||||
constexpr int64_t head_elems_per_group = 16; // AMX_TILE_ROW_BYTES/4
|
||||
const int64_t group_num = head_dim / head_elems_per_group;
|
||||
const int64_t group_size = block_size * head_elems_per_group;
|
||||
uint8_t* vc = reinterpret_cast<uint8_t*>(value_cache);
|
||||
|
||||
#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_mapping[tok];
|
||||
if (slot < 0) continue;
|
||||
const int64_t block_idx = slot / block_size;
|
||||
const int64_t block_offset = slot % block_size;
|
||||
const int64_t sub_group_idx = block_offset / token_num_per_sub_group;
|
||||
const int64_t sub_group_offset =
|
||||
block_offset % token_num_per_sub_group;
|
||||
|
||||
const scalar_t* vsrc =
|
||||
value + tok * value_token_num_stride + h * value_head_num_stride;
|
||||
uint8_t* vdst =
|
||||
vc + block_idx * num_blocks_stride + h * cache_head_num_stride +
|
||||
sub_group_idx * token_num_per_sub_group * head_elems_per_group +
|
||||
sub_group_offset;
|
||||
|
||||
for (int64_t i = 0; i < group_num; ++i) {
|
||||
for (int64_t j = 0; j < head_elems_per_group; ++j)
|
||||
vdst[j * token_num_per_sub_group] =
|
||||
qfn(static_cast<float>(vsrc[j]), v_inv);
|
||||
vsrc += head_elems_per_group;
|
||||
vdst += group_size;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private:
|
||||
alignas(64) __tilecfg amx_tile_config_;
|
||||
int32_t current_q_head_num_;
|
||||
};
|
||||
|
||||
} // namespace cpu_attention
|
||||
|
||||
#endif // CPU_CAPABILITY_AMXFP8
|
||||
#endif // CPU_ATTN_AMX_FP8_HPP
|
||||
+54
-36
@@ -12,7 +12,7 @@
|
||||
#include "cpu/utils.hpp"
|
||||
|
||||
namespace cpu_attention {
|
||||
enum class ISA { AMX, VEC, VEC16, NEON, VXE, RVV, VSX };
|
||||
enum class ISA { AMX, VEC, VEC16, NEON, VXE, RVV, VSX, AMX_FP8 };
|
||||
|
||||
// Mirrors csrc/attention/dtype_fp8.cuh Fp8KVCacheDataType exactly.
|
||||
enum class Fp8KVCacheDataType {
|
||||
@@ -152,6 +152,9 @@ struct AttentionMetadata {
|
||||
case ISA::AMX:
|
||||
ss << "AMX, ";
|
||||
break;
|
||||
case ISA::AMX_FP8:
|
||||
ss << "AMX_FP8, ";
|
||||
break;
|
||||
case ISA::VEC:
|
||||
ss << "VEC, ";
|
||||
break;
|
||||
@@ -885,6 +888,8 @@ struct AttentionInput {
|
||||
// FP8 KV cache scales (used by FP8 attention implementations)
|
||||
float k_scale_fp8 = 1.0f;
|
||||
float v_scale_fp8 = 1.0f;
|
||||
// FP8 query scale (used by AMX_FP8 which quantizes query inside the kernel)
|
||||
float q_scale_fp8 = 1.0f;
|
||||
};
|
||||
|
||||
#define DEFINE_CPU_ATTENTION_PARAMS \
|
||||
@@ -1092,11 +1097,6 @@ class AttentionMainLoop {
|
||||
|
||||
// process logits
|
||||
{
|
||||
// if (debug_info){
|
||||
// print_logits("raw logits", logits_buffer, q_head_num,
|
||||
// kv_tile_token_num, kv_tile_token_num);
|
||||
// }
|
||||
|
||||
if (softcap_scale != 0.0f) {
|
||||
apply_softcap(logits_buffer, kv_tile_token_num, q_head_num,
|
||||
kv_tile_token_num, softcap_scale);
|
||||
@@ -1127,30 +1127,33 @@ class AttentionMainLoop {
|
||||
apply_softmax(logits_buffer, partial_q_buffer, max_buffer, sum_buffer,
|
||||
kv_tile_token_num, q_head_num, kv_tile_token_num,
|
||||
is_first_iter, use_sink);
|
||||
|
||||
// if (debug_info){
|
||||
// print_logits("softmax logits",
|
||||
// reinterpret_cast<prob_buffer_t*>(logits_buffer), q_head_num,
|
||||
// kv_tile_token_num, kv_tile_token_num * sizeof(logits_buffer_t) /
|
||||
// sizeof(prob_buffer_t));
|
||||
// print_logits("new_max", max_buffer, 1, q_head_num, q_head_num);
|
||||
// print_logits("new_sum", sum_buffer, 1, q_head_num, q_head_num);
|
||||
// }
|
||||
}
|
||||
|
||||
// compute P@V
|
||||
{
|
||||
constexpr bool prequantize_probabilities = []() {
|
||||
if constexpr (requires { tile_gemm_t::prequantize_probabilities; }) {
|
||||
return tile_gemm_t::prequantize_probabilities;
|
||||
}
|
||||
return false;
|
||||
}();
|
||||
int32_t curr_group_offset =
|
||||
start_block_group_offset * v_cache_token_group_stride;
|
||||
int32_t curr_group_num_in_block =
|
||||
token_group_num_per_block - start_block_group_offset;
|
||||
int32_t remaining_group_num = token_group_num;
|
||||
int32_t head_dim_group_num = head_dim / headdim_alignment;
|
||||
prob_buffer_t* curr_prob_buffer =
|
||||
reinterpret_cast<prob_buffer_t*>(logits_buffer);
|
||||
int64_t prob_buffer_stride =
|
||||
kv_tile_token_num *
|
||||
(sizeof(logits_buffer_t) / sizeof(prob_buffer_t));
|
||||
using pv_prob_buffer_t = std::conditional_t<prequantize_probabilities,
|
||||
uint8_t, prob_buffer_t>;
|
||||
pv_prob_buffer_t* curr_prob_buffer =
|
||||
reinterpret_cast<pv_prob_buffer_t*>(logits_buffer);
|
||||
const int64_t prob_buffer_stride =
|
||||
prequantize_probabilities
|
||||
? kv_tile_token_num
|
||||
: kv_tile_token_num *
|
||||
(sizeof(logits_buffer_t) / sizeof(prob_buffer_t));
|
||||
constexpr int64_t prob_buffer_elem_size =
|
||||
prequantize_probabilities ? sizeof(uint8_t) : sizeof(prob_buffer_t);
|
||||
partial_output_buffer_t* curr_partial_q_buffer = partial_q_buffer;
|
||||
bool accum_c = !is_first_iter;
|
||||
for (int32_t block_idx = start_block_idx; block_idx < end_block_idx;
|
||||
@@ -1184,15 +1187,13 @@ class AttentionMainLoop {
|
||||
remaining_group_num -= curr_group_num_in_block;
|
||||
curr_group_offset = 0;
|
||||
curr_group_num_in_block = token_group_num_per_block;
|
||||
curr_prob_buffer += curr_token_num;
|
||||
curr_prob_buffer = reinterpret_cast<pv_prob_buffer_t*>(
|
||||
reinterpret_cast<uint8_t*>(curr_prob_buffer) +
|
||||
curr_token_num * prob_buffer_elem_size);
|
||||
curr_partial_q_buffer = partial_q_buffer;
|
||||
accum_c = true;
|
||||
}
|
||||
}
|
||||
// if (debug_info) {
|
||||
// print_logits("output", partial_q_buffer, q_head_num, head_dim,
|
||||
// head_dim);
|
||||
// }
|
||||
}
|
||||
|
||||
void apply_mask(logits_buffer_t* __restrict__ logits_buffer,
|
||||
@@ -1262,9 +1263,19 @@ class AttentionMainLoop {
|
||||
#endif
|
||||
|
||||
using prob_buffer_vec_t = typename VecTypeTrait<prob_buffer_t>::vec_t;
|
||||
constexpr bool prequantize_probabilities = []() {
|
||||
if constexpr (requires { tile_gemm_t::prequantize_probabilities; }) {
|
||||
return tile_gemm_t::prequantize_probabilities;
|
||||
}
|
||||
return false;
|
||||
}();
|
||||
using softmax_prob_buffer_t =
|
||||
std::conditional_t<prequantize_probabilities, uint8_t, prob_buffer_t>;
|
||||
static_assert(sizeof(prob_buffer_t) <= sizeof(logits_buffer_t));
|
||||
|
||||
logits_buffer_t* __restrict__ curr_logits_buffer = logits_buffer;
|
||||
softmax_prob_buffer_t* __restrict__ curr_prob_buffer =
|
||||
reinterpret_cast<softmax_prob_buffer_t*>(logits_buffer);
|
||||
float* __restrict__ curr_partial_q_buffer = partial_q_buffer;
|
||||
const int32_t vec_num = kv_tile_token_num / 16;
|
||||
const int32_t head_vec_num = head_dim / 16;
|
||||
@@ -1302,8 +1313,8 @@ class AttentionMainLoop {
|
||||
{
|
||||
logits_buffer_t* __restrict__ curr_logits_buffer_iter =
|
||||
curr_logits_buffer;
|
||||
prob_buffer_t* __restrict__ curr_prob_buffer_iter =
|
||||
reinterpret_cast<prob_buffer_t*>(curr_logits_buffer);
|
||||
softmax_prob_buffer_t* __restrict__ curr_prob_buffer_iter =
|
||||
curr_prob_buffer;
|
||||
for (int32_t j = 0; j < vec_num; ++j) {
|
||||
vec_op::FP32Vec16 vec(curr_logits_buffer_iter);
|
||||
vec = vec - max_vec;
|
||||
@@ -1320,8 +1331,12 @@ class AttentionMainLoop {
|
||||
vec = fast_exp(vec);
|
||||
}
|
||||
|
||||
prob_buffer_vec_t output_vec(vec);
|
||||
output_vec.save(curr_prob_buffer_iter);
|
||||
if constexpr (prequantize_probabilities) {
|
||||
vec.save(curr_logits_buffer_iter);
|
||||
} else {
|
||||
prob_buffer_vec_t output_vec(vec);
|
||||
output_vec.save(curr_prob_buffer_iter);
|
||||
}
|
||||
#else
|
||||
vec.save(curr_logits_buffer_iter);
|
||||
for (int32_t k = 0; k < 16; ++k) {
|
||||
@@ -1335,6 +1350,10 @@ class AttentionMainLoop {
|
||||
curr_logits_buffer_iter += 16;
|
||||
curr_prob_buffer_iter += 16;
|
||||
}
|
||||
if constexpr (prequantize_probabilities) {
|
||||
tile_gemm_t::quantize_probability_row(
|
||||
curr_logits_buffer, curr_prob_buffer, kv_tile_token_num);
|
||||
}
|
||||
}
|
||||
float new_sum_val = sum_vec.reduce_sum();
|
||||
|
||||
@@ -1366,6 +1385,12 @@ class AttentionMainLoop {
|
||||
sum_buffer[i] = new_sum_val;
|
||||
|
||||
curr_logits_buffer += logits_buffer_stride;
|
||||
if constexpr (prequantize_probabilities) {
|
||||
curr_prob_buffer += kv_tile_token_num;
|
||||
} else {
|
||||
curr_prob_buffer =
|
||||
reinterpret_cast<softmax_prob_buffer_t*>(curr_logits_buffer);
|
||||
}
|
||||
curr_partial_q_buffer += head_dim;
|
||||
}
|
||||
}
|
||||
@@ -1840,13 +1865,6 @@ class AttentionMainLoop {
|
||||
float* curr_sum_buffer = sum_buffer + q_tile_head_offset;
|
||||
|
||||
bool debug_info = false;
|
||||
// bool debug_info = (
|
||||
// q_head_start_idx == 4 &&
|
||||
// (q_token_start_idx + q_head_tile_token_offset) <=
|
||||
// 4
|
||||
// && (q_token_start_idx + q_head_tile_token_offset +
|
||||
// q_tile_token_num) > 4
|
||||
// );
|
||||
// if (debug_info) {
|
||||
// std::printf("\tq_iter_idx: %d, q_token_start: %d,"
|
||||
// "q_token_end: %d, q_token_num: %d, q_head_num: %d,"
|
||||
|
||||
@@ -350,6 +350,64 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
|
||||
};
|
||||
#endif
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Vectorized BF16 → FP8 E4M3 quantization helpers.
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// Quantize 16 FP32 values to 16 FP8-E4M3 bytes using AVX-512.
|
||||
// inv_scale = 1.0f / q_scale (pre-computed by caller).
|
||||
// Returns 16 packed uint8_t in the low 128-bits of an __m128i.
|
||||
#if defined(__AVX512F__)
|
||||
FORCE_INLINE __m128i quant_fp32x16_to_fp8e4m3_avx512(const float* src,
|
||||
float inv_scale) {
|
||||
__m512 v = _mm512_loadu_ps(src);
|
||||
v = _mm512_mul_ps(v, _mm512_set1_ps(inv_scale));
|
||||
// Clamp to FP8-E4M3 representable range [-448, 448]
|
||||
v = _mm512_min_ps(v, _mm512_set1_ps(448.0f));
|
||||
v = _mm512_max_ps(v, _mm512_set1_ps(-448.0f));
|
||||
// Shift exponent bias: FP32 bias=127 → E4M3 bias=7, delta=120
|
||||
// Multiply by 2^-120 to move the exponent bias from 127 to 7.
|
||||
v = _mm512_mul_ps(v, _mm512_set1_ps(0x1p-120f));
|
||||
__m512i vi = _mm512_castps_si512(v);
|
||||
// sign: bit31 → bit7
|
||||
__m512i sign = _mm512_srli_epi32(
|
||||
_mm512_and_si512(vi, _mm512_set1_epi32(0x80000000u)), 24);
|
||||
// payload: bits[26:20] -> bits[6:0] after the exponent-bias shift.
|
||||
__m512i payload = _mm512_srli_epi32(
|
||||
_mm512_and_si512(vi, _mm512_set1_epi32(0x07F00000u)), 20);
|
||||
// Keep 0x7F (all-ones) reserved as NaN encoding: clamp to 0x7E
|
||||
payload = _mm512_min_epu32(payload, _mm512_set1_epi32(0x7E));
|
||||
__m512i fp8 = _mm512_or_si512(sign, payload);
|
||||
// Pack 16 × int32 → 16 × uint8 (cvt truncates, so the int32 values must
|
||||
// fit in uint8; they do since max is 0xFF = sign|payload).
|
||||
return _mm512_cvtepi32_epi8(fp8);
|
||||
}
|
||||
|
||||
// Quantize 32 BF16 values to 32 FP8-E4M3 bytes using AVX-512.
|
||||
// Writes 32 bytes to dst. inv_scale = 1.0f / q_scale.
|
||||
FORCE_INLINE void quant_bf16x32_to_fp8e4m3_avx512(const c10::BFloat16* src,
|
||||
uint8_t* dst,
|
||||
float inv_scale) {
|
||||
// Convert 32 BF16 → 2x16 FP32 then quantize each half.
|
||||
const uint16_t* u16 = reinterpret_cast<const uint16_t*>(src);
|
||||
__m256i b16_lo = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(u16));
|
||||
__m256i b16_hi =
|
||||
_mm256_loadu_si256(reinterpret_cast<const __m256i*>(u16 + 16));
|
||||
// Zero-extend uint16 -> uint32, then shift left by 16 to form FP32 bits.
|
||||
__m512 lo_fp32 =
|
||||
_mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(b16_lo), 16));
|
||||
__m512 hi_fp32 =
|
||||
_mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(b16_hi), 16));
|
||||
alignas(64) float fp32_buf[32];
|
||||
_mm512_store_ps(fp32_buf, lo_fp32);
|
||||
_mm512_store_ps(fp32_buf + 16, hi_fp32);
|
||||
__m128i lo8 = quant_fp32x16_to_fp8e4m3_avx512(fp32_buf, inv_scale);
|
||||
__m128i hi8 = quant_fp32x16_to_fp8e4m3_avx512(fp32_buf + 16, inv_scale);
|
||||
_mm_storeu_si128(reinterpret_cast<__m128i*>(dst), lo8);
|
||||
_mm_storeu_si128(reinterpret_cast<__m128i*>(dst + 16), hi8);
|
||||
}
|
||||
#endif // __AVX512F__
|
||||
|
||||
struct FP32Vec4 : public Vec<FP32Vec4> {
|
||||
constexpr static int VEC_ELEM_NUM = 4;
|
||||
union AliasReg {
|
||||
|
||||
@@ -22,6 +22,7 @@ ISA_TYPES = {
|
||||
"VXE": 4,
|
||||
"RVV": 5,
|
||||
"VSX": 6,
|
||||
"AMX_FP8": 7,
|
||||
}
|
||||
|
||||
# KV cache index: 0 = auto (same as scalar_t), 1 = fp8_e4m3, 2 = fp8_e5m2
|
||||
@@ -39,13 +40,13 @@ KV_CACHE_CPP_TYPES = {
|
||||
}
|
||||
|
||||
# ISAs supported for head_dims divisible by 32
|
||||
ISA_FOR_32 = ["AMX", "NEON", "VEC", "VEC16", "VXE", "RVV", "VSX"]
|
||||
ISA_FOR_32 = ["AMX", "AMX_FP8", "NEON", "VEC", "VEC16", "VXE", "RVV", "VSX"]
|
||||
|
||||
# ISAs supported for head_dims divisible by 16 only
|
||||
ISA_FOR_16 = ["VEC16"]
|
||||
|
||||
# ISAs that support FP8 KV cache (x86 AVX2/AVX-512 required)
|
||||
ISA_FOR_FP8 = ["AMX", "VEC"]
|
||||
ISA_FOR_FP8 = ["AMX_FP8", "AMX", "VEC"]
|
||||
|
||||
|
||||
def encode_params(head_dim: int, isa_type: str, kv_cache: str = "auto") -> int:
|
||||
@@ -64,12 +65,21 @@ def _make_case(
|
||||
encoded = encode_params(head_dim, isa, kv_cache)
|
||||
actual_isa = isa_override if isa_override else isa
|
||||
cpp_type = KV_CACHE_CPP_TYPES[kv_cache]
|
||||
attn_impl = (
|
||||
f"cpu_attention::AttentionImpl<"
|
||||
f"cpu_attention::ISA::{actual_isa}, \\\n"
|
||||
f" "
|
||||
native_impl = (
|
||||
f"cpu_attention::AttentionImpl<cpu_attention::ISA::{actual_isa}, "
|
||||
f"scalar_t, head_dim, {cpp_type}>"
|
||||
)
|
||||
if actual_isa == "AMX_FP8":
|
||||
fallback_impl = (
|
||||
f"cpu_attention::AttentionImpl<cpu_attention::ISA::AMX, "
|
||||
f"scalar_t, head_dim, {cpp_type}>"
|
||||
)
|
||||
attn_impl = (
|
||||
"std::conditional_t<std::is_same_v<scalar_t, c10::BFloat16>, "
|
||||
f"{native_impl}, {fallback_impl}>"
|
||||
)
|
||||
else:
|
||||
attn_impl = native_impl
|
||||
comment = (
|
||||
f"head_dim={head_dim}, isa={isa}"
|
||||
if kv_cache == "auto"
|
||||
@@ -98,7 +108,8 @@ def generate_cases_for_isa_group(isa_list: list[str], include_fp8: bool = False)
|
||||
for isa in isa_list:
|
||||
if isa not in ISA_FOR_32:
|
||||
continue
|
||||
cases.append(_make_case(head_dim, isa, "auto"))
|
||||
isa_override = "AMX" if isa == "AMX_FP8" else None
|
||||
cases.append(_make_case(head_dim, isa, "auto", isa_override))
|
||||
|
||||
# Non-FP8 cases for head_dims divisible by 16 only
|
||||
for head_dim in HEAD_DIMS_16:
|
||||
@@ -112,7 +123,8 @@ def generate_cases_for_isa_group(isa_list: list[str], include_fp8: bool = False)
|
||||
for isa in isa_list:
|
||||
if isa not in ISA_FOR_FP8:
|
||||
continue
|
||||
cases.append(_make_case(head_dim, isa, fp8_type))
|
||||
isa_override = "AMX" if isa == "AMX_FP8" and head_dim % 64 else None
|
||||
cases.append(_make_case(head_dim, isa, fp8_type, isa_override))
|
||||
|
||||
return "\n".join(cases)
|
||||
|
||||
@@ -138,6 +150,10 @@ def generate_header_file() -> str:
|
||||
#include "cpu_attn_vec.hpp"
|
||||
#include "cpu_attn_vec16.hpp"
|
||||
|
||||
#ifdef CPU_CAPABILITY_AMXFP8
|
||||
#include "cpu_attn_amx_fp8.hpp"
|
||||
#endif
|
||||
|
||||
#ifdef CPU_CAPABILITY_AMXBF16
|
||||
#include "cpu_attn_amx.hpp"
|
||||
#endif
|
||||
@@ -206,7 +222,12 @@ def generate_header_file() -> str:
|
||||
)
|
||||
|
||||
header += _macro_block(
|
||||
"#if defined(CPU_CAPABILITY_AMXBF16)",
|
||||
"#if defined(CPU_CAPABILITY_AMXFP8)",
|
||||
["AMX_FP8", "AMX", "VEC", "VEC16"],
|
||||
fp8=True,
|
||||
)
|
||||
header += _macro_block(
|
||||
"#elif defined(CPU_CAPABILITY_AMXBF16)",
|
||||
["AMX", "VEC", "VEC16"],
|
||||
fp8=True,
|
||||
)
|
||||
|
||||
+14
-6
@@ -27,7 +27,7 @@
|
||||
# cache mounts (ccache, cargo registry) aren't exported by `--cache-to
|
||||
# type=registry`, so they don't survive a build landing on a different/fresh
|
||||
# builder. sccache's cache lives in S3 instead, so it does.
|
||||
FROM ubuntu:22.04 AS base
|
||||
FROM ubuntu:25.04 AS base
|
||||
|
||||
RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
|
||||
--mount=type=cache,target=/var/lib/apt,sharing=locked \
|
||||
@@ -66,9 +66,15 @@ ENV MAX_JOBS=${max_jobs}
|
||||
RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
|
||||
--mount=type=cache,target=/var/lib/apt,sharing=locked \
|
||||
apt-get update -y \
|
||||
&& apt-get install -y --no-install-recommends sudo ccache wget zlib1g-dev \
|
||||
gcc-12 g++-12 libtcmalloc-minimal4 libnuma-dev ffmpeg libsm6 libxext6 libgl1 jq lsof make xz-utils \
|
||||
&& update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-12 10 --slave /usr/bin/g++ g++ /usr/bin/g++-12 \
|
||||
&& apt-get install -y --no-install-recommends sudo ccache git curl wget ca-certificates zlib1g-dev \
|
||||
software-properties-common libtcmalloc-minimal4 libnuma-dev ffmpeg libsm6 libxext6 libgl1 jq lsof make xz-utils \
|
||||
&& for i in 1 2 3; do \
|
||||
add-apt-repository -y ppa:ubuntu-toolchain-r/test && break || \
|
||||
{ echo "Attempt $i failed, retrying toolchain PPA..."; sleep 5; }; \
|
||||
done \
|
||||
&& apt-get update -y \
|
||||
&& apt-get install -y --no-install-recommends gcc-15 g++-15 \
|
||||
&& update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-15 10 --slave /usr/bin/g++ g++ /usr/bin/g++-15 \
|
||||
&& curl -LsSf https://astral.sh/uv/install.sh | sh
|
||||
|
||||
ARG TARGETARCH
|
||||
@@ -83,7 +89,7 @@ ENV SCCACHE_S3_NO_CREDENTIALS=${USE_SCCACHE:+${SCCACHE_S3_NO_CREDENTIALS}}
|
||||
ENV SCCACHE_IDLE_TIMEOUT=${USE_SCCACHE:+0}
|
||||
|
||||
# Compiler and linker environment
|
||||
ENV CC=/usr/bin/gcc-12 CXX=/usr/bin/g++-12
|
||||
ENV CC=/usr/bin/gcc-15 CXX=/usr/bin/g++-15
|
||||
ENV CCACHE_DIR=/root/.cache/ccache
|
||||
ENV CMAKE_CXX_COMPILER_LAUNCHER=ccache
|
||||
|
||||
@@ -159,7 +165,7 @@ RUN apt-get update -y \
|
||||
WORKDIR /workspace
|
||||
|
||||
COPY requirements/build/rust.txt requirements/build/rust.txt
|
||||
RUN python3 -m pip install --no-cache-dir -r requirements/build/rust.txt
|
||||
RUN python3 -m pip install --break-system-packages --no-cache-dir -r requirements/build/rust.txt
|
||||
|
||||
ARG USE_SCCACHE
|
||||
ARG SCCACHE_ENDPOINT
|
||||
@@ -230,6 +236,7 @@ ENV VLLM_CPU_X86=${VLLM_CPU_X86}
|
||||
# Support for cross-compilation with ARM BF16 ISA: docker build --build-arg VLLM_CPU_ARM_BF16="true" ...
|
||||
ARG VLLM_CPU_ARM_BF16=0
|
||||
ENV VLLM_CPU_ARM_BF16=${VLLM_CPU_ARM_BF16}
|
||||
ARG ENABLE_AMX_FP8=1
|
||||
|
||||
WORKDIR /vllm-workspace
|
||||
|
||||
@@ -297,6 +304,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
fi; \
|
||||
sccache --show-stats; \
|
||||
fi && \
|
||||
CMAKE_ARGS="-DENABLE_AMX_FP8=${ENABLE_AMX_FP8}" \
|
||||
VLLM_TARGET_DEVICE=cpu python3 setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38 && \
|
||||
if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi
|
||||
|
||||
|
||||
@@ -608,6 +608,21 @@ def varlen_with_paged_kv(
|
||||
slot_mapping=slot_mapping,
|
||||
isa=isa,
|
||||
)
|
||||
ref_metadata = cpu_attn_get_scheduler_metadata(
|
||||
num_reqs=num_seqs,
|
||||
num_heads=num_query_heads,
|
||||
num_kv_heads=num_kv_heads,
|
||||
head_dim=head_size,
|
||||
seq_lens=kv_lens_tensor,
|
||||
dtype=dtype,
|
||||
query_start_loc=cu_query_lens,
|
||||
causal=dynamic_causal is None,
|
||||
sliding_window_size=(sliding_window if sliding_window is not None else -1),
|
||||
isa=isa,
|
||||
enable_kv_split=True,
|
||||
dynamic_causal=dynamic_causal_tensor,
|
||||
kv_cache_dtype="auto",
|
||||
)
|
||||
ref_output = torch.empty_like(query)
|
||||
cpu_attention_with_kv_cache(
|
||||
query=query,
|
||||
@@ -622,7 +637,7 @@ def varlen_with_paged_kv(
|
||||
sliding_window=sliding_window if sliding_window is not None else -1,
|
||||
block_table=block_tables,
|
||||
softcap=soft_cap if soft_cap is not None else 0,
|
||||
scheduler_metadata=metadata,
|
||||
scheduler_metadata=ref_metadata,
|
||||
s_aux=s_aux,
|
||||
dynamic_causal=dynamic_causal_tensor,
|
||||
)
|
||||
@@ -1166,3 +1181,135 @@ def test_varlen_with_paged_kv_dynamic_causal(
|
||||
kv_cache_dtype=kv_cache_dtype,
|
||||
dynamic_causal=dynamic_causal,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# AMX_FP8 (Diamond Rapids) tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _amx_fp8_available() -> bool:
|
||||
"""Return True iff the runtime reports AMX_FP8 capability."""
|
||||
return torch.cpu._is_amx_tile_supported() and torch.ops._C.cpu_attn_has_isa(
|
||||
"amx_fp8"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kv_cache_dtype", ["fp8_e4m3", "fp8_e5m2"])
|
||||
@pytest.mark.parametrize("seq_lens", SEQ_LENS)
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
@pytest.mark.parametrize("head_size", [64, 128, 256])
|
||||
@pytest.mark.parametrize("block_size", [32, 64])
|
||||
@pytest.mark.parametrize("sliding_window", [None])
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16])
|
||||
@pytest.mark.parametrize("soft_cap", [None])
|
||||
@pytest.mark.parametrize("num_blocks", NUM_BLOCKS)
|
||||
@pytest.mark.parametrize("use_alibi", [False])
|
||||
@pytest.mark.parametrize("use_sink", [False])
|
||||
@pytest.mark.parametrize("isa", ["amx_fp8"])
|
||||
@pytest.mark.skipif(
|
||||
not _amx_fp8_available(), reason="no AMX_FP8 support (requires Diamond Rapids)."
|
||||
)
|
||||
def test_varlen_with_paged_kv_amx_fp8(
|
||||
seq_lens: list[tuple[int, int]],
|
||||
num_heads: tuple[int, int],
|
||||
head_size: int,
|
||||
sliding_window: int | None,
|
||||
dtype: torch.dtype,
|
||||
block_size: int,
|
||||
soft_cap: float | None,
|
||||
num_blocks: int,
|
||||
use_alibi: bool,
|
||||
use_sink: bool,
|
||||
isa: str,
|
||||
kv_cache_dtype: str,
|
||||
) -> None:
|
||||
"""Test AMX_FP8 native FP8×FP8 attention (Diamond Rapids).
|
||||
|
||||
Verifies that:
|
||||
- QK uses _tile_dpfp8ps (native FP8 MMA, no K dequant).
|
||||
- PV uses native FP8 MMA for both E4M3 and E5M2.
|
||||
- Output cosine similarity vs fp32 reference > 0.99.
|
||||
"""
|
||||
if block_size % 64 != 0:
|
||||
pytest.skip("native FP8 PV requires block_size divisible by 64")
|
||||
|
||||
varlen_with_paged_kv(
|
||||
seq_lens=seq_lens,
|
||||
num_heads=num_heads,
|
||||
head_size=head_size,
|
||||
sliding_window=sliding_window,
|
||||
dtype=dtype,
|
||||
block_size=block_size,
|
||||
soft_cap=soft_cap,
|
||||
num_blocks=num_blocks,
|
||||
use_alibi=use_alibi,
|
||||
use_sink=use_sink,
|
||||
isa=isa,
|
||||
kv_cache_dtype=kv_cache_dtype,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not _amx_fp8_available(), reason="no AMX_FP8 support (requires Diamond Rapids)."
|
||||
)
|
||||
def test_amx_fp8_qk_covers_full_64_token_group() -> None:
|
||||
head_size = 64
|
||||
block_size = 64
|
||||
query = torch.ones((1, 1, head_size), dtype=torch.bfloat16)
|
||||
key = torch.empty((block_size, 1, head_size), dtype=torch.bfloat16)
|
||||
value = torch.empty_like(key)
|
||||
key[:32].fill_(-1)
|
||||
key[32:].fill_(1)
|
||||
value[:32].fill_(-1)
|
||||
value[32:].fill_(1)
|
||||
|
||||
key_cache = torch.empty((1, 1, block_size, head_size), dtype=torch.uint8)
|
||||
value_cache = torch.empty_like(key_cache)
|
||||
slot_mapping = torch.arange(block_size, dtype=torch.int64)
|
||||
cpu_attn_reshape_and_cache(
|
||||
key=key,
|
||||
value=value,
|
||||
key_cache=key_cache,
|
||||
value_cache=value_cache,
|
||||
slot_mapping=slot_mapping,
|
||||
isa="amx_fp8",
|
||||
kv_cache_dtype="fp8_e4m3",
|
||||
)
|
||||
|
||||
query_start_loc = torch.tensor([0, 1], dtype=torch.int32)
|
||||
seq_lens = torch.tensor([block_size], dtype=torch.int32)
|
||||
metadata = cpu_attn_get_scheduler_metadata(
|
||||
num_reqs=1,
|
||||
num_heads=1,
|
||||
num_kv_heads=1,
|
||||
head_dim=head_size,
|
||||
seq_lens=seq_lens,
|
||||
dtype=torch.bfloat16,
|
||||
query_start_loc=query_start_loc,
|
||||
causal=True,
|
||||
sliding_window_size=-1,
|
||||
isa="amx_fp8",
|
||||
enable_kv_split=False,
|
||||
kv_cache_dtype="fp8_e4m3",
|
||||
)
|
||||
output = torch.empty_like(query)
|
||||
cpu_attention_with_kv_cache(
|
||||
query=query,
|
||||
key_cache=key_cache,
|
||||
value_cache=value_cache,
|
||||
output=output,
|
||||
query_start_loc=query_start_loc,
|
||||
seq_lens=seq_lens,
|
||||
scale=head_size**-0.5,
|
||||
causal=True,
|
||||
alibi_slopes=None,
|
||||
sliding_window=-1,
|
||||
block_table=torch.tensor([[0]], dtype=torch.int32),
|
||||
softcap=0,
|
||||
scheduler_metadata=metadata,
|
||||
s_aux=None,
|
||||
kv_cache_dtype="fp8_e4m3",
|
||||
)
|
||||
|
||||
torch.testing.assert_close(output, torch.ones_like(output), atol=0.05, rtol=0)
|
||||
|
||||
Reference in New Issue
Block a user