@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 22524f7a92
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 4, 12:30 UTC

[Feat] CPU fp8 attn for AMX/AVX-512 (#39445)

Signed-off-by: Li, Tianmu <tianmu.li@intel.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Li, Jiang <jiang1.li@intel.com>

Tianmu Li committed Apr 29, 2026 at 05:43 UTC 22524f7a92b71c8e65eade20ef274fa3b4006d3e
19 files changed +1068 -284
csrc/cpu/cpu_attn.cpp
+78 -25
@@ -1,5 +1,16 @@
1 #include "cpu_attn_dispatch_generated.h"
2
3 +// Maps kv_cache_dtype string to Fp8KVCacheDataType enum.
4 +// "auto" -> kAuto(0); "fp8"/"fp8_e4m3" -> kFp8E4M3; "fp8_e5m2" -> kFp8E5M2.
5 +static inline cpu_attention::Fp8KVCacheDataType parse_fp8_kv_dtype(
6 + const std::string& kv_cache_dtype) {
7 + if (kv_cache_dtype == "fp8_e5m2")
8 + return cpu_attention::Fp8KVCacheDataType::kFp8E5M2;
9 + if (kv_cache_dtype == "fp8_e4m3" || kv_cache_dtype == "fp8")
10 + return cpu_attention::Fp8KVCacheDataType::kFp8E4M3;
11 + return cpu_attention::Fp8KVCacheDataType::kAuto;
12 +}
13 +
14 torch::Tensor get_scheduler_metadata(
15 const int64_t num_req, const int64_t num_heads_q,
16 const int64_t num_heads_kv, const int64_t head_dim,
@@ -49,7 +60,7 @@ torch::Tensor get_scheduler_metadata(
60 input.enable_kv_split = enable_kv_split;
61
62 VLLM_DISPATCH_FLOATING_TYPES(dtype, "get_scheduler_metadata", [&]() {
52 - CPU_ATTN_DISPATCH(head_dim, isa, [&]() {
63 + CPU_ATTN_DISPATCH(head_dim, isa, 0, [&]() {
64 input.elem_size = sizeof(scalar_t);
65 input.q_buffer_elem_size = sizeof(attn_impl::q_buffer_t);
66 input.logits_buffer_elem_size = sizeof(attn_impl::logits_buffer_t);
@@ -72,7 +83,9 @@ void cpu_attn_reshape_and_cache(
83 key_cache, // [num_blocks, num_kv_heads, block_size, head_size]
84 torch::Tensor&
85 value_cache, // [num_blocks, num_kv_heads, block_size, head_size]
75 - const torch::Tensor& slot_mapping, const std::string& isa) {
86 + const torch::Tensor& slot_mapping, const std::string& isa,
87 + const double k_scale = 1.0, const double v_scale = 1.0,
88 + const std::string& kv_cache_dtype = "auto") {
89 TORCH_CHECK_EQ(key.dim(), 3);
90 TORCH_CHECK_EQ(value.dim(), 3);
91 TORCH_CHECK_EQ(key_cache.dim(), 4);
@@ -80,18 +93,30 @@ void cpu_attn_reshape_and_cache(
93 TORCH_CHECK_EQ(key.stride(2), 1);
94 TORCH_CHECK_EQ(value.stride(2), 1);
95
96 + const int64_t kv_cache_idx =
97 + static_cast<int64_t>(parse_fp8_kv_dtype(kv_cache_dtype));
98 + const bool is_fp8 = (kv_cache_idx != 0);
99 +
100 + if (is_fp8) {
101 + TORCH_CHECK(key_cache.scalar_type() == at::ScalarType::Byte,
102 + "key_cache must be uint8 for FP8 path");
103 + TORCH_CHECK(value_cache.scalar_type() == at::ScalarType::Byte,
104 + "value_cache must be uint8 for FP8 path");
105 + TORCH_CHECK(k_scale > 0, "k_scale must be positive for FP8 path");
106 + TORCH_CHECK(v_scale > 0, "v_scale must be positive for FP8 path");
107 + }
108 +
109 + const float k_inv = is_fp8 ? 1.0f / static_cast<float>(k_scale) : 0.0f;
110 + const float v_inv = is_fp8 ? 1.0f / static_cast<float>(v_scale) : 0.0f;
111 +
112 const int64_t token_num = key.size(0);
84 - const int64_t key_token_num_stride = key.stride(0);
85 - const int64_t value_token_num_stride = value.stride(0);
86 - const int64_t head_num = value.size(1);
87 - const int64_t key_head_num_stride = key.stride(1);
88 - const int64_t value_head_num_stride = value.stride(1);
113 + const int64_t head_num = key.size(1);
114 + const int64_t head_dim = key.size(2);
115 const int64_t num_blocks = key_cache.size(0);
116 const int64_t num_blocks_stride = key_cache.stride(0);
117 const int64_t cache_head_num_stride = key_cache.stride(1);
118 const int64_t block_size = key_cache.size(2);
119 const int64_t block_size_stride = key_cache.stride(2);
94 - const int64_t head_dim = key.size(-1);
120
121 cpu_attention::ISA isa_tag = [&]() {
122 if (isa == "amx") {
@@ -109,16 +134,24 @@ void cpu_attn_reshape_and_cache(
134 }
135 }();
136
137 + if (is_fp8) {
138 + TORCH_CHECK(isa_tag == cpu_attention::ISA::AMX ||
139 + isa_tag == cpu_attention::ISA::VEC,
140 + "FP8 KV cache is only supported on x86 (AMX/VEC) ISA");
141 + }
142 +
143 VLLM_DISPATCH_FLOATING_TYPES(
144 key.scalar_type(), "cpu_attn_reshape_and_cache", [&]() {
114 - CPU_ATTN_DISPATCH(head_dim, isa_tag, [&]() {
145 + CPU_ATTN_DISPATCH(head_dim, isa_tag, kv_cache_idx, [&]() {
146 + using kv_t = typename attn_impl::kv_cache_t;
147 attn_impl::reshape_and_cache(
148 key.data_ptr<scalar_t>(), value.data_ptr<scalar_t>(),
117 - key_cache.data_ptr<scalar_t>(), value_cache.data_ptr<scalar_t>(),
118 - slot_mapping.data_ptr<int64_t>(), token_num, key_token_num_stride,
119 - value_token_num_stride, head_num, key_head_num_stride,
120 - value_head_num_stride, num_blocks, num_blocks_stride,
121 - cache_head_num_stride, block_size, block_size_stride);
149 + reinterpret_cast<kv_t*>(key_cache.data_ptr()),
150 + reinterpret_cast<kv_t*>(value_cache.data_ptr()),
151 + slot_mapping.data_ptr<int64_t>(), token_num, key.stride(0),
152 + value.stride(0), head_num, key.stride(1), value.stride(1),
153 + num_blocks, num_blocks_stride, cache_head_num_stride, block_size,
154 + block_size_stride, k_inv, v_inv);
155 });
156 });
157 }
@@ -137,13 +170,26 @@ void cpu_attention_with_kv_cache(
170 const int64_t sliding_window_left, const int64_t sliding_window_right,
171 const torch::Tensor& block_table, // [num_tokens, max_block_num]
172 const double softcap, const torch::Tensor& scheduler_metadata,
140 - const std::optional<torch::Tensor>& s_aux // [num_heads]
141 -) {
173 + const std::optional<torch::Tensor>& s_aux, // [num_heads]
174 + const double k_scale = 1.0, const double v_scale = 1.0,
175 + const std::string& kv_cache_dtype = "auto") {
176 TORCH_CHECK_EQ(query.dim(), 3);
177 TORCH_CHECK_EQ(query.stride(2), 1);
178 TORCH_CHECK_EQ(key_cache.dim(), 4);
179 TORCH_CHECK_EQ(value_cache.dim(), 4);
180
181 + const int64_t kv_cache_idx =
182 + static_cast<int64_t>(parse_fp8_kv_dtype(kv_cache_dtype));
183 + const bool is_fp8 = (kv_cache_idx != 0);
184 + if (is_fp8) {
185 + TORCH_CHECK(key_cache.scalar_type() == at::ScalarType::Byte,
186 + "key_cache must be uint8 for FP8 path");
187 + TORCH_CHECK(value_cache.scalar_type() == at::ScalarType::Byte,
188 + "value_cache must be uint8 for FP8 path");
189 + TORCH_CHECK(k_scale > 0, "k_scale must be positive for FP8 path");
190 + TORCH_CHECK(v_scale > 0, "v_scale must be positive for FP8 path");
191 + }
192 +
193 cpu_attention::AttentionInput input;
194 input.metadata = reinterpret_cast<cpu_attention::AttentionMetadata*>(
195 scheduler_metadata.data_ptr());
@@ -165,25 +211,32 @@ void cpu_attention_with_kv_cache(
211 input.block_table = block_table.data_ptr<int32_t>();
212 input.alibi_slopes =
213 alibi_slopes.has_value() ? alibi_slopes->data_ptr<float>() : nullptr;
168 - // For now sink must be bf16
214 input.s_aux = s_aux.has_value() ? s_aux->data_ptr<c10::BFloat16>() : nullptr;
215 input.scale = scale;
216 input.causal = causal;
217 input.sliding_window_left = sliding_window_left;
218 input.sliding_window_right = sliding_window_right;
219 if (input.causal) {
175 - // to make boundary calculation easier
220 input.sliding_window_right = 0;
221 }
178 - float softcap_fp32 = softcap;
179 - input.softcap = softcap_fp32;
222 + input.softcap = static_cast<float>(softcap);
223 +
224 + if (is_fp8) {
225 + input.k_scale_fp8 = static_cast<float>(k_scale);
226 + input.v_scale_fp8 = static_cast<float>(v_scale);
227 + TORCH_CHECK(input.metadata->isa == cpu_attention::ISA::AMX ||
228 + input.metadata->isa == cpu_attention::ISA::VEC,
229 + "FP8 KV cache is only supported on x86 (AMX/VEC) ISA");
230 + }
231
232 VLLM_DISPATCH_FLOATING_TYPES(
233 query.scalar_type(), "cpu_attention_with_kv_cache", [&]() {
183 - CPU_ATTN_DISPATCH(query.size(2), input.metadata->isa, [&]() {
184 - TORCH_CHECK_EQ(input.block_size % attn_impl::BlockSizeAlignment, 0);
185 - cpu_attention::AttentionMainLoop<attn_impl> mainloop;
186 - mainloop(&input);
187 - });
234 + CPU_ATTN_DISPATCH(
235 + query.size(2), input.metadata->isa, kv_cache_idx, [&]() {
236 + TORCH_CHECK_EQ(input.block_size % attn_impl::BlockSizeAlignment,
237 + 0);
238 + cpu_attention::AttentionMainLoop<attn_impl> mainloop;
239 + mainloop(&input);
240 + });
241 });
242 }
csrc/cpu/cpu_attn_amx.hpp
+171 -46
@@ -1,6 +1,7 @@
1 #ifndef CPU_ATTN_AMX_HPP
2 #define CPU_ATTN_AMX_HPP
3
4 +#include "cpu_attn_fp8.hpp"
5 #include "cpu_attn_impl.hpp"
6
7 namespace cpu_attention {
@@ -21,9 +22,10 @@ typedef struct __tile_config {
22 // 2-2-4 pattern, for 16 < m <= 32
23 // TILE 0, 1: load A matrix, row num should be 16, m - 16
24 // TILE 2, 3: load B matrix, row num should be 16
24 -// TILE 4, 5, 6, 7: store results C matrix, row num should be 16, 16, m - 16, m
25 -// - 16
26 -template <typename kv_cache_t>
25 +// TILE 4, 5, 6, 7: store results C matrix, row num should be 16, 16,
26 +// m - 16, m - 16
27 +// q_buffer_t: A (Q/P) tile type; kv_cache_t: B (K/V cache) tile type.
28 +template <typename q_buffer_t, typename kv_cache_t>
29 class TileGemm224 {
30 public:
31 template <AttentionGemmPhase phase, int32_t k_size>
@@ -42,13 +44,56 @@ class TileGemm224 {
44 }
45 };
46
45 -template <>
46 -class TileGemm224<c10::BFloat16> {
47 +// Dequantize one FP8 tile (AMX_TILE_ROW_NUM rows x 32 cols) to BF16.
48 +template <typename kv_cache_t>
49 +FORCE_INLINE void deq_tile_amx(const uint8_t* src, c10::BFloat16* dst) {
50 + for (int r = 0; r < AMX_TILE_ROW_NUM; ++r) {
51 + if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e4m3fn>) {
52 + vec_op::BF16Vec32(src + r * 32, vec_op::fp8_bf16_e4m3_tag{})
53 + .save(dst + r * 32);
54 + } else {
55 + vec_op::BF16Vec32(src + r * 32, vec_op::fp8_bf16_e5m2_tag{})
56 + .save(dst + r * 32);
57 + }
58 + }
59 +}
60 +
61 +// For FP8: dequant src into scratch and return scratch.
62 +// For BF16: return src directly (scratch is unused; the compiler elides it).
63 +template <typename kv_cache_t>
64 +FORCE_INLINE const c10::BFloat16* prepare_b_tile(const kv_cache_t* src,
65 + c10::BFloat16* scratch) {
66 + if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e4m3fn> ||
67 + std::is_same_v<kv_cache_t, c10::Float8_e5m2>) {
68 + deq_tile_amx<kv_cache_t>(reinterpret_cast<const uint8_t*>(src), scratch);
69 + return scratch;
70 + } else {
71 + return reinterpret_cast<const c10::BFloat16*>(src);
72 + }
73 +}
74 +
75 +// Handles both BF16 and FP8 KV cache (2-2-4 pattern).
76 +template <typename kv_cache_t>
77 +class TileGemm224<c10::BFloat16, kv_cache_t> {
78 + static_assert(std::is_same_v<kv_cache_t, c10::BFloat16> ||
79 + std::is_same_v<kv_cache_t, c10::Float8_e4m3fn> ||
80 + std::is_same_v<kv_cache_t, c10::Float8_e5m2>,
81 + "kv_cache_t must be BFloat16, Float8_e4m3fn, or Float8_e5m2");
82 +
83 + static constexpr bool fp8_kv =
84 + std::is_same_v<kv_cache_t, c10::Float8_e4m3fn> ||
85 + std::is_same_v<kv_cache_t, c10::Float8_e5m2>;
86 +
87 + static constexpr int64_t tile_elems = AMX_TILE_BYTES / sizeof(c10::BFloat16);
88 + // BF16 path: scratch_elems=1 so the scratch array is eliminated by the
89 + // compiler.
90 + static constexpr int64_t scratch_elems = fp8_kv ? tile_elems : 1;
91 +
92 public:
93 template <AttentionGemmPhase phase, int32_t k_size>
94 FORCE_INLINE static void gemm(const int32_t m_size,
95 c10::BFloat16* __restrict__ a_tile,
51 - c10::BFloat16* __restrict__ b_tile,
96 + kv_cache_t* __restrict__ b_tile,
97 float* __restrict__ c_tile, const int64_t lda,
98 const int64_t ldb, const int64_t ldc,
99 const int32_t block_size,
@@ -56,6 +101,7 @@ class TileGemm224<c10::BFloat16> {
101 const bool accum_c) {
102 const int32_t k_times =
103 dynamic_k_size / (AMX_TILE_ROW_NUM * 4 / sizeof(c10::BFloat16));
104 +
105 c10::BFloat16* __restrict__ a_tile_0 = a_tile;
106 c10::BFloat16* __restrict__ a_tile_1 = a_tile + lda * AMX_TILE_ROW_NUM;
107 const int64_t a_tile_stride = [&]() {
@@ -70,8 +116,8 @@ class TileGemm224<c10::BFloat16> {
116 }
117 }();
118
73 - c10::BFloat16* __restrict__ b_tile_2 = b_tile;
74 - c10::BFloat16* __restrict__ b_tile_3 = [&]() {
119 + kv_cache_t* __restrict__ b_tile_2 = b_tile;
120 + kv_cache_t* __restrict__ b_tile_3 = [&]() {
121 if constexpr (phase == AttentionGemmPhase::QK) {
122 // k_cache is prepacked
123 return b_tile + (k_size * AMX_TILE_ROW_BYTES / 4);
@@ -106,11 +152,16 @@ class TileGemm224<c10::BFloat16> {
152 _tile_zero(7);
153 }
154
155 + alignas(64) c10::BFloat16 scratch_2[scratch_elems];
156 + alignas(64) c10::BFloat16 scratch_3[scratch_elems];
157 for (int32_t k = 0; k < k_times; ++k) {
158 + const c10::BFloat16* load_2 = prepare_b_tile(b_tile_2, scratch_2);
159 + const c10::BFloat16* load_3 = prepare_b_tile(b_tile_3, scratch_3);
160 +
161 _tile_loadd(0, a_tile_0, a_tile_stride);
111 - _tile_stream_loadd(2, b_tile_2, b_tile_stride);
162 + _tile_stream_loadd(2, const_cast<c10::BFloat16*>(load_2), b_tile_stride);
163 _tile_dpbf16ps(4, 0, 2);
113 - _tile_stream_loadd(3, b_tile_3, b_tile_stride);
164 + _tile_stream_loadd(3, const_cast<c10::BFloat16*>(load_3), b_tile_stride);
165 _tile_dpbf16ps(5, 0, 3);
166 _tile_loadd(1, a_tile_1, a_tile_stride);
167 _tile_dpbf16ps(6, 1, 2);
@@ -154,13 +205,13 @@ class TileGemm224<c10::BFloat16> {
205 };
206
207 // 1-2-2 pattern, for 0 < m <= 16
157 -// TILE 0, (1): load A matrix, use extra 1 tile for prefetch, row num should be
158 -// m, m
159 -// TILE 2, 3, (4, 5): load B matrix, use extra 2 tiles for prefetch, row
160 -// num should be 16
161 -// TILE 6, 7, (6, 7): store results C matrix, row num should be
162 -// m
163 -template <typename kv_cache_t>
208 +// TILE 0, (1): load A matrix, use extra 1 tile for prefetch, row num should
209 +// be m, m
210 +// TILE 2, 3, (4, 5): load B matrix, use extra 2 tiles for prefetch, row num
211 +// should be 16
212 +// TILE 6, 7: store results C matrix, row num should be m
213 +// q_buffer_t: A (Q/P) tile type; kv_cache_t: B (K/V cache) tile type.
214 +template <typename q_buffer_t, typename kv_cache_t>
215 class TileGemm122 {
216 public:
217 template <AttentionGemmPhase phase, int32_t k_size>
@@ -179,13 +230,26 @@ class TileGemm122 {
230 }
231 };
232
182 -template <>
183 -class TileGemm122<c10::BFloat16> {
233 +// Handles both BF16 and FP8 KV cache (1-2-2 pattern).
234 +template <typename kv_cache_t>
235 +class TileGemm122<c10::BFloat16, kv_cache_t> {
236 + static_assert(std::is_same_v<kv_cache_t, c10::BFloat16> ||
237 + std::is_same_v<kv_cache_t, c10::Float8_e4m3fn> ||
238 + std::is_same_v<kv_cache_t, c10::Float8_e5m2>,
239 + "kv_cache_t must be BFloat16, Float8_e4m3fn, or Float8_e5m2");
240 +
241 + static constexpr bool fp8_kv =
242 + std::is_same_v<kv_cache_t, c10::Float8_e4m3fn> ||
243 + std::is_same_v<kv_cache_t, c10::Float8_e5m2>;
244 +
245 + static constexpr int64_t tile_elems = AMX_TILE_BYTES / sizeof(c10::BFloat16);
246 + static constexpr int64_t scratch_elems = fp8_kv ? tile_elems : 1;
247 +
248 public:
249 template <AttentionGemmPhase phase, int32_t k_size>
250 FORCE_INLINE static void gemm(const int32_t m_size,
251 c10::BFloat16* __restrict__ a_tile,
188 - c10::BFloat16* __restrict__ b_tile,
252 + kv_cache_t* __restrict__ b_tile,
253 float* __restrict__ c_tile, const int64_t lda,
254 const int64_t ldb, const int64_t ldc,
255 const int32_t block_size,
@@ -215,21 +279,19 @@ class TileGemm122<c10::BFloat16> {
279 }
280 }();
281
218 - c10::BFloat16* __restrict__ b_tile_2 = b_tile;
219 - c10::BFloat16* __restrict__ b_tile_3 = [&]() {
282 + kv_cache_t* __restrict__ b_tile_2 = b_tile;
283 + kv_cache_t* __restrict__ b_tile_3 = [&]() {
284 if constexpr (phase == AttentionGemmPhase::QK) {
221 - // k_cache is prepacked
285 return b_tile + (k_size * AMX_TILE_ROW_BYTES / 4);
286 } else if constexpr (phase == AttentionGemmPhase::PV) {
224 - // v_cache is prepacked
287 return b_tile + (block_size * AMX_TILE_ROW_BYTES / 4);
288 } else {
289 TORCH_CHECK(false, "Unreachable");
290 }
291 }();
230 - c10::BFloat16* __restrict__ b_tile_4 =
292 + kv_cache_t* __restrict__ b_tile_4 =
293 b_tile_2 + AMX_TILE_BYTES / sizeof(c10::BFloat16);
232 - c10::BFloat16* __restrict__ b_tile_5 =
294 + kv_cache_t* __restrict__ b_tile_5 =
295 b_tile_3 + AMX_TILE_BYTES / sizeof(c10::BFloat16);
296 int64_t b_stride = AMX_TILE_ROW_BYTES;
297
@@ -250,16 +312,25 @@ class TileGemm122<c10::BFloat16> {
312 _tile_zero(7);
313 }
314
315 + alignas(64) c10::BFloat16 scratch_2[scratch_elems];
316 + alignas(64) c10::BFloat16 scratch_3[scratch_elems];
317 + alignas(64) c10::BFloat16 scratch_4[scratch_elems];
318 + alignas(64) c10::BFloat16 scratch_5[scratch_elems];
319 for (int32_t k = 0; k < k_group_times; ++k) {
320 + const c10::BFloat16* load_2 = prepare_b_tile(b_tile_2, scratch_2);
321 + const c10::BFloat16* load_3 = prepare_b_tile(b_tile_3, scratch_3);
322 + const c10::BFloat16* load_4 = prepare_b_tile(b_tile_4, scratch_4);
323 + const c10::BFloat16* load_5 = prepare_b_tile(b_tile_5, scratch_5);
324 +
325 _tile_loadd(0, a_tile_0, a_tile_stride);
255 - _tile_stream_loadd(2, b_tile_2, b_stride);
326 + _tile_stream_loadd(2, const_cast<c10::BFloat16*>(load_2), b_stride);
327 _tile_dpbf16ps(6, 0, 2);
257 - _tile_stream_loadd(3, b_tile_3, b_stride);
328 + _tile_stream_loadd(3, const_cast<c10::BFloat16*>(load_3), b_stride);
329 _tile_dpbf16ps(7, 0, 3);
330 _tile_loadd(1, a_tile_1, a_tile_stride);
260 - _tile_stream_loadd(4, b_tile_4, b_stride);
331 + _tile_stream_loadd(4, const_cast<c10::BFloat16*>(load_4), b_stride);
332 _tile_dpbf16ps(6, 1, 4);
262 - _tile_stream_loadd(5, b_tile_5, b_stride);
333 + _tile_stream_loadd(5, const_cast<c10::BFloat16*>(load_5), b_stride);
334 _tile_dpbf16ps(7, 1, 5);
335
336 // update ptrs
@@ -279,10 +350,13 @@ class TileGemm122<c10::BFloat16> {
350 }
351
352 if (has_tail) {
353 + const c10::BFloat16* load_2 = prepare_b_tile(b_tile_2, scratch_2);
354 + const c10::BFloat16* load_3 = prepare_b_tile(b_tile_3, scratch_3);
355 +
356 _tile_loadd(0, a_tile_0, a_tile_stride);
283 - _tile_stream_loadd(2, b_tile_2, b_stride);
357 + _tile_stream_loadd(2, const_cast<c10::BFloat16*>(load_2), b_stride);
358 _tile_dpbf16ps(6, 0, 2);
285 - _tile_stream_loadd(3, b_tile_3, b_stride);
359 + _tile_stream_loadd(3, const_cast<c10::BFloat16*>(load_3), b_stride);
360 _tile_dpbf16ps(7, 0, 3);
361 }
362
@@ -302,21 +376,25 @@ class TileGemm122<c10::BFloat16> {
376 _tile_loadconfig(&config);
377 }
378 };
379 +
380 } // namespace
381
307 -template <typename scalar_t, int64_t head_dim>
308 -class AttentionImpl<ISA::AMX, scalar_t, head_dim> {
382 +template <typename scalar_t, int64_t head_dim, typename kv_cache_scalar_t>
383 +class AttentionImpl<ISA::AMX, scalar_t, head_dim, kv_cache_scalar_t> {
384 + static constexpr bool fp8_kv =
385 + std::is_same_v<kv_cache_scalar_t, c10::Float8_e4m3fn> ||
386 + std::is_same_v<kv_cache_scalar_t, c10::Float8_e5m2>;
387 +
388 public:
389 using query_t = scalar_t;
390 using q_buffer_t = scalar_t;
312 - using kv_cache_t = scalar_t;
391 + using kv_cache_t = kv_cache_scalar_t;
392 using logits_buffer_t = float;
393 using partial_output_buffer_t = float;
394 using prob_buffer_t = scalar_t;
395
396 constexpr static int64_t BlockSizeAlignment =
318 - AMX_TILE_ROW_BYTES /
319 - sizeof(kv_cache_t); // KV token num unit of QK and PV phases
397 + 32; // AMX_TILE_ROW_NUM = 16 tokens/tile; 32 = 2 tiles
398 constexpr static int64_t HeadDimAlignment =
399 2 * (AMX_TILE_ROW_BYTES / 4); // headdim num unit of PV phase
400 constexpr static int64_t MaxQHeadNumPerIteration = 32;
@@ -324,6 +402,9 @@ class AttentionImpl<ISA::AMX, scalar_t, head_dim> {
402 constexpr static ISA ISAType = ISA::AMX;
403 constexpr static bool scale_on_logits = true;
404
405 + float k_scale = 1.0f;
406 + float v_scale = 1.0f;
407 +
408 public:
409 AttentionImpl() : current_q_head_num_(0) {
410 // Use all columns in AMX tiles
@@ -332,21 +413,50 @@ class AttentionImpl<ISA::AMX, scalar_t, head_dim> {
413
414 ~AttentionImpl() { _tile_release(); }
415
416 + void init_from_input(const AttentionInput* input) {
417 + if constexpr (fp8_kv) {
418 + k_scale = input->k_scale_fp8;
419 + v_scale = input->v_scale_fp8;
420 + }
421 + }
422 +
423 + float get_output_v_scale() const noexcept {
424 + if constexpr (fp8_kv) {
425 + // AMX dequant places FP8 payload into a BF16 field (exponent bias 127).
426 + // Correction = 2^(127 - FP8_bias): E4M3 bias=7 → 2^120, E5M2 bias=15 →
427 + // 2^112.
428 + constexpr float bias =
429 + std::is_same_v<kv_cache_t, c10::Float8_e5m2> ? 0x1p112f : 0x1p120f;
430 + return v_scale * bias;
431 + }
432 + return 1.0f;
433 + }
434 +
435 template <template <typename tile_gemm_t> typename attention>
436 FORCE_INLINE void execute_attention(DEFINE_CPU_ATTENTION_PARAMS) {
437 + if constexpr (fp8_kv) {
438 + // Same bias correction as get_output_v_scale: AMX FP8→BF16 dequant
439 + // shifts the exponent bias from FP8 to BF16 (127), so we multiply by
440 + // 2^(127-FP8_bias) to recover the true value. E4M3: 2^120, E5M2: 2^112.
441 + const float bias =
442 + std::is_same_v<kv_cache_t, c10::Float8_e5m2> ? 0x1p112f : 0x1p120f;
443 + scale *= k_scale * bias;
444 + }
445 if (q_head_num > AMX_TILE_ROW_NUM) {
446 if (q_head_num != current_q_head_num_) {
447 current_q_head_num_ = q_head_num;
340 - TileGemm224<kv_cache_t>::init_tile_config(q_head_num, amx_tile_config_);
448 + TileGemm224<q_buffer_t, kv_cache_t>::init_tile_config(q_head_num,
449 + amx_tile_config_);
450 }
342 - attention<TileGemm224<kv_cache_t>> attention_iteration;
451 + attention<TileGemm224<q_buffer_t, kv_cache_t>> attention_iteration;
452 attention_iteration(CPU_ATTENTION_PARAMS);
453 } else {
454 if (q_head_num != current_q_head_num_) {
455 current_q_head_num_ = q_head_num;
347 - TileGemm122<kv_cache_t>::init_tile_config(q_head_num, amx_tile_config_);
456 + TileGemm122<q_buffer_t, kv_cache_t>::init_tile_config(q_head_num,
457 + amx_tile_config_);
458 }
349 - attention<TileGemm122<kv_cache_t>> attention_iteration;
459 + attention<TileGemm122<q_buffer_t, kv_cache_t>> attention_iteration;
460 attention_iteration(CPU_ATTENTION_PARAMS);
461 }
462 }
@@ -411,13 +521,26 @@ class AttentionImpl<ISA::AMX, scalar_t, head_dim> {
521 // reshape KV to AMX friendly layout
522 static void reshape_and_cache(
523 const scalar_t* __restrict__ key, const scalar_t* __restrict__ value,
414 - scalar_t* __restrict__ key_cache, scalar_t* __restrict__ value_cache,
524 + kv_cache_t* __restrict__ key_cache, kv_cache_t* __restrict__ value_cache,
525 const int64_t* __restrict__ slot_mapping, const int64_t token_num,
526 const int64_t key_token_num_stride, const int64_t value_token_num_stride,
527 const int64_t head_num, const int64_t key_head_num_stride,
528 const int64_t value_head_num_stride, const int64_t num_blocks,
529 const int64_t num_blocks_stride, const int64_t cache_head_num_stride,
420 - const int64_t block_size, const int64_t block_size_stride) {
530 + const int64_t block_size, const int64_t block_size_stride,
531 + const float k_inv = 0.0f, const float v_inv = 0.0f) {
532 + if constexpr (fp8_kv) {
533 + constexpr auto qfn = select_fp8_quant_fn<kv_cache_t>();
534 + reshape_and_cache_fp8_amx_impl<scalar_t, qfn>(
535 + key, value, reinterpret_cast<uint8_t*>(key_cache),
536 + reinterpret_cast<uint8_t*>(value_cache), slot_mapping, token_num,
537 + head_num, head_dim, block_size, key_token_num_stride,
538 + key_head_num_stride, value_token_num_stride, value_head_num_stride,
539 + num_blocks_stride, cache_head_num_stride, num_blocks_stride,
540 + cache_head_num_stride, k_inv, v_inv);
541 + return;
542 + }
543 +
544 // For AMX 2D tiles, size of each line is 64 bytes
545 constexpr int64_t amx_tile_row_size = AMX_TILE_ROW_BYTES;
546 // For AMX B matrix, N always is 16
@@ -426,6 +549,9 @@ class AttentionImpl<ISA::AMX, scalar_t, head_dim> {
549 // For now suppose block_size is divisible by amx_tile_column_num
550 TORCH_CHECK_EQ(block_size % amx_b_tile_k_size, 0);
551
552 + scalar_t* __restrict__ kc = reinterpret_cast<scalar_t*>(key_cache);
553 + scalar_t* __restrict__ vc = reinterpret_cast<scalar_t*>(value_cache);
554 +
555 #pragma omp parallel for collapse(2)
556 for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
557 for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
@@ -453,8 +579,7 @@ class AttentionImpl<ISA::AMX, scalar_t, head_dim> {
579 constexpr int64_t quadword_num_per_group =
580 token_num_per_group * quadword_num;
581 int32_t* key_cache_start_ptr =
456 - reinterpret_cast<int32_t*>(key_cache +
457 - block_idx * num_blocks_stride +
582 + reinterpret_cast<int32_t*>(kc + block_idx * num_blocks_stride +
583 head_idx * cache_head_num_stride) +
584 group_idx * quadword_num_per_group + group_offset;
585
@@ -483,7 +608,7 @@ class AttentionImpl<ISA::AMX, scalar_t, head_dim> {
608 token_idx * value_token_num_stride +
609 head_idx * value_head_num_stride;
610 scalar_t* value_cache_start_ptr =
486 - value_cache + block_idx * num_blocks_stride +
611 + vc + block_idx * num_blocks_stride +
612 head_idx * cache_head_num_stride +
613 sub_group_idx * token_num_per_sub_group * amx_b_tile_n_size +
614 sub_group_offset;
csrc/cpu/cpu_attn_fp8.hpp new
+214
@@ -0,0 +1,214 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +#pragma once
4 +#include <algorithm>
5 +#include <cmath>
6 +#include <cstdint>
7 +#include <limits>
8 +#include <type_traits>
9 +
10 +#include "cpu/utils.hpp"
11 +
12 +typedef uint32_t __attribute__((__may_alias__)) u32_alias_t;
13 +typedef uint16_t __attribute__((__may_alias__)) u16_alias_t;
14 +typedef float __attribute__((__may_alias__)) f32_alias_t;
15 +
16 +// Reference scalar dequant — used to verify vectorized AMX dequant.
17 +inline float fp8e4m3_to_float_scalar(uint8_t b, float scale) noexcept {
18 + // NaN encoding in E4M3
19 + if ((b & 0x7F) == 0x7F) return std::numeric_limits<float>::quiet_NaN();
20 + uint32_t b_u32 = static_cast<uint32_t>(b);
21 + uint32_t sign = (b_u32 & 0x80) << 24;
22 + uint32_t payload = (b_u32 & 0x7F) << 20;
23 + uint32_t bits = sign | payload;
24 + float b_f32_unscaled = *reinterpret_cast<const f32_alias_t*>(&bits);
25 + float b_f32_scaled = b_f32_unscaled * scale * 0x1p120f;
26 + return b_f32_scaled;
27 +}
28 +
29 +inline uint8_t float_to_fp8e4m3_scalar(float v, float inv_scale) noexcept {
30 + v *= inv_scale;
31 + constexpr float fp8_max = 448.0f;
32 + v = std::max(-fp8_max, std::min(fp8_max, v));
33 + if (v == 0.0f) return 0;
34 +
35 + // Inverse mapping of fp8e4m3_to_float_scalar: shift the effective exponent
36 + // bias from fp32 (127) back to fp8 e4m3 (7), then pack sign|payload.
37 + float v_f32_unscaled = v * 0x1p-120f;
38 + uint32_t bits = *reinterpret_cast<const u32_alias_t*>(&v_f32_unscaled);
39 + uint8_t sign = static_cast<uint8_t>((bits >> 24) & 0x80);
40 + uint8_t payload = static_cast<uint8_t>((bits >> 20) & 0x7F);
41 + if (payload == 0) return sign;
42 + payload = std::min<uint8_t>(payload, 0x7E); // keep 0x7F as NaN encoding
43 + return static_cast<uint8_t>(sign | payload);
44 +}
45 +
46 +// ---------------------------------------------------------------------------
47 +// AMX reshape impl — parameterised on the quantisation function.
48 +// Writes key/value into uint8 FP8 KV cache using the AMX tile-friendly layout.
49 +// K: halfword-packed (2 FP8 per uint16, token_num_per_group=16).
50 +// V: sub-group packing (token_num_per_sub_group=2, head_elems_per_group=16).
51 +// block_size must be divisible by 32.
52 +// ---------------------------------------------------------------------------
53 +template <typename scalar_t, uint8_t (*quant_fn)(float, float)>
54 +inline void reshape_and_cache_fp8_amx_impl(
55 + const scalar_t* key_ptr, const scalar_t* value_ptr, uint8_t* key_cache_ptr,
56 + uint8_t* value_cache_ptr, const int64_t* slot_ptr, int64_t token_num,
57 + int64_t head_num, int64_t head_dim, int64_t block_size, int64_t k_stride0,
58 + int64_t k_stride1, int64_t v_stride0, int64_t v_stride1, int64_t kc_stride0,
59 + int64_t kc_stride1, int64_t vc_stride0, int64_t vc_stride1, float k_inv,
60 + float v_inv) {
61 + constexpr int64_t token_num_per_group = 16; // AMX_TILE_ROW_NUM
62 + const int64_t halfword_num = head_dim / 2; // 2 FP8 per uint16
63 + const int64_t halfword_num_per_group = token_num_per_group * halfword_num;
64 + constexpr int64_t head_elems_per_group = 16;
65 + constexpr int64_t token_num_per_sub_group = 2; // = 4 / sizeof(BF16)
66 + const int64_t group_num = head_dim / head_elems_per_group;
67 + const int64_t group_size = block_size * head_elems_per_group;
68 +
69 +#pragma omp parallel for collapse(2) schedule(static)
70 + for (int64_t tok = 0; tok < token_num; ++tok) {
71 + for (int64_t h = 0; h < head_num; ++h) {
72 + const int64_t slot = slot_ptr[tok];
73 + if (slot < 0) continue;
74 + const int64_t block_idx = slot / block_size;
75 + const int64_t block_offset = slot % block_size;
76 +
77 + // Key: halfword-packed, 2 FP8 per uint16
78 + {
79 + const scalar_t* ksrc = key_ptr + tok * k_stride0 + h * k_stride1;
80 + const int64_t group_idx = block_offset / token_num_per_group;
81 + const int64_t group_offset = block_offset % token_num_per_group;
82 + uint16_t* kdst =
83 + reinterpret_cast<uint16_t*>(key_cache_ptr + block_idx * kc_stride0 +
84 + h * kc_stride1) +
85 + group_idx * halfword_num_per_group + group_offset;
86 + for (int64_t j = 0; j < halfword_num; ++j) {
87 + uint8_t fp8_0 = quant_fn(static_cast<float>(ksrc[j * 2]), k_inv);
88 + uint8_t fp8_1 = quant_fn(static_cast<float>(ksrc[j * 2 + 1]), k_inv);
89 + uint8_t bytes[2] = {fp8_0, fp8_1};
90 + uint16_t hw = *reinterpret_cast<const u16_alias_t*>(bytes);
91 + kdst[j * token_num_per_group] = hw;
92 + }
93 + }
94 +
95 + // Value: sub-group packing (token_num_per_sub_group = 2)
96 + {
97 + const scalar_t* vsrc = value_ptr + tok * v_stride0 + h * v_stride1;
98 + const int64_t sub_group_idx = block_offset / token_num_per_sub_group;
99 + const int64_t sub_group_offset = block_offset % token_num_per_sub_group;
100 + uint8_t* vdst =
101 + value_cache_ptr + block_idx * vc_stride0 + h * vc_stride1 +
102 + sub_group_idx * token_num_per_sub_group * head_elems_per_group +
103 + sub_group_offset;
104 + for (int64_t i = 0; i < group_num; ++i) {
105 + for (int64_t j = 0; j < head_elems_per_group; ++j)
106 + vdst[j * token_num_per_sub_group] =
107 + quant_fn(static_cast<float>(vsrc[j]), v_inv);
108 + vsrc += head_elems_per_group;
109 + vdst += group_size;
110 + }
111 + }
112 + }
113 + }
114 +}
115 +
116 +// ---------------------------------------------------------------------------
117 +// FP8 E5M2 scalar helpers
118 +// ---------------------------------------------------------------------------
119 +
120 +// Reference scalar dequant — used to verify vectorized AMX dequant.
121 +// FP8 E5M2: s[7] e[6:2] m[1:0], exponent bias = 15 (same as FP16).
122 +// Byte b → FP16 bits = b << 8 (no bias correction needed).
123 +inline float fp8e5m2_to_float_scalar(uint8_t b, float scale) noexcept {
124 + const uint8_t exp_bits = (b >> 2) & 0x1F;
125 + const uint8_t mant_bits = b & 0x03;
126 + // NaN: exp=11111, mant!=00
127 + if (exp_bits == 0x1F && mant_bits != 0)
128 + return std::numeric_limits<float>::quiet_NaN();
129 + const uint32_t sign = static_cast<uint32_t>(b & 0x80) << 24;
130 + if (exp_bits == 0x1F)
131 + return sign ? -std::numeric_limits<float>::infinity()
132 + : std::numeric_limits<float>::infinity();
133 + if (exp_bits == 0) { // subnormal: (-1)^s * 2^-14 * mant/4
134 + if (mant_bits == 0) return 0.0f;
135 + float v = mant_bits * 0x1p-16f;
136 + return (sign ? -v : v) * scale;
137 + }
138 + // Normal: FP32 exp = exp5 - 15 + 127, mantissa top 2 bits
139 + uint32_t fp32_bits = sign |
140 + ((static_cast<uint32_t>(exp_bits) - 15 + 127) << 23) |
141 + (static_cast<uint32_t>(mant_bits) << 21);
142 + float val = *reinterpret_cast<const f32_alias_t*>(&fp32_bits);
143 + return val * scale;
144 +}
145 +
146 +inline uint8_t float_to_fp8e5m2_scalar(float v, float inv_scale) noexcept {
147 + v *= inv_scale;
148 + constexpr float fp8_e5m2_max = 57344.0f;
149 + v = std::max(-fp8_e5m2_max, std::min(fp8_e5m2_max, v));
150 + if (v == 0.0f) return 0;
151 + uint32_t bits = *reinterpret_cast<const u32_alias_t*>(&v);
152 + const uint8_t sign = static_cast<uint8_t>((bits >> 24) & 0x80);
153 + const int32_t exp_fp32 = static_cast<int32_t>((bits >> 23) & 0xFF) - 127;
154 + const uint8_t mant2 = static_cast<uint8_t>((bits >> 21) & 0x03);
155 + if (exp_fp32 < -14) { // subnormal in E5M2
156 + const int shift = -14 - exp_fp32;
157 + if (shift + 21 >= 32)
158 + return sign; // underflow: too small for E5M2 subnormal
159 + const uint32_t m = (0x800000u | (bits & 0x7FFFFFu)) >> (shift + 21);
160 + return sign | static_cast<uint8_t>(std::min<uint32_t>(m, 3u));
161 + }
162 + const uint8_t exp5 = static_cast<uint8_t>(exp_fp32 + 15);
163 + return sign | (exp5 << 2) | mant2;
164 +}
165 +
166 +// ---------------------------------------------------------------------------
167 +// Select the FP8 quant function at compile time based on kv_cache_t.
168 +// ---------------------------------------------------------------------------
169 +template <typename kv_cache_t>
170 +constexpr auto select_fp8_quant_fn() {
171 + if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e5m2>)
172 + return float_to_fp8e5m2_scalar;
173 + else
174 + return float_to_fp8e4m3_scalar;
175 +}
176 +
177 +// ---------------------------------------------------------------------------
178 +// VEC reshape impl — parameterised on the quantisation function.
179 +// Writes key (column-major) and value (row-major) into uint8 FP8 KV cache.
180 +// The pragma omp must live outside VLLM_DISPATCH_FLOATING_TYPES because
181 +// #pragma cannot appear inside variadic macro arguments.
182 +// ---------------------------------------------------------------------------
183 +template <typename scalar_t, uint8_t (*quant_fn)(float, float)>
184 +inline void reshape_and_cache_fp8_vec_impl(
185 + const scalar_t* key_ptr, const scalar_t* value_ptr, uint8_t* key_cache_ptr,
186 + uint8_t* value_cache_ptr, const int64_t* slot_ptr, int64_t token_num,
187 + int64_t head_num, int64_t head_dim, int64_t block_size, int64_t k_stride0,
188 + int64_t k_stride1, int64_t v_stride0, int64_t v_stride1, int64_t kc_stride0,
189 + int64_t kc_stride1, int64_t vc_stride0, int64_t vc_stride1, float k_inv,
190 + float v_inv) {
191 +#pragma omp parallel for collapse(2) schedule(static)
192 + for (int64_t tok = 0; tok < token_num; ++tok) {
193 + for (int64_t h = 0; h < head_num; ++h) {
194 + const int64_t slot = slot_ptr[tok];
195 + if (slot < 0) continue;
196 + const int64_t block_idx = slot / block_size;
197 + const int64_t block_offset = slot % block_size;
198 +
199 + // Key layout: column-major within block
200 + const scalar_t* ksrc = key_ptr + tok * k_stride0 + h * k_stride1;
201 + uint8_t* kdst = key_cache_ptr + block_idx * kc_stride0 + h * kc_stride1 +
202 + block_offset;
203 + for (int64_t i = 0; i < head_dim; ++i)
204 + kdst[i * block_size] = quant_fn(static_cast<float>(ksrc[i]), k_inv);
205 +
206 + // Value layout: row-major within block (contiguous head_dim bytes)
207 + const scalar_t* vsrc = value_ptr + tok * v_stride0 + h * v_stride1;
208 + uint8_t* vdst = value_cache_ptr + block_idx * vc_stride0 +
209 + h * vc_stride1 + block_offset * head_dim;
210 + for (int64_t i = 0; i < head_dim; ++i)
211 + vdst[i] = quant_fn(static_cast<float>(vsrc[i]), v_inv);
212 + }
213 + }
214 +}
csrc/cpu/cpu_attn_impl.hpp
+31 -7
@@ -14,8 +14,22 @@
14 namespace cpu_attention {
15 enum class ISA { AMX, VEC, VEC16, NEON, VXE };
16
17 -template <ISA isa, typename scalar_t, int64_t head_dim>
18 -class AttentionImpl {};
17 +// Mirrors csrc/attention/dtype_fp8.cuh Fp8KVCacheDataType exactly.
18 +enum class Fp8KVCacheDataType {
19 + kAuto = 0,
20 + kFp8E4M3 = 1,
21 + kFp8E5M2 = 2,
22 +};
23 +
24 +struct AttentionInput;
25 +
26 +template <ISA isa, typename scalar_t, int64_t head_dim,
27 + typename kv_cache_scalar_t = scalar_t>
28 +class AttentionImpl {
29 + public:
30 + void init_from_input(const AttentionInput*) {}
31 + float get_output_v_scale() const noexcept { return 1.0f; }
32 +};
33
34 struct AttentionWorkItemGroup {
35 int32_t req_id;
@@ -780,6 +794,9 @@ struct AttentionInput {
794 int32_t sliding_window_left;
795 int32_t sliding_window_right;
796 float softcap;
797 + // FP8 KV cache scales (used by FP8 attention implementations)
798 + float k_scale_fp8 = 1.0f;
799 + float v_scale_fp8 = 1.0f;
800 };
801
802 #define DEFINE_CPU_ATTENTION_PARAMS \
@@ -1374,6 +1391,13 @@ class AttentionMainLoop {
1391 }
1392
1393 attention_impl_t attn_impl;
1394 + constexpr bool fp8_kv = std::is_same_v<kv_cache_t, c10::Float8_e4m3fn> ||
1395 + std::is_same_v<kv_cache_t, c10::Float8_e5m2>;
1396 + float output_v_scale = 1.0f;
1397 + if constexpr (fp8_kv) {
1398 + attn_impl.init_from_input(input);
1399 + output_v_scale = attn_impl.get_output_v_scale();
1400 + }
1401
1402 // general information
1403 const int32_t q_head_num = input->num_heads;
@@ -1753,7 +1777,7 @@ class AttentionMainLoop {
1777 reinterpret_cast<query_t*>(input->output) +
1778 output_buffer_offset,
1779 sum_buffer, actual_q_heads_per_kv,
1756 - actual_q_token_num, q_head_num);
1780 + actual_q_token_num, q_head_num, output_v_scale);
1781 } else {
1782 const int32_t stride =
1783 actual_q_heads_per_kv * split_kv_q_token_num_threshold;
@@ -1823,7 +1847,7 @@ class AttentionMainLoop {
1847 split_output_buffer,
1848 reinterpret_cast<query_t*>(input->output) + output_buffer_offset,
1849 split_sum_buffer, actual_q_heads_per_kv, curr_output_token_num,
1826 - q_head_num);
1850 + q_head_num, output_v_scale);
1851 }
1852 }
1853 }
@@ -1947,8 +1971,8 @@ class AttentionMainLoop {
1971 query_t* __restrict__ curr_output_buffer,
1972 float* __restrict__ sum_buffer,
1973 const int32_t q_heads_per_kv,
1950 - const int32_t actual_q_token_num,
1951 - const int32_t q_head_num) {
1974 + const int32_t actual_q_token_num, const int32_t q_head_num,
1975 + const float v_scale = 1.0f) {
1976 // final output
1977 using output_vec_t = typename VecTypeTrait<query_t>::vec_t;
1978
@@ -1962,7 +1986,7 @@ class AttentionMainLoop {
1986 curr_partial_output_buffer;
1987 query_t* __restrict__ curr_output_buffer_iter = curr_output_buffer;
1988 for (int32_t head_idx = 0; head_idx < q_heads_per_kv; ++head_idx) {
1965 - vec_op::FP32Vec16 inv_sum_scale_vec(1.0 / *curr_sum_buffer);
1989 + vec_op::FP32Vec16 inv_sum_scale_vec(v_scale / *curr_sum_buffer);
1990
1991 for (int32_t i = 0; i < group_num_per_head; ++i) {
1992 vec_op::FP32Vec16 vec(curr_partial_output_buffer_iter);
csrc/cpu/cpu_attn_neon.hpp
+5 -4
@@ -248,8 +248,8 @@ class TileGemmNeonFMLA {
248 } // namespace
249
250 // this is similar to "ISA::VEC" at the moment
251 -template <typename scalar_t, int64_t head_dim>
252 -class AttentionImpl<ISA::NEON, scalar_t, head_dim> {
251 +template <typename scalar_t, int64_t head_dim, typename kv_cache_scalar_t>
252 +class AttentionImpl<ISA::NEON, scalar_t, head_dim, kv_cache_scalar_t> {
253 public:
254 using query_t = scalar_t;
255 using q_buffer_t = float;
@@ -343,7 +343,8 @@ class AttentionImpl<ISA::NEON, scalar_t, head_dim> {
343 const int64_t head_num, const int64_t key_head_num_stride,
344 const int64_t value_head_num_stride, const int64_t num_blocks,
345 const int64_t num_blocks_stride, const int64_t cache_head_num_stride,
346 - const int64_t block_size, const int64_t block_size_stride) {
346 + const int64_t block_size, const int64_t block_size_stride,
347 + const float /*k_inv*/ = 0.0f, const float /*v_inv*/ = 0.0f) {
348 #pragma omp parallel for collapse(2)
349 for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
350 for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
@@ -388,7 +389,7 @@ class AttentionImpl<ISA::NEON, scalar_t, head_dim> {
389 #ifdef ARM_BF16_SUPPORT
390 // For BF16 on Arm, reuse the BFMMLA kernels with 32-token alignment.
391 template <int64_t head_dim>
391 -class AttentionImpl<ISA::NEON, c10::BFloat16, head_dim>
392 +class AttentionImpl<ISA::NEON, c10::BFloat16, head_dim, c10::BFloat16>
393 : public AttentionImplNEONBFMMLA<BLOCK_SIZE_ALIGNMENT, ISA::NEON,
394 head_dim> {};
395 #endif
csrc/cpu/cpu_attn_neon_bfmmla.hpp
+2 -1
@@ -602,7 +602,8 @@ class AttentionImplNEONBFMMLA {
602 [[maybe_unused]] const int64_t num_blocks,
603 const int64_t num_blocks_stride, const int64_t cache_head_num_stride,
604 const int64_t block_size,
605 - [[maybe_unused]] const int64_t block_size_stride) {
605 + [[maybe_unused]] const int64_t block_size_stride,
606 + const float /*k_inv*/ = 0.0f, const float /*v_inv*/ = 0.0f) {
607 const int64_t k_block_stride = (head_dim / TILE_K) * K_INNER_STRIDE;
608 const int64_t v_pair_stride =
609 (block_size / V_TOKENS_PER_ROW_BLOCK) * V_INNER_STRIDE;
csrc/cpu/cpu_attn_vec.hpp
+105 -28
@@ -1,11 +1,37 @@
1 #ifndef CPU_ATTN_VEC_HPP
2 #define CPU_ATTN_VEC_HPP
3
4 +#include "cpu_attn_fp8.hpp"
5 #include "cpu_attn_impl.hpp"
6
7 namespace cpu_attention {
8
9 namespace {
10 +
11 +// Load 32 kv_cache_t elements starting at ptr and return them as two FP32Vec16s
12 +// covering the lower 16 and upper 16 positions.
13 +// For FP8: both halves come from a single BF16Vec32 dequant of 32 bytes.
14 +// For BF16/FP16/FP32: two separate vector loads at ptr and ptr+16.
15 +template <typename kv_cache_t>
16 +FORCE_INLINE std::pair<vec_op::FP32Vec16, vec_op::FP32Vec16> load_b_pair_vec(
17 + const kv_cache_t* ptr) {
18 + if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e4m3fn>) {
19 + // BF16 container, but values are in the FP16 exponent range (bias 15 not
20 + // 127).
21 + vec_op::BF16Vec32 bf16_b_reg(reinterpret_cast<const uint8_t*>(ptr),
22 + vec_op::fp8_e4m3_tag{});
23 + return {vec_op::FP32Vec16(bf16_b_reg, 0), vec_op::FP32Vec16(bf16_b_reg, 1)};
24 + } else if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e5m2>) {
25 + vec_op::BF16Vec32 bf16_b_reg(reinterpret_cast<const uint8_t*>(ptr),
26 + vec_op::fp8_e5m2_tag{});
27 + return {vec_op::FP32Vec16(bf16_b_reg, 0), vec_op::FP32Vec16(bf16_b_reg, 1)};
28 + } else {
29 + using load_vec_t = typename VecTypeTrait<kv_cache_t>::vec_t;
30 + return {vec_op::FP32Vec16(load_vec_t(ptr)),
31 + vec_op::FP32Vec16(load_vec_t(ptr + 16))};
32 + }
33 +}
34 +
35 // 8-2-16 pattern, 8 regs for A, 2 regs for B, 16 regs for C, [8, K] @ [k, 32]
36 template <typename kv_cache_t>
37 class TileGemm82 {
@@ -54,10 +80,7 @@ class TileGemm82 {
80 const int32_t block_size, const int32_t dynamic_k_size,
81 const bool accum_c) {
82 static_assert(0 < M && M <= 8);
57 - using load_vec_t = typename VecTypeTrait<kv_cache_t>::vec_t;
83
59 - kv_cache_t* __restrict__ curr_b_0 = b_tile;
60 - kv_cache_t* __restrict__ curr_b_1 = b_tile + 16;
84 float* __restrict__ curr_c_0 = c_tile;
85 float* __restrict__ curr_c_1 = c_tile + 16;
86
@@ -76,16 +99,14 @@ class TileGemm82 {
99 }
100
101 float* __restrict__ curr_a = a_tile;
102 + kv_cache_t* __restrict__ curr_b = b_tile;
103 +
104 for (int32_t k = 0; k < dynamic_k_size; ++k) {
80 - load_vec_t b_0_reg(curr_b_0);
81 - vec_op::FP32Vec16 fp32_b_0_reg(b_0_reg);
82 - load_vec_t b_1_reg(curr_b_1);
83 - vec_op::FP32Vec16 fp32_b_1_reg(b_1_reg);
105 + auto [fp32_b_0_reg, fp32_b_1_reg] = load_b_pair_vec(curr_b);
106
107 float* __restrict__ curr_m_a = curr_a;
108 vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
87 - float v = *curr_m_a;
88 - vec_op::FP32Vec16 a_reg(v);
109 + vec_op::FP32Vec16 a_reg(*curr_m_a);
110 c_regs[i * 2] = c_regs[i * 2] + a_reg * fp32_b_0_reg;
111 c_regs[i * 2 + 1] = c_regs[i * 2 + 1] + a_reg * fp32_b_1_reg;
112
@@ -95,8 +116,7 @@ class TileGemm82 {
116
117 // update
118 curr_a += 1;
98 - curr_b_0 += ldb;
99 - curr_b_1 += ldb;
119 + curr_b += ldb;
120 }
121
122 vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
@@ -109,15 +129,20 @@ class TileGemm82 {
129 });
130 }
131 };
132 +
133 } // namespace
134
135 // This is a general but naive implementation based on vector instructions
115 -template <typename scalar_t, int64_t head_dim>
116 -class AttentionImpl<ISA::VEC, scalar_t, head_dim> {
136 +template <typename scalar_t, int64_t head_dim, typename kv_cache_scalar_t>
137 +class AttentionImpl<ISA::VEC, scalar_t, head_dim, kv_cache_scalar_t> {
138 + static constexpr bool fp8_kv =
139 + std::is_same_v<kv_cache_scalar_t, c10::Float8_e4m3fn> ||
140 + std::is_same_v<kv_cache_scalar_t, c10::Float8_e5m2>;
141 +
142 public:
143 using query_t = scalar_t;
144 using q_buffer_t = float;
120 - using kv_cache_t = scalar_t;
145 + using kv_cache_t = kv_cache_scalar_t;
146 using logits_buffer_t = float;
147 using partial_output_buffer_t = float;
148 using prob_buffer_t = float;
@@ -129,11 +154,45 @@ class AttentionImpl<ISA::VEC, scalar_t, head_dim> {
154 constexpr static int64_t MaxQHeadNumPerIteration = 8;
155 constexpr static int64_t HeadDim = head_dim;
156 constexpr static ISA ISAType = ISA::VEC;
132 - constexpr static bool scale_on_logits = false; // apply scale on q_buffer
157 + constexpr static bool scale_on_logits = fp8_kv;
158 +
159 + float k_scale = 1.0f;
160 + float v_scale = 1.0f;
161
162 public:
163 + void init_from_input(const AttentionInput* input) {
164 + if constexpr (fp8_kv) {
165 + k_scale = input->k_scale_fp8;
166 + v_scale = input->v_scale_fp8;
167 + }
168 + }
169 +
170 + float get_output_v_scale() const noexcept {
171 + if constexpr (fp8_kv) {
172 + // VEC dequant unpacks FP8 into a pseudo-FP16 layout (exponent bias 15).
173 + // E4M3 (bias=7) needs correction 2^(15-7) = 2^8; E5M2 bias matches FP16
174 + // so no correction.
175 + if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e5m2>) {
176 + return v_scale;
177 + } else {
178 + return v_scale * 0x1p8f;
179 + }
180 + }
181 + return 1.0f;
182 + }
183 +
184 template <template <typename tile_gemm_t> typename attention>
185 FORCE_INLINE void execute_attention(DEFINE_CPU_ATTENTION_PARAMS) {
186 + if constexpr (fp8_kv) {
187 + // Same bias correction as get_output_v_scale: VEC FP8→pseudo-FP16 dequant
188 + // uses bias 15; E4M3 (bias=7) needs ×2^8, E5M2 (bias=15) needs no
189 + // correction.
190 + if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e5m2>) {
191 + scale *= k_scale;
192 + } else {
193 + scale *= k_scale * 0x1p8f;
194 + }
195 + }
196 attention<TileGemm82<kv_cache_t>> attention_iteration;
197 attention_iteration(CPU_ATTENTION_PARAMS);
198 }
@@ -161,17 +220,19 @@ class AttentionImpl<ISA::VEC, scalar_t, head_dim> {
220 // row-major
221 }
222
164 - // Copy q to q_buffer and cast it to fp32
165 - static void copy_q_heads_tile(
166 - scalar_t* __restrict__ src, // [q_num, q_heads_per_kv, head_size]
167 - float* __restrict__ q_buffer, const int32_t q_num,
168 - const int32_t q_heads_per_kv, const int64_t q_num_stride,
169 - const int64_t q_head_stride, float scale) {
223 + // Copy q to q_buffer and cast it to fp32.
224 + // FP8: QK scale is folded into execute_attention; copy Q unscaled here.
225 + void copy_q_heads_tile(scalar_t* __restrict__ src,
226 + float* __restrict__ q_buffer, const int32_t q_num,
227 + const int32_t q_heads_per_kv,
228 + const int64_t q_num_stride,
229 + const int64_t q_head_stride, float scale) {
230 static_assert(head_dim % 16 == 0);
231 constexpr int32_t unroll_size = head_dim / 16;
232 using load_vec_t = typename VecTypeTrait<scalar_t>::vec_t;
233
174 - vec_op::FP32Vec16 scale_vec(scale);
234 + const float effective_scale = fp8_kv ? 1.0f : scale;
235 + vec_op::FP32Vec16 scale_vec(effective_scale);
236 for (int32_t q_num_idx = 0; q_num_idx < q_num; ++q_num_idx) {
237 for (int32_t q_head_idx = 0; q_head_idx < q_heads_per_kv; ++q_head_idx) {
238 scalar_t* __restrict__ curr_q =
@@ -196,13 +257,26 @@ class AttentionImpl<ISA::VEC, scalar_t, head_dim> {
257 // reshape K as column-major and V as row-major
258 static void reshape_and_cache(
259 const scalar_t* __restrict__ key, const scalar_t* __restrict__ value,
199 - scalar_t* __restrict__ key_cache, scalar_t* __restrict__ value_cache,
260 + kv_cache_t* __restrict__ key_cache, kv_cache_t* __restrict__ value_cache,
261 const int64_t* __restrict__ slot_mapping, const int64_t token_num,
262 const int64_t key_token_num_stride, const int64_t value_token_num_stride,
263 const int64_t head_num, const int64_t key_head_num_stride,
264 const int64_t value_head_num_stride, const int64_t num_blocks,
265 const int64_t num_blocks_stride, const int64_t cache_head_num_stride,
205 - const int64_t block_size, const int64_t block_size_stride) {
266 + const int64_t block_size, const int64_t block_size_stride,
267 + const float k_inv = 0.0f, const float v_inv = 0.0f) {
268 + if constexpr (fp8_kv) {
269 + constexpr auto qfn = select_fp8_quant_fn<kv_cache_t>();
270 + reshape_and_cache_fp8_vec_impl<scalar_t, qfn>(
271 + key, value, reinterpret_cast<uint8_t*>(key_cache),
272 + reinterpret_cast<uint8_t*>(value_cache), slot_mapping, token_num,
273 + head_num, head_dim, block_size, key_token_num_stride,
274 + key_head_num_stride, value_token_num_stride, value_head_num_stride,
275 + num_blocks_stride, cache_head_num_stride, num_blocks_stride,
276 + cache_head_num_stride, k_inv, v_inv);
277 + return;
278 + }
279 +
280 #pragma omp parallel for collapse(2)
281 for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
282 for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
@@ -220,8 +294,9 @@ class AttentionImpl<ISA::VEC, scalar_t, head_dim> {
294 token_idx * key_token_num_stride +
295 head_idx * key_head_num_stride;
296 scalar_t* key_cache_start_ptr =
223 - key_cache + block_idx * num_blocks_stride +
224 - head_idx * cache_head_num_stride + block_offset;
297 + reinterpret_cast<scalar_t*>(key_cache) +
298 + block_idx * num_blocks_stride + head_idx * cache_head_num_stride +
299 + block_offset;
300
301 #pragma GCC unroll 8
302 for (int64_t i = 0, j = 0; i < head_dim; ++i, j += block_size) {
@@ -234,8 +309,9 @@ class AttentionImpl<ISA::VEC, scalar_t, head_dim> {
309 token_idx * value_token_num_stride +
310 head_idx * value_head_num_stride;
311 scalar_t* value_cache_start_ptr =
237 - value_cache + block_idx * num_blocks_stride +
238 - head_idx * cache_head_num_stride + block_offset * head_dim;
312 + reinterpret_cast<scalar_t*>(value_cache) +
313 + block_idx * num_blocks_stride + head_idx * cache_head_num_stride +
314 + block_offset * head_dim;
315 std::memcpy(value_cache_start_ptr, value_start_ptr,
316 sizeof(scalar_t) * head_dim);
317 }
@@ -243,6 +319,7 @@ class AttentionImpl<ISA::VEC, scalar_t, head_dim> {
319 }
320 }
321 };
322 +
323 } // namespace cpu_attention
324
325 #endif
csrc/cpu/cpu_attn_vec16.hpp
+3 -3
@@ -116,9 +116,9 @@ class TileGemm161 {
116 } // namespace
117
118 // This is a general but naive implementation based on vector instructions
119 -template <typename scalar_t, int64_t head_dim>
120 -class AttentionImpl<ISA::VEC16, scalar_t, head_dim>
121 - : public AttentionImpl<ISA::VEC, scalar_t, head_dim> {
119 +template <typename scalar_t, int64_t head_dim, typename kv_cache_scalar_t>
120 +class AttentionImpl<ISA::VEC16, scalar_t, head_dim, kv_cache_scalar_t>
121 + : public AttentionImpl<ISA::VEC, scalar_t, head_dim, kv_cache_scalar_t> {
122 public:
123 using query_t = scalar_t;
124 using q_buffer_t = float;
csrc/cpu/cpu_attn_vxe.hpp
+4 -3
@@ -244,8 +244,8 @@ class TileGemmS390X {
244
245 } // namespace
246
247 -template <typename scalar_t, int64_t head_dim>
248 -class AttentionImpl<ISA::VXE, scalar_t, head_dim> {
247 +template <typename scalar_t, int64_t head_dim, typename kv_cache_scalar_t>
248 +class AttentionImpl<ISA::VXE, scalar_t, head_dim, kv_cache_scalar_t> {
249 public:
250 using query_t = scalar_t;
251 using q_buffer_t = float;
@@ -342,7 +342,8 @@ class AttentionImpl<ISA::VXE, scalar_t, head_dim> {
342 const int64_t head_num, const int64_t key_head_num_stride,
343 const int64_t value_head_num_stride, const int64_t num_blocks,
344 const int64_t num_blocks_stride, const int64_t cache_head_num_stride,
345 - const int64_t block_size, const int64_t block_size_stride) {
345 + const int64_t block_size, const int64_t block_size_stride,
346 + const float /*k_inv*/ = 0.0f, const float /*v_inv*/ = 0.0f) {
347 #pragma omp parallel for collapse(2)
348 for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
349 for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
csrc/cpu/cpu_types_arm.hpp
+6
@@ -15,6 +15,9 @@ using namespace at::vec;
15
16 namespace vec_op {
17
18 +struct fp8_e4m3_tag {};
19 +struct fp8_e5m2_tag {};
20 +
21 #define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \
22 AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
23 AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \
@@ -322,6 +325,9 @@ struct BF16Vec32 : public VectorizedRegWrapper<BF16Vec32, 4, c10::BFloat16> {
325 reg.val[2] = vec8_data.reg.val[0];
326 reg.val[3] = vec8_data.reg.val[0];
327 };
328 +
329 + explicit BF16Vec32(const uint8_t*, fp8_e4m3_tag) : Base() {}
330 + explicit BF16Vec32(const uint8_t*, fp8_e5m2_tag) : Base() {}
331 };
332
333 struct FP32Vec4 : public VectorizedRegWrapper<FP32Vec4, 1, float> {
csrc/cpu/cpu_types_vxe.hpp
+6
@@ -8,6 +8,9 @@
8 #include <torch/all.h>
9 namespace vec_op {
10
11 +struct fp8_e4m3_tag {};
12 +struct fp8_e5m2_tag {};
13 +
14 #define vec_neg(a) (-(a))
15 #define vec_add(a, b) ((a) + (b))
16 #define vec_sub(a, b) ((a) - (b))
@@ -241,6 +244,9 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
244 explicit BF16Vec32(const BF16Vec8& vec8_data)
245 : reg({vec8_data.reg, vec8_data.reg, vec8_data.reg, vec8_data.reg}) {}
246
247 + explicit BF16Vec32(const uint8_t*, fp8_e4m3_tag) : reg{} {}
248 + explicit BF16Vec32(const uint8_t*, fp8_e5m2_tag) : reg{} {}
249 +
250 void save(void* ptr) const { *reinterpret_cast<ss16x8x4_t*>(ptr) = reg; }
251 };
252
csrc/cpu/cpu_types_x86.hpp
+139
@@ -11,6 +11,17 @@ static_assert(false, "AVX2 must be supported for the current implementation.");
11
12 namespace vec_op {
13
14 +// Tags for FP8 BF16Vec32 constructors (avoid overload collision with
15 +// BF16Vec32(void*)).
16 +// VEC path (FP8 → pseudo-FP16 layout, scale correction applied later):
17 +struct fp8_e4m3_tag {}; // E4M3 → pseudo-FP16; BF16 value = true_E4M3 * 2^-8
18 +struct fp8_e5m2_tag {}; // E5M2 → FP16 bits directly (same exponent bias=15)
19 +// AMX path (FP8 → unscaled BF16, no FP32 round-trip):
20 +// BF16 value = true_E4M3 * 2^-120 (E4M3) or true_E5M2 * 2^-112 (E5M2).
21 +// Exponent rebiasing is folded into k/v scales by the caller.
22 +struct fp8_bf16_e4m3_tag {};
23 +struct fp8_bf16_e5m2_tag {};
24 +
25 #define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \
26 AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
27 AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) \
@@ -176,6 +187,50 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
187 (__m128i)vec8_data.reg, 2),
188 (__m128i)vec8_data.reg, 3)) {}
189
190 + // Decode 32 FP8-E4M3 bytes to pseudo-FP16 layout (stored in the BF16
191 + // register). Result = true_E4M3 * 2^-8; caller applies scale * 2^8.
192 + explicit BF16Vec32(const uint8_t* ptr, fp8_e4m3_tag) {
193 + __m256i b8 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(ptr));
194 + __m512i b16 = _mm512_cvtepu8_epi16(b8);
195 + __m512i sign =
196 + _mm512_slli_epi16(_mm512_and_si512(b16, _mm512_set1_epi16(0x80)), 8);
197 + __m512i payload =
198 + _mm512_slli_epi16(_mm512_and_si512(b16, _mm512_set1_epi16(0x7F)), 7);
199 + reg = _mm512_or_si512(sign, payload);
200 + }
201 +
202 + // Decode 32 FP8-E5M2 bytes to FP16 layout.
203 + // E5M2 and FP16 share the same 5-bit exponent bias (15), so FP8 byte b maps
204 + // directly to FP16 bits by shifting left 8 — no sign/payload reconstruction.
205 + explicit BF16Vec32(const uint8_t* ptr, fp8_e5m2_tag) {
206 + __m256i b8 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(ptr));
207 + reg = _mm512_slli_epi16(_mm512_cvtepu8_epi16(b8), 8);
208 + }
209 +
210 + // Direct FP8-E4M3 → unscaled BF16 for AMX (no FP32 round-trip).
211 + // BF16 value = true_E4M3 * 2^-120; exponent rebiasing folded into k/v scales.
212 + explicit BF16Vec32(const uint8_t* ptr, fp8_bf16_e4m3_tag) {
213 + __m256i b8 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(ptr));
214 + __m512i b16 = _mm512_cvtepu8_epi16(b8);
215 + __m512i sign =
216 + _mm512_slli_epi16(_mm512_and_si512(b16, _mm512_set1_epi16(0x80)), 8);
217 + __m512i payload =
218 + _mm512_slli_epi16(_mm512_and_si512(b16, _mm512_set1_epi16(0x7F)), 4);
219 + reg = _mm512_or_si512(sign, payload);
220 + }
221 +
222 + // Direct FP8-E5M2 → unscaled BF16 for AMX (no FP32 round-trip).
223 + // BF16 value = true_E5M2 * 2^-112; exponent rebiasing folded into k/v scales.
224 + explicit BF16Vec32(const uint8_t* ptr, fp8_bf16_e5m2_tag) {
225 + __m256i b8 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(ptr));
226 + __m512i b16 = _mm512_cvtepu8_epi16(b8);
227 + __m512i sign =
228 + _mm512_slli_epi16(_mm512_and_si512(b16, _mm512_set1_epi16(0x80)), 8);
229 + __m512i payload =
230 + _mm512_slli_epi16(_mm512_and_si512(b16, _mm512_set1_epi16(0x7F)), 5);
231 + reg = _mm512_or_si512(sign, payload);
232 + }
233 +
234 void save(void* ptr) const { *reinterpret_cast<__m512i*>(ptr) = reg; }
235 };
236 #else
@@ -200,6 +255,77 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
255 _mm256_castsi128_si256((__m128i)vec8_data.reg),
256 (__m128i)vec8_data.reg, 1)) {}
257
258 + // E4M3 decode (AVX2 path) — same bit-layout trick as the AVX512 variant
259 + // above. Result = true_E4M3 * 2^-8; caller applies scale * 2^8.
260 + explicit BF16Vec32(const uint8_t* ptr, fp8_e4m3_tag) {
261 + __m256i b8 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(ptr));
262 + __m128i b8_low = _mm256_extracti128_si256(b8, 0);
263 + __m128i b8_high = _mm256_extracti128_si256(b8, 1);
264 + __m256i b16_low = _mm256_cvtepu8_epi16(b8_low);
265 + __m256i b16_high = _mm256_cvtepu8_epi16(b8_high);
266 +
267 + __m256i sign_low = _mm256_slli_epi16(
268 + _mm256_and_si256(b16_low, _mm256_set1_epi16(0x80)), 8);
269 + __m256i payload_low = _mm256_slli_epi16(
270 + _mm256_and_si256(b16_low, _mm256_set1_epi16(0x7F)), 7);
271 + __m256i sign_high = _mm256_slli_epi16(
272 + _mm256_and_si256(b16_high, _mm256_set1_epi16(0x80)), 8);
273 + __m256i payload_high = _mm256_slli_epi16(
274 + _mm256_and_si256(b16_high, _mm256_set1_epi16(0x7F)), 7);
275 + reg_low = _mm256_or_si256(sign_low, payload_low);
276 + reg_high = _mm256_or_si256(sign_high, payload_high);
277 + }
278 +
279 + // E5M2 decode (AVX2 path) — b << 8 maps to FP16 bits; see AVX512 variant
280 + // above.
281 + explicit BF16Vec32(const uint8_t* ptr, fp8_e5m2_tag) {
282 + __m256i b8 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(ptr));
283 + __m128i b8_low = _mm256_extracti128_si256(b8, 0);
284 + __m128i b8_high = _mm256_extracti128_si256(b8, 1);
285 + reg_low = _mm256_slli_epi16(_mm256_cvtepu8_epi16(b8_low), 8);
286 + reg_high = _mm256_slli_epi16(_mm256_cvtepu8_epi16(b8_high), 8);
287 + }
288 +
289 + // Direct FP8-E4M3 → unscaled BF16 for AMX (AVX2 path, no FP32 round-trip).
290 + // BF16 value = true_E4M3 * 2^-120; exponent rebiasing folded into k/v scales.
291 + explicit BF16Vec32(const uint8_t* ptr, fp8_bf16_e4m3_tag) {
292 + __m256i b8 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(ptr));
293 + __m128i b8_low = _mm256_extracti128_si256(b8, 0);
294 + __m128i b8_high = _mm256_extracti128_si256(b8, 1);
295 + __m256i b16_low = _mm256_cvtepu8_epi16(b8_low);
296 + __m256i b16_high = _mm256_cvtepu8_epi16(b8_high);
297 + reg_low = _mm256_or_si256(
298 + _mm256_slli_epi16(_mm256_and_si256(b16_low, _mm256_set1_epi16(0x80)),
299 + 8),
300 + _mm256_slli_epi16(_mm256_and_si256(b16_low, _mm256_set1_epi16(0x7F)),
301 + 4));
302 + reg_high = _mm256_or_si256(
303 + _mm256_slli_epi16(_mm256_and_si256(b16_high, _mm256_set1_epi16(0x80)),
304 + 8),
305 + _mm256_slli_epi16(_mm256_and_si256(b16_high, _mm256_set1_epi16(0x7F)),
306 + 4));
307 + }
308 +
309 + // Direct FP8-E5M2 → unscaled BF16 for AMX (AVX2 path, no FP32 round-trip).
310 + // BF16 value = true_E5M2 * 2^-112; exponent rebiasing folded into k/v scales.
311 + explicit BF16Vec32(const uint8_t* ptr, fp8_bf16_e5m2_tag) {
312 + __m256i b8 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(ptr));
313 + __m128i b8_low = _mm256_extracti128_si256(b8, 0);
314 + __m128i b8_high = _mm256_extracti128_si256(b8, 1);
315 + __m256i b16_low = _mm256_cvtepu8_epi16(b8_low);
316 + __m256i b16_high = _mm256_cvtepu8_epi16(b8_high);
317 + reg_low = _mm256_or_si256(
318 + _mm256_slli_epi16(_mm256_and_si256(b16_low, _mm256_set1_epi16(0x80)),
319 + 8),
320 + _mm256_slli_epi16(_mm256_and_si256(b16_low, _mm256_set1_epi16(0x7F)),
321 + 5));
322 + reg_high = _mm256_or_si256(
323 + _mm256_slli_epi16(_mm256_and_si256(b16_high, _mm256_set1_epi16(0x80)),
324 + 8),
325 + _mm256_slli_epi16(_mm256_and_si256(b16_high, _mm256_set1_epi16(0x7F)),
326 + 5));
327 + }
328 +
329 void save(void* ptr) const {
330 _mm256_storeu_si256((__m256i*)ptr, reg_low);
331 _mm256_storeu_si256((__m256i*)ptr + 1, reg_high);
@@ -390,6 +516,11 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
516 : reg(_mm512_castsi512_ps(
517 _mm512_bslli_epi128(_mm512_cvtepu16_epi32(v.reg), 2))) {}
518
519 + explicit FP32Vec16(const BF16Vec32& v, int upper) {
520 + __m256i v_half_i = _mm512_extracti32x8_epi32(v.reg, upper);
521 + reg = _mm512_cvtph_ps(v_half_i);
522 + }
523 +
524 explicit FP32Vec16(const FP16Vec16& v) : reg(_mm512_cvtph_ps(v.reg)) {}
525
526 explicit FP32Vec16(const FP16Vec8& v) : FP32Vec16(FP32Vec8(v)) {}
@@ -494,6 +625,14 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
625 explicit FP32Vec16(const FP32Vec8& data)
626 : reg_low(data.reg), reg_high(data.reg) {}
627
628 + explicit FP32Vec16(const BF16Vec32& v, int upper) {
629 + const __m256i& half = upper ? v.reg_high : v.reg_low;
630 + __m128i lo = _mm256_extractf128_si256(half, 0);
631 + __m128i hi = _mm256_extractf128_si256(half, 1);
632 + reg_low = _mm256_cvtph_ps(lo);
633 + reg_high = _mm256_cvtph_ps(hi);
634 + }
635 +
636 explicit FP32Vec16(const FP16Vec16& v) {
637 __m128i low = _mm256_extractf128_si256(v.reg, 0);
638 __m128i high = _mm256_extractf128_si256(v.reg, 1);
csrc/cpu/generate_cpu_attn_dispatch.py
+142 -126
@@ -22,71 +22,95 @@ ISA_TYPES = {
22 "VXE": 4,
23 }
24
25 +# KV cache index: 0 = auto (same as scalar_t), 1 = fp8_e4m3, 2 = fp8_e5m2
26 +KV_CACHE_IDX = {
27 + "auto": 0,
28 + "fp8_e4m3": 1,
29 + "fp8_e5m2": 2,
30 +}
31 +
32 +# C++ type for each kv_cache index
33 +KV_CACHE_CPP_TYPES = {
34 + "auto": "scalar_t",
35 + "fp8_e4m3": "c10::Float8_e4m3fn",
36 + "fp8_e5m2": "c10::Float8_e5m2",
37 +}
38 +
39 # ISAs supported for head_dims divisible by 32
40 ISA_FOR_32 = ["AMX", "NEON", "VEC", "VEC16", "VXE"]
41
42 # ISAs supported for head_dims divisible by 16 only
43 ISA_FOR_16 = ["VEC16"]
44
45 +# ISAs that support FP8 KV cache (x86 AVX2/AVX-512 required)
46 +ISA_FOR_FP8 = ["AMX", "VEC"]
47 +
48
32 -def encode_params(head_dim: int, isa_type: str) -> int:
33 - """Encode head_dim and ISA type into a single int64_t."""
49 +def encode_params(head_dim: int, isa_type: str, kv_cache: str = "auto") -> int:
50 + """Encode head_dim, ISA type, and KV cache type into a single int64_t."""
51 isa_val = ISA_TYPES[isa_type]
35 - # Encoding: (head_dim << 8) | isa_type
36 - # This allows head_dim up to 2^56 - 1 and 256 ISA types
37 - return (head_dim << 8) | isa_val
52 + kv_val = KV_CACHE_IDX[kv_cache]
53 + # Encoding: (head_dim << 16) | (kv_cache_idx << 8) | isa_type
54 + # This allows head_dim up to 2^48 - 1, 256 KV cache types, and 256 ISA types
55 + return (head_dim << 16) | (kv_val << 8) | isa_val
56 +
57 +
58 +def _make_case(
59 + head_dim: int, isa: str, kv_cache: str = "auto", isa_override: str | None = None
60 +) -> str:
61 + """Generate a single switch case line."""
62 + encoded = encode_params(head_dim, isa, kv_cache)
63 + actual_isa = isa_override if isa_override else isa
64 + cpp_type = KV_CACHE_CPP_TYPES[kv_cache]
65 + attn_impl = (
66 + f"cpu_attention::AttentionImpl<"
67 + f"cpu_attention::ISA::{actual_isa}, \\\n"
68 + f" "
69 + f"scalar_t, head_dim, {cpp_type}>"
70 + )
71 + comment = (
72 + f"head_dim={head_dim}, isa={isa}"
73 + if kv_cache == "auto"
74 + else f"head_dim={head_dim}, isa={isa}, kv_cache={kv_cache}"
75 + )
76 + return (
77 + f""" case {encoded}LL: {{ """
78 + f"""/* {comment} */ \\"""
79 + f"""
80 + constexpr size_t head_dim = {head_dim}; \\"""
81 + f"""
82 + using attn_impl = {attn_impl}; \\"""
83 + f"""
84 + return __VA_ARGS__(); \\"""
85 + f"""
86 + }} \\"""
87 + )
88
89
40 -def generate_cases_for_isa_group(isa_list: list[str]) -> str:
90 +def generate_cases_for_isa_group(isa_list: list[str], include_fp8: bool = False) -> str:
91 """Generate switch cases for a specific ISA group."""
92 cases = []
93
44 - # Generate cases for head_dims divisible by 32
94 + # Non-FP8 cases for head_dims divisible by 32
95 for head_dim in HEAD_DIMS_32:
96 for isa in isa_list:
97 if isa not in ISA_FOR_32:
98 continue
49 - encoded = encode_params(head_dim, isa)
50 - case_str = (
51 - f""" case {encoded}LL: {{ """
52 - f"""/* head_dim={head_dim}, isa={isa} */ \\"""
53 - f"""
54 - constexpr size_t head_dim = {head_dim}; \\"""
55 - f"""
56 - using attn_impl = cpu_attention::AttentionImpl<"""
57 - f"""cpu_attention::ISA::{isa}, \\"""
58 - f"""
59 - """
60 - f"""scalar_t, head_dim>; \\"""
61 - f"""
62 - return __VA_ARGS__(); \\"""
63 - f"""
64 - }} \\"""
65 - )
66 - cases.append(case_str)
99 + cases.append(_make_case(head_dim, isa, "auto"))
100
68 - # Generate cases for head_dims divisible by 16 only
101 + # Non-FP8 cases for head_dims divisible by 16 only
102 for head_dim in HEAD_DIMS_16:
103 for isa in isa_list:
71 - encoded = encode_params(head_dim, isa)
72 - case_str = (
73 - f""" case {encoded}LL: {{ """
74 - f"""/* head_dim={head_dim}, isa={isa} """
75 - f"""(using VEC16) */ \\"""
76 - f"""
77 - constexpr size_t head_dim = {head_dim}; \\"""
78 - f"""
79 - using attn_impl = cpu_attention::AttentionImpl<"""
80 - f"""cpu_attention::ISA::VEC16, \\"""
81 - f"""
82 - """
83 - f"""scalar_t, head_dim>; \\"""
84 - f"""
85 - return __VA_ARGS__(); \\"""
86 - f"""
87 - }} \\"""
88 - )
89 - cases.append(case_str)
104 + cases.append(_make_case(head_dim, isa, "auto", isa_override="VEC16"))
105 +
106 + # FP8 cases: only AMX and VEC, only head_dims divisible by 32
107 + if include_fp8:
108 + for fp8_type in ("fp8_e4m3", "fp8_e5m2"):
109 + for head_dim in HEAD_DIMS_32:
110 + for isa in isa_list:
111 + if isa not in ISA_FOR_FP8:
112 + continue
113 + cases.append(_make_case(head_dim, isa, fp8_type))
114
115 return "\n".join(cases)
116
@@ -94,8 +118,9 @@ def generate_cases_for_isa_group(isa_list: list[str]) -> str:
118 def generate_helper_function() -> str:
119 """Generate helper function to encode parameters."""
120 return """
97 -inline int64_t encode_cpu_attn_params(int64_t head_dim, cpu_attention::ISA isa) {
98 - return (head_dim << 8) | static_cast<int64_t>(isa);
121 +inline int64_t encode_cpu_attn_params(int64_t head_dim, cpu_attention::ISA isa,
122 + int64_t kv_cache_idx = 0) {
123 + return (head_dim << 16) | (kv_cache_idx << 8) | static_cast<int64_t>(isa);
124 }
125 """
126
@@ -129,87 +154,78 @@ def generate_header_file() -> str:
154
155 # Generate dispatch macro with conditional compilation for different ISA sets
156 header += """
132 -// Dispatch macro using encoded parameters
133 -"""
134 -
135 - # x86_64 with AMX
136 - header += """#if defined(CPU_CAPABILITY_AMXBF16)
137 -#define CPU_ATTN_DISPATCH(HEAD_DIM, ISA_TYPE, ...) \\
138 - [&] { \\
139 - int64_t encoded_params = encode_cpu_attn_params(HEAD_DIM, ISA_TYPE); \\
140 - switch (encoded_params) { \\
141 -"""
142 - header += generate_cases_for_isa_group(["AMX", "VEC", "VEC16"])
143 - header += """
144 - default: { \\
145 - TORCH_CHECK(false, "Unsupported CPU attention configuration: head_dim=" + \\
146 - std::to_string(HEAD_DIM) + " isa=" + \\
147 - std::to_string(static_cast<int>(ISA_TYPE))); \\
148 - } \\
149 - } \\
150 - }()
151 -
157 +// Dispatch macro using encoded parameters.
158 +// KV_CACHE_IDX: Fp8KVCacheDataType enum value (kAuto=0, kFp8E4M3=1, kFp8E5M2=2).
159 +// FP8 cases (kv_cache_idx != 0) are generated on x86 platforms with AVX2 or
160 +// AVX-512: BF16Vec32 FP8 constructors have both AVX-512 and AVX2 implementations
161 +// in cpu_types_x86.hpp. Non-x86 platforms (#else fallback) have fp8=False.
162 """
163
154 - # ARM64 with NEON
155 - header += """#elif defined(__aarch64__)
156 -#define CPU_ATTN_DISPATCH(HEAD_DIM, ISA_TYPE, ...) \\
157 - [&] { \\
158 - int64_t encoded_params = encode_cpu_attn_params(HEAD_DIM, ISA_TYPE); \\
159 - switch (encoded_params) { \\
160 -"""
161 - header += generate_cases_for_isa_group(["NEON", "VEC", "VEC16"])
162 - header += """
163 - default: { \\
164 - TORCH_CHECK(false, "Unsupported CPU attention configuration: head_dim=" + \\
165 - std::to_string(HEAD_DIM) + " isa=" + \\
166 - std::to_string(static_cast<int>(ISA_TYPE))); \\
167 - } \\
168 - } \\
169 - }()
170 -
171 -"""
172 -
173 - # s390x with VXE
174 - header += """#elif defined(__s390x__)
175 -#define CPU_ATTN_DISPATCH(HEAD_DIM, ISA_TYPE, ...) \\
176 - [&] { \\
177 - int64_t encoded_params = encode_cpu_attn_params(HEAD_DIM, ISA_TYPE); \\
178 - switch (encoded_params) { \\
179 -"""
180 - header += generate_cases_for_isa_group(["VXE", "VEC", "VEC16"])
181 - header += """
182 - default: { \\
183 - TORCH_CHECK(false, "Unsupported CPU attention configuration: head_dim=" + \\
184 - std::to_string(HEAD_DIM) + " isa=" + \\
185 - std::to_string(static_cast<int>(ISA_TYPE))); \\
186 - } \\
187 - } \\
188 - }()
189 -
190 -"""
191 -
192 - # Fallback: VEC and VEC16 only
193 - header += """#else
194 -#define CPU_ATTN_DISPATCH(HEAD_DIM, ISA_TYPE, ...) \\
195 - [&] { \\
196 - int64_t encoded_params = encode_cpu_attn_params(HEAD_DIM, ISA_TYPE); \\
197 - switch (encoded_params) { \\
198 -"""
199 - header += generate_cases_for_isa_group(["VEC", "VEC16"])
200 - header += """
201 - default: { \\
202 - TORCH_CHECK(false, "Unsupported CPU attention configuration: head_dim=" + \\
203 - std::to_string(HEAD_DIM) + " isa=" + \\
204 - std::to_string(static_cast<int>(ISA_TYPE))); \\
205 - } \\
206 - } \\
207 - }()
208 -
209 -#endif /* CPU_CAPABILITY_AMXBF16 / __aarch64__ / __s390x__ */
210 -
211 -#endif // CPU_ATTN_DISPATCH_GENERATED_H
212 -"""
164 + def _macro_block(guard: str, isa_list: list[str], fp8: bool) -> str:
165 + """Return one CPU_ATTN_DISPATCH macro block for a given guard."""
166 + enc = (
167 + " int64_t encoded_params = encode_cpu_attn_params("
168 + "HEAD_DIM, ISA_TYPE, KV_CACHE_IDX); \\"
169 + )
170 + cases = generate_cases_for_isa_group(isa_list, include_fp8=fp8)
171 + tail = (
172 + "\n"
173 + " default: { \\\n"
174 + " TORCH_CHECK(false, "
175 + '"Unsupported CPU attention configuration: head_dim=" + \\\n'
176 + ' std::to_string(HEAD_DIM) + " isa=" + \\\n'
177 + " std::to_string(static_cast<int>(ISA_TYPE))"
178 + " + \\\n"
179 + ' " kv_cache_idx=" + '
180 + "std::to_string(KV_CACHE_IDX)); \\\n"
181 + " } \\\n"
182 + " } \\\n"
183 + " }()\n\n"
184 + )
185 + return (
186 + f"{guard}\n"
187 + "#define CPU_ATTN_DISPATCH(HEAD_DIM, ISA_TYPE, KV_CACHE_IDX, ...) \\\n"
188 + " [&] { \\\n"
189 + f"{enc}\n"
190 + " switch (encoded_params) { \\\n"
191 + f"{cases}"
192 + f"{tail}"
193 + )
194 +
195 + header += _macro_block(
196 + "#if defined(CPU_CAPABILITY_AMXBF16)",
197 + ["AMX", "VEC", "VEC16"],
198 + fp8=True,
199 + )
200 + header += _macro_block(
201 + "#elif defined(__aarch64__)",
202 + ["NEON", "VEC", "VEC16"],
203 + fp8=False,
204 + )
205 + header += _macro_block(
206 + "#elif defined(__s390x__)",
207 + ["VXE", "VEC", "VEC16"],
208 + fp8=False,
209 + )
210 + header += _macro_block(
211 + "#elif defined(__AVX512F__)",
212 + ["VEC", "VEC16"],
213 + fp8=True,
214 + )
215 + header += _macro_block(
216 + "#elif defined(__AVX2__)",
217 + ["VEC", "VEC16"],
218 + fp8=False,
219 + )
220 + header += _macro_block(
221 + "#else",
222 + ["VEC", "VEC16"],
223 + fp8=False,
224 + )
225 + header += (
226 + "#endif /* CPU_CAPABILITY_AMXBF16 / __aarch64__ / __s390x__ */\n\n"
227 + "#endif // CPU_ATTN_DISPATCH_GENERATED_H\n"
228 + )
229
230 return header
231
csrc/cpu/torch_bindings.cpp
+11 -5
@@ -101,7 +101,9 @@ void cpu_attn_reshape_and_cache(const torch::Tensor& key,
101 torch::Tensor& key_cache,
102 torch::Tensor& value_cache,
103 const torch::Tensor& slot_mapping,
104 - const std::string& isa);
104 + const std::string& isa, const double k_scale,
105 + const double v_scale,
106 + const std::string& kv_cache_dtype);
107
108 void cpu_attention_with_kv_cache(
109 const torch::Tensor& query, const torch::Tensor& key_cache,
@@ -112,7 +114,8 @@ void cpu_attention_with_kv_cache(
114 const int64_t sliding_window_left, const int64_t sliding_window_right,
115 const torch::Tensor& block_table, const double softcap,
116 const torch::Tensor& scheduler_metadata,
115 - const std::optional<torch::Tensor>& s_aux);
117 + const std::optional<torch::Tensor>& s_aux, const double k_scale,
118 + const double v_scale, const std::string& kv_cache_dtype);
119
120 // Note: just for avoiding importing errors
121 void placeholder_op() { TORCH_CHECK(false, "Unimplemented"); }
@@ -384,15 +387,18 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
387 &get_scheduler_metadata);
388 ops.def(
389 "cpu_attn_reshape_and_cache(Tensor key, Tensor value, Tensor(a2!) "
387 - "key_cache, Tensor(a3!) value_cache, Tensor slot_mapping, str "
388 - "isa) -> ()",
390 + "key_cache, Tensor(a3!) value_cache, Tensor slot_mapping, str isa, "
391 + "float k_scale=1.0, float v_scale=1.0, str kv_cache_dtype=\"auto\") -> "
392 + "()",
393 &cpu_attn_reshape_and_cache);
394 ops.def(
395 "cpu_attention_with_kv_cache(Tensor query, Tensor key_cache, Tensor "
396 "value_cache, Tensor(a3!) output, Tensor query_start_loc, Tensor "
397 "seq_lens, float scale, bool causal, Tensor? alibi_slopes, SymInt "
398 "sliding_window_left, SymInt sliding_window_right, Tensor block_table, "
395 - "float softcap, Tensor scheduler_metadata, Tensor? s_aux) -> ()",
399 + "float softcap, Tensor scheduler_metadata, Tensor? s_aux, "
400 + "float k_scale=1.0, float v_scale=1.0, str kv_cache_dtype=\"auto\") -> "
401 + "()",
402 &cpu_attention_with_kv_cache);
403
404 // placeholders
docs/design/attention_backends.md
+1 -1
@@ -167,7 +167,7 @@ Priority is **1 = highest** (tried first).
167
168 | Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | MM Prefix | DCP | Attention Types | Compute Cap. |
169 | ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | --------- | --- | --------------- | ------------ |
170 -| `CPU_ATTN` | | fp16, bf16, fp32 | `auto` | Any | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ❌ | ❌ | All | N/A |
170 +| `CPU_ATTN` | | fp16, bf16, fp32 | `auto`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | Any | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ❌ | ❌ | All | N/A |
171 | `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ❌ | ❌ | ✅ | Decoder | 7.x-9.x |
172 | `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ✅ | Decoder | 10.x |
173 | `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥8.0 |
tests/kernels/attention/test_cpu_attn.py
+96 -15
@@ -20,6 +20,12 @@ from vllm._custom_ops import (
20 cpu_attn_reshape_and_cache,
21 )
22
23 +# Enable AMX tile data registers so isolated runs (e.g. -k fp8_amx) don't rely
24 +# on ref_paged_attn's einsum to trigger oneDNN's _init_amx() first.
25 +if torch.cpu._is_amx_tile_supported():
26 + torch.cpu._init_amx()
27 +
28 +
29 NUM_HEADS = [
30 (4, 4),
31 (8, 2),
@@ -178,6 +184,10 @@ def ref_paged_attn(
184 return torch.cat(outputs, dim=0)
185
186
187 +_FP8_ATOL = {"fp8_e4m3": 0.2, "fp8_e5m2": 0.3}
188 +_FP8_RTOL = 0.1
189 +
190 +
191 @torch.inference_mode()
192 def varlen_with_paged_kv(
193 seq_lens: list[tuple[int, int]],
@@ -191,6 +201,9 @@ def varlen_with_paged_kv(
201 use_alibi: bool,
202 use_sink: bool,
203 isa: str,
204 + kv_cache_dtype: str = "auto",
205 + k_scale: float = 1.0,
206 + v_scale: float = 1.0,
207 ) -> None:
208 set_random_seed(0)
209 num_seqs = len(seq_lens)
@@ -212,6 +225,10 @@ def varlen_with_paged_kv(
225 15 * torch.rand((num_query_heads,), dtype=torch.bfloat16) if use_sink else None
226 )
227
228 + is_fp8 = kv_cache_dtype != "auto"
229 + if is_fp8 and current_platform.get_cpu_architecture() != CpuArchEnum.X86:
230 + pytest.skip("FP8 KV cache only supported on x86")
231 +
232 query = tensor_cache(
233 elem_num=token_num * num_query_heads * head_size,
234 dtype=dtype,
@@ -233,11 +250,17 @@ def varlen_with_paged_kv(
250 num_kv_heads,
251 head_size,
252 )
253 + if is_fp8:
254 + # Clamp KV to [-1, 1] so FP8 quantization error (<=12.5% for E4M3,
255 + # <=25% for E5M2) stays within the test tolerances regardless of
256 + # which tensor_cache values happen to be in use.
257 + key_value = key_value.clamp(-1, 1)
258 key_cache, value_cache = key_value.unbind(0)
259
260 # KV cache for CPU attention
261 + cache_dtype = torch.uint8 if is_fp8 else dtype
262 packed_key_cache = torch.empty(
240 - num_blocks, num_kv_heads, block_size, head_size, dtype=dtype
263 + num_blocks, num_kv_heads, block_size, head_size, dtype=cache_dtype
264 )
265 packed_value_cache = torch.empty_like(packed_key_cache)
266
@@ -252,6 +275,11 @@ def varlen_with_paged_kv(
275
276 # use reshape_and_cache to pack key_cache and value_cache
277 slot_mapping = torch.arange(0, num_blocks * block_size, dtype=torch.int64)
278 + fp8_kwargs: dict = (
279 + dict(k_scale=k_scale, v_scale=v_scale, kv_cache_dtype=kv_cache_dtype)
280 + if is_fp8
281 + else {}
282 + )
283 cpu_attn_reshape_and_cache(
284 key=key_cache.view(-1, num_kv_heads, head_size),
285 value=value_cache.view(-1, num_kv_heads, head_size),
@@ -259,6 +287,7 @@ def varlen_with_paged_kv(
287 value_cache=packed_value_cache,
288 slot_mapping=slot_mapping,
289 isa=isa,
290 + **fp8_kwargs,
291 )
292
293 metadata = cpu_attn_get_scheduler_metadata(
@@ -291,6 +320,7 @@ def varlen_with_paged_kv(
320 softcap=soft_cap if soft_cap is not None else 0,
321 scheduler_metadata=metadata,
322 s_aux=s_aux,
323 + **fp8_kwargs,
324 )
325
326 metadata = cpu_attn_get_scheduler_metadata(
@@ -323,23 +353,59 @@ def varlen_with_paged_kv(
353 softcap=soft_cap if soft_cap is not None else 0,
354 scheduler_metadata=metadata,
355 s_aux=s_aux,
356 + **fp8_kwargs,
357 )
358
328 - ref_output = ref_paged_attn(
329 - query=query,
330 - key_cache=key_cache,
331 - value_cache=value_cache,
332 - query_lens=query_lens,
333 - kv_lens=kv_lens,
334 - block_tables=block_tables,
335 - scale=scale,
336 - sliding_window=sliding_window,
337 - soft_cap=soft_cap,
338 - alibi_slopes=alibi_slopes,
339 - s_aux=s_aux,
340 - )
359 + if is_fp8:
360 + # Build a float KV cache via the non-FP8 path and run float attention
361 + # to use as the reference.
362 + ref_key_cache = torch.empty(
363 + num_blocks, num_kv_heads, block_size, head_size, dtype=dtype
364 + )
365 + ref_value_cache = torch.empty_like(ref_key_cache)
366 + cpu_attn_reshape_and_cache(
367 + key=key_cache.view(-1, num_kv_heads, head_size),
368 + value=value_cache.view(-1, num_kv_heads, head_size),
369 + key_cache=ref_key_cache,
370 + value_cache=ref_value_cache,
371 + slot_mapping=slot_mapping,
372 + isa=isa,
373 + )
374 + ref_output = torch.empty_like(query)
375 + cpu_attention_with_kv_cache(
376 + query=query,
377 + key_cache=ref_key_cache,
378 + value_cache=ref_value_cache,
379 + output=ref_output,
380 + query_start_loc=cu_query_lens,
381 + seq_lens=kv_lens_tensor,
382 + scale=scale,
383 + causal=True,
384 + alibi_slopes=alibi_slopes,
385 + sliding_window=window_size,
386 + block_table=block_tables,
387 + softcap=soft_cap if soft_cap is not None else 0,
388 + scheduler_metadata=metadata,
389 + s_aux=s_aux,
390 + )
391 + atol = _FP8_ATOL[kv_cache_dtype]
392 + rtol = _FP8_RTOL
393 + else:
394 + ref_output = ref_paged_attn(
395 + query=query,
396 + key_cache=key_cache,
397 + value_cache=value_cache,
398 + query_lens=query_lens,
399 + kv_lens=kv_lens,
400 + block_tables=block_tables,
401 + scale=scale,
402 + sliding_window=sliding_window,
403 + soft_cap=soft_cap,
404 + alibi_slopes=alibi_slopes,
405 + s_aux=s_aux,
406 + )
407 + atol, rtol = 1.5e-2, 1e-2
408
342 - atol, rtol = 1.5e-2, 1e-2
409 (
410 torch.testing.assert_close(out_with_split, ref_output, atol=atol, rtol=rtol),
411 f"{torch.max(torch.abs(out_with_split - ref_output))}",
@@ -350,6 +416,7 @@ def varlen_with_paged_kv(
416 )
417
418
419 +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8_e4m3", "fp8_e5m2"])
420 @pytest.mark.parametrize("seq_lens", SEQ_LENS)
421 @pytest.mark.parametrize("num_heads", NUM_HEADS)
422 @pytest.mark.parametrize("head_size", HEAD_SIZES)
@@ -373,6 +440,7 @@ def test_varlen_with_paged_kv_normal_vec(
440 use_alibi: bool,
441 use_sink: bool,
442 isa: str,
443 + kv_cache_dtype: str,
444 ) -> None:
445 varlen_with_paged_kv(
446 seq_lens=seq_lens,
@@ -386,9 +454,11 @@ def test_varlen_with_paged_kv_normal_vec(
454 use_alibi=use_alibi,
455 use_sink=use_sink,
456 isa=isa,
457 + kv_cache_dtype=kv_cache_dtype,
458 )
459
460
461 +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8_e4m3", "fp8_e5m2"])
462 @pytest.mark.parametrize("seq_lens", SEQ_LENS)
463 @pytest.mark.parametrize("num_heads", NUM_HEADS)
464 @pytest.mark.parametrize("head_size", HEAD_SIZES)
@@ -413,6 +483,7 @@ def test_varlen_with_paged_kv_normal_amx(
483 use_alibi: bool,
484 use_sink: bool,
485 isa: str,
486 + kv_cache_dtype: str,
487 ) -> None:
488 varlen_with_paged_kv(
489 seq_lens=seq_lens,
@@ -426,6 +497,7 @@ def test_varlen_with_paged_kv_normal_amx(
497 use_alibi=use_alibi,
498 use_sink=use_sink,
499 isa=isa,
500 + kv_cache_dtype=kv_cache_dtype,
501 )
502
503
@@ -511,6 +583,7 @@ def test_varlen_with_paged_kv_normal_neon(
583 )
584
585
586 +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8_e4m3"])
587 @pytest.mark.parametrize("seq_lens", SEQ_LENS)
588 @pytest.mark.parametrize("num_heads", NUM_HEADS)
589 @pytest.mark.parametrize("head_size", [96])
@@ -534,6 +607,7 @@ def test_varlen_with_paged_kv_softcap(
607 use_alibi: bool,
608 use_sink: bool,
609 isa: str,
610 + kv_cache_dtype: str,
611 ) -> None:
612 varlen_with_paged_kv(
613 seq_lens=seq_lens,
@@ -547,9 +621,11 @@ def test_varlen_with_paged_kv_softcap(
621 use_alibi=use_alibi,
622 use_sink=use_sink,
623 isa=isa,
624 + kv_cache_dtype=kv_cache_dtype,
625 )
626
627
628 +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8_e4m3"])
629 @pytest.mark.parametrize("seq_lens", SEQ_LENS)
630 @pytest.mark.parametrize("num_heads", NUM_HEADS)
631 @pytest.mark.parametrize("head_size", [96])
@@ -573,6 +649,7 @@ def test_varlen_with_paged_kv_alibi(
649 use_alibi: bool,
650 use_sink: bool,
651 isa: str,
652 + kv_cache_dtype: str,
653 ) -> None:
654 varlen_with_paged_kv(
655 seq_lens=seq_lens,
@@ -586,9 +663,11 @@ def test_varlen_with_paged_kv_alibi(
663 use_alibi=use_alibi,
664 use_sink=use_sink,
665 isa=isa,
666 + kv_cache_dtype=kv_cache_dtype,
667 )
668
669
670 +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8_e4m3"])
671 @pytest.mark.parametrize("seq_lens", SEQ_LENS)
672 @pytest.mark.parametrize("num_heads", NUM_HEADS)
673 @pytest.mark.parametrize("head_size", [96])
@@ -612,6 +691,7 @@ def test_varlen_with_paged_kv_sink(
691 use_alibi: bool,
692 use_sink: bool,
693 isa: str,
694 + kv_cache_dtype: str,
695 ) -> None:
696 varlen_with_paged_kv(
697 seq_lens=seq_lens,
@@ -625,4 +705,5 @@ def test_varlen_with_paged_kv_sink(
705 use_alibi=use_alibi,
706 use_sink=use_sink,
707 isa=isa,
708 + kv_cache_dtype=kv_cache_dtype,
709 )
vllm/_custom_ops.py
+12
@@ -3403,6 +3403,9 @@ def cpu_attn_reshape_and_cache(
3403 value_cache: torch.Tensor,
3404 slot_mapping: torch.Tensor,
3405 isa: str,
3406 + k_scale: float = 1.0,
3407 + v_scale: float = 1.0,
3408 + kv_cache_dtype: str = "auto",
3409 ) -> None:
3410 torch.ops._C.cpu_attn_reshape_and_cache(
3411 key,
@@ -3411,6 +3414,9 @@ def cpu_attn_reshape_and_cache(
3414 value_cache,
3415 slot_mapping,
3416 isa,
3417 + k_scale,
3418 + v_scale,
3419 + kv_cache_dtype,
3420 )
3421
3422
@@ -3429,6 +3435,9 @@ def cpu_attention_with_kv_cache(
3435 softcap: float,
3436 scheduler_metadata: torch.Tensor,
3437 s_aux: torch.Tensor | None,
3438 + k_scale: float = 1.0,
3439 + v_scale: float = 1.0,
3440 + kv_cache_dtype: str = "auto",
3441 ) -> None:
3442 torch.ops._C.cpu_attention_with_kv_cache(
3443 query,
@@ -3446,6 +3455,9 @@ def cpu_attention_with_kv_cache(
3455 softcap,
3456 scheduler_metadata,
3457 s_aux,
3458 + k_scale,
3459 + v_scale,
3460 + kv_cache_dtype,
3461 )
3462
3463
vllm/platforms/cpu.py
+4 -15
@@ -16,7 +16,6 @@ from vllm.utils.cpu_resource_utils import (
16 get_memory_node_info,
17 )
18 from vllm.utils.mem_constants import GiB_bytes
19 -from vllm.utils.torch_utils import is_quantized_kv_cache
19 from vllm.v1.attention.backends.registry import AttentionBackendEnum
20
21 from .interface import CpuArchEnum, Platform, PlatformEnum
@@ -134,20 +133,6 @@ class CpuPlatform(Platform):
133 scheduler_config = vllm_config.scheduler_config
134 # async scheduling is not required on CPU
135 scheduler_config.async_scheduling = False
137 - if (
138 - scheduler_config.enable_chunked_prefill
139 - or cache_config.enable_prefix_caching
140 - ) and is_quantized_kv_cache(cache_config.cache_dtype):
141 - raise RuntimeError(
142 - "Chunked-prefill and prefix-cache on the CPU "
143 - "backend is not compatible with FP8 KV cache."
144 - )
145 -
146 - if is_quantized_kv_cache(cache_config.cache_dtype):
147 - logger.warning(
148 - "CPU backend doesn't support KV cache quantization fallback to auto."
149 - )
150 - cache_config.cache_dtype = "auto"
136
137 parallel_config = vllm_config.parallel_config
138 # OMP requires the MP executor to function correctly, UniProc is not
@@ -458,6 +443,10 @@ class CpuPlatform(Platform):
443 block_offsets.reshape(1, block_size)
444 + indices.reshape(num_blocks, 1) * block_size
445 ).flatten()
446 + if key_cache.dtype == torch.uint8:
447 + raise NotImplementedError(
448 + "FP8 KV cache is not yet supported with KV transfer on CPU"
449 + )
450 cpu_attn_reshape_and_cache(
451 key,
452 value,
vllm/v1/attention/backends/cpu_attn.py
+38 -5
@@ -1,7 +1,10 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from dataclasses import dataclass
4 -from typing import ClassVar
4 +from typing import TYPE_CHECKING, ClassVar
5 +
6 +if TYPE_CHECKING:
7 + from vllm.config.cache import CacheDType
8
9 import torch
10
@@ -35,6 +38,12 @@ class CPUAttentionBackend(AttentionBackend):
38 torch.bfloat16,
39 torch.float32,
40 ]
41 + supported_kv_cache_dtypes: ClassVar[list["CacheDType"]] = [
42 + "auto",
43 + "fp8",
44 + "fp8_e4m3",
45 + "fp8_e5m2",
46 + ]
47
48 @classmethod
49 def get_supported_head_sizes(cls) -> list[int]:
@@ -133,7 +142,13 @@ class CPUAttentionMetadataBuilder(AttentionMetadataBuilder[CPUAttentionMetadata]
142 if self.window_size is None:
143 self.window_size = -1
144 self.block_size = vllm_config.cache_config.block_size
136 - self.isa = _get_attn_isa(self.dtype, self.block_size, self.head_dim)
145 + kv_cache_dtype_str = vllm_config.cache_config.cache_dtype
146 + self.isa = _get_attn_isa(
147 + self.dtype,
148 + self.block_size,
149 + self.head_dim,
150 + kv_cache_dtype_str,
151 + )
152 self.is_cross_attention = isinstance(kv_cache_spec, CrossAttentionSpec)
153
154 def build(
@@ -247,8 +262,7 @@ class CPUAttentionBackendImpl(AttentionImpl):
262 self.kv_cache_dtype = kv_cache_dtype
263 self.num_queries_per_kv = self.num_heads // self.num_kv_heads
264
250 - if is_quantized_kv_cache(kv_cache_dtype):
251 - raise NotImplementedError("FP8 KV cache is unsupported in CPU_ATTN")
265 + self.is_fp8_kv_cache = is_quantized_kv_cache(kv_cache_dtype)
266 self.attn_type = attn_type
267
268 self.sinks = sinks
@@ -325,6 +339,9 @@ class CPUAttentionBackendImpl(AttentionImpl):
339 value_cache,
340 attn_metadata.slot_mapping,
341 attn_metadata.isa,
342 + k_scale=layer._k_scale_float,
343 + v_scale=layer._v_scale_float,
344 + kv_cache_dtype=self.kv_cache_dtype,
345 )
346
347 if attn_metadata.use_sdpa_prefill:
@@ -356,6 +373,9 @@ class CPUAttentionBackendImpl(AttentionImpl):
373 softcap=self.logits_soft_cap,
374 scheduler_metadata=attn_metadata.scheduler_metadata,
375 s_aux=self.sinks,
376 + k_scale=layer._k_scale_float,
377 + v_scale=layer._v_scale_float,
378 + kv_cache_dtype=self.kv_cache_dtype,
379 )
380
381 return output
@@ -477,13 +497,26 @@ def _make_sliding_window_bias(
497
498
499 def _get_attn_isa(
480 - dtype: torch.dtype, block_size: int, head_size: int | None = None
500 + dtype: torch.dtype,
501 + block_size: int,
502 + head_size: int | None = None,
503 + kv_cache_dtype: str | None = None,
504 ) -> str:
505 + fp8_kv = is_quantized_kv_cache(kv_cache_dtype) if kv_cache_dtype else False
506 if head_size is not None and head_size % 32 != 0 and head_size % 16 == 0:
507 + if fp8_kv:
508 + raise NotImplementedError(
509 + "FP8 KV cache requires head_size divisible by 32 on CPU."
510 + )
511 return "vec16"
512 supports_amx = torch.cpu._is_amx_tile_supported()
513 supports_arm = current_platform.get_cpu_architecture() == CpuArchEnum.ARM
514 supports_vxe = current_platform.get_cpu_architecture() == CpuArchEnum.S390X
515 + supports_avx512 = torch.cpu._is_avx512_supported()
516 + if fp8_kv and not supports_amx and not supports_avx512:
517 + raise NotImplementedError(
518 + "FP8 KV cache on CPU requires x86 with AVX-512 or AMX."
519 + )
520 if supports_amx and dtype in (torch.bfloat16,) and block_size % 32 == 0:
521 return "amx"
522 elif block_size % 32 == 0: