@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 1c344ed41e

[CPU] [Feat] Add native AMX-FP8 attention impl for Diamond Rapids (#49410)

Signed-off-by: zhejiangxiaomai <zhenhui.zhao@intel.com>

zhao, zhenhui committed Sep 6, 2026 at 11:06 UTC 1c344ed41e1e3bd44d5836fbcf5ddceb56ef5ab9

10 files changed +1211 -61

.buildkite/scripts/hardware_ci/run-cpu-test.sh

+1 -1

@@ -57,7 +57,7 @@ prune_if_disk_pressure

57 # omits `--pull` so it already prefers a cached image over the network. A

58 # single attempt only -- the retry loop below covers a miss here too.

59 echo "--- :docker: Pre-fetching base image"

60 -docker pull ubuntu:22.04 || true

60 +docker pull ubuntu:25.04 || true

62 # building the docker image

63 echo "--- :docker: Building Docker image"

benchmarks/kernels/cpu/benchmark_cpu_attn.py

+52 -6

@@ -13,13 +13,26 @@ from vllm._custom_ops import (

13 cpu_attn_reshape_and_cache,

14 )

15 from vllm.utils.argparse_utils import FlexibleArgumentParser

16 -from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, set_random_seed

16 +from vllm.utils.torch_utils import (

17 + STR_DTYPE_TO_TORCH_DTYPE,

18 + is_quantized_kv_cache,

19 + set_random_seed,

20 +)

21 from vllm.v1.attention.backends.cpu_attn import CPUAttentionBackend, _get_attn_isa

23 +# Enable AMX tile data registers so isolated runs don't rely on other operations

24 +# to trigger oneDNN's _init_amx() first (which can cause hangs otherwise).

25 +if torch.cpu._is_amx_tile_supported():

26 + torch.cpu._init_amx()

27 +

28 +KV_CACHE_DTYPE_CHOICES = ["auto", "fp8_e4m3", "fp8_e5m2"]

29 +

31 def get_attn_isa(

32 block_size: int | None = None,

33 dtype: torch.dtype | None = None,

34 + head_size: int | None = None,

35 + kv_cache_dtype: str = "auto",

36 ):

37 # Delegate to _get_attn_isa so the fallback path applies the same arch

38 # gating (e.g. RISC-V RVV is only chosen when the build's hardcoded

@@ -28,6 +41,8 @@ def get_attn_isa(

41 return _get_attn_isa(

42 dtype if dtype is not None else torch.bfloat16,

43 block_size if block_size else 32,

44 + head_size=head_size,

45 + kv_cache_dtype=kv_cache_dtype if kv_cache_dtype != "auto" else None,

46 )

@@ -53,6 +68,7 @@ def main(

68 use_sink: bool = False,

69 enable_kv_split: bool = False,

70 isa: str | None = None,

71 + kv_cache_dtype: str = "auto",

72 seed: int = 0,

73 iters: int = 20,

74 ) -> None:

@@ -68,8 +84,18 @@ def main(

84 scale = head_size**-0.5

85 token_num = sum(query_lens)

87 + # Resolve kv_cache_dtype: "auto" means same as compute dtype (no quantization)

88 + effective_kv_cache_dtype = kv_cache_dtype if kv_cache_dtype != "auto" else None

89 + is_fp8_kv = (

90 + is_quantized_kv_cache(effective_kv_cache_dtype)

91 + if effective_kv_cache_dtype

92 + else False

93 + )

94 + # FP8 KV cache is stored as uint8 (byte-level view of fp8 values)

95 + kv_cache_torch_dtype = torch.uint8 if is_fp8_kv else dtype

96 +

97 if isa is None:

72 - isa = get_attn_isa(block_size, dtype)

98 + isa = get_attn_isa(block_size, dtype, kv_cache_dtype)

100 s_aux = (

101 15 * torch.rand((num_query_heads,), dtype=torch.bfloat16) if use_sink else None

@@ -98,9 +124,9 @@ def main(

124 )

125 key_cache, value_cache = key_value.unbind(0)

126

101 - # KV cache for CPU attention

127 + # KV cache for CPU attention; FP8 variants store quantized values as uint8

128 packed_key_cache = torch.empty(

103 - num_blocks, num_kv_heads, block_size, head_size, dtype=dtype

129 + num_blocks, num_kv_heads, block_size, head_size, dtype=kv_cache_torch_dtype

130 )

131 packed_value_cache = torch.empty_like(packed_key_cache)

132

@@ -122,6 +148,7 @@ def main(

148 value_cache=packed_value_cache,

149 slot_mapping=slot_mapping,

150 isa=isa,

151 + kv_cache_dtype=kv_cache_dtype,

152 )

153

154 metadata = cpu_attn_get_scheduler_metadata(

@@ -159,6 +186,7 @@ def main(

186 softcap=0,

187 scheduler_metadata=metadata,

188 s_aux=s_aux,

189 + kv_cache_dtype=kv_cache_dtype,

190 )

191 end_time = time.perf_counter_ns()

192 times.append((end_time - start_time) / 1e6)

@@ -233,7 +261,19 @@ if __name__ == "__main__":

261 )

262 parser.add_argument("--use-sink", action="store_true")

263 parser.add_argument(

236 - "--isa", type=str, choices=["vec", "neon", "amx", "vec16", "rvv"], default=None

264 + "--isa",

265 + type=str,

266 + choices=["vec", "neon", "amx", "amx_fp8", "vec16", "rvv"],

267 + default=None,

268 + )

269 + parser.add_argument(

270 + "--kv-cache-dtype",

271 + type=str,

272 + choices=KV_CACHE_DTYPE_CHOICES,

273 + default="auto",

274 + help="KV cache dtype: 'auto' uses compute dtype (bfloat16), "

275 + "'fp8_e4m3'/'fp8_e5m2' enables FP8 quantized KV cache "

276 + "(requires AMX-FP8 capable hardware for amx_fp8 ISA).",

277 )

278 parser.add_argument("--seed", type=int, default=0)

279 parser.add_argument("--iters", type=int, default=20)

@@ -264,7 +304,13 @@ if __name__ == "__main__":

304 enable_kv_split=args.enable_kv_split,

305 isa=args.isa

306 if args.isa is not None

267 - else get_attn_isa(args.block_size, STR_DTYPE_TO_TORCH_DTYPE[args.dtype]),

307 + else get_attn_isa(

308 + args.block_size,

309 + STR_DTYPE_TO_TORCH_DTYPE[args.dtype],

310 + args.head_size,

311 + args.kv_cache_dtype,

312 + ),

313 + kv_cache_dtype=args.kv_cache_dtype,

314 seed=args.seed,

315 iters=args.iters,

316 )

cmake/cpu_extension.cmake

+29

@@ -1,4 +1,5 @@

1 include(FetchContent)

2 +include(CheckCXXCompilerFlag)

4 set(CMAKE_CXX_STANDARD_REQUIRED ON)

5 set(CMAKE_CXX_STANDARD 20)

@@ -134,6 +135,12 @@ if (CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|amd64" OR ENABLE_X86_ISA)

135 CMAKE_CXX_COMPILER_VERSION VERSION_GREATER_EQUAL 12.3))

136 message(FATAL_ERROR "X86 backend requires gcc/g++ >= 12.3")

137 endif()

138 + # Work around a GCC 15 optimizer crash in oneDNN (tree-ssa-pre pass).

139 + # Keep this scoped to GCC 15+ so older toolchains are unaffected.

140 + if (CMAKE_CXX_COMPILER_ID STREQUAL "GNU" AND

141 + CMAKE_CXX_COMPILER_VERSION VERSION_GREATER_EQUAL 15)

142 + list(APPEND CXX_COMPILE_FLAGS "-fno-tree-pre")

143 + endif()

144 list(APPEND CXX_COMPILE_FLAGS "-mf16c")

145 list(APPEND CXX_COMPILE_FLAGS_AVX512 ${CXX_COMPILE_FLAGS})

146 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)

155 "-mamx-tile"

156 "-mavx512bf16"

157 "-mavx512vnni")

158 + # Some packaged GCC 14 builds still do not accept -mamx-fp8, so gate on

159 + # actual flag support rather than compiler version alone.

160 + # Option to enable/disable AMX-FP8 support (default: ON)

161 + option(ENABLE_AMX_FP8 "Enable AMX-FP8 support" ON)

162 +

163 + if (ENABLE_AMX_FP8)

164 + check_cxx_compiler_flag("-mamx-fp8" COMPILER_SUPPORTS_AMX_FP8_FLAG)

165 + if (COMPILER_SUPPORTS_AMX_FP8_FLAG)

166 + list(APPEND CXX_COMPILE_FLAGS_AVX512_AMX "-mamx-fp8")

167 + set(AMX_FP8_SUPPORTED TRUE)

168 + else()

169 + set(AMX_FP8_SUPPORTED FALSE)

170 + message(STATUS "AMX-FP8 disabled: compiler does not support -mamx-fp8 (compiler: ${CMAKE_CXX_COMPILER_ID} ${CMAKE_CXX_COMPILER_VERSION})")

171 + endif()

172 + else()

173 + set(AMX_FP8_SUPPORTED FALSE)

174 + message(STATUS "AMX-FP8 disabled by ENABLE_AMX_FP8=OFF")

175 + endif()

176 list(APPEND CXX_COMPILE_FLAGS_AVX2

177 "-mavx2")

178 elseif (POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND)

@@ -581,6 +606,10 @@ if (ENABLE_X86_ISA)

606

607 # For AMX kernels

608 target_compile_definitions(_C PRIVATE "-DCPU_CAPABILITY_AMXBF16")

609 + if (AMX_FP8_SUPPORTED)

610 + target_compile_definitions(_C PRIVATE "-DCPU_CAPABILITY_AMXFP8")

611 + message(STATUS "AMX-FP8 (Diamond Rapids) enabled")

612 + endif()

613

614 # AVX512F

615 define_extension_target(

csrc/cpu/cpu_attn.cpp

+30 -2

@@ -1,5 +1,17 @@

1 #include "cpu_attn_dispatch_generated.h"

3 +#if defined(__x86_64__) || defined(_M_X64)

4 + #include <cpuid.h>

5 +// Runtime check for AMX-FP8 (Diamond Rapids): CPUID leaf 7, subleaf 1, EAX[21].

6 +static bool runtime_has_amx_fp8() {

7 + unsigned int eax = 0, ebx = 0, ecx = 0, edx = 0;

8 + if (__get_cpuid_count(7, 1, &eax, &ebx, &ecx, &edx)) {

9 + return (eax >> 21) & 1u;

10 + }

11 + return false;

12 +}

13 +#endif // x86_64

14 +

15 // Maps kv_cache_dtype string to Fp8KVCacheDataType enum.

16 // "auto" -> kAuto(0); "fp8"/"fp8_e4m3" -> kFp8E4M3; "fp8_e5m2" -> kFp8E5M2.

17 static inline cpu_attention::Fp8KVCacheDataType parse_fp8_kv_dtype(

@@ -18,6 +30,16 @@ bool cpu_attn_has_isa(const std::string& isa) {

30 return true;

31 #else

32 return false;

33 +#endif

34 + }

35 + if (isa == "amx_fp8") {

36 +#ifdef CPU_CAPABILITY_AMXFP8

37 + // Guard with runtime CPUID check: the binary may be compiled with

38 + // -mamx-fp8 (CPU_CAPABILITY_AMXFP8 defined) but run on a CPU that

39 + // has AMX-BF16 without AMX-FP8 (e.g. Sapphire Rapids, Emerald Rapids).

40 + return runtime_has_amx_fp8();

41 +#else

42 + return false;

43 #endif

44 }

45 return false;

@@ -35,6 +57,8 @@ torch::Tensor get_scheduler_metadata(

57 cpu_attention::ISA isa;

58 if (isa_hint == "amx") {

59 isa = cpu_attention::ISA::AMX;

60 + } else if (isa_hint == "amx_fp8") {

61 + isa = cpu_attention::ISA::AMX_FP8;

62 } else if (isa_hint == "vec") {

63 isa = cpu_attention::ISA::VEC;

64 } else if (isa_hint == "vec16") {

@@ -130,6 +154,8 @@ void cpu_attn_reshape_and_cache(

154 cpu_attention::ISA isa_tag = [&]() {

155 if (isa == "amx") {

156 return cpu_attention::ISA::AMX;

157 + } else if (isa == "amx_fp8") {

158 + return cpu_attention::ISA::AMX_FP8;

159 } else if (isa == "vec") {

160 return cpu_attention::ISA::VEC;

161 } else if (isa == "vec16") {

@@ -149,8 +175,9 @@ void cpu_attn_reshape_and_cache(

175

176 if (is_fp8) {

177 TORCH_CHECK(isa_tag == cpu_attention::ISA::AMX ||

178 + isa_tag == cpu_attention::ISA::AMX_FP8 ||

179 isa_tag == cpu_attention::ISA::VEC,

153 - "FP8 KV cache is only supported on x86 (AMX/VEC) ISA");

180 + "FP8 KV cache is only supported on x86 (AMX_FP8/AMX/VEC) ISA");

181 }

182

183 VLLM_DISPATCH_FLOATING_TYPES(

@@ -237,8 +264,9 @@ void cpu_attention_with_kv_cache(

264 input.k_scale_fp8 = static_cast<float>(k_scale);

265 input.v_scale_fp8 = static_cast<float>(v_scale);

266 TORCH_CHECK(input.metadata->isa == cpu_attention::ISA::AMX ||

267 + input.metadata->isa == cpu_attention::ISA::AMX_FP8 ||

268 input.metadata->isa == cpu_attention::ISA::VEC,

241 - "FP8 KV cache is only supported on x86 (AMX/VEC) ISA");

269 + "FP8 KV cache is only supported on x86 (AMX_FP8/AMX/VEC) ISA");

270 }

271

272 VLLM_DISPATCH_FLOATING_TYPES(

csrc/cpu/cpu_attn_amx_fp8.hpp new

+795

@@ -0,0 +1,795 @@

1 +// SPDX-License-Identifier: Apache-2.0

2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +//

4 +// AMX_FP8 attention implementation for Intel Diamond Rapids (DMR).

5 +//

6 +// Architecture:

7 +// QK phase: native FP8×FP8 MMA via _tile_dpfp8ps / _tile_dpbf8ps

8 +// Q is quantized from BF16 to the KV cache FP8 type.

9 +// K cache is stored in the native AMX-FP8 layout (1 byte/element).

10 +// PV phase: native FP8 MMA with 64-byte VNNI4 tiles.

11 +// Softmax writes P once in compact FP8 form for reuse across all

12 +// output-dimension groups.

13 +//

14 +// Scale bookkeeping for QK:

15 +// true_score = Q_fp32 · K_fp32

16 +// = (Q_fp8 * q_scale) · (K_fp8 * k_scale)

17 +// = dot(Q_fp8, K_fp8) * q_scale * k_scale

18 +// execute_attention() applies: scale *= q_dynamic_scale * k_scale

19 +// where q_dynamic_scale is computed per-tile in copy_q_heads_tile().

20 +//

21 +// Scale bookkeeping for PV:

22 +// raw_C = dot(P_fp8, V_fp8)

23 +// ≈ dot(P/p_scale, V_real/v_scale)

24 +// = dot(P, V_real) / (p_scale * v_scale)

25 +// In final_output: multiply by output_v_scale = p_scale * v_scale.

26 +//

27 +#ifndef CPU_ATTN_AMX_FP8_HPP

28 +#define CPU_ATTN_AMX_FP8_HPP

29 +

30 +#ifdef CPU_CAPABILITY_AMXFP8

31 +

32 + #include "cpu_attn_amx.hpp" // pulls in TileGemm224/TileGemm122 BF16 specialisations

33 + #include "cpu_attn_fp8.hpp"

34 + #include "cpu_attn_impl.hpp"

35 + #include "cpu_types.hpp"

36 +

37 + #include <amxfp8intrin.h> // requires GCC 14+ / Clang 17+, -mamx-fp8

38 +

39 +namespace cpu_attention {

40 +

41 +// ---------------------------------------------------------------------------

42 +// Native-FP8 K cache layout reshape helper.

43 +//

44 +// Layout per 16-token group (token_num_per_group = AMX_TILE_ROW_NUM = 16):

45 +// k-slice s (s = 0 .. head_dim/64 - 1):

46 +// row k (k = 0..15): four consecutive K values for each of 16 tokens

47 +// (VNNI4 B-operand layout)

48 +// One k-slice = 16 rows × 64 bytes = 1024 bytes = 1 AMX FP8 tile.

49 +// Full group = (head_dim / 64) slices × 1024 bytes.

50 +//

51 +// For head_dim=128: group = 2 × 1024 = 2048 bytes.

52 +// Consecutive 16-token groups are contiguous within each cache block.

53 +// ---------------------------------------------------------------------------

54 +template <typename scalar_t, uint8_t (*quant_fn)(float, float)>

55 +inline void reshape_and_cache_k_amx_fp8_impl(

56 + const scalar_t* key_ptr, uint8_t* key_cache_ptr, const int64_t* slot_ptr,

57 + int64_t token_num, int64_t head_num, int64_t head_dim, int64_t block_size,

58 + int64_t k_stride0, int64_t k_stride1, int64_t kc_stride0,

59 + int64_t kc_stride1, float k_inv) {

60 + constexpr int64_t token_num_per_group = AMX_TILE_ROW_NUM; // 16

61 + constexpr int64_t fp8_per_row = AMX_TILE_ROW_BYTES; // 64 FP8/row

62 + const int64_t slice_num = head_dim / fp8_per_row;

63 + // Bytes per 16-token group = slice_num * (16 * 64) = head_dim * 16

64 + const int64_t group_bytes = token_num_per_group * head_dim;

65 +

66 + #pragma omp parallel for collapse(2) schedule(static)

67 + for (int64_t tok = 0; tok < token_num; ++tok) {

68 + for (int64_t h = 0; h < head_num; ++h) {

69 + const int64_t slot = slot_ptr[tok];

70 + if (slot < 0) continue;

71 + const int64_t block_idx = slot / block_size;

72 + const int64_t block_offset = slot % block_size;

73 + const int64_t group_idx = block_offset / token_num_per_group;

74 + const int64_t group_offset = block_offset % token_num_per_group;

75 +

76 + const scalar_t* ksrc = key_ptr + tok * k_stride0 + h * k_stride1;

77 + uint8_t* group_base = key_cache_ptr + block_idx * kc_stride0 +

78 + h * kc_stride1 + group_idx * group_bytes;

79 +

80 + // Each k-slice is one 16x64-byte VNNI4 B tile. Row k stores dimensions

81 + // [4*k, 4*k+4) for every token; each token occupies four bytes.

82 + for (int64_t s = 0; s < slice_num; ++s) {

83 + uint8_t* tile_base = group_base + s * token_num_per_group * fp8_per_row;

84 + for (int64_t k = 0; k < fp8_per_row / 4; ++k) {

85 + uint8_t* packed_values =

86 + tile_base + k * fp8_per_row + group_offset * 4;

87 + for (int64_t lane = 0; lane < 4; ++lane) {

88 + packed_values[lane] = quant_fn(

89 + static_cast<float>(ksrc[s * fp8_per_row + k * 4 + lane]),

90 + k_inv);

91 + }

92 + }

93 + }

94 + }

95 + }

96 +}

97 +

98 +template <typename fp8_t>

99 +__attribute__((target("avx10.2"), noinline)) void quantize_probability_row(

100 + const float* __restrict__ src, uint8_t* __restrict__ dst, int32_t cols) {

101 + constexpr float fp8_max =

102 + std::is_same_v<fp8_t, c10::Float8_e4m3fn> ? 448.0f : 57344.0f;

103 + const __m512 scale = _mm512_set1_ps(fp8_max);

104 + for (int32_t col = 0; col < cols; col += 16) {

105 + const __m512 scaled = _mm512_mul_ps(_mm512_loadu_ps(src + col), scale);

106 + const __m256h fp16 = (__m256h)_mm512_cvtps_ph(

107 + scaled, _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC);

108 + if constexpr (std::is_same_v<fp8_t, c10::Float8_e4m3fn>) {

109 + _mm_storeu_si128(reinterpret_cast<__m128i*>(dst + col),

110 + _mm256_cvtph_hf8(fp16));

111 + } else {

112 + _mm_storeu_si128(reinterpret_cast<__m128i*>(dst + col),

113 + _mm256_cvtph_bf8(fp16));

114 + }

115 + }

116 +}

117 +

118 +template <typename fp8_t>

119 +__attribute__((target("avx10.2"), noinline)) void quantize_q_tile(

120 + const c10::BFloat16* __restrict__ src, uint8_t* __restrict__ dst,

121 + int32_t q_num, int32_t q_heads_per_kv, int64_t q_num_stride,

122 + int64_t q_head_stride, int64_t head_dim, float inv_scale) {

123 + const __m512 scale = _mm512_set1_ps(inv_scale);

124 + constexpr float fp8_max =

125 + std::is_same_v<fp8_t, c10::Float8_e4m3fn> ? 448.0f : 57344.0f;

126 + const __m512 max_value = _mm512_set1_ps(fp8_max);

127 + const __m512 min_value = _mm512_set1_ps(-fp8_max);

128 + const int64_t head_size_block_num = head_dim / AMX_TILE_ROW_BYTES;

129 + int32_t packed_row = 0;

130 +

131 + for (int32_t q = 0; q < q_num; ++q) {

132 + const c10::BFloat16* head = src + q * q_num_stride;

133 + for (int32_t h = 0; h < q_heads_per_kv;

134 + ++h, head += q_head_stride, ++packed_row) {

135 + uint8_t* packed_head =

136 + dst +

137 + (packed_row / AMX_TILE_ROW_NUM) * head_size_block_num *

138 + AMX_TILE_BYTES +

139 + (packed_row % AMX_TILE_ROW_NUM) * AMX_TILE_ROW_BYTES;

140 + for (int64_t e = 0; e < head_dim; e += 16) {

141 + const __m256i bf16 =

142 + _mm256_loadu_si256(reinterpret_cast<const __m256i*>(head + e));

143 + __m512 values = _mm512_castsi512_ps(

144 + _mm512_slli_epi32(_mm512_cvtepu16_epi32(bf16), 16));

145 + values = _mm512_mul_ps(values, scale);

146 + values = _mm512_min_ps(values, max_value);

147 + values = _mm512_max_ps(values, min_value);

148 + const __m256h fp16 = (__m256h)_mm512_cvtps_ph(

149 + values, _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC);

150 + const int64_t block = e / AMX_TILE_ROW_BYTES;

151 + const int64_t offset = e % AMX_TILE_ROW_BYTES;

152 + __m128i fp8;

153 + if constexpr (std::is_same_v<fp8_t, c10::Float8_e4m3fn>) {

154 + fp8 = _mm256_cvtph_hf8(fp16);

155 + } else {

156 + fp8 = _mm256_cvtph_bf8(fp16);

157 + }

158 + _mm_storeu_si128(reinterpret_cast<__m128i*>(

159 + packed_head + block * AMX_TILE_BYTES + offset),

160 + fp8);

161 + }

162 + }

163 + }

164 +}

165 +

166 +// ---------------------------------------------------------------------------

167 +// ---------------------------------------------------------------------------

168 +// TileGemm224<uint8_t, kv_cache_t>: 2-2-4 pattern for AMX_FP8.

169 +//

170 +// QK phase: native FP8 × FP8 MMA (_tile_dphf8ps / _tile_dpbf8ps).

171 +// PV phase: native wide FP8 MMA (colsb=64 tiles).

172 +// P is quantized once per KV tile and reused by every output-dimension group.

173 +// V (FP8 E4M3) loaded directly from cache (no dequant).

174 +// Result correction is folded into the final output scale.

175 +// ---------------------------------------------------------------------------

176 +template <typename kv_cache_t>

177 +class TileGemm224<uint8_t, kv_cache_t> {

178 + static_assert(std::is_same_v<kv_cache_t, c10::Float8_e4m3fn> ||

179 + std::is_same_v<kv_cache_t, c10::Float8_e5m2>,

180 + "kv_cache_t must be Float8_e4m3fn or Float8_e5m2");

181 +

182 + // FP8 tile size in uint8_t elements (= AMX_TILE_BYTES, one per byte).

183 + static constexpr int64_t fp8_tile_elems = AMX_TILE_BYTES; // 1024

184 +

185 + public:

186 + static constexpr bool prequantize_probabilities = true;

187 +

188 + FORCE_INLINE static void quantize_probability_row(const float* src,

189 + uint8_t* dst,

190 + int32_t cols) {

191 + cpu_attention::quantize_probability_row<kv_cache_t>(src, dst, cols);

192 + }

193 +

194 + // ------------------------------------------------------------------

195 + // QK: k_times = head_dim / 64 (each FP8 tile covers 64 k-elements).

196 + // PV: native wide FP8 MMA.

197 + // ------------------------------------------------------------------

198 + template <AttentionGemmPhase phase, int32_t k_size>

199 + FORCE_INLINE static void gemm(const int32_t m_size, void* __restrict__ a_tile,

200 + void* __restrict__ b_tile,

201 + float* __restrict__ c_tile, const int64_t lda,

202 + const int64_t ldb, const int64_t ldc,

203 + const int32_t block_size,

204 + const int32_t dynamic_k_size,

205 + const bool accum_c) {

206 + if constexpr (phase == AttentionGemmPhase::PV) {

207 + // -------- Native wide FP8 PV (FP8 × FP8 → FP32) --------

208 + //

209 + // A fixed FP8-dependent P scale lets the tile accumulate all k-steps

210 + // before get_output_v_scale() applies the correction.

211 +

212 + const int32_t m_0 =

213 + std::min(m_size, static_cast<int32_t>(AMX_TILE_ROW_NUM));

214 + const int32_t m_1 = m_size - m_0;

215 +

216 + constexpr int32_t v_stride = AMX_TILE_ROW_BYTES;

217 + const int32_t k_times =

218 + dynamic_k_size / (AMX_TILE_ROW_NUM * 4 / sizeof(uint8_t));

219 + constexpr int32_t a_advance = AMX_TILE_ROW_BYTES;

220 + constexpr int32_t v_advance = AMX_TILE_BYTES;

221 + const int32_t c_stride = ldc * sizeof(float);

222 +

223 + for (int32_t output_group = 0; output_group < 2; ++output_group) {

224 + uint8_t* a0 = reinterpret_cast<uint8_t*>(a_tile);

225 + uint8_t* a1 = a0 + lda * AMX_TILE_ROW_NUM;

226 + uint8_t* v2 = reinterpret_cast<uint8_t*>(b_tile) +

227 + output_group * 2 * block_size * AMX_TILE_ROW_BYTES / 4;

228 + uint8_t* v3 = v2 + block_size * AMX_TILE_ROW_BYTES / 4;

229 + float* c4 = c_tile + output_group * 2 * AMX_TILE_ROW_NUM;

230 + float* c5 = c4 + AMX_TILE_ROW_NUM;

231 + float* c6 = c4 + AMX_TILE_ROW_NUM * ldc;

232 + float* c7 = c6 + AMX_TILE_ROW_NUM;

233 +

234 + if (accum_c) {

235 + _tile_loadd(4, c4, c_stride);

236 + _tile_loadd(5, c5, c_stride);

237 + if (m_1 > 0) {

238 + _tile_loadd(6, c6, c_stride);

239 + _tile_loadd(7, c7, c_stride);

240 + }

241 + } else {

242 + _tile_zero(4);

243 + _tile_zero(5);

244 + if (m_1 > 0) {

245 + _tile_zero(6);

246 + _tile_zero(7);

247 + }

248 + }

249 +

250 + for (int32_t k = 0; k < k_times; ++k) {

251 + _tile_loadd(0, a0, lda);

252 + _tile_stream_loadd(2, v2, v_stride);

253 + _tile_stream_loadd(3, v3, v_stride);

254 + if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e4m3fn>) {

255 + _tile_dphf8ps(4, 0, 2);

256 + _tile_dphf8ps(5, 0, 3);

257 + } else {

258 + _tile_dpbf8ps(4, 0, 2);

259 + _tile_dpbf8ps(5, 0, 3);

260 + }

261 + if (m_1 > 0) {

262 + _tile_loadd(1, a1, lda);

263 + if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e4m3fn>) {

264 + _tile_dphf8ps(6, 1, 2);

265 + _tile_dphf8ps(7, 1, 3);

266 + } else {

267 + _tile_dpbf8ps(6, 1, 2);

268 + _tile_dpbf8ps(7, 1, 3);

269 + }

270 + }

271 + a0 += a_advance;

272 + a1 += a_advance;

273 + v2 += v_advance;

274 + v3 += v_advance;

275 + }

276 +

277 + _tile_stored(4, c4, c_stride);

278 + _tile_stored(5, c5, c_stride);

279 + if (m_1 > 0) {

280 + _tile_stored(6, c6, c_stride);

281 + _tile_stored(7, c7, c_stride);

282 + }

283 + }

284 + return;

285 + }

286 +

287 + // -------- QK phase: native FP8 MMA --------

288 + // k_times: each FP8 AMX tile covers (AMX_TILE_ROW_NUM * 4) FP8 = 64

289 + // inner-product elements.

290 + const int32_t k_times =

291 + dynamic_k_size / (AMX_TILE_ROW_NUM * 4 / sizeof(uint8_t));

292 +

293 + const int32_t b_tile_stride = AMX_TILE_ROW_BYTES; // 64 bytes

294 +

295 + const int32_t m_0 =

296 + std::min(m_size, static_cast<int32_t>(AMX_TILE_ROW_NUM));

297 + const int32_t m_1 = m_size - m_0;

298 + const int32_t c_tile_stride = ldc * sizeof(float);

299 + constexpr int32_t output_group_num = 2;

300 + for (int32_t output_group = 0; output_group < output_group_num;

301 + ++output_group) {

302 + uint8_t* __restrict__ a_tile_0 = static_cast<uint8_t*>(a_tile);

303 + uint8_t* __restrict__ a_tile_1 = a_tile_0 + k_times * fp8_tile_elems;

304 + uint8_t* __restrict__ b_tile_2 =

305 + static_cast<uint8_t*>(b_tile) +

306 + output_group * 2 * k_size * AMX_TILE_ROW_NUM;

307 + uint8_t* __restrict__ b_tile_3 = b_tile_2 + (k_size * AMX_TILE_ROW_NUM);

308 + float* __restrict__ c_tile_4 = c_tile + output_group * 32;

309 + float* __restrict__ c_tile_5 = c_tile_4 + AMX_TILE_ROW_NUM;

310 + float* __restrict__ c_tile_6 = c_tile_4 + AMX_TILE_ROW_NUM * ldc;

311 + float* __restrict__ c_tile_7 = c_tile_6 + AMX_TILE_ROW_NUM;

312 +

313 + if (accum_c) {

314 + _tile_loadd(4, c_tile_4, c_tile_stride);

315 + _tile_loadd(5, c_tile_5, c_tile_stride);

316 + _tile_loadd(6, c_tile_6, c_tile_stride);

317 + _tile_loadd(7, c_tile_7, c_tile_stride);

318 + } else {

319 + _tile_zero(4);

320 + _tile_zero(5);

321 + _tile_zero(6);

322 + _tile_zero(7);

323 + }

324 +

325 + for (int32_t k = 0; k < k_times; ++k) {

326 + _tile_loadd(0, a_tile_0, AMX_TILE_ROW_BYTES);

327 + _tile_stream_loadd(2, b_tile_2, b_tile_stride);

328 + if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e4m3fn>) {

329 + _tile_dphf8ps(4, 0, 2);

330 + } else {

331 + _tile_dpbf8ps(4, 0, 2);

332 + }

333 + _tile_stream_loadd(3, b_tile_3, b_tile_stride);

334 + if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e4m3fn>) {

335 + _tile_dphf8ps(5, 0, 3);

336 + } else {

337 + _tile_dpbf8ps(5, 0, 3);

338 + }

339 + _tile_loadd(1, a_tile_1, AMX_TILE_ROW_BYTES);

340 + if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e4m3fn>) {

341 + _tile_dphf8ps(6, 1, 2);

342 + _tile_dphf8ps(7, 1, 3);

343 + } else {

344 + _tile_dpbf8ps(6, 1, 2);

345 + _tile_dpbf8ps(7, 1, 3);

346 + }

347 +

348 + // Q is packed as [row group][K slice].

349 + a_tile_0 += fp8_tile_elems;

350 + a_tile_1 += fp8_tile_elems;

351 + // Advance K cache (native layout): 1 FP8 tile per k-step.

352 + b_tile_2 += fp8_tile_elems;

353 + b_tile_3 += fp8_tile_elems;

354 + }

355 +

356 + _tile_stored(4, c_tile_4, c_tile_stride);

357 + _tile_stored(5, c_tile_5, c_tile_stride);

358 + _tile_stored(6, c_tile_6, c_tile_stride);

359 + _tile_stored(7, c_tile_7, c_tile_stride);

360 + }

361 + }

362 +

363 + FORCE_INLINE static void init_tile_config(int32_t m, __tilecfg& config) {

364 + // QK and PV share one 64-byte tile configuration.

365 + const int32_t m_0 = std::min(m, static_cast<int32_t>(AMX_TILE_ROW_NUM));

366 + const int32_t m_1 = m - m_0;

367 + // Use all 64 bytes per row for every tile (same as BF16 AMX).

368 + // This works for both FP8 (64 FP8/row) and BF16 (32 BF16/row = 64B).

369 + for (int i = 0; i < 8; ++i) config.colsb[i] = AMX_TILE_ROW_BYTES;

370 + config.rows[0] = m_0;

371 + config.rows[1] = m_1;

372 + config.rows[2] = AMX_TILE_ROW_NUM;

373 + config.rows[3] = AMX_TILE_ROW_NUM;

374 + config.rows[4] = m_0;

375 + config.rows[5] = m_0;

376 + config.rows[6] = m_1;

377 + config.rows[7] = m_1;

378 + _tile_loadconfig(&config);

379 + }

380 +};

381 +

382 +// ---------------------------------------------------------------------------

383 +// TileGemm122<uint8_t, kv_cache_t>: 1-2-2 pattern for AMX_FP8.

384 +// ---------------------------------------------------------------------------

385 +template <typename kv_cache_t>

386 +class TileGemm122<uint8_t, kv_cache_t> {

387 + static_assert(std::is_same_v<kv_cache_t, c10::Float8_e4m3fn> ||

388 + std::is_same_v<kv_cache_t, c10::Float8_e5m2>,

389 + "kv_cache_t must be Float8_e4m3fn or Float8_e5m2");

390 +

391 + static constexpr int64_t fp8_tile_elems = AMX_TILE_BYTES;

392 +

393 + public:

394 + static constexpr bool prequantize_probabilities = true;

395 +

396 + FORCE_INLINE static void quantize_probability_row(const float* src,

397 + uint8_t* dst,

398 + int32_t cols) {

399 + cpu_attention::quantize_probability_row<kv_cache_t>(src, dst, cols);

400 + }

401 +

402 + template <AttentionGemmPhase phase, int32_t k_size>

403 + FORCE_INLINE static void gemm(const int32_t m_size, void* __restrict__ a_tile,

404 + void* __restrict__ b_tile,

405 + float* __restrict__ c_tile, const int64_t lda,

406 + const int64_t ldb, const int64_t ldc,

407 + const int32_t block_size,

408 + const int32_t dynamic_k_size,

409 + const bool accum_c) {

410 + if constexpr (phase == AttentionGemmPhase::PV) {

411 + // Native wide FP8 PV (1-2-2 pattern)

412 + // A fixed FP8-dependent P scale enables accumulation across all k-steps.

413 + constexpr int32_t v_stride = AMX_TILE_ROW_BYTES;

414 + constexpr int32_t a_advance = AMX_TILE_ROW_BYTES;

415 + constexpr int32_t v_advance = AMX_TILE_BYTES;

416 + const int32_t k_times =

417 + dynamic_k_size / (AMX_TILE_ROW_NUM * 4 / sizeof(uint8_t));

418 +

419 + uint8_t* a = reinterpret_cast<uint8_t*>(a_tile);

420 + uint8_t* v2 = reinterpret_cast<uint8_t*>(b_tile);

421 + uint8_t* v3 = v2 + (block_size * AMX_TILE_ROW_BYTES / 4);

422 + uint8_t* v4 = v3 + (block_size * AMX_TILE_ROW_BYTES / 4);

423 + uint8_t* v5 = v4 + (block_size * AMX_TILE_ROW_BYTES / 4);

424 + float* c4 = c_tile + 2 * AMX_TILE_ROW_NUM;

425 + float* c5 = c_tile + 3 * AMX_TILE_ROW_NUM;

426 + float* c6 = c_tile;

427 + float* c7 = c6 + AMX_TILE_ROW_BYTES / sizeof(float);

428 +

429 + const int32_t c_stride = ldc * sizeof(float);

430 + if (accum_c) {

431 + _tile_loadd(4, c4, c_stride);

432 + _tile_loadd(5, c5, c_stride);

433 + _tile_loadd(6, c6, c_stride);

434 + _tile_loadd(7, c7, c_stride);

435 + } else {

436 + _tile_zero(4);

437 + _tile_zero(5);

438 + _tile_zero(6);

439 + _tile_zero(7);

440 + }

441 +

442 + for (int32_t k = 0; k < k_times; ++k) {

443 + _tile_loadd(0, a, lda);

444 + _tile_stream_loadd(2, v2, v_stride);

445 + _tile_stream_loadd(3, v3, v_stride);

446 + if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e4m3fn>) {

447 + _tile_dphf8ps(6, 0, 2);

448 + _tile_dphf8ps(7, 0, 3);

449 + } else {

450 + _tile_dpbf8ps(6, 0, 2);

451 + _tile_dpbf8ps(7, 0, 3);

452 + }

453 + _tile_stream_loadd(2, v4, v_stride);

454 + _tile_stream_loadd(3, v5, v_stride);

455 + if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e4m3fn>) {

456 + _tile_dphf8ps(4, 0, 2);

457 + _tile_dphf8ps(5, 0, 3);

458 + } else {

459 + _tile_dpbf8ps(4, 0, 2);

460 + _tile_dpbf8ps(5, 0, 3);

461 + }

462 + a += a_advance;

463 + v2 += v_advance;

464 + v3 += v_advance;

465 + v4 += v_advance;

466 + v5 += v_advance;

467 + }

468 +

469 + _tile_stored(4, c4, c_stride);

470 + _tile_stored(5, c5, c_stride);

471 + _tile_stored(6, c6, c_stride);

472 + _tile_stored(7, c7, c_stride);

473 + return;

474 + }

475 +

476 + // -------- QK phase: native FP8 MMA (1-2-2 pattern) --------

477 + const int32_t k_times =

478 + dynamic_k_size / (AMX_TILE_ROW_NUM * 4 / sizeof(uint8_t));

479 + const int32_t k_group_times = k_times / 2;

480 + const bool has_tail = (k_times % 2 == 1);

481 +

482 + const int32_t b_stride = AMX_TILE_ROW_BYTES;

483 + const int32_t c_stride = ldc * sizeof(float);

484 + constexpr int32_t output_group_num = 2;

485 + for (int32_t output_group = 0; output_group < output_group_num;

486 + ++output_group) {

487 + uint8_t* __restrict__ a_tile_0 = static_cast<uint8_t*>(a_tile);

488 + uint8_t* __restrict__ a_tile_1 = a_tile_0 + fp8_tile_elems;

489 + uint8_t* __restrict__ b_tile_2 =

490 + static_cast<uint8_t*>(b_tile) +

491 + output_group * 2 * k_size * AMX_TILE_ROW_NUM;

492 + uint8_t* __restrict__ b_tile_3 = b_tile_2 + (k_size * AMX_TILE_ROW_NUM);

493 + uint8_t* __restrict__ b_tile_4 = b_tile_2 + fp8_tile_elems;

494 + uint8_t* __restrict__ b_tile_5 = b_tile_3 + fp8_tile_elems;

495 + float* __restrict__ c_tile_6 = c_tile + output_group * 32;

496 + float* __restrict__ c_tile_7 = c_tile_6 + AMX_TILE_ROW_NUM;

497 +

498 + if (accum_c) {

499 + _tile_loadd(6, c_tile_6, c_stride);

500 + _tile_loadd(7, c_tile_7, c_stride);

501 + } else {

502 + _tile_zero(6);

503 + _tile_zero(7);

504 + }

505 +

506 + for (int32_t k = 0; k < k_group_times; ++k) {

507 + _tile_loadd(0, a_tile_0, AMX_TILE_ROW_BYTES);

508 + _tile_stream_loadd(2, b_tile_2, b_stride);

509 + if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e4m3fn>) {

510 + _tile_dphf8ps(6, 0, 2);

511 + } else {

512 + _tile_dpbf8ps(6, 0, 2);

513 + }

514 + _tile_stream_loadd(3, b_tile_3, b_stride);

515 + if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e4m3fn>) {

516 + _tile_dphf8ps(7, 0, 3);

517 + } else {

518 + _tile_dpbf8ps(7, 0, 3);

519 + }

520 + _tile_loadd(1, a_tile_1, AMX_TILE_ROW_BYTES);

521 + _tile_stream_loadd(2, b_tile_4, b_stride);

522 + if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e4m3fn>) {

523 + _tile_dphf8ps(6, 1, 2);

524 + } else {

525 + _tile_dpbf8ps(6, 1, 2);

526 + }

527 + _tile_stream_loadd(3, b_tile_5, b_stride);

528 + if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e4m3fn>) {

529 + _tile_dphf8ps(7, 1, 3);

530 + } else {

531 + _tile_dpbf8ps(7, 1, 3);

532 + }

533 +

534 + a_tile_0 += 2 * fp8_tile_elems;

535 + a_tile_1 += 2 * fp8_tile_elems;

536 + b_tile_2 += 2 * fp8_tile_elems;

537 + b_tile_3 += 2 * fp8_tile_elems;

538 + b_tile_4 += 2 * fp8_tile_elems;

539 + b_tile_5 += 2 * fp8_tile_elems;

540 + }

541 +

542 + if (has_tail) {

543 + _tile_loadd(0, a_tile_0, AMX_TILE_ROW_BYTES);

544 + _tile_stream_loadd(2, b_tile_2, b_stride);

545 + if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e4m3fn>) {

546 + _tile_dphf8ps(6, 0, 2);

547 + } else {

548 + _tile_dpbf8ps(6, 0, 2);

549 + }

550 + _tile_stream_loadd(3, b_tile_3, b_stride);

551 + if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e4m3fn>) {

552 + _tile_dphf8ps(7, 0, 3);

553 + } else {

554 + _tile_dpbf8ps(7, 0, 3);

555 + }

556 + }

557 +

558 + _tile_stored(6, c_tile_6, c_stride);

559 + _tile_stored(7, c_tile_7, c_stride);

560 + }

561 + }

562 +

563 + FORCE_INLINE static void init_tile_config(int32_t m, __tilecfg& config) {

564 + for (int i = 0; i < 8; ++i) config.colsb[i] = AMX_TILE_ROW_BYTES;

565 + config.rows[0] = m;

566 + config.rows[1] = m;

567 + config.rows[2] = AMX_TILE_ROW_NUM;

568 + config.rows[3] = AMX_TILE_ROW_NUM;

569 + config.rows[4] = m;

570 + config.rows[5] = m;

571 + config.rows[6] = m;

572 + config.rows[7] = m;

573 + _tile_loadconfig(&config);

574 + }

575 +};

576 +

577 +// ---------------------------------------------------------------------------

578 +// AttentionImpl<ISA::AMX_FP8, scalar_t, head_dim, kv_cache_scalar_t>

579 +//

580 +// scalar_t = BF16 (input query/output type)

581 +// kv_cache_scalar_t = Float8_e4m3fn or Float8_e5m2

582 +// ---------------------------------------------------------------------------

583 +template <typename scalar_t, int64_t head_dim, typename kv_cache_scalar_t>

584 + requires(std::is_same_v<scalar_t, c10::BFloat16> &&

585 + (std::is_same_v<kv_cache_scalar_t, c10::Float8_e4m3fn> ||

586 + std::is_same_v<kv_cache_scalar_t, c10::Float8_e5m2>) &&

587 + (head_dim % 64 == 0))

588 +class AttentionImpl<ISA::AMX_FP8, scalar_t, head_dim, kv_cache_scalar_t> {

589 + public:

590 + // Q buffer holds FP8-quantized query.

591 + using query_t = scalar_t; // BF16 input

592 + using q_buffer_t = uint8_t; // FP8 quantized, stored in Q buffer

593 + using kv_cache_t = kv_cache_scalar_t;

594 + using logits_buffer_t = float;

595 + using partial_output_buffer_t = float;

596 + using prob_buffer_t = scalar_t; // BF16 softmax probs (PV phase)

597 +

598 + constexpr static int64_t BlockSizeAlignment = 64;

599 + constexpr static int64_t HeadDimAlignment = 64;

600 + constexpr static int64_t MaxQHeadNumPerIteration = 32;

601 + constexpr static int64_t HeadDim = head_dim;

602 + constexpr static ISA ISAType = ISA::AMX_FP8;

603 + constexpr static bool scale_on_logits = true;

604 +

605 + float k_scale = 1.0f;

606 + float v_scale = 1.0f;

607 + // Dynamically computed per copy_q_heads_tile() call.

608 + float q_dynamic_scale = 1.0f;

609 +

610 + public:

611 + AttentionImpl() : current_q_head_num_(0) {

612 + vec_op::unroll_loop<int, 8>([&](int i) { amx_tile_config_.colsb[i] = 64; });

613 + }

614 +

615 + ~AttentionImpl() { _tile_release(); }

616 +

617 + void init_from_input(const AttentionInput* input) {

618 + k_scale = input->k_scale_fp8;

619 + v_scale = input->v_scale_fp8;

620 + // q_scale_fp8 from input is used as a hint; dynamic scale is recomputed

621 + // per-tile in copy_q_heads_tile().

622 + }

623 +

624 + // Output scale for the PV result.

625 + float get_output_v_scale() const noexcept {

626 + constexpr float fp8_max =

627 + std::is_same_v<kv_cache_t, c10::Float8_e4m3fn> ? 448.0f : 57344.0f;

628 + return v_scale / fp8_max;

629 + }

630 +

631 + template <template <typename tile_gemm_t> typename attention>

632 + FORCE_INLINE void execute_attention(DEFINE_CPU_ATTENTION_PARAMS) {

633 + // Apply QK scale correction: score_fp32 = dot(Q_fp8, K_fp8) * q_dyn * k

634 + scale *= q_dynamic_scale * k_scale;

635 +

636 + if (q_head_num > AMX_TILE_ROW_NUM) {

637 + if (q_head_num != current_q_head_num_) {

638 + current_q_head_num_ = q_head_num;

639 + TileGemm224<q_buffer_t, kv_cache_t>::init_tile_config(q_head_num,

640 + amx_tile_config_);

641 + }

642 + attention<TileGemm224<q_buffer_t, kv_cache_t>> attention_iteration;

643 + attention_iteration(CPU_ATTENTION_PARAMS);

644 + } else {

645 + if (q_head_num != current_q_head_num_) {

646 + current_q_head_num_ = q_head_num;

647 + TileGemm122<q_buffer_t, kv_cache_t>::init_tile_config(q_head_num,

648 + amx_tile_config_);

649 + }

650 + attention<TileGemm122<q_buffer_t, kv_cache_t>> attention_iteration;

651 + attention_iteration(CPU_ATTENTION_PARAMS);

652 + }

653 + }

654 +

655 + // K cache stride: same element-count formula as AMX_BF16.

656 + // (BlockSizeAlignment * head_dim) uint8_t elements = same byte count as

657 + // before since kv_cache_t is 1-byte FP8.

658 + constexpr static int64_t k_cache_token_group_stride(

659 + const int32_t /*block_size*/) {

660 + return BlockSizeAlignment * head_dim;

661 + }

662 +

663 + // Both native FP8 formats use VNNI4 packing.

664 + constexpr static int64_t v_cache_token_group_stride(

665 + const int32_t /*block_size*/) {

666 + return BlockSizeAlignment * (AMX_TILE_ROW_BYTES / 4);

667 + }

668 +

669 + constexpr static int64_t v_cache_head_group_stride(const int32_t block_size) {

670 + return block_size * HeadDimAlignment;

671 + }

672 +

673 + // Copy Q heads tile: quantize BF16 to the KV cache FP8 type and pack it.

674 + //

675 + // Q buffer layout (FP8):

676 + // Same tile-prepacking as AMX_BF16 copy_q_heads_tile but with uint8_t

677 + // (FP8) elements. Each FP8 tile row = 64 bytes = 64 FP8 elements.

678 + // head_size_block_num = head_dim / 64 (vs head_dim/32 for BF16).

679 + //

680 + // Dynamic scale: max(abs(Q)) / fp8_max is computed over all elements,

681 + // stored in q_dynamic_scale for use in execute_attention().

682 + void copy_q_heads_tile(

683 + scalar_t* __restrict__ src, // [q_num, q_heads_per_kv, head_size] BF16

684 + q_buffer_t* __restrict__ q_buffer, const int32_t q_num,

685 + const int32_t q_heads_per_kv, const int64_t q_num_stride,

686 + const int64_t q_head_stride, const float /*scale*/) {

687 + // ---- Step 1: compute dynamic per-tile Q scale ----

688 + constexpr float fp8_max =

689 + std::is_same_v<kv_cache_t, c10::Float8_e4m3fn> ? 448.0f : 57344.0f;

690 + float max_abs = 0.0f;

691 + {

692 + scalar_t* src_iter = src;

693 + __m512 max_abs_vec = _mm512_setzero_ps();

694 + const __m512i abs_mask = _mm512_set1_epi32(0x7FFFFFFF);

695 + for (int32_t q = 0; q < q_num; ++q, src_iter += q_num_stride) {

696 + scalar_t* head_iter = src_iter;

697 + for (int32_t h = 0; h < q_heads_per_kv;

698 + ++h, head_iter += q_head_stride) {

699 + for (int64_t e = 0; e < head_dim; e += 16) {

700 + const __m256i bf16 = _mm256_loadu_si256(

701 + reinterpret_cast<const __m256i*>(head_iter + e));

702 + const __m512i fp32_bits =

703 + _mm512_slli_epi32(_mm512_cvtepu16_epi32(bf16), 16);

704 + const __m512 abs_values =

705 + _mm512_castsi512_ps(_mm512_and_si512(fp32_bits, abs_mask));

706 + max_abs_vec = _mm512_max_ps(max_abs_vec, abs_values);

707 + }

708 + }

709 + }

710 + max_abs = _mm512_reduce_max_ps(max_abs_vec);

711 + }

712 + if (max_abs == 0.0f) max_abs = 1.0f; // avoid div-by-zero

713 + const float inv_scale = fp8_max / max_abs;

714 + // Store true scale for score correction in execute_attention().

715 + q_dynamic_scale = max_abs / fp8_max;

716 +

717 + // ---- Step 2: quantize BF16 Q to FP8 and pack into q_buffer ----

718 + // Layout: same tile-prepacking as AMX_BF16, but with FP8 (1 byte/elem).

719 + // head_size_block_num = head_dim * sizeof(uint8_t) / AMX_TILE_ROW_BYTES

720 + // = head_dim / 64

721 + constexpr int64_t bytes_per_head = head_dim * sizeof(uint8_t);

722 + static_assert(bytes_per_head % AMX_TILE_ROW_BYTES == 0,

723 + "head_dim must be divisible by 64 for AMX_FP8");

724 + quantize_q_tile<kv_cache_t>(src, reinterpret_cast<uint8_t*>(q_buffer),

725 + q_num, q_heads_per_kv, q_num_stride,

726 + q_head_stride, head_dim, inv_scale);

727 + }

728 +

729 + // reshape KV to AMX-FP8 friendly layout.

730 + // K: native FP8 layout (1 byte/element, 16-token groups × k-slices).

731 + // V: native VNNI4 packing.

732 + static void reshape_and_cache(

733 + const scalar_t* __restrict__ key, const scalar_t* __restrict__ value,

734 + kv_cache_t* __restrict__ key_cache, kv_cache_t* __restrict__ value_cache,

735 + const int64_t* __restrict__ slot_mapping, const int64_t token_num,

736 + const int64_t key_token_num_stride, const int64_t value_token_num_stride,

737 + const int64_t head_num, const int64_t key_head_num_stride,

738 + const int64_t value_head_num_stride, const int64_t num_blocks,

739 + const int64_t num_blocks_stride, const int64_t cache_head_num_stride,

740 + const int64_t block_size, const int64_t /*block_size_stride*/,

741 + const float k_inv = 0.0f, const float v_inv = 0.0f) {

742 + constexpr auto qfn = select_fp8_quant_fn<kv_cache_t>();

743 +

744 + // K: native AMX-FP8 layout

745 + reshape_and_cache_k_amx_fp8_impl<scalar_t, qfn>(

746 + key, reinterpret_cast<uint8_t*>(key_cache), slot_mapping, token_num,

747 + head_num, head_dim, block_size, key_token_num_stride,

748 + key_head_num_stride, num_blocks_stride, cache_head_num_stride, k_inv);

749 +

750 + {

751 + constexpr int64_t token_num_per_sub_group = 4;

752 + constexpr int64_t head_elems_per_group = 16; // AMX_TILE_ROW_BYTES/4

753 + const int64_t group_num = head_dim / head_elems_per_group;

754 + const int64_t group_size = block_size * head_elems_per_group;

755 + uint8_t* vc = reinterpret_cast<uint8_t*>(value_cache);

756 +

757 + #pragma omp parallel for collapse(2) schedule(static)

758 + for (int64_t tok = 0; tok < token_num; ++tok) {

759 + for (int64_t h = 0; h < head_num; ++h) {

760 + const int64_t slot = slot_mapping[tok];

761 + if (slot < 0) continue;

762 + const int64_t block_idx = slot / block_size;

763 + const int64_t block_offset = slot % block_size;

764 + const int64_t sub_group_idx = block_offset / token_num_per_sub_group;

765 + const int64_t sub_group_offset =

766 + block_offset % token_num_per_sub_group;

767 +

768 + const scalar_t* vsrc =

769 + value + tok * value_token_num_stride + h * value_head_num_stride;

770 + uint8_t* vdst =

771 + vc + block_idx * num_blocks_stride + h * cache_head_num_stride +

772 + sub_group_idx * token_num_per_sub_group * head_elems_per_group +

773 + sub_group_offset;

774 +

775 + for (int64_t i = 0; i < group_num; ++i) {

776 + for (int64_t j = 0; j < head_elems_per_group; ++j)

777 + vdst[j * token_num_per_sub_group] =

778 + qfn(static_cast<float>(vsrc[j]), v_inv);

779 + vsrc += head_elems_per_group;

780 + vdst += group_size;

781 + }

782 + }

783 + }

784 + }

785 + }

786 +

787 + private:

788 + alignas(64) __tilecfg amx_tile_config_;

789 + int32_t current_q_head_num_;

790 +};

791 +

792 +} // namespace cpu_attention

793 +

794 +#endif // CPU_CAPABILITY_AMXFP8

795 +#endif // CPU_ATTN_AMX_FP8_HPP

csrc/cpu/cpu_attn_impl.hpp

+54 -36

@@ -12,7 +12,7 @@

12 #include "cpu/utils.hpp"

14 namespace cpu_attention {

15 -enum class ISA { AMX, VEC, VEC16, NEON, VXE, RVV, VSX };

15 +enum class ISA { AMX, VEC, VEC16, NEON, VXE, RVV, VSX, AMX_FP8 };

17 // Mirrors csrc/attention/dtype_fp8.cuh Fp8KVCacheDataType exactly.

18 enum class Fp8KVCacheDataType {

@@ -152,6 +152,9 @@ struct AttentionMetadata {

152 case ISA::AMX:

153 ss << "AMX, ";

154 break;

155 + case ISA::AMX_FP8:

156 + ss << "AMX_FP8, ";

157 + break;

158 case ISA::VEC:

159 ss << "VEC, ";

160 break;

@@ -885,6 +888,8 @@ struct AttentionInput {

888 // FP8 KV cache scales (used by FP8 attention implementations)

889 float k_scale_fp8 = 1.0f;

890 float v_scale_fp8 = 1.0f;

891 + // FP8 query scale (used by AMX_FP8 which quantizes query inside the kernel)

892 + float q_scale_fp8 = 1.0f;

893 };

894

895 #define DEFINE_CPU_ATTENTION_PARAMS \

@@ -1092,11 +1097,6 @@ class AttentionMainLoop {

1097

1098 // process logits

1099 {

1095 - // if (debug_info){

1096 - // print_logits("raw logits", logits_buffer, q_head_num,

1097 - // kv_tile_token_num, kv_tile_token_num);

1098 - // }

1099 -

1100 if (softcap_scale != 0.0f) {

1101 apply_softcap(logits_buffer, kv_tile_token_num, q_head_num,

1102 kv_tile_token_num, softcap_scale);

@@ -1127,30 +1127,33 @@ class AttentionMainLoop {

1127 apply_softmax(logits_buffer, partial_q_buffer, max_buffer, sum_buffer,

1128 kv_tile_token_num, q_head_num, kv_tile_token_num,

1129 is_first_iter, use_sink);

1130 -

1131 - // if (debug_info){

1132 - // print_logits("softmax logits",

1133 - // reinterpret_cast<prob_buffer_t*>(logits_buffer), q_head_num,

1134 - // kv_tile_token_num, kv_tile_token_num * sizeof(logits_buffer_t) /

1135 - // sizeof(prob_buffer_t));

1136 - // print_logits("new_max", max_buffer, 1, q_head_num, q_head_num);

1137 - // print_logits("new_sum", sum_buffer, 1, q_head_num, q_head_num);

1138 - // }

1130 }

1131

1132 // compute P@V

1133 {

1134 + constexpr bool prequantize_probabilities = []() {

1135 + if constexpr (requires { tile_gemm_t::prequantize_probabilities; }) {

1136 + return tile_gemm_t::prequantize_probabilities;

1137 + }

1138 + return false;

1139 + }();

1140 int32_t curr_group_offset =

1141 start_block_group_offset * v_cache_token_group_stride;

1142 int32_t curr_group_num_in_block =

1143 token_group_num_per_block - start_block_group_offset;

1144 int32_t remaining_group_num = token_group_num;

1145 int32_t head_dim_group_num = head_dim / headdim_alignment;

1149 - prob_buffer_t* curr_prob_buffer =

1150 - reinterpret_cast<prob_buffer_t*>(logits_buffer);

1151 - int64_t prob_buffer_stride =

1152 - kv_tile_token_num *

1153 - (sizeof(logits_buffer_t) / sizeof(prob_buffer_t));

1146 + using pv_prob_buffer_t = std::conditional_t<prequantize_probabilities,

1147 + uint8_t, prob_buffer_t>;

1148 + pv_prob_buffer_t* curr_prob_buffer =

1149 + reinterpret_cast<pv_prob_buffer_t*>(logits_buffer);

1150 + const int64_t prob_buffer_stride =

1151 + prequantize_probabilities

1152 + ? kv_tile_token_num

1153 + : kv_tile_token_num *

1154 + (sizeof(logits_buffer_t) / sizeof(prob_buffer_t));

1155 + constexpr int64_t prob_buffer_elem_size =

1156 + prequantize_probabilities ? sizeof(uint8_t) : sizeof(prob_buffer_t);

1157 partial_output_buffer_t* curr_partial_q_buffer = partial_q_buffer;

1158 bool accum_c = !is_first_iter;

1159 for (int32_t block_idx = start_block_idx; block_idx < end_block_idx;

@@ -1184,15 +1187,13 @@ class AttentionMainLoop {

1187 remaining_group_num -= curr_group_num_in_block;

1188 curr_group_offset = 0;

1189 curr_group_num_in_block = token_group_num_per_block;

1187 - curr_prob_buffer += curr_token_num;

1190 + curr_prob_buffer = reinterpret_cast<pv_prob_buffer_t*>(

1191 + reinterpret_cast<uint8_t*>(curr_prob_buffer) +

1192 + curr_token_num * prob_buffer_elem_size);

1193 curr_partial_q_buffer = partial_q_buffer;

1194 accum_c = true;

1195 }

1196 }

1192 - // if (debug_info) {

1193 - // print_logits("output", partial_q_buffer, q_head_num, head_dim,

1194 - // head_dim);

1195 - // }

1197 }

1198

1199 void apply_mask(logits_buffer_t* __restrict__ logits_buffer,

@@ -1262,9 +1263,19 @@ class AttentionMainLoop {

1263 #endif

1264

1265 using prob_buffer_vec_t = typename VecTypeTrait<prob_buffer_t>::vec_t;

1266 + constexpr bool prequantize_probabilities = []() {

1267 + if constexpr (requires { tile_gemm_t::prequantize_probabilities; }) {

1268 + return tile_gemm_t::prequantize_probabilities;

1269 + }

1270 + return false;

1271 + }();

1272 + using softmax_prob_buffer_t =

1273 + std::conditional_t<prequantize_probabilities, uint8_t, prob_buffer_t>;

1274 static_assert(sizeof(prob_buffer_t) <= sizeof(logits_buffer_t));

1275

1276 logits_buffer_t* __restrict__ curr_logits_buffer = logits_buffer;

1277 + softmax_prob_buffer_t* __restrict__ curr_prob_buffer =

1278 + reinterpret_cast<softmax_prob_buffer_t*>(logits_buffer);

1279 float* __restrict__ curr_partial_q_buffer = partial_q_buffer;

1280 const int32_t vec_num = kv_tile_token_num / 16;

1281 const int32_t head_vec_num = head_dim / 16;

@@ -1302,8 +1313,8 @@ class AttentionMainLoop {

1313 {

1314 logits_buffer_t* __restrict__ curr_logits_buffer_iter =

1315 curr_logits_buffer;

1305 - prob_buffer_t* __restrict__ curr_prob_buffer_iter =

1306 - reinterpret_cast<prob_buffer_t*>(curr_logits_buffer);

1316 + softmax_prob_buffer_t* __restrict__ curr_prob_buffer_iter =

1317 + curr_prob_buffer;

1318 for (int32_t j = 0; j < vec_num; ++j) {

1319 vec_op::FP32Vec16 vec(curr_logits_buffer_iter);

1320 vec = vec - max_vec;

@@ -1320,8 +1331,12 @@ class AttentionMainLoop {

1331 vec = fast_exp(vec);

1332 }

1333

1323 - prob_buffer_vec_t output_vec(vec);

1324 - output_vec.save(curr_prob_buffer_iter);

1334 + if constexpr (prequantize_probabilities) {

1335 + vec.save(curr_logits_buffer_iter);

1336 + } else {

1337 + prob_buffer_vec_t output_vec(vec);

1338 + output_vec.save(curr_prob_buffer_iter);

1339 + }

1340 #else

1341 vec.save(curr_logits_buffer_iter);

1342 for (int32_t k = 0; k < 16; ++k) {

@@ -1335,6 +1350,10 @@ class AttentionMainLoop {

1350 curr_logits_buffer_iter += 16;

1351 curr_prob_buffer_iter += 16;

1352 }

1353 + if constexpr (prequantize_probabilities) {

1354 + tile_gemm_t::quantize_probability_row(

1355 + curr_logits_buffer, curr_prob_buffer, kv_tile_token_num);

1356 + }

1357 }

1358 float new_sum_val = sum_vec.reduce_sum();

1359

@@ -1366,6 +1385,12 @@ class AttentionMainLoop {

1385 sum_buffer[i] = new_sum_val;

1386

1387 curr_logits_buffer += logits_buffer_stride;

1388 + if constexpr (prequantize_probabilities) {

1389 + curr_prob_buffer += kv_tile_token_num;

1390 + } else {

1391 + curr_prob_buffer =

1392 + reinterpret_cast<softmax_prob_buffer_t*>(curr_logits_buffer);

1393 + }

1394 curr_partial_q_buffer += head_dim;

1395 }

1396 }

@@ -1840,13 +1865,6 @@ class AttentionMainLoop {

1865 float* curr_sum_buffer = sum_buffer + q_tile_head_offset;

1866

1867 bool debug_info = false;

1843 - // bool debug_info = (

1844 - // q_head_start_idx == 4 &&

1845 - // (q_token_start_idx + q_head_tile_token_offset) <=

1846 - // 4

1847 - // && (q_token_start_idx + q_head_tile_token_offset +

1848 - // q_tile_token_num) > 4

1849 - // );

1868 // if (debug_info) {

1869 // std::printf("\tq_iter_idx: %d, q_token_start: %d,"

1870 // "q_token_end: %d, q_token_num: %d, q_head_num: %d,"

csrc/cpu/cpu_types_x86.hpp

+58

@@ -350,6 +350,64 @@ struct BF16Vec32 : public Vec<BF16Vec32> {

350 };

351 #endif

352

353 +// ---------------------------------------------------------------------------

354 +// Vectorized BF16 → FP8 E4M3 quantization helpers.

355 +// ---------------------------------------------------------------------------

356 +

357 +// Quantize 16 FP32 values to 16 FP8-E4M3 bytes using AVX-512.

358 +// inv_scale = 1.0f / q_scale (pre-computed by caller).

359 +// Returns 16 packed uint8_t in the low 128-bits of an __m128i.

360 +#if defined(__AVX512F__)

361 +FORCE_INLINE __m128i quant_fp32x16_to_fp8e4m3_avx512(const float* src,

362 + float inv_scale) {

363 + __m512 v = _mm512_loadu_ps(src);

364 + v = _mm512_mul_ps(v, _mm512_set1_ps(inv_scale));

365 + // Clamp to FP8-E4M3 representable range [-448, 448]

366 + v = _mm512_min_ps(v, _mm512_set1_ps(448.0f));

367 + v = _mm512_max_ps(v, _mm512_set1_ps(-448.0f));

368 + // Shift exponent bias: FP32 bias=127 → E4M3 bias=7, delta=120

369 + // Multiply by 2^-120 to move the exponent bias from 127 to 7.

370 + v = _mm512_mul_ps(v, _mm512_set1_ps(0x1p-120f));

371 + __m512i vi = _mm512_castps_si512(v);

372 + // sign: bit31 → bit7

373 + __m512i sign = _mm512_srli_epi32(

374 + _mm512_and_si512(vi, _mm512_set1_epi32(0x80000000u)), 24);

375 + // payload: bits[26:20] -> bits[6:0] after the exponent-bias shift.

376 + __m512i payload = _mm512_srli_epi32(

377 + _mm512_and_si512(vi, _mm512_set1_epi32(0x07F00000u)), 20);

378 + // Keep 0x7F (all-ones) reserved as NaN encoding: clamp to 0x7E

379 + payload = _mm512_min_epu32(payload, _mm512_set1_epi32(0x7E));

380 + __m512i fp8 = _mm512_or_si512(sign, payload);

381 + // Pack 16 × int32 → 16 × uint8 (cvt truncates, so the int32 values must

382 + // fit in uint8; they do since max is 0xFF = sign|payload).

383 + return _mm512_cvtepi32_epi8(fp8);

384 +}

385 +

386 +// Quantize 32 BF16 values to 32 FP8-E4M3 bytes using AVX-512.

387 +// Writes 32 bytes to dst. inv_scale = 1.0f / q_scale.

388 +FORCE_INLINE void quant_bf16x32_to_fp8e4m3_avx512(const c10::BFloat16* src,

389 + uint8_t* dst,

390 + float inv_scale) {

391 + // Convert 32 BF16 → 2x16 FP32 then quantize each half.

392 + const uint16_t* u16 = reinterpret_cast<const uint16_t*>(src);

393 + __m256i b16_lo = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(u16));

394 + __m256i b16_hi =

395 + _mm256_loadu_si256(reinterpret_cast<const __m256i*>(u16 + 16));

396 + // Zero-extend uint16 -> uint32, then shift left by 16 to form FP32 bits.

397 + __m512 lo_fp32 =

398 + _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(b16_lo), 16));

399 + __m512 hi_fp32 =

400 + _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(b16_hi), 16));

401 + alignas(64) float fp32_buf[32];

402 + _mm512_store_ps(fp32_buf, lo_fp32);

403 + _mm512_store_ps(fp32_buf + 16, hi_fp32);

404 + __m128i lo8 = quant_fp32x16_to_fp8e4m3_avx512(fp32_buf, inv_scale);

405 + __m128i hi8 = quant_fp32x16_to_fp8e4m3_avx512(fp32_buf + 16, inv_scale);

406 + _mm_storeu_si128(reinterpret_cast<__m128i*>(dst), lo8);

407 + _mm_storeu_si128(reinterpret_cast<__m128i*>(dst + 16), hi8);

408 +}

409 +#endif // __AVX512F__

410 +

411 struct FP32Vec4 : public Vec<FP32Vec4> {

412 constexpr static int VEC_ELEM_NUM = 4;

413 union AliasReg {

csrc/cpu/generate_cpu_attn_dispatch.py

+30 -9

@@ -22,6 +22,7 @@ ISA_TYPES = {

22 "VXE": 4,

23 "RVV": 5,

24 "VSX": 6,

25 + "AMX_FP8": 7,

26 }

28 # KV cache index: 0 = auto (same as scalar_t), 1 = fp8_e4m3, 2 = fp8_e5m2

@@ -39,13 +40,13 @@ KV_CACHE_CPP_TYPES = {

40 }

42 # ISAs supported for head_dims divisible by 32

42 -ISA_FOR_32 = ["AMX", "NEON", "VEC", "VEC16", "VXE", "RVV", "VSX"]

43 +ISA_FOR_32 = ["AMX", "AMX_FP8", "NEON", "VEC", "VEC16", "VXE", "RVV", "VSX"]

45 # ISAs supported for head_dims divisible by 16 only

46 ISA_FOR_16 = ["VEC16"]

48 # ISAs that support FP8 KV cache (x86 AVX2/AVX-512 required)

48 -ISA_FOR_FP8 = ["AMX", "VEC"]

49 +ISA_FOR_FP8 = ["AMX_FP8", "AMX", "VEC"]

52 def encode_params(head_dim: int, isa_type: str, kv_cache: str = "auto") -> int:

@@ -64,12 +65,21 @@ def _make_case(

65 encoded = encode_params(head_dim, isa, kv_cache)

66 actual_isa = isa_override if isa_override else isa

67 cpp_type = KV_CACHE_CPP_TYPES[kv_cache]

67 - attn_impl = (

68 - f"cpu_attention::AttentionImpl<"

69 - f"cpu_attention::ISA::{actual_isa}, \\\n"

70 - f" "

68 + native_impl = (

69 + f"cpu_attention::AttentionImpl<cpu_attention::ISA::{actual_isa}, "

70 f"scalar_t, head_dim, {cpp_type}>"

71 )

72 + if actual_isa == "AMX_FP8":

73 + fallback_impl = (

74 + f"cpu_attention::AttentionImpl<cpu_attention::ISA::AMX, "

75 + f"scalar_t, head_dim, {cpp_type}>"

76 + )

77 + attn_impl = (

78 + "std::conditional_t<std::is_same_v<scalar_t, c10::BFloat16>, "

79 + f"{native_impl}, {fallback_impl}>"

80 + )

81 + else:

82 + attn_impl = native_impl

83 comment = (

84 f"head_dim={head_dim}, isa={isa}"

85 if kv_cache == "auto"

@@ -98,7 +108,8 @@ def generate_cases_for_isa_group(isa_list: list[str], include_fp8: bool = False)

108 for isa in isa_list:

109 if isa not in ISA_FOR_32:

110 continue

101 - cases.append(_make_case(head_dim, isa, "auto"))

111 + isa_override = "AMX" if isa == "AMX_FP8" else None

112 + cases.append(_make_case(head_dim, isa, "auto", isa_override))

113

114 # Non-FP8 cases for head_dims divisible by 16 only

115 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)

123 for isa in isa_list:

124 if isa not in ISA_FOR_FP8:

125 continue

115 - cases.append(_make_case(head_dim, isa, fp8_type))

126 + isa_override = "AMX" if isa == "AMX_FP8" and head_dim % 64 else None

127 + cases.append(_make_case(head_dim, isa, fp8_type, isa_override))

128

129 return "\n".join(cases)

130

@@ -138,6 +150,10 @@ def generate_header_file() -> str:

150 #include "cpu_attn_vec.hpp"

151 #include "cpu_attn_vec16.hpp"

152

153 +#ifdef CPU_CAPABILITY_AMXFP8

154 + #include "cpu_attn_amx_fp8.hpp"

155 +#endif

156 +

157 #ifdef CPU_CAPABILITY_AMXBF16

158 #include "cpu_attn_amx.hpp"

159 #endif

@@ -206,7 +222,12 @@ def generate_header_file() -> str:

222 )

223

224 header += _macro_block(

209 - "#if defined(CPU_CAPABILITY_AMXBF16)",

225 + "#if defined(CPU_CAPABILITY_AMXFP8)",

226 + ["AMX_FP8", "AMX", "VEC", "VEC16"],

227 + fp8=True,

228 + )

229 + header += _macro_block(

230 + "#elif defined(CPU_CAPABILITY_AMXBF16)",

231 ["AMX", "VEC", "VEC16"],

232 fp8=True,

233 )

docker/Dockerfile.cpu

+14 -6

@@ -27,7 +27,7 @@

27 # cache mounts (ccache, cargo registry) aren't exported by `--cache-to

28 # type=registry`, so they don't survive a build landing on a different/fresh

29 # builder. sccache's cache lives in S3 instead, so it does.

30 -FROM ubuntu:22.04 AS base

30 +FROM ubuntu:25.04 AS base

32 RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \

33 --mount=type=cache,target=/var/lib/apt,sharing=locked \

@@ -66,9 +66,15 @@ ENV MAX_JOBS=${max_jobs}

66 RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \

67 --mount=type=cache,target=/var/lib/apt,sharing=locked \

68 apt-get update -y \

69 - && apt-get install -y --no-install-recommends sudo ccache wget zlib1g-dev \

70 - gcc-12 g++-12 libtcmalloc-minimal4 libnuma-dev ffmpeg libsm6 libxext6 libgl1 jq lsof make xz-utils \

71 - && update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-12 10 --slave /usr/bin/g++ g++ /usr/bin/g++-12 \

69 + && apt-get install -y --no-install-recommends sudo ccache git curl wget ca-certificates zlib1g-dev \

70 + software-properties-common libtcmalloc-minimal4 libnuma-dev ffmpeg libsm6 libxext6 libgl1 jq lsof make xz-utils \

71 + && for i in 1 2 3; do \

72 + add-apt-repository -y ppa:ubuntu-toolchain-r/test && break || \

73 + { echo "Attempt $i failed, retrying toolchain PPA..."; sleep 5; }; \

74 + done \

75 + && apt-get update -y \

76 + && apt-get install -y --no-install-recommends gcc-15 g++-15 \

77 + && update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-15 10 --slave /usr/bin/g++ g++ /usr/bin/g++-15 \

78 && curl -LsSf https://astral.sh/uv/install.sh | sh

80 ARG TARGETARCH

@@ -83,7 +89,7 @@ ENV SCCACHE_S3_NO_CREDENTIALS=${USE_SCCACHE:+${SCCACHE_S3_NO_CREDENTIALS}}

89 ENV SCCACHE_IDLE_TIMEOUT=${USE_SCCACHE:+0}

91 # Compiler and linker environment

86 -ENV CC=/usr/bin/gcc-12 CXX=/usr/bin/g++-12

92 +ENV CC=/usr/bin/gcc-15 CXX=/usr/bin/g++-15

93 ENV CCACHE_DIR=/root/.cache/ccache

94 ENV CMAKE_CXX_COMPILER_LAUNCHER=ccache

@@ -159,7 +165,7 @@ RUN apt-get update -y \

165 WORKDIR /workspace

166

167 COPY requirements/build/rust.txt requirements/build/rust.txt

162 -RUN python3 -m pip install --no-cache-dir -r requirements/build/rust.txt

168 +RUN python3 -m pip install --break-system-packages --no-cache-dir -r requirements/build/rust.txt

169

170 ARG USE_SCCACHE

171 ARG SCCACHE_ENDPOINT

@@ -230,6 +236,7 @@ ENV VLLM_CPU_X86=${VLLM_CPU_X86}

236 # Support for cross-compilation with ARM BF16 ISA: docker build --build-arg VLLM_CPU_ARM_BF16="true" ...

237 ARG VLLM_CPU_ARM_BF16=0

238 ENV VLLM_CPU_ARM_BF16=${VLLM_CPU_ARM_BF16}

239 +ARG ENABLE_AMX_FP8=1

240

241 WORKDIR /vllm-workspace

242

@@ -297,6 +304,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \

304 fi; \

305 sccache --show-stats; \

306 fi && \

307 + CMAKE_ARGS="-DENABLE_AMX_FP8=${ENABLE_AMX_FP8}" \

308 VLLM_TARGET_DEVICE=cpu python3 setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38 && \

309 if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi

310

tests/kernels/attention/test_cpu_attn.py

+148 -1

@@ -608,6 +608,21 @@ def varlen_with_paged_kv(

608 slot_mapping=slot_mapping,

609 isa=isa,

610 )

611 + ref_metadata = cpu_attn_get_scheduler_metadata(

612 + num_reqs=num_seqs,

613 + num_heads=num_query_heads,

614 + num_kv_heads=num_kv_heads,

615 + head_dim=head_size,

616 + seq_lens=kv_lens_tensor,

617 + dtype=dtype,

618 + query_start_loc=cu_query_lens,

619 + causal=dynamic_causal is None,

620 + sliding_window_size=(sliding_window if sliding_window is not None else -1),

621 + isa=isa,

622 + enable_kv_split=True,

623 + dynamic_causal=dynamic_causal_tensor,

624 + kv_cache_dtype="auto",

625 + )

626 ref_output = torch.empty_like(query)

627 cpu_attention_with_kv_cache(

628 query=query,

@@ -622,7 +637,7 @@ def varlen_with_paged_kv(

637 sliding_window=sliding_window if sliding_window is not None else -1,

638 block_table=block_tables,

639 softcap=soft_cap if soft_cap is not None else 0,

625 - scheduler_metadata=metadata,

640 + scheduler_metadata=ref_metadata,

641 s_aux=s_aux,

642 dynamic_causal=dynamic_causal_tensor,

643 )

@@ -1166,3 +1181,135 @@ def test_varlen_with_paged_kv_dynamic_causal(

1181 kv_cache_dtype=kv_cache_dtype,

1182 dynamic_causal=dynamic_causal,

1183 )

1184 +

1185 +

1186 +# ---------------------------------------------------------------------------

1187 +# AMX_FP8 (Diamond Rapids) tests

1188 +# ---------------------------------------------------------------------------

1189 +

1190 +

1191 +def _amx_fp8_available() -> bool:

1192 + """Return True iff the runtime reports AMX_FP8 capability."""

1193 + return torch.cpu._is_amx_tile_supported() and torch.ops._C.cpu_attn_has_isa(

1194 + "amx_fp8"

1195 + )

1196 +

1197 +

1198 +@pytest.mark.parametrize("kv_cache_dtype", ["fp8_e4m3", "fp8_e5m2"])

1199 +@pytest.mark.parametrize("seq_lens", SEQ_LENS)

1200 +@pytest.mark.parametrize("num_heads", NUM_HEADS)

1201 +@pytest.mark.parametrize("head_size", [64, 128, 256])

1202 +@pytest.mark.parametrize("block_size", [32, 64])

1203 +@pytest.mark.parametrize("sliding_window", [None])

1204 +@pytest.mark.parametrize("dtype", [torch.bfloat16])

1205 +@pytest.mark.parametrize("soft_cap", [None])

1206 +@pytest.mark.parametrize("num_blocks", NUM_BLOCKS)

1207 +@pytest.mark.parametrize("use_alibi", [False])

1208 +@pytest.mark.parametrize("use_sink", [False])

1209 +@pytest.mark.parametrize("isa", ["amx_fp8"])

1210 +@pytest.mark.skipif(

1211 + not _amx_fp8_available(), reason="no AMX_FP8 support (requires Diamond Rapids)."

1212 +)

1213 +def test_varlen_with_paged_kv_amx_fp8(

1214 + seq_lens: list[tuple[int, int]],

1215 + num_heads: tuple[int, int],

1216 + head_size: int,

1217 + sliding_window: int | None,

1218 + dtype: torch.dtype,

1219 + block_size: int,

1220 + soft_cap: float | None,

1221 + num_blocks: int,

1222 + use_alibi: bool,

1223 + use_sink: bool,

1224 + isa: str,

1225 + kv_cache_dtype: str,

1226 +) -> None:

1227 + """Test AMX_FP8 native FP8×FP8 attention (Diamond Rapids).

1228 +

1229 + Verifies that:

1230 + - QK uses _tile_dpfp8ps (native FP8 MMA, no K dequant).

1231 + - PV uses native FP8 MMA for both E4M3 and E5M2.

1232 + - Output cosine similarity vs fp32 reference > 0.99.

1233 + """

1234 + if block_size % 64 != 0:

1235 + pytest.skip("native FP8 PV requires block_size divisible by 64")

1236 +

1237 + varlen_with_paged_kv(

1238 + seq_lens=seq_lens,

1239 + num_heads=num_heads,

1240 + head_size=head_size,

1241 + sliding_window=sliding_window,

1242 + dtype=dtype,

1243 + block_size=block_size,

1244 + soft_cap=soft_cap,

1245 + num_blocks=num_blocks,

1246 + use_alibi=use_alibi,

1247 + use_sink=use_sink,

1248 + isa=isa,

1249 + kv_cache_dtype=kv_cache_dtype,

1250 + )

1251 +

1252 +

1253 +@pytest.mark.skipif(

1254 + not _amx_fp8_available(), reason="no AMX_FP8 support (requires Diamond Rapids)."

1255 +)

1256 +def test_amx_fp8_qk_covers_full_64_token_group() -> None:

1257 + head_size = 64

1258 + block_size = 64

1259 + query = torch.ones((1, 1, head_size), dtype=torch.bfloat16)

1260 + key = torch.empty((block_size, 1, head_size), dtype=torch.bfloat16)

1261 + value = torch.empty_like(key)

1262 + key[:32].fill_(-1)

1263 + key[32:].fill_(1)

1264 + value[:32].fill_(-1)

1265 + value[32:].fill_(1)

1266 +

1267 + key_cache = torch.empty((1, 1, block_size, head_size), dtype=torch.uint8)

1268 + value_cache = torch.empty_like(key_cache)

1269 + slot_mapping = torch.arange(block_size, dtype=torch.int64)

1270 + cpu_attn_reshape_and_cache(

1271 + key=key,

1272 + value=value,

1273 + key_cache=key_cache,

1274 + value_cache=value_cache,

1275 + slot_mapping=slot_mapping,

1276 + isa="amx_fp8",

1277 + kv_cache_dtype="fp8_e4m3",

1278 + )

1279 +

1280 + query_start_loc = torch.tensor([0, 1], dtype=torch.int32)

1281 + seq_lens = torch.tensor([block_size], dtype=torch.int32)

1282 + metadata = cpu_attn_get_scheduler_metadata(

1283 + num_reqs=1,

1284 + num_heads=1,

1285 + num_kv_heads=1,

1286 + head_dim=head_size,

1287 + seq_lens=seq_lens,

1288 + dtype=torch.bfloat16,

1289 + query_start_loc=query_start_loc,

1290 + causal=True,

1291 + sliding_window_size=-1,

1292 + isa="amx_fp8",

1293 + enable_kv_split=False,

1294 + kv_cache_dtype="fp8_e4m3",

1295 + )

1296 + output = torch.empty_like(query)

1297 + cpu_attention_with_kv_cache(

1298 + query=query,

1299 + key_cache=key_cache,

1300 + value_cache=value_cache,

1301 + output=output,

1302 + query_start_loc=query_start_loc,

1303 + seq_lens=seq_lens,

1304 + scale=head_size**-0.5,

1305 + causal=True,

1306 + alibi_slopes=None,

1307 + sliding_window=-1,

1308 + block_table=torch.tensor([[0]], dtype=torch.int32),

1309 + softcap=0,

1310 + scheduler_metadata=metadata,

1311 + s_aux=None,

1312 + kv_cache_dtype="fp8_e4m3",

1313 + )

1314 +

1315 + torch.testing.assert_close(output, torch.ones_like(output), atol=0.05, rtol=0)