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 13, 10:45 UTC
[CPU] Refactor CPU attention backend (#27954)
Signed-off-by: jiang1.li <jiang1.li@intel.com>
Li, Jiang committed
Nov 12, 2025 at 09:43 UTC
7f829be7d3d734020606fcca520f3c500581beb8
34 files changed
+4352
-1900
.buildkite/release-pipeline.yaml
+1
-1
@@ -132,7 +132,7 @@ steps:
132
queue: cpu_queue_postmerge
133
commands:
134
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
135
- - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --build-arg VLLM_CPU_AVX512BF16=true --build-arg VLLM_CPU_AVX512VNNI=true --tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:$(buildkite-agent meta-data get release-version) --tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:latest --progress plain --target vllm-openai -f docker/Dockerfile.cpu ."
135
+ - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --build-arg VLLM_CPU_AVX512BF16=true --build-arg VLLM_CPU_AVX512VNNI=true --build-arg VLLM_CPU_AMXBF16=true --tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:$(buildkite-agent meta-data get release-version) --tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:latest --progress plain --target vllm-openai -f docker/Dockerfile.cpu ."
136
- "docker push public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:latest"
137
- "docker push public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:$(buildkite-agent meta-data get release-version)"
138
env:
.buildkite/scripts/hardware_ci/run-cpu-test.sh
+2
-1
@@ -49,6 +49,7 @@ function cpu_tests() {
49
# Run kernel tests
50
docker exec cpu-test-"$NUMA_NODE" bash -c "
51
set -e
52
+ pytest -x -v -s tests/kernels/attention/test_cpu_attn.py
53
pytest -x -v -s tests/kernels/test_onednn.py"
54
55
# Run basic model test
@@ -116,4 +117,4 @@ function cpu_tests() {
117
118
# All of CPU tests are expected to be finished less than 40 mins.
119
export -f cpu_tests
119
-timeout 2h bash -c "cpu_tests $CORE_RANGE $NUMA_NODE"
120
+timeout 2.5h bash -c "cpu_tests $CORE_RANGE $NUMA_NODE"
cmake/cpu_extension.cmake
+24
-4
@@ -15,6 +15,7 @@ endif()
15
#
16
set(ENABLE_AVX512BF16 $ENV{VLLM_CPU_AVX512BF16})
17
set(ENABLE_AVX512VNNI $ENV{VLLM_CPU_AVX512VNNI})
18
+set(ENABLE_AMXBF16 $ENV{VLLM_CPU_AMXBF16})
19
20
include_directories("${CMAKE_SOURCE_DIR}/csrc")
21
@@ -140,6 +141,22 @@ if (AVX512_FOUND AND NOT AVX512_DISABLED)
141
set(ENABLE_AVX512VNNI OFF)
142
message(WARNING "Disable AVX512-VNNI ISA support, no avx512_vnni found in local CPU flags." " If cross-compilation is required, please set env VLLM_CPU_AVX512VNNI=1.")
143
endif()
144
+
145
+ find_isa(${CPUINFO} "amx_bf16" AMXBF16_FOUND)
146
+ if (AMXBF16_FOUND OR ENABLE_AMXBF16)
147
+ if (CMAKE_CXX_COMPILER_ID STREQUAL "GNU" AND
148
+ CMAKE_CXX_COMPILER_VERSION VERSION_GREATER_EQUAL 12.3)
149
+ list(APPEND CXX_COMPILE_FLAGS "-mamx-bf16" "-mamx-tile")
150
+ set(ENABLE_AMXBF16 ON)
151
+ add_compile_definitions(-DCPU_CAPABILITY_AMXBF16)
152
+ else()
153
+ set(ENABLE_AMXBF16 OFF)
154
+ message(WARNING "Disable AMX_BF16 ISA support, requires gcc/g++ >= 12.3")
155
+ endif()
156
+ else()
157
+ set(ENABLE_AMXBF16 OFF)
158
+ message(WARNING "Disable AMX_BF16 ISA support, no amx_bf16 found in local CPU flags." " If cross-compilation is required, please set env VLLM_CPU_AMXBF16=1.")
159
+ endif()
160
161
elseif (AVX2_FOUND)
162
list(APPEND CXX_COMPILE_FLAGS "-mavx2")
@@ -275,7 +292,10 @@ if ((AVX512_FOUND AND NOT AVX512_DISABLED) OR (ASIMD_FOUND AND NOT APPLE_SILICON
292
set(ONEDNN_VERBOSE "OFF")
293
set(CMAKE_POLICY_DEFAULT_CMP0077 NEW)
294
295
+ set(VLLM_BUILD_TYPE ${CMAKE_BUILD_TYPE})
296
+ set(CMAKE_BUILD_TYPE "Release") # remove oneDNN debug symbols to reduce size
297
FetchContent_MakeAvailable(oneDNN)
298
+ set(CMAKE_BUILD_TYPE ${VLLM_BUILD_TYPE})
299
add_library(dnnl_ext OBJECT "csrc/cpu/dnnl_helper.cpp")
300
target_include_directories(
301
dnnl_ext
@@ -305,14 +325,14 @@ endif()
325
#
326
set(VLLM_EXT_SRC
327
"csrc/cpu/activation.cpp"
308
- "csrc/cpu/attention.cpp"
309
- "csrc/cpu/cache.cpp"
328
"csrc/cpu/utils.cpp"
329
"csrc/cpu/layernorm.cpp"
330
"csrc/cpu/mla_decode.cpp"
331
"csrc/cpu/pos_encoding.cpp"
314
- "csrc/cpu/torch_bindings.cpp"
315
- "csrc/moe/dynamic_4bit_int_moe_cpu.cpp")
332
+ "csrc/moe/dynamic_4bit_int_moe_cpu.cpp"
333
+ "csrc/cpu/cpu_attn.cpp"
334
+ "csrc/cpu/scratchpad_manager.cpp"
335
+ "csrc/cpu/torch_bindings.cpp")
336
337
if (AVX512_FOUND AND NOT AVX512_DISABLED)
338
set(VLLM_EXT_SRC
csrc/cpu/attention.cpp
deleted
-798
@@ -1,798 +0,0 @@
1
-#include "cpu_types.hpp"
2
-
3
-namespace {
4
-
5
-template <typename scalar_t>
6
-struct KernelVecType {
7
- using q_load_vec_type = void;
8
- using q_vec_type = void;
9
- using k_load_vec_type = void;
10
- using k_vec_type = void;
11
- using qk_acc_vec_type = void;
12
- using v_load_vec_type = void;
13
-};
14
-
15
-template <>
16
-struct KernelVecType<float> {
17
- using q_load_vec_type = vec_op::FP32Vec4;
18
- using q_vec_type = vec_op::FP32Vec16;
19
- using k_load_vec_type = vec_op::FP32Vec16;
20
- using k_vec_type = vec_op::FP32Vec16;
21
- using qk_acc_vec_type = vec_op::FP32Vec16;
22
- using v_load_vec_type = vec_op::FP32Vec16;
23
-};
24
-
25
-template <>
26
-struct KernelVecType<c10::Half> {
27
-#if defined(__powerpc64__) || defined(__s390x__)
28
- // Power and s390x architecture-specific vector types
29
- using q_load_vec_type = vec_op::FP32Vec8;
30
- using k_load_vec_type = vec_op::FP32Vec16;
31
- using v_load_vec_type = vec_op::FP32Vec16;
32
-#else
33
- // Fallback for other architectures, including x86
34
- using q_load_vec_type = vec_op::FP16Vec8;
35
- using k_load_vec_type = vec_op::FP16Vec16;
36
- using v_load_vec_type = vec_op::FP16Vec16;
37
-#endif
38
- using q_vec_type = vec_op::FP32Vec16;
39
- using k_vec_type = vec_op::FP32Vec16;
40
- using qk_acc_vec_type = vec_op::FP32Vec16;
41
-};
42
-
43
-#ifdef __AVX512BF16__
44
-template <>
45
-struct KernelVecType<c10::BFloat16> {
46
- using q_load_vec_type = vec_op::BF16Vec8;
47
- using q_vec_type = vec_op::BF16Vec32;
48
- using k_load_vec_type = vec_op::BF16Vec32;
49
- using k_vec_type = vec_op::BF16Vec32;
50
- using qk_acc_vec_type = vec_op::FP32Vec16;
51
- using v_load_vec_type = vec_op::BF16Vec16;
52
-};
53
-#else
54
- #ifdef __aarch64__
55
- #ifndef ARM_BF16_SUPPORT
56
- // pass
57
- #else
58
-template <>
59
-struct KernelVecType<c10::BFloat16> {
60
- using q_load_vec_type = vec_op::BF16Vec8;
61
- using q_vec_type = vec_op::FP32Vec16;
62
- using k_load_vec_type = vec_op::BF16Vec16;
63
- using k_vec_type = vec_op::FP32Vec16;
64
- using qk_acc_vec_type = vec_op::FP32Vec16;
65
- using v_load_vec_type = vec_op::BF16Vec16;
66
-};
67
- #endif
68
- #else
69
-template <>
70
-struct KernelVecType<c10::BFloat16> {
71
- using q_load_vec_type = vec_op::BF16Vec8;
72
- using q_vec_type = vec_op::FP32Vec16;
73
- using k_load_vec_type = vec_op::BF16Vec16;
74
- using k_vec_type = vec_op::FP32Vec16;
75
- using qk_acc_vec_type = vec_op::FP32Vec16;
76
- using v_load_vec_type = vec_op::BF16Vec16;
77
-};
78
- #endif
79
-#endif
80
-
81
-template <typename T>
82
-FORCE_INLINE std::pair<T, T> reduceSoftmax(T* data, const int size,
83
- const int capacity) {
84
- T max = data[0];
85
- for (int i = 1; i < size; ++i) {
86
- max = max >= data[i] ? max : data[i];
87
- }
88
-
89
- T sum = 0;
90
- for (int i = 0; i < size; ++i) {
91
- data[i] = std::exp(data[i] - max);
92
- sum += data[i];
93
- }
94
-
95
- int i = 0;
96
- for (; i < size; ++i) {
97
- data[i] /= sum;
98
- }
99
-
100
- for (; i < capacity; ++i) {
101
- data[i] = 0;
102
- }
103
-
104
- return {max, sum};
105
-}
106
-
107
-template <typename T>
108
-FORCE_INLINE std::pair<T, T> reduceSoftmaxAlibi(T* data, const int size,
109
- const int capacity,
110
- const float alibi_slope,
111
- const int start_index,
112
- const int seq_len) {
113
- data[0] += alibi_slope * (start_index - seq_len + 1);
114
- T max = data[0];
115
- for (int i = 1; i < size; ++i) {
116
- T qk = data[i] + alibi_slope * (start_index + i - seq_len + 1);
117
- data[i] = qk;
118
- max = max >= qk ? max : qk;
119
- }
120
-
121
- T sum = 0;
122
- for (int i = 0; i < size; ++i) {
123
- data[i] = std::exp(data[i] - max);
124
- sum += data[i];
125
- }
126
-
127
- int i = 0;
128
- for (; i < size; ++i) {
129
- data[i] /= sum;
130
- }
131
-
132
- for (; i < capacity; ++i) {
133
- data[i] = 0;
134
- }
135
-
136
- return {max, sum};
137
-}
138
-
139
-template <typename T>
140
-FORCE_INLINE void reducePartitionSoftmax(const T* max_data, T* sum_data,
141
- const int size) {
142
- T max = max_data[0];
143
- for (int i = 1; i < size; ++i) {
144
- max = max >= max_data[i] ? max : max_data[i];
145
- }
146
-
147
- T rescaled_sum = 0;
148
- for (int i = 0; i < size; ++i) {
149
- T rescale_factor = std::exp(max_data[i] - max);
150
- rescaled_sum += rescale_factor * sum_data[i];
151
- sum_data[i] *= rescale_factor;
152
- }
153
- for (int i = 0; i < size; ++i) {
154
- sum_data[i] /= rescaled_sum + 1e-8;
155
- }
156
-}
157
-
158
-template <typename scalar_t, int HEAD_SIZE, int BLOCK_SIZE, int x>
159
-struct reduceQKBlockKernel {
160
- using q_load_vec_type = typename KernelVecType<scalar_t>::q_load_vec_type;
161
- using q_vec_type = typename KernelVecType<scalar_t>::q_vec_type;
162
- using k_load_vec_type = typename KernelVecType<scalar_t>::k_load_vec_type;
163
- using k_vec_type = typename KernelVecType<scalar_t>::k_vec_type;
164
- using qk_acc_vec_type = typename KernelVecType<scalar_t>::qk_acc_vec_type;
165
-
166
- constexpr static int TOKEN_PER_GROUP = k_load_vec_type::get_elem_num() / x;
167
- constexpr static int MAX_GROUP_NUM = 16 / TOKEN_PER_GROUP;
168
- constexpr static int UNROLL_GROUP_NUM = MAX_GROUP_NUM / 4;
169
-
170
- static_assert(MAX_GROUP_NUM == 8 || MAX_GROUP_NUM == 4);
171
- static_assert(k_load_vec_type::get_elem_num() % x == 0);
172
- static_assert(q_load_vec_type::get_elem_num() * sizeof(scalar_t) == 16);
173
-
174
- FORCE_INLINE static void call(const scalar_t* __restrict__ q,
175
- const scalar_t* __restrict__ k_block,
176
- float* __restrict__ logits, float scale,
177
- const int token_num) {
178
- const int group_num = (token_num + TOKEN_PER_GROUP - 1) / TOKEN_PER_GROUP;
179
-
180
- qk_acc_vec_type group_accums[MAX_GROUP_NUM];
181
- if (token_num == BLOCK_SIZE) {
182
- for (int q_offset = 0; q_offset < HEAD_SIZE;
183
- q_offset += x, k_block += x * BLOCK_SIZE) {
184
- q_load_vec_type q_load_group_vec(q + q_offset);
185
- q_vec_type q_group_vec(q_load_group_vec);
186
-
187
- vec_op::unroll_loop<int, MAX_GROUP_NUM>(
188
- [k_block, &q_group_vec, &group_accums](int token_group_idx) {
189
- k_load_vec_type k_load_group_vec(k_block + token_group_idx * x *
190
- TOKEN_PER_GROUP);
191
- k_vec_type k_group_vec(k_load_group_vec);
192
- vec_op::fma(group_accums[token_group_idx], q_group_vec,
193
- k_group_vec);
194
- vec_op::prefetch(k_block + x * BLOCK_SIZE +
195
- token_group_idx * x * TOKEN_PER_GROUP);
196
- });
197
- }
198
- } else {
199
- for (int q_offset = 0; q_offset < HEAD_SIZE;
200
- q_offset += x, k_block += x * BLOCK_SIZE) {
201
- q_load_vec_type q_load_group_vec(q + q_offset);
202
- q_vec_type q_group_vec(q_load_group_vec);
203
- for (int token_group_start = 0; token_group_start < group_num;
204
- token_group_start += UNROLL_GROUP_NUM) {
205
- vec_op::unroll_loop<int, UNROLL_GROUP_NUM>(
206
- [token_group_start, k_block, &q_group_vec,
207
- &group_accums](int token_group_idx) {
208
- token_group_idx += token_group_start;
209
- k_load_vec_type k_load_group_vec(k_block + token_group_idx * x *
210
- TOKEN_PER_GROUP);
211
- k_vec_type k_group_vec(k_load_group_vec);
212
- vec_op::fma(group_accums[token_group_idx], q_group_vec,
213
- k_group_vec);
214
- vec_op::prefetch(k_block + x * BLOCK_SIZE +
215
- token_group_idx * x * TOKEN_PER_GROUP);
216
- });
217
- }
218
- }
219
- }
220
-
221
- for (int token_group_idx = 0; token_group_idx < group_num;
222
- ++token_group_idx) {
223
- vec_op::unroll_loop<int, TOKEN_PER_GROUP>(
224
- [&group_accums, logits, scale, token_group_idx](int token_idx) {
225
- float dot_v =
226
- group_accums[token_group_idx]
227
- .template reduce_sub_sum<qk_acc_vec_type::get_elem_num() /
228
- TOKEN_PER_GROUP>(token_idx);
229
- logits[token_group_idx * TOKEN_PER_GROUP + token_idx] =
230
- dot_v * scale;
231
- });
232
- }
233
- }
234
-};
235
-
236
-template <typename scalar_t, int HEAD_SIZE, int BLOCK_SIZE,
237
- int HEAD_PARTITION_SIZE, typename acc_t>
238
-FORCE_INLINE void reduceValueBlock(const float* prob, const scalar_t* v_block,
239
- acc_t&& acc) {
240
- using v_load_vec_type = typename KernelVecType<scalar_t>::v_load_vec_type;
241
- constexpr int ELEM_NUM = v_load_vec_type::get_elem_num();
242
- static_assert(BLOCK_SIZE == ELEM_NUM);
243
- vec_op::FP32Vec16 prob_vec(prob);
244
-
245
- vec_op::unroll_loop<int, HEAD_PARTITION_SIZE>([&](int head_elem_idx) {
246
- v_load_vec_type v_vec(v_block + BLOCK_SIZE * head_elem_idx);
247
- vec_op::FP32Vec16 fp32_v_vec(v_vec);
248
- acc[head_elem_idx] = acc[head_elem_idx] + prob_vec * fp32_v_vec;
249
- });
250
-}
251
-}; // namespace
252
-
253
-// Paged attention v1
254
-namespace {
255
-template <typename scalar_t, int HEAD_SIZE, int BLOCK_SIZE>
256
-struct paged_attention_v1_impl {
257
- static void call(
258
- scalar_t* __restrict__ out, // [num_seqs, num_heads, head_size]
259
- const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size]
260
- const scalar_t* __restrict__ k_cache, // [num_blocks, num_kv_heads,
261
- // head_size/x, block_size, x]
262
- const scalar_t* __restrict__ v_cache, // [num_blocks, num_kv_heads,
263
- // head_size, block_size]
264
- const int num_kv_heads, const float scale,
265
- const int* __restrict__ block_tables, // [num_seqs,
266
- // max_num_blocks_per_seq]
267
- const int* __restrict__ seq_lens, // [num_seqs]
268
- const int max_num_blocks_per_seq,
269
- const float* __restrict__ alibi_slopes, // [num_heads]
270
- const int q_stride, const int kv_block_stride, const int kv_head_stride,
271
- const int num_seqs, const int num_heads) {
272
- constexpr int x = 16 / sizeof(scalar_t);
273
- const int num_queries_per_kv = num_heads / num_kv_heads;
274
-
275
- static_assert(BLOCK_SIZE == 16);
276
-
277
- int max_seq_len = max_num_blocks_per_seq * BLOCK_SIZE;
278
- int max_seq_len_padded = (max_seq_len + 15) & 0xFFFFFFF0;
279
- TORCH_CHECK((max_seq_len_padded * sizeof(float)) % 64 == 0);
280
-
281
- const int parallel_work_item_num = omp_get_max_threads();
282
-
283
- size_t logits_bytes =
284
- parallel_work_item_num * max_seq_len_padded * sizeof(float);
285
- float* logits = (float*)std::aligned_alloc(
286
- 64, logits_bytes); // Cacheline alignment for each context token.
287
- // [parallel_work_item_num, max_seq_len_padded]
288
-
289
-#pragma omp parallel for collapse(2) schedule(dynamic, 1)
290
- for (int seq_idx = 0; seq_idx < num_seqs; ++seq_idx) {
291
- for (int head_idx = 0; head_idx < num_heads; ++head_idx) {
292
- int seq_len = seq_lens[seq_idx];
293
- const int* seq_block_table =
294
- block_tables + max_num_blocks_per_seq * seq_idx;
295
- const int block_num = (seq_len + BLOCK_SIZE - 1) / BLOCK_SIZE;
296
- const int64_t kv_head_idx = head_idx / num_queries_per_kv;
297
- const scalar_t* __restrict__ q_vec_ptr =
298
- q + seq_idx * q_stride + head_idx * HEAD_SIZE;
299
- const int last_block_token_num = seq_len - (block_num - 1) * BLOCK_SIZE;
300
- float* __restrict__ thread_block_logits =
301
- logits + omp_get_thread_num() * max_seq_len_padded;
302
-
303
- // Compute logits
304
- for (int block_idx = 0; block_idx < block_num; ++block_idx) {
305
- const int64_t physical_block_idx = seq_block_table[block_idx];
306
- const scalar_t* __restrict__ k_block_cache_ptr =
307
- k_cache + physical_block_idx * kv_block_stride +
308
- kv_head_idx * kv_head_stride;
309
- float* __restrict__ head_block_logits =
310
- thread_block_logits + block_idx * BLOCK_SIZE;
311
-
312
- reduceQKBlockKernel<scalar_t, HEAD_SIZE, BLOCK_SIZE, x>::call(
313
- q_vec_ptr, k_block_cache_ptr, head_block_logits, scale,
314
- block_idx == block_num - 1 ? last_block_token_num : BLOCK_SIZE);
315
- }
316
-
317
- // Compute softmax
318
- if (alibi_slopes) {
319
- reduceSoftmaxAlibi(thread_block_logits, seq_len,
320
- block_num * BLOCK_SIZE, alibi_slopes[head_idx], 0,
321
- seq_len);
322
- } else {
323
- reduceSoftmax(thread_block_logits, seq_len, block_num * BLOCK_SIZE);
324
- }
325
-
326
- // Compute value
327
- constexpr int head_elem_num_per_partition = 16;
328
- constexpr int head_partition_num =
329
- HEAD_SIZE / head_elem_num_per_partition;
330
- for (int head_part_idx = 0; head_part_idx < head_partition_num;
331
- ++head_part_idx) {
332
- vec_op::FP32Vec16 accums[head_elem_num_per_partition];
333
- scalar_t* __restrict__ out_ptr =
334
- out + seq_idx * num_heads * HEAD_SIZE + head_idx * HEAD_SIZE +
335
- head_part_idx * head_elem_num_per_partition;
336
- for (int block_idx = 0; block_idx < block_num; ++block_idx) {
337
- const int64_t physical_block_idx = seq_block_table[block_idx];
338
- const float* __restrict__ prob_vec_ptr =
339
- thread_block_logits + block_idx * BLOCK_SIZE;
340
- const scalar_t* __restrict__ v_block_cache_ptr =
341
- v_cache + physical_block_idx * kv_block_stride +
342
- kv_head_idx * kv_head_stride +
343
- BLOCK_SIZE * head_part_idx * head_elem_num_per_partition;
344
- reduceValueBlock<scalar_t, HEAD_SIZE, BLOCK_SIZE,
345
- head_elem_num_per_partition>(
346
- prob_vec_ptr, v_block_cache_ptr, accums);
347
-
348
- if (block_idx != block_num - 1) {
349
- const int64_t next_physical_block_idx =
350
- seq_block_table[block_idx + 1];
351
- const scalar_t* __restrict__ next_v_block_cache_ptr =
352
- v_cache + next_physical_block_idx * kv_block_stride +
353
- kv_head_idx * kv_head_stride +
354
- BLOCK_SIZE * head_part_idx * head_elem_num_per_partition;
355
- vec_op::unroll_loop<int, head_elem_num_per_partition>(
356
- [&](int head_elem_idx) {
357
- if (head_elem_idx % 2 == 0) {
358
- vec_op::prefetch(next_v_block_cache_ptr +
359
- BLOCK_SIZE * head_elem_idx);
360
- }
361
- });
362
- }
363
- }
364
-
365
- vec_op::unroll_loop<int, head_elem_num_per_partition>(
366
- [&](int head_elem_idx) {
367
- float value = accums[head_elem_idx].reduce_sum();
368
- vec_op::storeFP32(value, out_ptr + head_elem_idx);
369
- });
370
- }
371
- }
372
- }
373
- std::free(logits);
374
- }
375
-};
376
-
377
-#define LAUNCH_V1_ATTENTION_KERNEL(T, HEAD_SIZE, BLOCK_SIZE) \
378
- paged_attention_v1_impl<T, HEAD_SIZE, BLOCK_SIZE>::call( \
379
- out_ptr, query_ptr, key_cache_ptr, value_cache_ptr, num_kv_heads, scale, \
380
- block_tables_ptr, seq_lens_ptr, max_num_blocks_per_seq, \
381
- alibi_slopes_ptr, q_stride, kv_block_stride, kv_head_stride, num_seqs, \
382
- num_heads);
383
-
384
-template <typename T, int BLOCK_SIZE>
385
-void paged_attention_v1_impl_launcher(
386
- torch::Tensor& out, torch::Tensor& query, torch::Tensor& key_cache,
387
- torch::Tensor& value_cache, int num_kv_heads, float scale,
388
- torch::Tensor& block_tables, torch::Tensor& seq_lens, int max_seq_len,
389
- const std::optional<torch::Tensor>& alibi_slopes) {
390
- int num_seqs = query.size(0);
391
- int num_heads = query.size(1);
392
- int head_size = query.size(2);
393
- int max_num_blocks_per_seq = block_tables.size(1);
394
- int q_stride = query.stride(0);
395
- int kv_block_stride = key_cache.stride(0);
396
- int kv_head_stride = key_cache.stride(1);
397
-
398
- // NOTE: alibi_slopes is optional.
399
- const float* alibi_slopes_ptr =
400
- alibi_slopes
401
- ? reinterpret_cast<const float*>(alibi_slopes.value().data_ptr())
402
- : nullptr;
403
-
404
- T* out_ptr = reinterpret_cast<T*>(out.data_ptr());
405
- T* query_ptr = reinterpret_cast<T*>(query.data_ptr());
406
- T* key_cache_ptr = reinterpret_cast<T*>(key_cache.data_ptr());
407
- T* value_cache_ptr = reinterpret_cast<T*>(value_cache.data_ptr());
408
- int* block_tables_ptr = block_tables.data_ptr<int>();
409
- int* seq_lens_ptr = seq_lens.data_ptr<int>();
410
-
411
- switch (head_size) {
412
- case 32:
413
- LAUNCH_V1_ATTENTION_KERNEL(T, 32, BLOCK_SIZE);
414
- break;
415
- case 64:
416
- LAUNCH_V1_ATTENTION_KERNEL(T, 64, BLOCK_SIZE);
417
- break;
418
- case 80:
419
- LAUNCH_V1_ATTENTION_KERNEL(T, 80, BLOCK_SIZE);
420
- break;
421
- case 96:
422
- LAUNCH_V1_ATTENTION_KERNEL(T, 96, BLOCK_SIZE);
423
- break;
424
- case 112:
425
- LAUNCH_V1_ATTENTION_KERNEL(T, 112, BLOCK_SIZE);
426
- break;
427
- case 128:
428
- LAUNCH_V1_ATTENTION_KERNEL(T, 128, BLOCK_SIZE);
429
- break;
430
- case 192:
431
- LAUNCH_V1_ATTENTION_KERNEL(T, 192, BLOCK_SIZE);
432
- break;
433
- case 256:
434
- LAUNCH_V1_ATTENTION_KERNEL(T, 256, BLOCK_SIZE);
435
- break;
436
- default:
437
- TORCH_CHECK(false, "Unsupported head size: ", head_size);
438
- break;
439
- }
440
-}
441
-
442
-#define CALL_V1_KERNEL_LAUNCHER(T, BLOCK_SIZE) \
443
- paged_attention_v1_impl_launcher<T, BLOCK_SIZE>( \
444
- out, query, key_cache, value_cache, num_kv_heads, scale, block_tables, \
445
- seq_lens, max_seq_len, alibi_slopes);
446
-
447
-#define CALL_V1_KERNEL_LAUNCHER_BLOCK_SIZE(T) \
448
- switch (block_size) { \
449
- case 16: \
450
- CALL_V1_KERNEL_LAUNCHER(T, 16); \
451
- break; \
452
- default: \
453
- TORCH_CHECK(false, "Unsupported block size: ", block_size); \
454
- break; \
455
- }
456
-} // namespace
457
-
458
-void paged_attention_v1(
459
- torch::Tensor& out, torch::Tensor& query, torch::Tensor& key_cache,
460
- torch::Tensor& value_cache, int64_t num_kv_heads, double scale,
461
- torch::Tensor& block_tables, torch::Tensor& seq_lens, int64_t block_size,
462
- int64_t max_seq_len, const std::optional<torch::Tensor>& alibi_slopes,
463
- const std::string& kv_cache_dtype, torch::Tensor& k_scale,
464
- torch::Tensor& v_scale, const int64_t tp_rank,
465
- const int64_t blocksparse_local_blocks,
466
- const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
467
- const int64_t blocksparse_head_sliding_step) {
468
- TORCH_CHECK(blocksparse_vert_stride <= 1,
469
- "CPU backend does not support blocksparse attention yet.");
470
- VLLM_DISPATCH_FLOATING_TYPES(query.scalar_type(), "paged_attention_v1_impl",
471
- [&] {
472
- CPU_KERNEL_GUARD_IN(paged_attention_v1_impl)
473
- CALL_V1_KERNEL_LAUNCHER_BLOCK_SIZE(scalar_t);
474
- CPU_KERNEL_GUARD_OUT(paged_attention_v1_impl)
475
- });
476
-}
477
-
478
-// Paged attention v2
479
-namespace {
480
-template <typename scalar_t, int HEAD_SIZE, int BLOCK_SIZE, int PARTITION_SIZE>
481
-struct paged_attention_v2_impl {
482
- static void call(
483
- scalar_t* __restrict__ out, // [num_seqs, num_heads, head_size]
484
- float* __restrict__ exp_sums, // [num_seqs, num_heads,
485
- // max_num_partitions]
486
- float* __restrict__ max_logits, // [num_seqs, num_heads,
487
- // max_num_partitions]
488
- scalar_t* __restrict__ tmp_out, // [num_seqs, num_heads,
489
- // max_num_partitions, head_size]
490
- const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size]
491
- const scalar_t* __restrict__ k_cache, // [num_blocks, num_kv_heads,
492
- // head_size/x, block_size, x]
493
- const scalar_t* __restrict__ v_cache, // [num_blocks, num_kv_heads,
494
- // head_size, block_size]
495
- const int num_kv_heads, const float scale,
496
- const int* __restrict__ block_tables, // [num_seqs,
497
- // max_num_blocks_per_seq]
498
- const int* __restrict__ seq_lens, // [num_seqs]
499
- const int max_num_blocks_per_seq,
500
- const float* __restrict__ alibi_slopes, // [num_heads]
501
- const int q_stride, const int kv_block_stride, const int kv_head_stride,
502
- const int num_seqs, const int num_heads, const int max_num_partitions) {
503
- constexpr int x = 16 / sizeof(scalar_t);
504
- const int num_queries_per_kv = num_heads / num_kv_heads;
505
-
506
- static_assert(BLOCK_SIZE == 16);
507
- static_assert(PARTITION_SIZE * sizeof(float) % 64 == 0);
508
- static_assert(PARTITION_SIZE % BLOCK_SIZE == 0);
509
-
510
-#pragma omp parallel for collapse(3) schedule(static, 1)
511
- for (int seq_idx = 0; seq_idx < num_seqs; ++seq_idx) {
512
- for (int partition_idx = 0; partition_idx < max_num_partitions;
513
- ++partition_idx) {
514
- for (int head_idx = 0; head_idx < num_heads; ++head_idx) {
515
- const int seq_len = seq_lens[seq_idx];
516
- const int start_token_idx = partition_idx * PARTITION_SIZE;
517
-
518
- if (start_token_idx >= seq_len) continue;
519
-
520
- const int partition_num =
521
- (seq_len + PARTITION_SIZE - 1) / PARTITION_SIZE;
522
- const bool no_reduce = (partition_num == 1);
523
- const int token_num =
524
- (std::min(seq_len, start_token_idx + PARTITION_SIZE) -
525
- start_token_idx);
526
- const int block_num = (token_num + BLOCK_SIZE - 1) / BLOCK_SIZE;
527
- const int last_block_token_num =
528
- token_num - (block_num - 1) * BLOCK_SIZE;
529
- const int* seq_block_table = block_tables +
530
- max_num_blocks_per_seq * seq_idx +
531
- start_token_idx / BLOCK_SIZE;
532
- const int64_t kv_head_idx = head_idx / num_queries_per_kv;
533
- const scalar_t* __restrict__ q_vec_ptr =
534
- q + seq_idx * q_stride + head_idx * HEAD_SIZE;
535
-
536
- float logits[PARTITION_SIZE] __attribute__((aligned(64))) = {0};
537
-
538
- // Compute logits
539
- for (int block_idx = 0; block_idx < block_num; ++block_idx) {
540
- const int64_t physical_block_idx = seq_block_table[block_idx];
541
- const scalar_t* __restrict__ k_block_cache_ptr =
542
- k_cache + physical_block_idx * kv_block_stride +
543
- kv_head_idx * kv_head_stride;
544
- float* __restrict__ head_block_logits =
545
- logits + block_idx * BLOCK_SIZE;
546
-
547
- reduceQKBlockKernel<scalar_t, HEAD_SIZE, BLOCK_SIZE, x>::call(
548
- q_vec_ptr, k_block_cache_ptr, head_block_logits, scale,
549
- block_idx == block_num - 1 ? last_block_token_num : BLOCK_SIZE);
550
- }
551
-
552
- std::pair<float, float> max_and_sum;
553
- if (alibi_slopes) {
554
- max_and_sum = reduceSoftmaxAlibi(
555
- logits, token_num, block_num * BLOCK_SIZE,
556
- alibi_slopes[head_idx], start_token_idx, seq_len);
557
- } else {
558
- max_and_sum =
559
- reduceSoftmax(logits, token_num, block_num * BLOCK_SIZE);
560
- }
561
-
562
- auto&& [max_logit, exp_sum] = max_and_sum;
563
-
564
- scalar_t* __restrict__ output_buffer = nullptr;
565
- if (!no_reduce) {
566
- auto idx = seq_idx * num_heads * max_num_partitions +
567
- head_idx * max_num_partitions + partition_idx;
568
- max_logits[idx] = max_logit;
569
- exp_sums[idx] = exp_sum;
570
- output_buffer =
571
- tmp_out + seq_idx * num_heads * max_num_partitions * HEAD_SIZE +
572
- head_idx * max_num_partitions * HEAD_SIZE +
573
- partition_idx * HEAD_SIZE;
574
- } else {
575
- output_buffer =
576
- out + seq_idx * num_heads * HEAD_SIZE + head_idx * HEAD_SIZE;
577
- }
578
-
579
- // Compute value
580
- constexpr int head_elem_num_per_partition = 16;
581
- constexpr int head_partition_num =
582
- HEAD_SIZE / head_elem_num_per_partition;
583
- for (int head_part_idx = 0; head_part_idx < head_partition_num;
584
- ++head_part_idx) {
585
- vec_op::FP32Vec16 accums[head_elem_num_per_partition];
586
- scalar_t* __restrict__ out_ptr =
587
- output_buffer + head_part_idx * head_elem_num_per_partition;
588
- for (int block_idx = 0; block_idx < block_num; ++block_idx) {
589
- const int64_t physical_block_idx = seq_block_table[block_idx];
590
- const float* __restrict__ prob_vec_ptr =
591
- logits + block_idx * BLOCK_SIZE;
592
- const scalar_t* __restrict__ v_block_cache_ptr =
593
- v_cache + physical_block_idx * kv_block_stride +
594
- kv_head_idx * kv_head_stride +
595
- BLOCK_SIZE * head_part_idx * head_elem_num_per_partition;
596
- reduceValueBlock<scalar_t, HEAD_SIZE, BLOCK_SIZE,
597
- head_elem_num_per_partition>(
598
- prob_vec_ptr, v_block_cache_ptr, accums);
599
-
600
- if (block_idx != block_num - 1) {
601
- const int64_t next_physical_block_idx =
602
- seq_block_table[block_idx + 1];
603
- const scalar_t* __restrict__ next_v_block_cache_ptr =
604
- v_cache + next_physical_block_idx * kv_block_stride +
605
- kv_head_idx * kv_head_stride +
606
- BLOCK_SIZE * head_part_idx * head_elem_num_per_partition;
607
- vec_op::unroll_loop<int, head_elem_num_per_partition>(
608
- [&](int head_elem_idx) {
609
- if (head_elem_idx % 2 == 0) {
610
- vec_op::prefetch(next_v_block_cache_ptr +
611
- BLOCK_SIZE * head_elem_idx);
612
- }
613
- });
614
- }
615
- }
616
-
617
- vec_op::unroll_loop<int, head_elem_num_per_partition>(
618
- [&](int head_elem_idx) {
619
- float value = accums[head_elem_idx].reduce_sum();
620
- vec_op::storeFP32(value, out_ptr + head_elem_idx);
621
- });
622
- }
623
- }
624
- }
625
- }
626
-
627
- // Rescale partition softmax and store the factors to exp_sums
628
-#pragma omp parallel for collapse(2) schedule(static, 1)
629
- for (int seq_idx = 0; seq_idx < num_seqs; ++seq_idx) {
630
- for (int head_idx = 0; head_idx < num_heads; ++head_idx) {
631
- const int seq_len = seq_lens[seq_idx];
632
- const int partition_num =
633
- (seq_len + PARTITION_SIZE - 1) / PARTITION_SIZE;
634
-
635
- if (partition_num == 1) continue;
636
-
637
- reducePartitionSoftmax(
638
- max_logits + seq_idx * num_heads * max_num_partitions +
639
- head_idx * max_num_partitions,
640
- exp_sums + seq_idx * num_heads * max_num_partitions +
641
- head_idx * max_num_partitions,
642
- partition_num);
643
- }
644
- }
645
-
646
- // Reduce values
647
- using v_load_vec_type = typename KernelVecType<scalar_t>::v_load_vec_type;
648
- static_assert(v_load_vec_type::get_elem_num() == BLOCK_SIZE);
649
- constexpr int head_elem_num_per_group =
650
- 16; // Note: didn't align with the cacheline size, due to some
651
- // HEAD_SIZE didn't align with 64 bytes
652
- static_assert(HEAD_SIZE % head_elem_num_per_group == 0);
653
- constexpr int head_group_num = HEAD_SIZE / head_elem_num_per_group;
654
- const float* __restrict__ rescale_factors = exp_sums;
655
-#pragma omp parallel for collapse(3) schedule(static, 1)
656
- for (int seq_idx = 0; seq_idx < num_seqs; ++seq_idx) {
657
- for (int head_idx = 0; head_idx < num_heads; ++head_idx) {
658
- for (int group_idx = 0; group_idx < head_group_num; ++group_idx) {
659
- const int seq_len = seq_lens[seq_idx];
660
- const int partition_num =
661
- (seq_len + PARTITION_SIZE - 1) / PARTITION_SIZE;
662
-
663
- if (partition_num == 1) continue;
664
-
665
- const float* __restrict__ seq_head_rescale_factors =
666
- rescale_factors + seq_idx * num_heads * max_num_partitions +
667
- head_idx * max_num_partitions;
668
- const scalar_t* __restrict__ seq_head_tmp_out =
669
- tmp_out + seq_idx * num_heads * max_num_partitions * HEAD_SIZE +
670
- head_idx * max_num_partitions * HEAD_SIZE +
671
- group_idx * head_elem_num_per_group;
672
- scalar_t* __restrict__ seq_head_output =
673
- out + seq_idx * num_heads * HEAD_SIZE + head_idx * HEAD_SIZE +
674
- group_idx * head_elem_num_per_group;
675
-
676
- vec_op::FP32Vec16 acc;
677
- for (int i = 0; i < partition_num; ++i) {
678
- vec_op::FP32Vec16 rescale_factor(seq_head_rescale_factors[i]);
679
- v_load_vec_type value(seq_head_tmp_out + i * HEAD_SIZE);
680
- vec_op::FP32Vec16 fp32_value(value);
681
- acc = acc + fp32_value * rescale_factor;
682
- }
683
- v_load_vec_type cast_acc(acc);
684
- cast_acc.save(seq_head_output);
685
- }
686
- }
687
- }
688
- }
689
-};
690
-
691
-#define LAUNCH_V2_ATTENTION_KERNEL(T, HEAD_SIZE, BLOCK_SIZE) \
692
- paged_attention_v2_impl<T, HEAD_SIZE, BLOCK_SIZE, PARTITION_SIZE>::call( \
693
- out_ptr, exp_sums_ptr, max_logits_ptr, tmp_out_ptr, query_ptr, \
694
- key_cache_ptr, value_cache_ptr, num_kv_heads, scale, block_tables_ptr, \
695
- seq_lens_ptr, max_num_blocks_per_seq, alibi_slopes_ptr, q_stride, \
696
- kv_block_stride, kv_head_stride, num_seqs, num_heads, \
697
- max_num_partitions);
698
-
699
-template <typename T, int BLOCK_SIZE, int PARTITION_SIZE = 512>
700
-void paged_attention_v2_impl_launcher(
701
- torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits,
702
- torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache,
703
- torch::Tensor& value_cache, int num_kv_heads, float scale,
704
- torch::Tensor& block_tables, torch::Tensor& seq_lens, int block_size,
705
- int max_seq_len, const std::optional<torch::Tensor>& alibi_slopes) {
706
- int num_seqs = query.size(0);
707
- int num_heads = query.size(1);
708
- int head_size = query.size(2);
709
- int max_num_blocks_per_seq = block_tables.size(1);
710
- int q_stride = query.stride(0);
711
- int kv_block_stride = key_cache.stride(0);
712
- int kv_head_stride = key_cache.stride(1);
713
- int max_num_partitions = exp_sums.size(-1);
714
-
715
- // NOTE: alibi_slopes is optional.
716
- const float* alibi_slopes_ptr =
717
- alibi_slopes
718
- ? reinterpret_cast<const float*>(alibi_slopes.value().data_ptr())
719
- : nullptr;
720
-
721
- T* out_ptr = reinterpret_cast<T*>(out.data_ptr());
722
- float* exp_sums_ptr = reinterpret_cast<float*>(exp_sums.data_ptr());
723
- float* max_logits_ptr = reinterpret_cast<float*>(max_logits.data_ptr());
724
- T* tmp_out_ptr = reinterpret_cast<T*>(tmp_out.data_ptr());
725
- T* query_ptr = reinterpret_cast<T*>(query.data_ptr());
726
- T* key_cache_ptr = reinterpret_cast<T*>(key_cache.data_ptr());
727
- T* value_cache_ptr = reinterpret_cast<T*>(value_cache.data_ptr());
728
- int* block_tables_ptr = block_tables.data_ptr<int>();
729
- int* seq_lens_ptr = seq_lens.data_ptr<int>();
730
-
731
- switch (head_size) {
732
- case 32:
733
- LAUNCH_V2_ATTENTION_KERNEL(T, 32, BLOCK_SIZE);
734
- break;
735
- case 64:
736
- LAUNCH_V2_ATTENTION_KERNEL(T, 64, BLOCK_SIZE);
737
- break;
738
- case 80:
739
- LAUNCH_V2_ATTENTION_KERNEL(T, 80, BLOCK_SIZE);
740
- break;
741
- case 96:
742
- LAUNCH_V2_ATTENTION_KERNEL(T, 96, BLOCK_SIZE);
743
- break;
744
- case 112:
745
- LAUNCH_V2_ATTENTION_KERNEL(T, 112, BLOCK_SIZE);
746
- break;
747
- case 128:
748
- LAUNCH_V2_ATTENTION_KERNEL(T, 128, BLOCK_SIZE);
749
- break;
750
- case 192:
751
- LAUNCH_V2_ATTENTION_KERNEL(T, 192, BLOCK_SIZE);
752
- break;
753
- case 256:
754
- LAUNCH_V2_ATTENTION_KERNEL(T, 256, BLOCK_SIZE);
755
- break;
756
- default:
757
- TORCH_CHECK(false, "Unsupported head size: ", head_size);
758
- break;
759
- }
760
-}
761
-
762
-#define CALL_V2_KERNEL_LAUNCHER(T, BLOCK_SIZE) \
763
- paged_attention_v2_impl_launcher<T, BLOCK_SIZE>( \
764
- out, exp_sums, max_logits, tmp_out, query, key_cache, value_cache, \
765
- num_kv_heads, scale, block_tables, seq_lens, block_size, max_seq_len, \
766
- alibi_slopes);
767
-
768
-#define CALL_V2_KERNEL_LAUNCHER_BLOCK_SIZE(T) \
769
- switch (block_size) { \
770
- case 16: \
771
- CALL_V2_KERNEL_LAUNCHER(T, 16); \
772
- break; \
773
- default: \
774
- TORCH_CHECK(false, "Unsupported block size: ", block_size); \
775
- break; \
776
- }
777
-} // namespace
778
-
779
-void paged_attention_v2(
780
- torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits,
781
- torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache,
782
- torch::Tensor& value_cache, int64_t num_kv_heads, double scale,
783
- torch::Tensor& block_tables, torch::Tensor& seq_lens, int64_t block_size,
784
- int64_t max_seq_len, const std::optional<torch::Tensor>& alibi_slopes,
785
- const std::string& kv_cache_dtype, torch::Tensor& k_scale,
786
- torch::Tensor& v_scale, const int64_t tp_rank,
787
- const int64_t blocksparse_local_blocks,
788
- const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
789
- const int64_t blocksparse_head_sliding_step) {
790
- TORCH_CHECK(blocksparse_vert_stride <= 1,
791
- "CPU backend does not support blocksparse attention yet.");
792
- VLLM_DISPATCH_FLOATING_TYPES(query.scalar_type(), "paged_attention_v2_impl",
793
- [&] {
794
- CPU_KERNEL_GUARD_IN(paged_attention_v2_impl)
795
- CALL_V2_KERNEL_LAUNCHER_BLOCK_SIZE(scalar_t);
796
- CPU_KERNEL_GUARD_OUT(paged_attention_v2_impl)
797
- });
798
-}
\ No newline at end of file
csrc/cpu/cache.cpp
deleted
-214
@@ -1,214 +0,0 @@
1
-#include <map>
2
-#include <vector>
3
-
4
-#include "cpu_types.hpp"
5
-
6
-#if defined(__x86_64__)
7
- #define DISPATCH_MACRO VLLM_DISPATCH_FLOATING_TYPES_WITH_E5M2
8
-#else
9
- #define DISPATCH_MACRO VLLM_DISPATCH_FLOATING_TYPES
10
-#endif
11
-
12
-namespace {
13
-template <typename scalar_t>
14
-void copy_blocks_cpu_impl(std::vector<torch::Tensor> const& key_caches,
15
- std::vector<torch::Tensor> const& value_caches,
16
- const torch::Tensor& mapping_pairs,
17
- const int element_num_per_block,
18
- const int layer_num) {
19
- const size_t pair_num = mapping_pairs.size(0);
20
- const size_t block_bytes = sizeof(scalar_t) * element_num_per_block;
21
-#pragma omp parallel for collapse(2)
22
- for (int layer = 0; layer < layer_num; ++layer) {
23
- for (size_t pair = 0; pair < pair_num; ++pair) {
24
- int64_t source_offset =
25
- element_num_per_block * mapping_pairs[pair][0].item<int64_t>();
26
- int64_t target_offset =
27
- element_num_per_block * mapping_pairs[pair][1].item<int64_t>();
28
- scalar_t* key_cache_ptr = key_caches[layer].data_ptr<scalar_t>();
29
- scalar_t* source_ptr = key_cache_ptr + source_offset;
30
- scalar_t* target_ptr = key_cache_ptr + target_offset;
31
- std::memcpy(target_ptr, source_ptr, block_bytes);
32
-
33
- scalar_t* value_cache_ptr = value_caches[layer].data_ptr<scalar_t>();
34
- source_ptr = value_cache_ptr + source_offset;
35
- target_ptr = value_cache_ptr + target_offset;
36
- std::memcpy(target_ptr, source_ptr, block_bytes);
37
- }
38
- }
39
-}
40
-
41
-template <typename scalar_t>
42
-void reshape_and_cache_cpu_impl(
43
- const scalar_t* __restrict__ key, const scalar_t* __restrict__ value,
44
- scalar_t* __restrict__ key_cache, scalar_t* __restrict__ value_cache,
45
- const int64_t* __restrict__ slot_mapping, const int num_tokens,
46
- const int key_stride, const int value_stride, const int num_heads,
47
- const int head_size, const int block_size, const int x) {
48
- const int block_elem_num = num_heads * head_size * block_size;
49
-
50
-#pragma omp parallel for collapse(2)
51
- for (int token_idx = 0; token_idx < num_tokens; ++token_idx) {
52
- for (int head_idx = 0; head_idx < num_heads; ++head_idx) {
53
- const int64_t slot_idx = slot_mapping[token_idx];
54
- if (slot_idx >= 0) {
55
- int src_key_head_idx = token_idx * key_stride + head_idx * head_size;
56
- int src_value_head_idx =
57
- token_idx * value_stride + head_idx * head_size;
58
- const scalar_t* src_key_head_ptr = key + src_key_head_idx;
59
- const scalar_t* src_value_head_ptr = value + src_value_head_idx;
60
- const int64_t block_index = slot_idx / block_size;
61
- const int64_t block_offset = slot_idx % block_size;
62
- scalar_t* target_key_head_ptr = key_cache +
63
- block_elem_num * block_index +
64
- head_idx * block_size * head_size;
65
- scalar_t* target_value_head_ptr = value_cache +
66
- block_elem_num * block_index +
67
- head_idx * block_size * head_size;
68
-
69
- for (int src_key_idx = 0; src_key_idx < head_size; src_key_idx += x) {
70
- const int64_t target_offset =
71
- src_key_idx * block_size + block_offset * x;
72
- for (int i = 0; i < x; ++i) {
73
- target_key_head_ptr[target_offset + i] =
74
- src_key_head_ptr[src_key_idx + i];
75
- }
76
- }
77
-
78
- for (int src_value_idx = 0; src_value_idx < head_size;
79
- ++src_value_idx) {
80
- const int64_t target_offset =
81
- src_value_idx * block_size + block_offset;
82
- target_value_head_ptr[target_offset] =
83
- src_value_head_ptr[src_value_idx];
84
- }
85
- }
86
- }
87
- }
88
-}
89
-}; // namespace
90
-
91
-template <typename scalar_t>
92
-void concat_and_cache_mla_cpu_impl(
93
- const scalar_t* __restrict__ kv_c, // [num_tokens, kv_lora_rank]
94
- const scalar_t* __restrict__ k_pe, // [num_tokens, pe_dim]
95
- scalar_t* __restrict__ kv_cache, // [num_blocks, block_size, (kv_lora_rank
96
- // + pe_dim)]
97
- const int64_t* __restrict__ slot_mapping, // [num_tokens]
98
- const int num_tokens, //
99
- const int block_stride, //
100
- const int entry_stride, //
101
- const int kv_c_stride, //
102
- const int k_pe_stride, //
103
- const int kv_lora_rank, //
104
- const int pe_dim, //
105
- const int block_size //
106
-) {
107
-#pragma omp parallel for
108
- for (int token_idx = 0; token_idx < num_tokens; ++token_idx) {
109
- const int64_t slot_idx = slot_mapping[token_idx];
110
- // NOTE: slot_idx can be -1 if the token is padded
111
- if (slot_idx < 0) {
112
- continue;
113
- }
114
- const int64_t block_idx = slot_idx / block_size;
115
- const int64_t block_offset = slot_idx % block_size;
116
-
117
- auto copy = [&](const scalar_t* __restrict__ src,
118
- scalar_t* __restrict__ dst, int src_stride, int dst_stride,
119
- int size, int offset) {
120
- for (int i = 0; i < size; i++) {
121
- const int64_t src_idx = token_idx * src_stride + i;
122
- const int64_t dst_idx =
123
- block_idx * block_stride + block_offset * entry_stride + i + offset;
124
- dst[dst_idx] = src[src_idx];
125
- }
126
- };
127
-
128
- copy(kv_c, kv_cache, kv_c_stride, block_stride, kv_lora_rank, 0);
129
- copy(k_pe, kv_cache, k_pe_stride, block_stride, pe_dim, kv_lora_rank);
130
- }
131
-}
132
-
133
-// Note: the key_caches and value_caches vectors are constant but
134
-// not the Tensors they contain. The vectors need to be const refs
135
-// in order to satisfy pytorch's C++ operator registration code.
136
-void copy_blocks(std::vector<torch::Tensor> const& key_caches,
137
- std::vector<torch::Tensor> const& value_caches,
138
- const torch::Tensor& block_mapping) {
139
- unsigned num_layers = key_caches.size();
140
- TORCH_CHECK(num_layers == value_caches.size());
141
- if (num_layers == 0) {
142
- return;
143
- }
144
-
145
- const int element_num_per_block = key_caches[0][0].numel();
146
- DISPATCH_MACRO(key_caches[0].scalar_type(), "copy_blocks_cpu_impl", [&] {
147
- CPU_KERNEL_GUARD_IN(copy_blocks_cpu_impl)
148
- copy_blocks_cpu_impl<scalar_t>(key_caches, value_caches, block_mapping,
149
- element_num_per_block, num_layers);
150
- CPU_KERNEL_GUARD_OUT(copy_blocks_cpu_impl)
151
- });
152
-}
153
-
154
-void reshape_and_cache(torch::Tensor& key, torch::Tensor& value,
155
- torch::Tensor& key_cache, torch::Tensor& value_cache,
156
- torch::Tensor& slot_mapping,
157
- const std::string& kv_cache_dtype,
158
- torch::Tensor& k_scale, torch::Tensor& v_scale) {
159
- int num_tokens = key.size(0);
160
- int num_heads = key.size(1);
161
- int head_size = key.size(2);
162
- int block_size = key_cache.size(3);
163
- int x = key_cache.size(4);
164
-
165
- int key_stride = key.stride(0);
166
- int value_stride = value.stride(0);
167
-
168
- DISPATCH_MACRO(key.scalar_type(), "reshape_and_cache_cpu_impl", [&] {
169
- CPU_KERNEL_GUARD_IN(reshape_and_cache_cpu_impl)
170
- reshape_and_cache_cpu_impl<scalar_t>(
171
- key.data_ptr<scalar_t>(), value.data_ptr<scalar_t>(),
172
- key_cache.data_ptr<scalar_t>(), value_cache.data_ptr<scalar_t>(),
173
- slot_mapping.data_ptr<int64_t>(), num_tokens, key_stride, value_stride,
174
- num_heads, head_size, block_size, x);
175
- CPU_KERNEL_GUARD_OUT(reshape_and_cache_cpu_impl)
176
- });
177
-}
178
-
179
-void concat_and_cache_mla(
180
- torch::Tensor& kv_c, // [num_tokens, kv_lora_rank]
181
- torch::Tensor& k_pe, // [num_tokens, pe_dim]
182
- torch::Tensor& kv_cache, // [num_blocks, block_size, (kv_lora_rank +
183
- // pe_dim)]
184
- torch::Tensor& slot_mapping, // [num_tokens] or [num_actual_tokens]
185
- const std::string& kv_cache_dtype, torch::Tensor& scale) {
186
- int num_tokens = slot_mapping.size(0);
187
- int kv_lora_rank = kv_c.size(1);
188
- int pe_dim = k_pe.size(1);
189
- int block_size = kv_cache.size(1);
190
-
191
- TORCH_CHECK(kv_cache.size(2) == kv_lora_rank + pe_dim);
192
- TORCH_CHECK(kv_cache_dtype != "fp8");
193
-
194
- int kv_c_stride = kv_c.stride(0);
195
- int k_pe_stride = k_pe.stride(0);
196
- int block_stride = kv_cache.stride(0);
197
- int entry_stride = kv_cache.stride(1);
198
-
199
- VLLM_DISPATCH_FLOATING_TYPES(
200
- kv_c.scalar_type(), "concat_and_cache_mla_cpu_impl", [&] {
201
- CPU_KERNEL_GUARD_IN(concat_and_cache_mla_cpu_impl)
202
- concat_and_cache_mla_cpu_impl<scalar_t>(
203
- kv_c.data_ptr<scalar_t>(), k_pe.data_ptr<scalar_t>(),
204
- kv_cache.data_ptr<scalar_t>(), slot_mapping.data_ptr<int64_t>(),
205
- num_tokens, block_stride, entry_stride, kv_c_stride, k_pe_stride,
206
- kv_lora_rank, pe_dim, block_size);
207
- CPU_KERNEL_GUARD_OUT(concat_and_cache_mla_cpu_impl)
208
- });
209
-}
210
-
211
-void swap_blocks(torch::Tensor& src, torch::Tensor& dst,
212
- const torch::Tensor& block_mapping) {
213
- TORCH_CHECK(false, "swap_blocks is unsupported on CPU.")
214
-}
csrc/cpu/cpu_attn.cpp
new
+249
@@ -0,0 +1,249 @@
1
+#include "cpu_attn_vec.hpp"
2
+#include "cpu_attn_vec16.hpp"
3
+
4
+#ifdef CPU_CAPABILITY_AMXBF16
5
+ #include "cpu_attn_amx.hpp"
6
+ #define AMX_DISPATCH(...) \
7
+ case cpu_attention::ISA::AMX: { \
8
+ using attn_impl = cpu_attention::AttentionImpl<cpu_attention::ISA::AMX, \
9
+ scalar_t, head_dim>; \
10
+ return __VA_ARGS__(); \
11
+ }
12
+#else
13
+ #define AMX_DISPATCH(...) case cpu_attention::ISA::AMX:
14
+#endif
15
+
16
+#define CPU_ATTN_DISPATCH_CASE(HEAD_DIM, ...) \
17
+ case HEAD_DIM: { \
18
+ constexpr size_t head_dim = HEAD_DIM; \
19
+ return __VA_ARGS__(); \
20
+ }
21
+
22
+#define CPU_ATTN_DISPATCH_CASE_HEADDIM(HEAD_DIM, ...) \
23
+ [&] { \
24
+ switch (HEAD_DIM) { \
25
+ CPU_ATTN_DISPATCH_CASE(32, __VA_ARGS__) \
26
+ CPU_ATTN_DISPATCH_CASE(64, __VA_ARGS__) \
27
+ CPU_ATTN_DISPATCH_CASE(96, __VA_ARGS__) \
28
+ CPU_ATTN_DISPATCH_CASE(128, __VA_ARGS__) \
29
+ CPU_ATTN_DISPATCH_CASE(160, __VA_ARGS__) \
30
+ CPU_ATTN_DISPATCH_CASE(192, __VA_ARGS__) \
31
+ CPU_ATTN_DISPATCH_CASE(224, __VA_ARGS__) \
32
+ CPU_ATTN_DISPATCH_CASE(256, __VA_ARGS__) \
33
+ default: { \
34
+ TORCH_CHECK(false, "Invalid CPU attention head_dim: " + \
35
+ std::to_string(HEAD_DIM)); \
36
+ } \
37
+ } \
38
+ }()
39
+
40
+#define CPU_ATTN_DISPATCH_IMPL(ISA_TYPE, ...) \
41
+ [&] { \
42
+ switch (ISA_TYPE) { \
43
+ AMX_DISPATCH(__VA_ARGS__) \
44
+ case cpu_attention::ISA::VEC: { \
45
+ using attn_impl = \
46
+ cpu_attention::AttentionImpl<cpu_attention::ISA::VEC, scalar_t, \
47
+ head_dim>; \
48
+ return __VA_ARGS__(); \
49
+ } \
50
+ case cpu_attention::ISA::VEC16: { \
51
+ using attn_impl = \
52
+ cpu_attention::AttentionImpl<cpu_attention::ISA::VEC16, scalar_t, \
53
+ head_dim>; \
54
+ return __VA_ARGS__(); \
55
+ } \
56
+ default: { \
57
+ TORCH_CHECK(false, "Invalid CPU attention ISA type."); \
58
+ } \
59
+ } \
60
+ }()
61
+
62
+torch::Tensor get_scheduler_metadata(
63
+ const int64_t num_req, const int64_t num_heads_q,
64
+ const int64_t num_heads_kv, const int64_t head_dim,
65
+ const torch::Tensor& seq_lens, at::ScalarType dtype,
66
+ const torch::Tensor& query_start_loc, const bool casual,
67
+ const int64_t window_size, const std::string& isa_hint,
68
+ const bool enable_kv_split) {
69
+ cpu_attention::ISA isa;
70
+ if (isa_hint == "amx") {
71
+ isa = cpu_attention::ISA::AMX;
72
+ } else if (isa_hint == "vec") {
73
+ isa = cpu_attention::ISA::VEC;
74
+ } else if (isa_hint == "vec16") {
75
+ isa = cpu_attention::ISA::VEC16;
76
+ } else {
77
+ TORCH_CHECK(false, "Unsupported CPU attention ISA hint: " + isa_hint);
78
+ }
79
+
80
+ cpu_attention::AttentionScheduler::ScheduleInput input;
81
+ input.num_reqs = num_req;
82
+ input.num_heads_q = num_heads_q;
83
+ input.num_heads_kv = num_heads_kv;
84
+ input.head_dim = head_dim;
85
+ input.query_start_loc = query_start_loc.data_ptr<int32_t>();
86
+ input.seq_lens = seq_lens.data_ptr<int32_t>();
87
+ if (window_size != -1) {
88
+ input.left_sliding_window_size = window_size - 1;
89
+ if (casual) {
90
+ input.right_sliding_window_size = 0;
91
+ } else {
92
+ input.right_sliding_window_size = window_size - 1;
93
+ }
94
+ } else {
95
+ input.left_sliding_window_size = -1;
96
+ if (casual) {
97
+ input.right_sliding_window_size = 0;
98
+ } else {
99
+ input.right_sliding_window_size = -1;
100
+ }
101
+ }
102
+ input.casual = casual;
103
+ input.isa = isa;
104
+ input.enable_kv_split = enable_kv_split;
105
+ TORCH_CHECK(casual, "Only supports casual mask for now.");
106
+
107
+ VLLM_DISPATCH_FLOATING_TYPES(dtype, "get_scheduler_metadata", [&]() {
108
+ CPU_ATTN_DISPATCH_CASE_HEADDIM(head_dim, [&] {
109
+ CPU_ATTN_DISPATCH_IMPL(isa, [&]() {
110
+ input.elem_size = sizeof(scalar_t);
111
+ input.q_buffer_elem_size = sizeof(attn_impl::q_buffer_t);
112
+ input.logits_buffer_elem_size = sizeof(attn_impl::logits_buffer_t);
113
+ input.output_buffer_elem_size =
114
+ sizeof(attn_impl::partial_output_buffer_t);
115
+ input.max_num_q_per_iter = attn_impl::MaxQHeadNumPerIteration;
116
+ input.kv_block_alignment = attn_impl::BlockSizeAlignment;
117
+ });
118
+ });
119
+ });
120
+
121
+ cpu_attention::AttentionScheduler scheduler;
122
+ torch::Tensor metadata = scheduler.schedule(input);
123
+ return metadata;
124
+}
125
+
126
+void cpu_attn_reshape_and_cache(
127
+ const torch::Tensor& key, // [token_num, head_num, head_size]
128
+ const torch::Tensor& value, // [token_num, head_num, head_size]
129
+ torch::Tensor&
130
+ key_cache, // [num_blocks, num_kv_heads, block_size, head_size]
131
+ torch::Tensor&
132
+ value_cache, // [num_blocks, num_kv_heads, block_size, head_size]
133
+ const torch::Tensor& slot_mapping, const std::string& isa) {
134
+ TORCH_CHECK_EQ(key.dim(), 3);
135
+ TORCH_CHECK_EQ(value.dim(), 3);
136
+ TORCH_CHECK_EQ(key_cache.dim(), 4);
137
+ TORCH_CHECK_EQ(value_cache.dim(), 4);
138
+ TORCH_CHECK_EQ(key.stride(2), 1);
139
+ TORCH_CHECK_EQ(value.stride(2), 1);
140
+
141
+ const int64_t token_num = key.size(0);
142
+ const int64_t key_token_num_stride = key.stride(0);
143
+ const int64_t value_token_num_stride = value.stride(0);
144
+ const int64_t head_num = value.size(1);
145
+ const int64_t key_head_num_stride = key.stride(1);
146
+ const int64_t value_head_num_stride = value.stride(1);
147
+ const int64_t num_blocks = key_cache.size(0);
148
+ const int64_t num_blocks_stride = key_cache.stride(0);
149
+ const int64_t cache_head_num_stride = key_cache.stride(1);
150
+ const int64_t block_size = key_cache.size(2);
151
+ const int64_t block_size_stride = key_cache.stride(2);
152
+ const int64_t head_dim = key.size(-1);
153
+
154
+ cpu_attention::ISA isa_tag = [&]() {
155
+ if (isa == "amx") {
156
+ return cpu_attention::ISA::AMX;
157
+ } else if (isa == "vec") {
158
+ return cpu_attention::ISA::VEC;
159
+ } else if (isa == "vec16") {
160
+ return cpu_attention::ISA::VEC16;
161
+ } else {
162
+ TORCH_CHECK(false, "Invalid ISA type: " + isa);
163
+ }
164
+ }();
165
+
166
+ VLLM_DISPATCH_FLOATING_TYPES(
167
+ key.scalar_type(), "cpu_attn_reshape_and_cache", [&]() {
168
+ CPU_ATTN_DISPATCH_CASE_HEADDIM(head_dim, [&] {
169
+ CPU_ATTN_DISPATCH_IMPL(isa_tag, [&]() {
170
+ attn_impl::reshape_and_cache(
171
+ key.data_ptr<scalar_t>(), value.data_ptr<scalar_t>(),
172
+ key_cache.data_ptr<scalar_t>(),
173
+ value_cache.data_ptr<scalar_t>(),
174
+ slot_mapping.data_ptr<int64_t>(), token_num,
175
+ key_token_num_stride, value_token_num_stride, head_num,
176
+ key_head_num_stride, value_head_num_stride, num_blocks,
177
+ num_blocks_stride, cache_head_num_stride, block_size,
178
+ block_size_stride);
179
+ });
180
+ });
181
+ });
182
+}
183
+
184
+void cpu_attention_with_kv_cache(
185
+ const torch::Tensor& query, // [num_tokens, num_heads, head_size]
186
+ const torch::Tensor&
187
+ key_cache, // [num_blocks, num_kv_heads, block_size, head_size]
188
+ const torch::Tensor&
189
+ value_cache, // [num_blocks, num_kv_heads, block_size, head_size]
190
+ torch::Tensor& output, // [num_tokens, num_heads, head_size]
191
+ const torch::Tensor& query_start_loc, // [num_tokens + 1]
192
+ const torch::Tensor& seq_lens, // [num_tokens]
193
+ const double scale, const bool causal,
194
+ const std::optional<torch::Tensor>& alibi_slopes, // [num_heads]
195
+ const int64_t sliding_window_left, const int64_t sliding_window_right,
196
+ const torch::Tensor& block_table, // [num_tokens, max_block_num]
197
+ const double softcap, const torch::Tensor& scheduler_metadata,
198
+ const std::optional<torch::Tensor>& s_aux // [num_heads]
199
+) {
200
+ TORCH_CHECK_EQ(query.dim(), 3);
201
+ TORCH_CHECK_EQ(query.stride(2), 1);
202
+ TORCH_CHECK_EQ(key_cache.dim(), 4);
203
+ TORCH_CHECK_EQ(value_cache.dim(), 4);
204
+
205
+ cpu_attention::AttentionInput input;
206
+ input.metadata = reinterpret_cast<cpu_attention::AttentionMetadata*>(
207
+ scheduler_metadata.data_ptr());
208
+ input.num_tokens = query.size(0);
209
+ input.num_heads = query.size(1);
210
+ input.num_kv_heads = key_cache.size(1);
211
+ input.block_size = key_cache.size(2);
212
+ input.query = query.data_ptr();
213
+ input.query_num_tokens_stride = query.stride(0);
214
+ input.query_num_heads_stride = query.stride(1);
215
+ input.cache_num_blocks_stride = key_cache.stride(0);
216
+ input.cache_num_kv_heads_stride = key_cache.stride(1);
217
+ input.blt_num_tokens_stride = block_table.stride(0);
218
+ input.key_cache = key_cache.data_ptr();
219
+ input.value_cache = value_cache.data_ptr();
220
+ input.output = output.data_ptr();
221
+ input.query_start_loc = query_start_loc.data_ptr<int32_t>();
222
+ input.seq_lens = seq_lens.data_ptr<int32_t>();
223
+ input.block_table = block_table.data_ptr<int32_t>();
224
+ input.alibi_slopes =
225
+ alibi_slopes.has_value() ? alibi_slopes->data_ptr<float>() : nullptr;
226
+ // For now sink must be bf16
227
+ input.s_aux = s_aux.has_value() ? s_aux->data_ptr<c10::BFloat16>() : nullptr;
228
+ input.scale = scale;
229
+ input.causal = causal;
230
+ input.sliding_window_left = sliding_window_left;
231
+ input.sliding_window_right = sliding_window_right;
232
+ if (input.causal) {
233
+ // to make boundary calculation easier
234
+ input.sliding_window_right = 0;
235
+ }
236
+ float softcap_fp32 = softcap;
237
+ input.softcap = softcap_fp32;
238
+
239
+ VLLM_DISPATCH_FLOATING_TYPES(
240
+ query.scalar_type(), "cpu_attention_with_kv_cache", [&]() {
241
+ CPU_ATTN_DISPATCH_CASE_HEADDIM(query.size(2), [&] {
242
+ CPU_ATTN_DISPATCH_IMPL(input.metadata->isa, [&]() {
243
+ TORCH_CHECK_EQ(input.block_size % attn_impl::BlockSizeAlignment, 0);
244
+ cpu_attention::AttentionMainLoop<attn_impl> mainloop;
245
+ mainloop(&input);
246
+ });
247
+ });
248
+ });
249
+}
csrc/cpu/cpu_attn_amx.hpp
new
+511
@@ -0,0 +1,511 @@
1
+#ifndef CPU_ATTN_AMX_HPP
2
+#define CPU_ATTN_AMX_HPP
3
+
4
+#include "cpu_attn_impl.hpp"
5
+
6
+namespace cpu_attention {
7
+namespace {
8
+// AMX specific
9
+constexpr static int64_t AMX_TILE_ROW_BYTES = 64;
10
+constexpr static int64_t AMX_TILE_ROW_NUM = 16;
11
+constexpr static int64_t AMX_TILE_BYTES = AMX_TILE_ROW_BYTES * AMX_TILE_ROW_NUM;
12
+
13
+typedef struct __tile_config {
14
+ uint8_t palette_id = 1;
15
+ uint8_t start_row = 0;
16
+ uint8_t reserved_0[14] = {0};
17
+ uint16_t colsb[16] = {0};
18
+ uint8_t rows[16] = {0};
19
+} __tilecfg;
20
+
21
+// 2-2-4 pattern, for 16 < m <= 32
22
+// TILE 0, 1: load A matrix, row num should be 16, m - 16
23
+// 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>
27
+class TileGemm224 {
28
+ public:
29
+ template <AttentionGemmPhase phase, int32_t k_size>
30
+ FORCE_INLINE static void gemm(const int32_t m_size, void* __restrict__ a_tile,
31
+ void* __restrict__ b_tile,
32
+ float* __restrict__ c_tile, const int64_t lda,
33
+ const int64_t ldb, const int64_t ldc,
34
+ const int32_t block_size,
35
+ const int32_t dynamic_k_size,
36
+ const bool accum_c) {
37
+ TORCH_CHECK(false, "Unsupported kv cache type for TileGemm224");
38
+ }
39
+
40
+ FORCE_INLINE static void init_tile_config(int32_t m, __tilecfg& config) {
41
+ TORCH_CHECK(false, "Unsupported kv cache type for TileGemm224");
42
+ }
43
+};
44
+
45
+template <>
46
+class TileGemm224<c10::BFloat16> {
47
+ public:
48
+ template <AttentionGemmPhase phase, int32_t k_size>
49
+ FORCE_INLINE static void gemm(const int32_t m_size,
50
+ c10::BFloat16* __restrict__ a_tile,
51
+ c10::BFloat16* __restrict__ b_tile,
52
+ float* __restrict__ c_tile, const int64_t lda,
53
+ const int64_t ldb, const int64_t ldc,
54
+ const int32_t block_size,
55
+ const int32_t dynamic_k_size,
56
+ const bool accum_c) {
57
+ const int32_t k_times =
58
+ dynamic_k_size / (AMX_TILE_ROW_NUM * 4 / sizeof(c10::BFloat16));
59
+ c10::BFloat16* __restrict__ a_tile_0 = a_tile;
60
+ c10::BFloat16* __restrict__ a_tile_1 = a_tile + lda * AMX_TILE_ROW_NUM;
61
+ const int64_t a_tile_stride = [&]() {
62
+ if constexpr (phase == AttentionGemmPhase::QK) {
63
+ // q_buffer is prepacked
64
+ return AMX_TILE_ROW_BYTES;
65
+ } else if constexpr (phase == AttentionGemmPhase::PV) {
66
+ // logits_buffer is row-major
67
+ return lda * sizeof(c10::BFloat16);
68
+ } else {
69
+ TORCH_CHECK(false, "Unreachable");
70
+ }
71
+ }();
72
+
73
+ c10::BFloat16* __restrict__ b_tile_2 = b_tile;
74
+ c10::BFloat16* __restrict__ b_tile_3 = [&]() {
75
+ if constexpr (phase == AttentionGemmPhase::QK) {
76
+ // k_cache is prepacked
77
+ return b_tile + (k_size * AMX_TILE_ROW_BYTES / 4);
78
+ } else if constexpr (phase == AttentionGemmPhase::PV) {
79
+ // v_cache is prepacked
80
+ return b_tile + (block_size * AMX_TILE_ROW_BYTES / 4);
81
+ } else {
82
+ TORCH_CHECK(false, "Unreachable");
83
+ }
84
+ }();
85
+ // k_cache, v_cache are prepacked
86
+ const int32_t b_tile_stride = AMX_TILE_ROW_BYTES;
87
+
88
+ // logits_buffer, output_buffer are not prepacked
89
+ float* __restrict__ c_tile_4 = c_tile;
90
+ float* __restrict__ c_tile_5 =
91
+ c_tile_4 + AMX_TILE_ROW_BYTES / sizeof(float);
92
+ float* __restrict__ c_tile_6 = c_tile + AMX_TILE_ROW_NUM * ldc;
93
+ float* __restrict__ c_tile_7 =
94
+ c_tile_6 + AMX_TILE_ROW_BYTES / sizeof(float);
95
+ const int32_t c_tile_stride = ldc * sizeof(float);
96
+
97
+ if (accum_c) {
98
+ _tile_loadd(4, c_tile_4, c_tile_stride);
99
+ _tile_loadd(5, c_tile_5, c_tile_stride);
100
+ _tile_loadd(6, c_tile_6, c_tile_stride);
101
+ _tile_loadd(7, c_tile_7, c_tile_stride);
102
+ } else {
103
+ _tile_zero(4);
104
+ _tile_zero(5);
105
+ _tile_zero(6);
106
+ _tile_zero(7);
107
+ }
108
+
109
+ for (int32_t k = 0; k < k_times; ++k) {
110
+ _tile_loadd(0, a_tile_0, a_tile_stride);
111
+ _tile_stream_loadd(2, b_tile_2, b_tile_stride);
112
+ _tile_dpbf16ps(4, 0, 2);
113
+ _tile_stream_loadd(3, b_tile_3, b_tile_stride);
114
+ _tile_dpbf16ps(5, 0, 3);
115
+ _tile_loadd(1, a_tile_1, a_tile_stride);
116
+ _tile_dpbf16ps(6, 1, 2);
117
+ _tile_dpbf16ps(7, 1, 3);
118
+
119
+ // update ptrs
120
+ if constexpr (phase == AttentionGemmPhase::QK) {
121
+ // Q buffer is prepacked
122
+ a_tile_0 += AMX_TILE_BYTES / sizeof(c10::BFloat16);
123
+ a_tile_1 += AMX_TILE_BYTES / sizeof(c10::BFloat16);
124
+ } else if constexpr (phase == AttentionGemmPhase::PV) {
125
+ // P buffer is not prepacked
126
+ a_tile_0 += AMX_TILE_ROW_BYTES / sizeof(c10::BFloat16);
127
+ a_tile_1 += AMX_TILE_ROW_BYTES / sizeof(c10::BFloat16);
128
+ } else {
129
+ TORCH_CHECK(false, "Unreachable");
130
+ }
131
+ b_tile_2 += AMX_TILE_BYTES / sizeof(c10::BFloat16);
132
+ b_tile_3 += AMX_TILE_BYTES / sizeof(c10::BFloat16);
133
+ }
134
+
135
+ _tile_stored(4, c_tile_4, c_tile_stride);
136
+ _tile_stored(5, c_tile_5, c_tile_stride);
137
+ _tile_stored(6, c_tile_6, c_tile_stride);
138
+ _tile_stored(7, c_tile_7, c_tile_stride);
139
+ }
140
+
141
+ FORCE_INLINE static void init_tile_config(int32_t m, __tilecfg& config) {
142
+ const int32_t m_0 = AMX_TILE_ROW_NUM;
143
+ const int32_t m_1 = m - AMX_TILE_ROW_NUM;
144
+ config.rows[0] = m_0;
145
+ config.rows[1] = m_1;
146
+ config.rows[2] = AMX_TILE_ROW_NUM;
147
+ config.rows[3] = AMX_TILE_ROW_NUM;
148
+ config.rows[4] = m_0;
149
+ config.rows[5] = m_0;
150
+ config.rows[6] = m_1;
151
+ config.rows[7] = m_1;
152
+ _tile_loadconfig(&config);
153
+ }
154
+};
155
+
156
+// 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>
164
+class TileGemm122 {
165
+ public:
166
+ template <AttentionGemmPhase phase, int32_t k_size>
167
+ FORCE_INLINE static void gemm(const int32_t m_size, void* __restrict__ a_tile,
168
+ void* __restrict__ b_tile,
169
+ float* __restrict__ c_tile, const int64_t lda,
170
+ const int64_t ldb, const int64_t ldc,
171
+ const int32_t block_size,
172
+ const int32_t dynamic_k_size,
173
+ const bool accum_c) {
174
+ TORCH_CHECK(false, "Unsupported kv cache type for TileGemm122");
175
+ }
176
+
177
+ FORCE_INLINE static void init_tile_config(int32_t m, __tilecfg& config) {
178
+ TORCH_CHECK(false, "Unsupported kv cache type for TileGemm122");
179
+ }
180
+};
181
+
182
+template <>
183
+class TileGemm122<c10::BFloat16> {
184
+ public:
185
+ template <AttentionGemmPhase phase, int32_t k_size>
186
+ FORCE_INLINE static void gemm(const int32_t m_size,
187
+ c10::BFloat16* __restrict__ a_tile,
188
+ c10::BFloat16* __restrict__ b_tile,
189
+ float* __restrict__ c_tile, const int64_t lda,
190
+ const int64_t ldb, const int64_t ldc,
191
+ const int32_t block_size,
192
+ const int32_t dynamic_k_size,
193
+ const bool accum_c) {
194
+ c10::BFloat16* __restrict__ a_tile_0 = a_tile;
195
+ c10::BFloat16* __restrict__ a_tile_1 = [&]() {
196
+ if constexpr (phase == AttentionGemmPhase::QK) {
197
+ // q_buffer is prepacked
198
+ return a_tile + AMX_TILE_BYTES / sizeof(c10::BFloat16);
199
+ } else if constexpr (phase == AttentionGemmPhase::PV) {
200
+ // logits_buffer is row-major
201
+ return a_tile + AMX_TILE_ROW_BYTES / sizeof(c10::BFloat16);
202
+ } else {
203
+ TORCH_CHECK(false, "Unreachable");
204
+ }
205
+ }();
206
+ const int64_t a_tile_stride = [&]() {
207
+ if constexpr (phase == AttentionGemmPhase::QK) {
208
+ // q_buffer is prepacked
209
+ return AMX_TILE_ROW_BYTES;
210
+ } else if constexpr (phase == AttentionGemmPhase::PV) {
211
+ // logits_buffer is row-major
212
+ return lda * sizeof(c10::BFloat16);
213
+ } else {
214
+ TORCH_CHECK(false, "Unreachable");
215
+ }
216
+ }();
217
+
218
+ c10::BFloat16* __restrict__ b_tile_2 = b_tile;
219
+ c10::BFloat16* __restrict__ b_tile_3 = [&]() {
220
+ if constexpr (phase == AttentionGemmPhase::QK) {
221
+ // k_cache is prepacked
222
+ return b_tile + (k_size * AMX_TILE_ROW_BYTES / 4);
223
+ } else if constexpr (phase == AttentionGemmPhase::PV) {
224
+ // v_cache is prepacked
225
+ return b_tile + (block_size * AMX_TILE_ROW_BYTES / 4);
226
+ } else {
227
+ TORCH_CHECK(false, "Unreachable");
228
+ }
229
+ }();
230
+ c10::BFloat16* __restrict__ b_tile_4 =
231
+ b_tile_2 + AMX_TILE_BYTES / sizeof(c10::BFloat16);
232
+ c10::BFloat16* __restrict__ b_tile_5 =
233
+ b_tile_3 + AMX_TILE_BYTES / sizeof(c10::BFloat16);
234
+ int64_t b_stride = AMX_TILE_ROW_BYTES;
235
+
236
+ float* __restrict__ c_tile_6 = c_tile;
237
+ float* __restrict__ c_tile_7 = c_tile + AMX_TILE_ROW_BYTES / sizeof(float);
238
+ int64_t c_stride = ldc * sizeof(float);
239
+
240
+ const int32_t k_times =
241
+ dynamic_k_size / (AMX_TILE_ROW_NUM * 4 / sizeof(c10::BFloat16));
242
+ const int32_t k_group_times = k_times / 2;
243
+ const bool has_tail = (k_times % 2 == 1);
244
+
245
+ if (accum_c) {
246
+ _tile_loadd(6, c_tile_6, c_stride);
247
+ _tile_loadd(7, c_tile_7, c_stride);
248
+ } else {
249
+ _tile_zero(6);
250
+ _tile_zero(7);
251
+ }
252
+
253
+ for (int32_t k = 0; k < k_group_times; ++k) {
254
+ _tile_loadd(0, a_tile_0, a_tile_stride);
255
+ _tile_stream_loadd(2, b_tile_2, b_stride);
256
+ _tile_dpbf16ps(6, 0, 2);
257
+ _tile_stream_loadd(3, b_tile_3, b_stride);
258
+ _tile_dpbf16ps(7, 0, 3);
259
+ _tile_loadd(1, a_tile_1, a_tile_stride);
260
+ _tile_stream_loadd(4, b_tile_4, b_stride);
261
+ _tile_dpbf16ps(6, 1, 4);
262
+ _tile_stream_loadd(5, b_tile_5, b_stride);
263
+ _tile_dpbf16ps(7, 1, 5);
264
+
265
+ // update ptrs
266
+ if constexpr (phase == AttentionGemmPhase::QK) {
267
+ // Q buffer is prepacked
268
+ a_tile_0 += 2 * AMX_TILE_BYTES / sizeof(c10::BFloat16);
269
+ a_tile_1 += 2 * AMX_TILE_BYTES / sizeof(c10::BFloat16);
270
+ } else if constexpr (phase == AttentionGemmPhase::PV) {
271
+ // P buffer is not prepacked
272
+ a_tile_0 += 2 * AMX_TILE_ROW_BYTES / sizeof(c10::BFloat16);
273
+ a_tile_1 += 2 * AMX_TILE_ROW_BYTES / sizeof(c10::BFloat16);
274
+ }
275
+ b_tile_2 += 2 * AMX_TILE_BYTES / sizeof(c10::BFloat16);
276
+ b_tile_3 += 2 * AMX_TILE_BYTES / sizeof(c10::BFloat16);
277
+ b_tile_4 += 2 * AMX_TILE_BYTES / sizeof(c10::BFloat16);
278
+ b_tile_5 += 2 * AMX_TILE_BYTES / sizeof(c10::BFloat16);
279
+ }
280
+
281
+ if (has_tail) {
282
+ _tile_loadd(0, a_tile_0, a_tile_stride);
283
+ _tile_stream_loadd(2, b_tile_2, b_stride);
284
+ _tile_dpbf16ps(6, 0, 2);
285
+ _tile_stream_loadd(3, b_tile_3, b_stride);
286
+ _tile_dpbf16ps(7, 0, 3);
287
+ }
288
+
289
+ _tile_stored(6, c_tile_6, c_stride);
290
+ _tile_stored(7, c_tile_7, c_stride);
291
+ }
292
+
293
+ FORCE_INLINE static void init_tile_config(int32_t m, __tilecfg& config) {
294
+ config.rows[0] = m;
295
+ config.rows[1] = m;
296
+ config.rows[2] = AMX_TILE_ROW_NUM;
297
+ config.rows[3] = AMX_TILE_ROW_NUM;
298
+ config.rows[4] = AMX_TILE_ROW_NUM;
299
+ config.rows[5] = AMX_TILE_ROW_NUM;
300
+ config.rows[6] = m;
301
+ config.rows[7] = m;
302
+ _tile_loadconfig(&config);
303
+ }
304
+};
305
+} // namespace
306
+
307
+template <typename scalar_t, int64_t head_dim>
308
+class AttentionImpl<ISA::AMX, scalar_t, head_dim> {
309
+ public:
310
+ using query_t = scalar_t;
311
+ using q_buffer_t = scalar_t;
312
+ using kv_cache_t = scalar_t;
313
+ using logits_buffer_t = float;
314
+ using partial_output_buffer_t = float;
315
+ using prob_buffer_t = scalar_t;
316
+
317
+ constexpr static int64_t BlockSizeAlignment =
318
+ AMX_TILE_ROW_BYTES /
319
+ sizeof(kv_cache_t); // KV token num unit of QK and PV phases
320
+ constexpr static int64_t HeadDimAlignment =
321
+ 2 * (AMX_TILE_ROW_BYTES / 4); // headdim num unit of PV phase
322
+ constexpr static int64_t MaxQHeadNumPerIteration = 32;
323
+ constexpr static int64_t HeadDim = head_dim;
324
+ constexpr static ISA ISAType = ISA::AMX;
325
+ constexpr static bool scale_on_logits = true;
326
+
327
+ public:
328
+ AttentionImpl() : current_q_head_num_(0) {
329
+ // Use all columns in AMX tiles
330
+ vec_op::unroll_loop<int, 8>([&](int i) { amx_tile_config_.colsb[i] = 64; });
331
+ }
332
+
333
+ ~AttentionImpl() { _tile_release(); }
334
+
335
+ template <template <typename tile_gemm_t> typename attention>
336
+ FORCE_INLINE void execute_attention(DEFINE_CPU_ATTENTION_PARAMS) {
337
+ if (q_head_num > AMX_TILE_ROW_NUM) {
338
+ if (q_head_num != current_q_head_num_) {
339
+ current_q_head_num_ = q_head_num;
340
+ TileGemm224<kv_cache_t>::init_tile_config(q_head_num, amx_tile_config_);
341
+ }
342
+ attention<TileGemm224<kv_cache_t>> attention_iteration;
343
+ attention_iteration(CPU_ATTENTION_PARAMS);
344
+ } else {
345
+ if (q_head_num != current_q_head_num_) {
346
+ current_q_head_num_ = q_head_num;
347
+ TileGemm122<kv_cache_t>::init_tile_config(q_head_num, amx_tile_config_);
348
+ }
349
+ attention<TileGemm122<kv_cache_t>> attention_iteration;
350
+ attention_iteration(CPU_ATTENTION_PARAMS);
351
+ }
352
+ }
353
+
354
+ // k_cache_token_group_stride: stride of K cache when move to next
355
+ // BlockSizeAlignment tokens in a block
356
+ constexpr static int64_t k_cache_token_group_stride(
357
+ const int32_t block_size) {
358
+ return BlockSizeAlignment * head_dim;
359
+ }
360
+
361
+ // v_cache_token_group_stride: stride of V cache when move to next
362
+ // BlockSizeAlignment tokens in a block
363
+ constexpr static int64_t v_cache_token_group_stride(
364
+ const int32_t block_size) {
365
+ return BlockSizeAlignment * (AMX_TILE_ROW_BYTES / 4);
366
+ }
367
+
368
+ // v_cache_head_group_stride: stride of V cache when move to next
369
+ // HeadDimAlignment head dims in a block
370
+ constexpr static int64_t v_cache_head_group_stride(const int32_t block_size) {
371
+ return block_size * HeadDimAlignment;
372
+ }
373
+
374
+ static void copy_q_heads_tile(
375
+ scalar_t* __restrict__ src, // [q_num, q_heads_per_kv, head_size]
376
+ scalar_t* __restrict__ q_buffer, const int32_t q_num,
377
+ const int32_t q_heads_per_kv, const int64_t q_num_stride,
378
+ const int64_t q_head_stride, const float scale) {
379
+ constexpr int64_t bytes_per_head = head_dim * sizeof(scalar_t);
380
+ static_assert(bytes_per_head % AMX_TILE_ROW_BYTES == 0);
381
+ constexpr int64_t head_size_block_num = bytes_per_head / AMX_TILE_ROW_BYTES;
382
+ constexpr int64_t head_elem_num_pre_block =
383
+ AMX_TILE_ROW_BYTES / sizeof(scalar_t);
384
+
385
+ int32_t idx = 0;
386
+ int8_t* __restrict__ q_buffer_iter = reinterpret_cast<int8_t*>(q_buffer);
387
+ for (int32_t q_num_idx = 0; q_num_idx < q_num;
388
+ ++q_num_idx, src += q_num_stride) {
389
+ scalar_t* __restrict__ src_iter = src;
390
+ for (int32_t q_head_idx = 0; q_head_idx < q_heads_per_kv;
391
+ ++q_head_idx, src_iter += q_head_stride) {
392
+ vec_op::unroll_loop<int32_t, head_size_block_num>(
393
+ [&](int32_t head_size_block_idx) {
394
+ // Use INT8Vec64 for 64 bytes block
395
+ vec_op::INT8Vec64 vec(src_iter + head_size_block_idx *
396
+ head_elem_num_pre_block);
397
+ vec.save(q_buffer_iter + head_size_block_idx * AMX_TILE_BYTES);
398
+ });
399
+
400
+ ++idx;
401
+ q_buffer_iter += AMX_TILE_ROW_BYTES;
402
+ if ((idx & (AMX_TILE_ROW_NUM - 1)) == 0) {
403
+ // head is in another amx tile
404
+ q_buffer_iter -= AMX_TILE_ROW_NUM * AMX_TILE_ROW_BYTES;
405
+ q_buffer_iter += head_size_block_num * AMX_TILE_BYTES;
406
+ }
407
+ }
408
+ }
409
+ }
410
+
411
+ // reshape KV to AMX friendly layout
412
+ static void reshape_and_cache(
413
+ const scalar_t* __restrict__ key, const scalar_t* __restrict__ value,
414
+ scalar_t* __restrict__ key_cache, scalar_t* __restrict__ value_cache,
415
+ const int64_t* __restrict__ slot_mapping, const int64_t token_num,
416
+ const int64_t key_token_num_stride, const int64_t value_token_num_stride,
417
+ const int64_t head_num, const int64_t key_head_num_stride,
418
+ const int64_t value_head_num_stride, const int64_t num_blocks,
419
+ 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) {
421
+ // For AMX 2D tiles, size of each line is 64 bytes
422
+ constexpr int64_t amx_tile_row_size = AMX_TILE_ROW_BYTES;
423
+ // For AMX B martix, N always is 16
424
+ constexpr int64_t amx_b_tile_n_size = AMX_TILE_ROW_BYTES / 4;
425
+ constexpr int64_t amx_b_tile_k_size = amx_tile_row_size / sizeof(scalar_t);
426
+ // For now suppose block_size is divisible by amx_tile_column_num
427
+ TORCH_CHECK_EQ(block_size % amx_b_tile_k_size, 0);
428
+
429
+#pragma omp parallel for collapse(2)
430
+ for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
431
+ for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
432
+ const int64_t pos = slot_mapping[token_idx];
433
+ if (pos < 0) {
434
+ // skip
435
+ continue;
436
+ }
437
+
438
+ const int64_t block_idx = pos / block_size;
439
+ const int64_t block_offset = pos % block_size;
440
+ {
441
+ // Write Key
442
+ // Head elements should be packed as quand-words and stored in token
443
+ // groups with (quadword_stride/4) tokens
444
+ constexpr int64_t token_num_per_group = amx_tile_row_size / 4;
445
+ static_assert(head_dim % (4 / sizeof(scalar_t)) == 0);
446
+ constexpr int64_t quadword_num = head_dim / (4 / sizeof(scalar_t));
447
+ const int32_t* key_start_quadword_ptr =
448
+ reinterpret_cast<const int32_t*>(
449
+ key + token_idx * key_token_num_stride +
450
+ head_idx * key_head_num_stride);
451
+ const int64_t group_idx = block_offset / token_num_per_group;
452
+ const int64_t group_offset = block_offset % token_num_per_group;
453
+ constexpr int64_t quadword_num_per_group =
454
+ token_num_per_group * quadword_num;
455
+ int32_t* key_cache_start_ptr =
456
+ reinterpret_cast<int32_t*>(key_cache +
457
+ block_idx * num_blocks_stride +
458
+ head_idx * cache_head_num_stride) +
459
+ group_idx * quadword_num_per_group + group_offset;
460
+
461
+#pragma GCC unroll 8
462
+ for (int64_t i = 0, j = 0; j < quadword_num;
463
+ i += token_num_per_group, ++j) {
464
+ key_cache_start_ptr[i] = key_start_quadword_ptr[j];
465
+ }
466
+ }
467
+ {
468
+ // Write Value
469
+ // Different from Key, block_size dimension is packed rather than
470
+ // head_size dimension block_size dimension is packed as quand-words;
471
+ constexpr int64_t token_num_per_sub_group = 4 / sizeof(scalar_t);
472
+ const int64_t token_num_per_group = block_size;
473
+ constexpr int64_t head_elems_per_group = amx_b_tile_n_size;
474
+ const int64_t group_size = token_num_per_group * head_elems_per_group;
475
+ // For now suppose head_dim is divisible by amx_b_tile_n_size
476
+ static_assert(head_dim % head_elems_per_group == 0);
477
+ constexpr int64_t group_num = head_dim / head_elems_per_group;
478
+ const int64_t sub_group_idx = block_offset / token_num_per_sub_group;
479
+ const int64_t sub_group_offset =
480
+ block_offset % token_num_per_sub_group;
481
+
482
+ const scalar_t* value_start_ptr = value +
483
+ token_idx * value_token_num_stride +
484
+ head_idx * value_head_num_stride;
485
+ scalar_t* value_cache_start_ptr =
486
+ value_cache + block_idx * num_blocks_stride +
487
+ head_idx * cache_head_num_stride +
488
+ sub_group_idx * token_num_per_sub_group * amx_b_tile_n_size +
489
+ sub_group_offset;
490
+
491
+ for (int64_t i = 0; i < group_num; ++i) {
492
+#pragma GCC unroll head_elems_per_group
493
+ for (int64_t j = 0, k = 0; j < head_elems_per_group;
494
+ ++j, k += token_num_per_sub_group) {
495
+ value_cache_start_ptr[k] = value_start_ptr[j];
496
+ }
497
+ value_start_ptr += head_elems_per_group;
498
+ value_cache_start_ptr += group_size;
499
+ }
500
+ }
501
+ }
502
+ }
503
+ }
504
+
505
+ private:
506
+ alignas(64) __tilecfg amx_tile_config_;
507
+ int32_t current_q_head_num_;
508
+};
509
+} // namespace cpu_attention
510
+
511
+#endif
csrc/cpu/cpu_attn_impl.hpp
new
+1977
@@ -0,0 +1,1977 @@
1
+#ifndef CPU_ATTN_HPP
2
+#define CPU_ATTN_HPP
3
+
4
+#include <unistd.h>
5
+#include <type_traits>
6
+#include <cstddef>
7
+
8
+#include "cpu_types.hpp"
9
+#include "scratchpad_manager.h"
10
+#include "cpu_attn_macros.h"
11
+
12
+namespace cpu_attention {
13
+enum class ISA { AMX, VEC, VEC16 };
14
+
15
+template <ISA isa, typename scalar_t, int64_t head_dim>
16
+class AttentionImpl {};
17
+
18
+struct AttentionWorkItemGroup {
19
+ int32_t req_id;
20
+ int32_t q_token_id_start;
21
+ int32_t q_token_num;
22
+ int32_t kv_split_pos_start;
23
+ int32_t kv_split_pos_end;
24
+
25
+ int64_t total_kv_len;
26
+ int32_t split_id;
27
+ int32_t local_split_id;
28
+
29
+ AttentionWorkItemGroup(const int32_t req_id, const int32_t q_token_id_start,
30
+ const int32_t kv_split_pos_start,
31
+ const int32_t kv_split_pos_end)
32
+ : req_id(req_id),
33
+ q_token_id_start(q_token_id_start),
34
+ q_token_num(0),
35
+ kv_split_pos_start(kv_split_pos_start),
36
+ kv_split_pos_end(kv_split_pos_end),
37
+ total_kv_len(0),
38
+ split_id(-1),
39
+ local_split_id(0) {}
40
+
41
+ std::string to_string() const {
42
+ std::stringstream ss;
43
+ ss << '[' << "req_id: " << req_id << ",\n";
44
+ ss << "q_token_id_start: " << q_token_id_start << ",\n";
45
+ ss << "q_token_num: " << q_token_num << ",\n";
46
+ ss << "kv_split_pos_start: " << kv_split_pos_start << ",\n";
47
+ ss << "kv_split_pos_end: " << kv_split_pos_end << ",\n";
48
+ ss << "total_kv_len: " << total_kv_len << ",\n";
49
+ ss << "split_id: " << split_id << ",\n";
50
+ ss << "local_split_id: " << local_split_id << ",\n";
51
+ ss << ']';
52
+
53
+ return ss.str();
54
+ }
55
+};
56
+
57
+struct ReductionWorkItemGroup {
58
+ int32_t req_id;
59
+ int32_t q_token_id_start;
60
+ int32_t q_token_id_num;
61
+ int32_t split_start_id;
62
+ int32_t split_num;
63
+
64
+ ReductionWorkItemGroup(const int32_t req_id, const int32_t q_token_id_start,
65
+ const int32_t q_token_id_num,
66
+ const int32_t split_start_id)
67
+ : req_id(req_id),
68
+ q_token_id_start(q_token_id_start),
69
+ q_token_id_num(q_token_id_num),
70
+ split_start_id(split_start_id),
71
+ split_num(0) {}
72
+
73
+ std::string to_string() const {
74
+ std::stringstream ss;
75
+ ss << '[' << "req_id: " << req_id << ",\n";
76
+ ss << "q_token_id_start: " << q_token_id_start << ",\n";
77
+ ss << "q_token_id_num: " << q_token_id_num << ",\n";
78
+ ss << "split_start_id: " << split_start_id << ",\n";
79
+ ss << "split_num: " << split_num << ",\n";
80
+ ss << ']';
81
+
82
+ return ss.str();
83
+ }
84
+};
85
+
86
+struct AttentionMetadata {
87
+ std::atomic_int64_t counter;
88
+ char _padding1[56];
89
+ ISA isa;
90
+ int32_t workitem_group_num;
91
+ int32_t reduction_item_num;
92
+ int32_t reduction_split_num;
93
+ int32_t thread_num;
94
+ int32_t effective_thread_num; // non-zero item num in workitem_num_per_thread
95
+ int32_t split_kv_q_token_num_threshold;
96
+ int64_t attention_scratchpad_size_per_thread;
97
+ int64_t reduction_scratchpad_size_per_kv_head;
98
+ AttentionWorkItemGroup* workitem_groups_ptr;
99
+ ReductionWorkItemGroup* reduction_items_ptr;
100
+ int32_t cu_workitem_num_per_thread[1025] = {
101
+ 0}; // prefix sum of workitem_num_per_thread
102
+ char _padding2[56];
103
+
104
+ AttentionMetadata(ISA isa, int32_t workitem_group_num,
105
+ int32_t reduction_item_num, int32_t reduction_split_num,
106
+ int32_t split_kv_q_token_num_threshold)
107
+ : isa(isa),
108
+ workitem_group_num(workitem_group_num),
109
+ reduction_item_num(reduction_item_num),
110
+ reduction_split_num(reduction_split_num),
111
+ thread_num(omp_get_max_threads()),
112
+ effective_thread_num(thread_num),
113
+ split_kv_q_token_num_threshold(split_kv_q_token_num_threshold),
114
+ attention_scratchpad_size_per_thread(0),
115
+ reduction_scratchpad_size_per_kv_head(0),
116
+ workitem_groups_ptr(
117
+ (AttentionWorkItemGroup*)((char*)this + sizeof(AttentionMetadata))),
118
+ reduction_items_ptr(
119
+ (ReductionWorkItemGroup*)((char*)this + sizeof(AttentionMetadata) +
120
+ workitem_group_num *
121
+ sizeof(AttentionWorkItemGroup))),
122
+ counter(0) {
123
+ TORCH_CHECK_LE(thread_num, 1024);
124
+ static_assert(sizeof(AttentionMetadata) % 64 == 0);
125
+ TORCH_CHECK(reinterpret_cast<size_t>(this) % 64 == 0);
126
+ }
127
+
128
+ void reset_counter() { counter.store(0); }
129
+
130
+ int64_t acquire_counter() { return counter++; }
131
+
132
+ void print() const {
133
+ std::stringstream ss;
134
+ ss << "ISA: ";
135
+ switch (isa) {
136
+ case ISA::AMX:
137
+ ss << "AMX, ";
138
+ break;
139
+ case ISA::VEC:
140
+ ss << "VEC, ";
141
+ break;
142
+ }
143
+ ss << "workitem_group_num: " << workitem_group_num
144
+ << ", reduction_item_num: " << reduction_item_num
145
+ << ", reduction_split_num: " << reduction_split_num
146
+ << ", thread_num: " << thread_num
147
+ << ", effective_thread_num: " << effective_thread_num
148
+ << ", attention_scratchpad_size_per_thread: "
149
+ << attention_scratchpad_size_per_thread
150
+ << ", reduction_scratchpad_size_per_kv_head: "
151
+ << reduction_scratchpad_size_per_kv_head << ", workitem groups:\n";
152
+ for (int32_t i = 0; i < workitem_group_num; ++i) {
153
+ ss << (workitem_groups_ptr + i)->to_string() << ",\n";
154
+ }
155
+
156
+ ss << "cu_workitem_num_per_thread: [";
157
+ for (int32_t i = 0; i < thread_num + 1; ++i) {
158
+ ss << cu_workitem_num_per_thread[i] << ", ";
159
+ }
160
+ ss << "]\n";
161
+
162
+ ss << "reduction items: \n";
163
+
164
+ for (int32_t i = 0; i < reduction_item_num; ++i) {
165
+ ss << (reduction_items_ptr + i)->to_string() << ",\n";
166
+ }
167
+
168
+ std::printf("%s", ss.str().c_str());
169
+ }
170
+};
171
+
172
+// Thread attention scratchpad contains:
173
+// - Q: q_tile_size * head_dim * q_buffer_elem_size, gather Q heads, especially
174
+// for GQA
175
+// - Q@K^T: max_num_q_per_iter * k_tile_size * logits_buffer_elem_size, logits
176
+// - Intermediate outputs: q_tile_size * head_dim * output_buffer_elem_size + 2
177
+// * q_tile_size * 4, partial output, max + sum (float)
178
+// Reduction scratchpad contains:
179
+// - flags: bool array to indicate wether the split is finished
180
+// - outputs: split_num * q_tile_size * head_dim * output_buffer_elem_size
181
+// - max, sum: 2 * split_num * q_tile_size * 4
182
+class AttentionScratchPad {
183
+ public:
184
+ AttentionScratchPad(int64_t thread_id,
185
+ const AttentionMetadata& attention_metadata,
186
+ void* scratchpad_ptr)
187
+ : thread_scratchpad_ptr(
188
+ static_cast<int8_t*>(scratchpad_ptr) +
189
+ thread_id *
190
+ attention_metadata.attention_scratchpad_size_per_thread),
191
+ reduction_scratchpad_ptr(
192
+ static_cast<int8_t*>(scratchpad_ptr) +
193
+ attention_metadata.thread_num *
194
+ attention_metadata.attention_scratchpad_size_per_thread),
195
+ reduction_scratchpad_size_per_kv_head(
196
+ attention_metadata.reduction_scratchpad_size_per_kv_head) {}
197
+
198
+ // for attention
199
+ void update(const int64_t head_dim, const int64_t q_buffer_elem_size,
200
+ const int64_t logits_buffer_elem_size,
201
+ const int64_t output_buffer_elem_size,
202
+ const int64_t max_num_q_per_iter, const int64_t q_head_tile_size,
203
+ const int64_t kv_tile_size) {
204
+ int64_t buffer_offset = 0;
205
+ q_buffer_offset_ = buffer_offset;
206
+ buffer_offset +=
207
+ calcu_q_buffer_size(q_head_tile_size, head_dim, q_buffer_elem_size);
208
+ logits_buffer_offset_ = buffer_offset;
209
+ buffer_offset += calcu_logits_buffer_size(max_num_q_per_iter, kv_tile_size,
210
+ logits_buffer_elem_size);
211
+ output_buffer_offset_ = buffer_offset;
212
+ buffer_offset += calcu_partial_output_buffer_size(
213
+ q_head_tile_size, head_dim, output_buffer_elem_size);
214
+ max_buffer_offset_ = buffer_offset;
215
+ buffer_offset += calcu_partial_output_max_sum_buffer_size(q_head_tile_size);
216
+ sum_buffer_offset_ = buffer_offset;
217
+ }
218
+
219
+ // for reduction
220
+ void update(const int32_t kv_head_idx, const int32_t total_split_num,
221
+ const int64_t head_dim, const int64_t q_head_tile_size,
222
+ const int64_t output_buffer_elem_size) {
223
+ int64_t buffer_offset = kv_head_idx * reduction_scratchpad_size_per_kv_head;
224
+ reduce_flag_buffer_offset_ = buffer_offset;
225
+ buffer_offset += calcu_reduce_flag_buffer_size(total_split_num);
226
+ reduce_output_buffer_offset_ = buffer_offset;
227
+ buffer_offset += calcu_reduce_output_buffer_size(
228
+ total_split_num, q_head_tile_size, head_dim, output_buffer_elem_size);
229
+ reduce_max_buffer_offset_ = buffer_offset;
230
+ buffer_offset +=
231
+ calcu_reduce_max_sum_buffer_size(total_split_num, q_head_tile_size);
232
+ reduce_sum_buffer_offset_ = buffer_offset;
233
+ }
234
+
235
+ template <typename T>
236
+ T* get_q_buffer() {
237
+ return reinterpret_cast<T*>(thread_scratchpad_ptr + q_buffer_offset_);
238
+ }
239
+
240
+ float* get_logits_buffer() {
241
+ return reinterpret_cast<float*>(thread_scratchpad_ptr +
242
+ logits_buffer_offset_);
243
+ }
244
+
245
+ float* get_output_buffer() {
246
+ return reinterpret_cast<float*>(thread_scratchpad_ptr +
247
+ output_buffer_offset_);
248
+ }
249
+
250
+ float* get_max_buffer() {
251
+ return reinterpret_cast<float*>(thread_scratchpad_ptr + max_buffer_offset_);
252
+ }
253
+
254
+ float* get_sum_buffer() {
255
+ return reinterpret_cast<float*>(thread_scratchpad_ptr + sum_buffer_offset_);
256
+ }
257
+
258
+ volatile bool* get_reduce_flag_buffer() {
259
+ return reinterpret_cast<volatile bool*>(reduction_scratchpad_ptr +
260
+ reduce_flag_buffer_offset_);
261
+ }
262
+
263
+ float* get_reduce_output_buffer() {
264
+ return reinterpret_cast<float*>(reduction_scratchpad_ptr +
265
+ reduce_output_buffer_offset_);
266
+ }
267
+
268
+ float* get_reduce_max_buffer() {
269
+ return reinterpret_cast<float*>(reduction_scratchpad_ptr +
270
+ reduce_max_buffer_offset_);
271
+ }
272
+
273
+ float* get_reduce_sum_buffer() {
274
+ return reinterpret_cast<float*>(reduction_scratchpad_ptr +
275
+ reduce_sum_buffer_offset_);
276
+ }
277
+
278
+ int64_t get_thread_scratchpad_size() const {
279
+ return 2 * sum_buffer_offset_ - max_buffer_offset_;
280
+ }
281
+
282
+ int64_t get_reduction_scratchpad_size() const {
283
+ return 2 * reduce_sum_buffer_offset_ - reduce_max_buffer_offset_;
284
+ }
285
+
286
+ private:
287
+ static int64_t round_to_64(const int64_t num) {
288
+ return ((num + 63) >> 6) << 6;
289
+ }
290
+
291
+ static int64_t calcu_q_buffer_size(const int64_t q_tile_size,
292
+ const int64_t head_dim,
293
+ const int64_t elem_size) {
294
+ return round_to_64(q_tile_size * head_dim * elem_size);
295
+ }
296
+
297
+ static int64_t calcu_logits_buffer_size(const int64_t max_num_q_per_iter,
298
+ const int64_t k_tile_size,
299
+ const int64_t elem_size) {
300
+ return round_to_64(elem_size * max_num_q_per_iter * k_tile_size);
301
+ }
302
+
303
+ static int64_t calcu_partial_output_buffer_size(const int64_t q_tile_size,
304
+ const int64_t head_dim,
305
+ const int64_t elem_size) {
306
+ return round_to_64(q_tile_size * head_dim * elem_size);
307
+ }
308
+
309
+ static int64_t calcu_partial_output_max_sum_buffer_size(
310
+ const int64_t q_tile_size) {
311
+ return round_to_64(q_tile_size * sizeof(float));
312
+ }
313
+
314
+ static int64_t calcu_reduce_flag_buffer_size(const int64_t total_split_num) {
315
+ return round_to_64(total_split_num * sizeof(bool));
316
+ }
317
+
318
+ static int64_t calcu_reduce_max_sum_buffer_size(
319
+ const int64_t total_split_num, const int32_t q_head_tile_size) {
320
+ return round_to_64(total_split_num * q_head_tile_size * sizeof(float));
321
+ }
322
+
323
+ static int64_t calcu_reduce_output_buffer_size(
324
+ const int64_t total_split_num, const int64_t q_head_tile_size,
325
+ const int64_t head_dim, const int64_t output_buffer_elem_size) {
326
+ return round_to_64(total_split_num * q_head_tile_size * head_dim *
327
+ output_buffer_elem_size);
328
+ }
329
+
330
+ private:
331
+ int8_t* thread_scratchpad_ptr;
332
+ int8_t* reduction_scratchpad_ptr;
333
+ int64_t reduction_scratchpad_size_per_kv_head;
334
+ // attention buffers
335
+ int64_t q_buffer_offset_;
336
+ int64_t logits_buffer_offset_;
337
+ int64_t output_buffer_offset_;
338
+ int64_t max_buffer_offset_;
339
+ int64_t sum_buffer_offset_;
340
+ // reduction buffers
341
+ int64_t reduce_flag_buffer_offset_;
342
+ int64_t reduce_output_buffer_offset_;
343
+ int64_t reduce_max_buffer_offset_;
344
+ int64_t reduce_sum_buffer_offset_;
345
+};
346
+
347
+class AttentionScheduler {
348
+ public:
349
+ struct ScheduleInput {
350
+ int32_t num_reqs;
351
+ int32_t elem_size;
352
+ int32_t q_buffer_elem_size;
353
+ int32_t logits_buffer_elem_size;
354
+ int32_t output_buffer_elem_size;
355
+ int32_t num_heads_q;
356
+ int32_t num_heads_kv;
357
+ int32_t head_dim;
358
+ int32_t* query_start_loc;
359
+ int32_t* seq_lens;
360
+ int32_t left_sliding_window_size;
361
+ int32_t right_sliding_window_size;
362
+ bool casual;
363
+ cpu_attention::ISA isa;
364
+ int32_t max_num_q_per_iter; // max Q head num can be hold in registers
365
+ int32_t kv_block_alignment; // context length alignment requirement
366
+ bool enable_kv_split;
367
+ };
368
+
369
+ static constexpr int32_t MaxQTileIterNum = 128;
370
+
371
+ AttentionScheduler() : available_cache_size_(get_available_l2_size()) {}
372
+
373
+ torch::Tensor schedule(const ScheduleInput& input) const {
374
+ const bool casual = input.casual;
375
+ const int32_t thread_num = omp_get_max_threads();
376
+ const int64_t cache_size = get_available_l2_size();
377
+ const int32_t max_num_q_per_iter = input.max_num_q_per_iter;
378
+ const int32_t kv_len_alignment = input.kv_block_alignment;
379
+ int32_t q_head_per_kv = input.num_heads_q / input.num_heads_kv;
380
+ const bool use_gqa = (max_num_q_per_iter % q_head_per_kv == 0);
381
+ if (!use_gqa) {
382
+ q_head_per_kv = 1; // fallback to MHA
383
+ }
384
+ const int32_t min_split_kv_len =
385
+ ((max_num_q_per_iter * 4 + kv_len_alignment - 1) / kv_len_alignment) *
386
+ kv_len_alignment;
387
+ const int32_t max_num_q_token_per_iter = max_num_q_per_iter / q_head_per_kv;
388
+ const int64_t default_tile_size = calcu_default_tile_size(
389
+ cache_size, input.head_dim, input.elem_size, input.q_buffer_elem_size,
390
+ input.logits_buffer_elem_size, input.output_buffer_elem_size,
391
+ max_num_q_per_iter, max_num_q_per_iter);
392
+ const int32_t default_tile_token_num = default_tile_size / q_head_per_kv;
393
+ const int32_t split_kv_q_token_num_threshold =
394
+ input.enable_kv_split ? 1 : 0;
395
+ const int32_t left_sliding_window_size = input.left_sliding_window_size;
396
+ const int32_t right_sliding_window_size = input.right_sliding_window_size;
397
+ TORCH_CHECK_LE(split_kv_q_token_num_threshold * q_head_per_kv, 16);
398
+
399
+ // get total kv len
400
+ int64_t total_kv_len = 0;
401
+ for (int32_t req_id = 0; req_id < input.num_reqs; ++req_id) {
402
+ const int32_t seq_len = input.seq_lens[req_id];
403
+ const int32_t q_token_num =
404
+ input.query_start_loc[req_id + 1] - input.query_start_loc[req_id];
405
+ const int32_t q_start_pos = (casual ? (seq_len - q_token_num) : 0);
406
+ const int32_t kv_start_pos = 0;
407
+ const int32_t kv_end_pos = seq_len;
408
+
409
+ for (int32_t token_id = 0; token_id < q_token_num;
410
+ token_id += max_num_q_token_per_iter) {
411
+ const int32_t q_tile_token_num =
412
+ std::min(max_num_q_token_per_iter, q_token_num - token_id);
413
+ const int32_t q_tile_pos_left = q_start_pos + token_id;
414
+ const int32_t q_tile_pos_right = q_tile_pos_left + q_tile_token_num;
415
+ const auto [kv_tile_pos_left, kv_tile_pos_right] = calcu_kv_tile_pos(
416
+ kv_start_pos, kv_end_pos, q_tile_pos_left, q_tile_pos_right,
417
+ left_sliding_window_size, right_sliding_window_size);
418
+ const auto [aligned_kv_tile_pos_left, aligned_kv_tile_pos_right] =
419
+ align_kv_tile_pos(kv_tile_pos_left, kv_tile_pos_right,
420
+ kv_len_alignment);
421
+
422
+ int32_t curr_kv_len =
423
+ aligned_kv_tile_pos_right - aligned_kv_tile_pos_left;
424
+ total_kv_len += curr_kv_len;
425
+ }
426
+ }
427
+ const int64_t kv_len_per_thread =
428
+ (((total_kv_len / thread_num) + kv_len_alignment - 1) /
429
+ kv_len_alignment) *
430
+ kv_len_alignment * (use_gqa ? input.num_heads_kv : input.num_heads_q);
431
+ std::vector<AttentionWorkItemGroup> workitems;
432
+ std::vector<ReductionWorkItemGroup> reduce_workitems;
433
+ workitems.reserve(1024);
434
+ reduce_workitems.reserve(1024);
435
+ std::vector<int32_t> workitem_num_per_thread(thread_num, 0);
436
+
437
+ // split tasks
438
+ int32_t curr_thread_id = 0;
439
+ int64_t remaining_kv_len = kv_len_per_thread;
440
+ int32_t cum_split_num = 0;
441
+ for (int32_t req_id = 0; req_id < input.num_reqs; ++req_id) {
442
+ const int32_t seq_len = input.seq_lens[req_id];
443
+ const int32_t q_token_num =
444
+ input.query_start_loc[req_id + 1] - input.query_start_loc[req_id];
445
+ const int32_t q_start_pos = (casual ? (seq_len - q_token_num) : 0);
446
+ const int32_t kv_start_pos = 0;
447
+ const int32_t kv_end_pos = seq_len;
448
+ int32_t local_split_id = 0;
449
+
450
+ AttentionWorkItemGroup curr_workitem(req_id, 0, 0, seq_len);
451
+ for (int32_t token_id = 0; token_id < q_token_num;
452
+ token_id += max_num_q_token_per_iter) {
453
+ const int32_t q_tile_token_num =
454
+ std::min(max_num_q_token_per_iter, q_token_num - token_id);
455
+ const int32_t q_tile_pos_left = q_start_pos + token_id;
456
+ const int32_t q_tile_pos_right = q_tile_pos_left + q_tile_token_num;
457
+ const auto [kv_tile_pos_left, kv_tile_pos_right] = calcu_kv_tile_pos(
458
+ kv_start_pos, kv_end_pos, q_tile_pos_left, q_tile_pos_right,
459
+ left_sliding_window_size, right_sliding_window_size);
460
+ const auto [aligned_kv_tile_pos_left, aligned_kv_tile_pos_right] =
461
+ align_kv_tile_pos(kv_tile_pos_left, kv_tile_pos_right,
462
+ kv_len_alignment);
463
+ int32_t curr_kv_len =
464
+ aligned_kv_tile_pos_right - aligned_kv_tile_pos_left;
465
+ int32_t kv_token_pos_start = aligned_kv_tile_pos_left;
466
+
467
+ while (curr_kv_len > 0) {
468
+ if (curr_kv_len <= (remaining_kv_len + min_split_kv_len) ||
469
+ curr_thread_id == (thread_num - 1)) {
470
+ curr_workitem.q_token_num += q_tile_token_num;
471
+ curr_workitem.total_kv_len += curr_kv_len;
472
+ remaining_kv_len -= curr_kv_len;
473
+ curr_kv_len = 0;
474
+
475
+ if (remaining_kv_len < 0) {
476
+ // stop to accept more workitems
477
+ remaining_kv_len -= min_split_kv_len;
478
+ }
479
+
480
+ if (curr_workitem.kv_split_pos_start != 0) {
481
+ // got a partial kv spilt, need to create a single workitem
482
+ curr_workitem.split_id = cum_split_num;
483
+ curr_workitem.local_split_id = local_split_id;
484
+ workitems.emplace_back(curr_workitem);
485
+ ++workitem_num_per_thread[curr_thread_id];
486
+ ++reduce_workitems.back().split_num;
487
+ ++cum_split_num;
488
+
489
+ curr_workitem = AttentionWorkItemGroup(
490
+ req_id, token_id + max_num_q_token_per_iter, 0, seq_len);
491
+ }
492
+
493
+ break;
494
+ }
495
+
496
+ if (remaining_kv_len < min_split_kv_len &&
497
+ (curr_workitem.total_kv_len > 0 ||
498
+ workitem_num_per_thread[curr_thread_id] > 0)) {
499
+ // remaining_kv_len is too short, and have allocated workitems, just
500
+ // leave to next thread
501
+ if (curr_workitem.total_kv_len > 0) {
502
+ workitems.emplace_back(curr_workitem);
503
+ ++workitem_num_per_thread[curr_thread_id];
504
+ curr_workitem =
505
+ AttentionWorkItemGroup(req_id, token_id, 0, seq_len);
506
+ }
507
+
508
+ // switch to next thread
509
+ ++curr_thread_id;
510
+ remaining_kv_len = kv_len_per_thread;
511
+
512
+ // retry this iteration
513
+ continue;
514
+ }
515
+
516
+ // only split tail splits with q_tile_token_num <=
517
+ // split_kv_q_token_num_threshold
518
+ if (token_id + max_num_q_token_per_iter < q_token_num ||
519
+ q_tile_token_num > split_kv_q_token_num_threshold) {
520
+ // if requires a new q tile iteration and already has workitems,
521
+ // leave this workitem to next thread
522
+ if (curr_workitem.q_token_num % default_tile_token_num == 0 &&
523
+ (curr_workitem.total_kv_len > 0 ||
524
+ workitem_num_per_thread[curr_thread_id] > 0)) {
525
+ if (curr_workitem.total_kv_len > 0) {
526
+ workitems.emplace_back(curr_workitem);
527
+ ++workitem_num_per_thread[curr_thread_id];
528
+ }
529
+ curr_workitem =
530
+ AttentionWorkItemGroup(req_id, token_id, 0, seq_len);
531
+
532
+ // switch to next thread
533
+ ++curr_thread_id;
534
+ remaining_kv_len = kv_len_per_thread;
535
+ }
536
+
537
+ curr_workitem.q_token_num += q_tile_token_num;
538
+ curr_workitem.total_kv_len += curr_kv_len;
539
+ remaining_kv_len -= curr_kv_len;
540
+ curr_kv_len = 0;
541
+ break;
542
+ }
543
+
544
+ // split kv
545
+ if (curr_workitem.total_kv_len > 0) {
546
+ // write back curr workitem
547
+ workitems.emplace_back(curr_workitem);
548
+ ++workitem_num_per_thread[curr_thread_id];
549
+ }
550
+
551
+ if (kv_token_pos_start == aligned_kv_tile_pos_left) {
552
+ // first split, init the workitem
553
+ reduce_workitems.emplace_back(ReductionWorkItemGroup(
554
+ req_id, token_id, q_tile_token_num, cum_split_num));
555
+ }
556
+
557
+ int32_t spilt_size =
558
+ std::min(std::max(remaining_kv_len, (int64_t)min_split_kv_len),
559
+ (int64_t)curr_kv_len);
560
+ curr_workitem =
561
+ AttentionWorkItemGroup(req_id, token_id, kv_token_pos_start,
562
+ kv_token_pos_start + spilt_size);
563
+ curr_workitem.q_token_num += q_tile_token_num;
564
+ curr_workitem.total_kv_len += spilt_size;
565
+ curr_workitem.split_id = cum_split_num;
566
+ curr_workitem.local_split_id = local_split_id;
567
+ workitems.emplace_back(curr_workitem);
568
+ ++workitem_num_per_thread[curr_thread_id];
569
+ ++reduce_workitems.back().split_num;
570
+ ++cum_split_num;
571
+ ++local_split_id;
572
+
573
+ kv_token_pos_start += spilt_size;
574
+ curr_kv_len -= spilt_size;
575
+ curr_workitem = AttentionWorkItemGroup(req_id, token_id,
576
+ kv_token_pos_start, seq_len);
577
+
578
+ // switch to next thread
579
+ ++curr_thread_id;
580
+ remaining_kv_len = kv_len_per_thread;
581
+ }
582
+ }
583
+
584
+ if (curr_workitem.total_kv_len > 0) {
585
+ // write back curr workitem
586
+ workitems.emplace_back(curr_workitem);
587
+ ++workitem_num_per_thread[curr_thread_id];
588
+ }
589
+ }
590
+
591
+ int64_t metadata_tensor_size =
592
+ sizeof(AttentionMetadata) +
593
+ workitems.size() * sizeof(AttentionWorkItemGroup) +
594
+ reduce_workitems.size() * sizeof(ReductionWorkItemGroup);
595
+ auto options =
596
+ torch::TensorOptions().dtype(torch::kInt8).device(torch::kCPU);
597
+ torch::Tensor metadata_tensor =
598
+ torch::empty({metadata_tensor_size}, options);
599
+ AttentionMetadata* metadata_ptr = new (metadata_tensor.data_ptr())
600
+ AttentionMetadata(input.isa, workitems.size(), reduce_workitems.size(),
601
+ cum_split_num, split_kv_q_token_num_threshold);
602
+ AttentionWorkItemGroup* workitem_groups_ptr =
603
+ metadata_ptr->workitem_groups_ptr;
604
+ ReductionWorkItemGroup* reduction_items_ptr =
605
+ metadata_ptr->reduction_items_ptr;
606
+ std::memcpy(workitem_groups_ptr, workitems.data(),
607
+ workitems.size() * sizeof(AttentionWorkItemGroup));
608
+ std::memcpy(reduction_items_ptr, reduce_workitems.data(),
609
+ reduce_workitems.size() * sizeof(ReductionWorkItemGroup));
610
+
611
+ int32_t effective_thread_num = 0;
612
+ for (; effective_thread_num < thread_num; ++effective_thread_num) {
613
+ if (workitem_num_per_thread[effective_thread_num] == 0) {
614
+ break;
615
+ }
616
+ }
617
+
618
+ std::memcpy(metadata_ptr->cu_workitem_num_per_thread + 1,
619
+ workitem_num_per_thread.data(),
620
+ workitem_num_per_thread.size() * sizeof(int32_t));
621
+ for (int32_t i = 1; i <= thread_num; ++i) {
622
+ metadata_ptr->cu_workitem_num_per_thread[i] +=
623
+ metadata_ptr->cu_workitem_num_per_thread[i - 1];
624
+ }
625
+ metadata_ptr->effective_thread_num = effective_thread_num;
626
+
627
+ {
628
+ // when q_tile_size = max_num_q_per_iter, requires max
629
+ // attention_scratchpad_size
630
+ AttentionScratchPad sc(0, *metadata_ptr, 0x0);
631
+ int64_t n = AttentionScheduler::calcu_tile_size_with_constant_q(
632
+ cache_size, input.head_dim, input.elem_size, input.q_buffer_elem_size,
633
+ input.logits_buffer_elem_size, input.output_buffer_elem_size,
634
+ max_num_q_per_iter, kv_len_alignment, max_num_q_per_iter, true);
635
+ sc.update(input.head_dim, input.q_buffer_elem_size,
636
+ input.logits_buffer_elem_size, input.output_buffer_elem_size,
637
+ max_num_q_per_iter, max_num_q_per_iter, n);
638
+ metadata_ptr->attention_scratchpad_size_per_thread =
639
+ ((sc.get_thread_scratchpad_size() + 63) / 64) * 64;
640
+
641
+ sc.update(0, metadata_ptr->reduction_split_num, input.head_dim,
642
+ q_head_per_kv * split_kv_q_token_num_threshold,
643
+ input.output_buffer_elem_size);
644
+ metadata_ptr->reduction_scratchpad_size_per_kv_head =
645
+ ((sc.get_reduction_scratchpad_size() + 63) / 64) * 64;
646
+ }
647
+ int64_t scratchpad_size =
648
+ metadata_ptr->attention_scratchpad_size_per_thread *
649
+ metadata_ptr->thread_num +
650
+ metadata_ptr->reduction_scratchpad_size_per_kv_head *
651
+ (use_gqa ? input.num_heads_kv : input.num_heads_q);
652
+ DNNLScratchPadManager::get_dnnl_scratchpad_manager()->realloc(
653
+ scratchpad_size);
654
+
655
+ // metadata_ptr->print();
656
+
657
+ // test out of boundary access
658
+ // {
659
+ // float* cache_ptr =
660
+ // DNNLScratchPadManager::get_dnnl_scratchpad_manager()->get_data<float>();
661
+ // for (int64_t i = 0; i < scratchpad_size / sizeof(float); ++i) {
662
+ // cache_ptr[i] = std::numeric_limits<float>::quiet_NaN();
663
+ // }
664
+ // }
665
+
666
+ return metadata_tensor;
667
+ }
668
+
669
+ FORCE_INLINE static std::pair<int32_t, int32_t> calcu_kv_tile_pos(
670
+ int32_t kv_left_pos, int32_t kv_right_pos, int32_t q_left_pos,
671
+ int32_t q_right_pos, int32_t sliding_window_left,
672
+ int32_t sliding_window_right) {
673
+ if (sliding_window_left != -1) {
674
+ kv_left_pos = std::max(kv_left_pos, q_left_pos - sliding_window_left);
675
+ }
676
+ if (sliding_window_right != -1) {
677
+ kv_right_pos = std::min(kv_right_pos, q_right_pos + sliding_window_right);
678
+ }
679
+ return {kv_left_pos, kv_right_pos};
680
+ }
681
+
682
+ FORCE_INLINE static std::pair<int32_t, int32_t> align_kv_tile_pos(
683
+ int32_t kv_left_pos, int32_t kv_right_pos, int32_t align_factor) {
684
+ kv_left_pos = (kv_left_pos / align_factor) * align_factor;
685
+ kv_right_pos =
686
+ ((kv_right_pos + align_factor - 1) / align_factor) * align_factor;
687
+ return {kv_left_pos, kv_right_pos};
688
+ }
689
+
690
+ static int64_t calcu_default_tile_size(int64_t cache_size, int64_t head_dim,
691
+ int64_t elem_size,
692
+ int64_t q_buffer_elem_size,
693
+ int64_t logits_buffer_elem_size,
694
+ int64_t output_buffer_elem_size,
695
+ int64_t max_num_q_per_iter,
696
+ int64_t round_size) {
697
+ // For CPU, different from CUDA, Q@K^T results should also be hold in cache,
698
+ // using float32. Intermediate outputs should be float32 to be compatible
699
+ // with AMX Then the cache includes:
700
+ // - Q: q_tile_size * head_dim * q_buffer_elem_size
701
+ // - K, V: 2 * k_tile_size * head_dim * elem_size
702
+ // - Q@K^T: max_num_q_per_iter * k_tile_size * logits_buffer_elem_size
703
+ // - Intermediate outputs: q_tile_size * head_dim * output_buffer_elem_size
704
+ // By default, let tile_size = q_tile_size = k_tile_size. To record
705
+ // is_first_iter states in a static array, require the default tile <= 128 *
706
+ // max_num_q_per_iter
707
+
708
+ int64_t tile_size =
709
+ cache_size / (head_dim * (q_buffer_elem_size + 2 * elem_size +
710
+ output_buffer_elem_size) +
711
+ max_num_q_per_iter * logits_buffer_elem_size);
712
+ tile_size = std::min(tile_size, MaxQTileIterNum * max_num_q_per_iter);
713
+ int64_t rounded_tile_size = (tile_size / round_size) * round_size;
714
+ return std::max(rounded_tile_size, round_size);
715
+ }
716
+
717
+ static int64_t calcu_tile_size_with_constant_q(
718
+ int64_t cache_size, int64_t head_dim, int64_t elem_size,
719
+ int64_t q_buffer_elem_size, int64_t logits_buffer_elem_size,
720
+ int64_t output_buffer_elem_size, int64_t max_num_q_per_iter,
721
+ int64_t round_size, int64_t q_tile_size, bool one_round) {
722
+ // calculate tile_size with known q_tile_size
723
+ // If one_round is True, the outer Q tile loop time is 1, then the K,V will
724
+ // not be included in the cache
725
+ int64_t tile_size;
726
+ if (one_round) {
727
+ tile_size =
728
+ (cache_size - q_tile_size * head_dim *
729
+ (q_buffer_elem_size + output_buffer_elem_size)) /
730
+ (logits_buffer_elem_size * max_num_q_per_iter);
731
+ } else {
732
+ tile_size =
733
+ (cache_size - q_tile_size * head_dim *
734
+ (q_buffer_elem_size + output_buffer_elem_size)) /
735
+ (logits_buffer_elem_size * max_num_q_per_iter +
736
+ 2 * head_dim * elem_size);
737
+ }
738
+ int64_t rounded_tile_size = (tile_size / round_size) * round_size;
739
+ return std::max(rounded_tile_size, round_size);
740
+ }
741
+
742
+ static int64_t get_available_l2_size() {
743
+ static int64_t size = []() {
744
+ long l2_cache_size = sysconf(_SC_LEVEL2_CACHE_SIZE);
745
+ TORCH_CHECK_NE(l2_cache_size, -1);
746
+ return l2_cache_size >> 1; // use 50% of L2 cache
747
+ }();
748
+ return size;
749
+ }
750
+
751
+ private:
752
+ int64_t available_cache_size_;
753
+};
754
+
755
+struct AttentionInput {
756
+ AttentionMetadata* metadata;
757
+ int32_t num_tokens;
758
+ int32_t num_heads;
759
+ int32_t num_kv_heads;
760
+ int32_t block_size;
761
+ void* query;
762
+ int64_t query_num_tokens_stride;
763
+ int64_t query_num_heads_stride;
764
+ int64_t cache_num_blocks_stride;
765
+ int64_t cache_num_kv_heads_stride;
766
+ int64_t blt_num_tokens_stride;
767
+ void* key_cache;
768
+ void* value_cache;
769
+ void* output;
770
+ int32_t* query_start_loc;
771
+ int32_t* seq_lens;
772
+ int32_t* block_table;
773
+ float* alibi_slopes;
774
+ c10::BFloat16* s_aux;
775
+ float scale;
776
+ bool causal;
777
+ int32_t sliding_window_left;
778
+ int32_t sliding_window_right;
779
+ float softcap;
780
+};
781
+
782
+#define DEFINE_CPU_ATTENTION_PARAMS \
783
+ q_buffer_t *__restrict__ q_heads_buffer, \
784
+ kv_cache_t *__restrict__ k_head_cache_ptr, \
785
+ kv_cache_t *__restrict__ v_head_cache_ptr, \
786
+ logits_buffer_t *__restrict__ logits_buffer, \
787
+ float *__restrict__ partial_q_buffer, float *__restrict__ max_buffer, \
788
+ float *__restrict__ sum_buffer, int32_t *__restrict__ block_table, \
789
+ const int32_t kv_tile_start_pos, const int32_t kv_tile_end_pos, \
790
+ const int32_t kv_tile_token_num, \
791
+ const int64_t kv_cache_num_blocks_stride, const int32_t q_head_num, \
792
+ const int32_t q_token_num, const int32_t q_tile_start_pos, \
793
+ const int32_t q_heads_per_kv, const int32_t block_size, \
794
+ const int32_t left_window_size, const int32_t right_window_size, \
795
+ float scale, const float softcap_scale, \
796
+ const float *__restrict__ alibi_slopes, const bool is_first_iter, \
797
+ const bool use_sink, const bool debug_info
798
+
799
+#define CPU_ATTENTION_PARAMS \
800
+ q_heads_buffer, k_head_cache_ptr, v_head_cache_ptr, logits_buffer, \
801
+ partial_q_buffer, max_buffer, sum_buffer, block_table, \
802
+ kv_tile_start_pos, kv_tile_end_pos, kv_tile_token_num, \
803
+ kv_cache_num_blocks_stride, q_head_num, q_token_num, q_tile_start_pos, \
804
+ q_heads_per_kv, block_size, left_window_size, right_window_size, scale, \
805
+ softcap_scale, alibi_slopes, is_first_iter, use_sink, debug_info
806
+
807
+enum class AttentionGemmPhase { QK, PV };
808
+
809
+template <typename T>
810
+struct VecTypeTrait {
811
+ using vec_t = void;
812
+};
813
+
814
+template <>
815
+struct VecTypeTrait<float> {
816
+ using vec_t = vec_op::FP32Vec16;
817
+};
818
+
819
+template <>
820
+struct VecTypeTrait<c10::BFloat16> {
821
+ using vec_t = vec_op::BF16Vec16;
822
+};
823
+
824
+template <>
825
+struct VecTypeTrait<c10::Half> {
826
+ using vec_t = vec_op::FP16Vec16;
827
+};
828
+
829
+template <typename T>
830
+void print_logits(const char* name, T* ptr, int32_t row, int32_t col,
831
+ int32_t stride) {
832
+ std::stringstream ss;
833
+ ss << std::fixed << std::setprecision(5) << name << ": [\n";
834
+ auto* curr_logits_buffer = ptr;
835
+ for (int32_t m = 0; m < row; ++m) {
836
+ for (int32_t n = 0; n < col; ++n) {
837
+ ss << curr_logits_buffer[n] << ", ";
838
+ }
839
+ ss << "\n";
840
+ curr_logits_buffer += stride;
841
+ }
842
+ ss << "]\n";
843
+ std::printf("%s", ss.str().c_str());
844
+}
845
+
846
+template <typename attention_impl_t>
847
+class AttentionMainLoop {
848
+ public:
849
+ using query_t = typename attention_impl_t::query_t;
850
+ using q_buffer_t = typename attention_impl_t::q_buffer_t;
851
+ using kv_cache_t = typename attention_impl_t::kv_cache_t;
852
+ using logits_buffer_t = typename attention_impl_t::logits_buffer_t;
853
+ using partial_output_buffer_t =
854
+ typename attention_impl_t::partial_output_buffer_t;
855
+ using prob_buffer_t = typename attention_impl_t::prob_buffer_t;
856
+
857
+ static constexpr int64_t max_q_head_num_per_iter =
858
+ attention_impl_t::MaxQHeadNumPerIteration;
859
+ static constexpr int64_t blocksize_alignment =
860
+ attention_impl_t::BlockSizeAlignment;
861
+ static constexpr int64_t headdim_alignment =
862
+ attention_impl_t::HeadDimAlignment;
863
+ static constexpr int64_t head_dim = attention_impl_t::HeadDim;
864
+ static constexpr ISA ISAType = attention_impl_t::ISAType;
865
+ static constexpr bool scale_on_logits =
866
+ attention_impl_t::scale_on_logits; // apply scale on logits, otherwise
867
+ // apply scale on q_buffer
868
+
869
+ template <typename tile_gemm_t>
870
+ class Attention {
871
+ public:
872
+ // Args:
873
+ // - q_heads_buffer: [MaxQHeadNumPerIteration, head_dim]
874
+ // - k_head_cache_ptr: [num_blocks, block_size * head_dim]
875
+ // - v_head_cache_ptr: [num_blocks, block_size * head_dim]
876
+ // - logits_buffer: [MaxQHeadNumPerIteration, kv_tile_token_num], store Q@K
877
+ // - logits partial_q_buffer: [MaxQHeadNumPerIteration, head_dim], store
878
+ // partial output
879
+ // - max_buffer: [MaxQHeadNumPerIteration, 1], store max logits
880
+ // - sum_buffer: [MaxQHeadNumPerIteration, 1], store sum of exp
881
+ // - block_table
882
+ // - kv_tile_start_pos: start position of KV cache, aligned to
883
+ // BlockSizeAlignment
884
+ // - kv_tile_end_pos: end position of KV cache, aligned to
885
+ // BlockSizeAlignment
886
+ // - kv_tile_token_num: KV token num, aligned to BlockSizeAlignment
887
+ // - kv_cache_num_blocks_stride
888
+ // - q_head_num: head num of q_tile
889
+ // - q_token_num: token num of q_tile, should be q_head_num /
890
+ // q_heads_per_kv
891
+ // - q_tile_start_pos: start pos of the first token in q_heads_buffer
892
+ // - q_heads_per_kv
893
+ // - block_size
894
+ // - left_window_size
895
+ // - right_window_size
896
+ // - scale
897
+ // - softcap_scale
898
+ // - alibi_slopes
899
+ // - is_first_iter
900
+ // - use_sink
901
+ // - debug_info
902
+ void operator()(DEFINE_CPU_ATTENTION_PARAMS) {
903
+ // k_cache_token_group_stride: stride of K cache when move to next
904
+ // BlockSizeAlignment tokens in a block
905
+ const int64_t k_cache_token_group_stride =
906
+ attention_impl_t::k_cache_token_group_stride(block_size);
907
+ // v_cache_token_group_stride: stride of V cache when move to next
908
+ // BlockSizeAlignment tokens in a block
909
+ const int64_t v_cache_token_group_stride =
910
+ attention_impl_t::v_cache_token_group_stride(block_size);
911
+ // v_cache_head_group_stride: stride of V cache when move to next
912
+ // HeadDimAlignment head dims in a block
913
+ const int64_t v_cache_head_group_stride =
914
+ attention_impl_t::v_cache_head_group_stride(block_size);
915
+ const int32_t token_group_num = kv_tile_token_num / blocksize_alignment;
916
+ const int32_t token_group_num_per_block =
917
+ block_size / blocksize_alignment;
918
+ const int32_t start_block_idx = kv_tile_start_pos / block_size;
919
+ const int32_t start_block_offset = kv_tile_start_pos % block_size;
920
+ const int32_t start_block_group_offset =
921
+ start_block_offset / blocksize_alignment;
922
+ const int32_t end_block_idx =
923
+ (kv_tile_start_pos + kv_tile_token_num - 1) / block_size + 1;
924
+
925
+ // compute Q@K logits
926
+ {
927
+ int32_t curr_group_offset =
928
+ start_block_group_offset * k_cache_token_group_stride;
929
+ int32_t curr_group_num_in_block =
930
+ token_group_num_per_block - start_block_group_offset;
931
+ int32_t remaining_group_num = token_group_num;
932
+ logits_buffer_t* curr_logits_buffer = logits_buffer;
933
+ for (int32_t block_idx = start_block_idx; block_idx < end_block_idx;
934
+ ++block_idx) {
935
+ int32_t physical_block_idx = block_table[block_idx];
936
+ kv_cache_t* k_cache_block_ptr =
937
+ k_head_cache_ptr +
938
+ physical_block_idx * kv_cache_num_blocks_stride +
939
+ curr_group_offset;
940
+ curr_group_num_in_block =
941
+ std::min(remaining_group_num, curr_group_num_in_block);
942
+
943
+ for (int32_t block_group_idx = 0;
944
+ block_group_idx < curr_group_num_in_block; ++block_group_idx) {
945
+ // logits_tile = q_tile @ k_tile, [MaxQHeadNumPerIteration,
946
+ // BlockSizeAlignment] = [MaxQHeadNumPerIteration, head_dim] @
947
+ // [head_dim, BlockSizeAlignment]
948
+
949
+ // By default, logits_buffer, q_buffer and k_cache are row-major,
950
+ // but may be packed by ISA implementation.
951
+ tile_gemm_t::template gemm<AttentionGemmPhase::QK, head_dim>(
952
+ q_head_num, q_heads_buffer, k_cache_block_ptr,
953
+ curr_logits_buffer, head_dim, block_size, kv_tile_token_num,
954
+ block_size, head_dim, false);
955
+
956
+ if constexpr (scale_on_logits) {
957
+ float* __restrict__ scale_curr_logits_buffer = curr_logits_buffer;
958
+ vec_op::FP32Vec16 scale_vec(scale);
959
+ for (int32_t i = 0; i < q_head_num; ++i) {
960
+ static_assert(blocksize_alignment % 16 == 0);
961
+ constexpr int32_t vec_num = blocksize_alignment / 16;
962
+ vec_op::unroll_loop<int32_t, vec_num>([&](int32_t vec_idx) {
963
+ vec_op::FP32Vec16 vec(scale_curr_logits_buffer +
964
+ vec_idx * 16);
965
+ vec = vec * scale_vec;
966
+ vec.save(scale_curr_logits_buffer + vec_idx * 16);
967
+ });
968
+ scale_curr_logits_buffer += kv_tile_token_num;
969
+ }
970
+ }
971
+
972
+ // Move buffer ptrs
973
+ k_cache_block_ptr += k_cache_token_group_stride;
974
+ curr_logits_buffer += blocksize_alignment;
975
+ }
976
+
977
+ // Update
978
+ remaining_group_num -= curr_group_num_in_block;
979
+ curr_group_offset = 0;
980
+ curr_group_num_in_block = token_group_num_per_block;
981
+ }
982
+ }
983
+
984
+ // process logits
985
+ {
986
+ // if (debug_info){
987
+ // print_logits("raw logits", logits_buffer, q_head_num,
988
+ // kv_tile_token_num, kv_tile_token_num);
989
+ // }
990
+
991
+ if (softcap_scale != 0.0f) {
992
+ apply_softcap(logits_buffer, kv_tile_token_num, q_head_num,
993
+ kv_tile_token_num, softcap_scale);
994
+ // print_logits("softcap raw logits", logits_buffer, q_head_num,
995
+ // kv_tile_token_num, kv_tile_token_num);
996
+ }
997
+
998
+ if (alibi_slopes != nullptr) {
999
+ apply_alibi_slopes(logits_buffer, alibi_slopes, kv_tile_token_num,
1000
+ q_tile_start_pos, kv_tile_start_pos, q_token_num,
1001
+ kv_tile_token_num, q_heads_per_kv);
1002
+
1003
+ // print_logits("alibi raw logits", logits_buffer, q_head_num,
1004
+ // kv_tile_token_num, kv_tile_token_num);
1005
+ }
1006
+
1007
+ apply_mask(logits_buffer, kv_tile_token_num, q_tile_start_pos,
1008
+ kv_tile_start_pos, kv_tile_end_pos, q_token_num,
1009
+ q_heads_per_kv, left_window_size, right_window_size);
1010
+
1011
+ // if (debug_info){
1012
+ // print_logits("masked logits", logits_buffer, q_head_num,
1013
+ // kv_tile_token_num, kv_tile_token_num);
1014
+ // print_logits("old_max", max_buffer, 1, q_head_num, q_head_num);
1015
+ // print_logits("old_sum", sum_buffer, 1, q_head_num, q_head_num);
1016
+ // }
1017
+
1018
+ apply_softmax(logits_buffer, partial_q_buffer, max_buffer, sum_buffer,
1019
+ kv_tile_token_num, q_head_num, kv_tile_token_num,
1020
+ is_first_iter, use_sink);
1021
+
1022
+ // if (debug_info){
1023
+ // print_logits("softmax logits",
1024
+ // reinterpret_cast<prob_buffer_t*>(logits_buffer), q_head_num,
1025
+ // kv_tile_token_num, kv_tile_token_num * sizeof(logits_buffer_t) /
1026
+ // sizeof(prob_buffer_t));
1027
+ // print_logits("new_max", max_buffer, 1, q_head_num, q_head_num);
1028
+ // print_logits("new_sum", sum_buffer, 1, q_head_num, q_head_num);
1029
+ // }
1030
+ }
1031
+
1032
+ // compute P@V
1033
+ {
1034
+ int32_t curr_group_offset =
1035
+ start_block_group_offset * v_cache_token_group_stride;
1036
+ int32_t curr_group_num_in_block =
1037
+ token_group_num_per_block - start_block_group_offset;
1038
+ int32_t remaining_group_num = token_group_num;
1039
+ int32_t head_dim_group_num = head_dim / headdim_alignment;
1040
+ prob_buffer_t* curr_prob_buffer =
1041
+ reinterpret_cast<prob_buffer_t*>(logits_buffer);
1042
+ int64_t prob_buffer_stride =
1043
+ kv_tile_token_num *
1044
+ (sizeof(logits_buffer_t) / sizeof(prob_buffer_t));
1045
+ partial_output_buffer_t* curr_partial_q_buffer = partial_q_buffer;
1046
+ bool accum_c = !is_first_iter;
1047
+ for (int32_t block_idx = start_block_idx; block_idx < end_block_idx;
1048
+ ++block_idx) {
1049
+ int32_t physical_block_idx = block_table[block_idx];
1050
+ kv_cache_t* v_cache_block_ptr =
1051
+ v_head_cache_ptr +
1052
+ physical_block_idx * kv_cache_num_blocks_stride +
1053
+ curr_group_offset;
1054
+ curr_group_num_in_block =
1055
+ std::min(remaining_group_num, curr_group_num_in_block);
1056
+ int32_t curr_token_num =
1057
+ curr_group_num_in_block * blocksize_alignment;
1058
+
1059
+ for (int32_t head_dim_group_idx = 0;
1060
+ head_dim_group_idx < head_dim_group_num; ++head_dim_group_idx) {
1061
+ // output_tile = p_tile @ v_tile, [MaxQHeadNumPerIteration,
1062
+ // HeadDimAlignment] = [MaxQHeadNumPerIteration, block_size] @
1063
+ // [block_size, HeadDimAlignment]
1064
+ tile_gemm_t::template gemm<AttentionGemmPhase::PV, -1>(
1065
+ q_head_num, curr_prob_buffer, v_cache_block_ptr,
1066
+ curr_partial_q_buffer, prob_buffer_stride, head_dim, head_dim,
1067
+ block_size, curr_token_num, accum_c);
1068
+
1069
+ // Update
1070
+ curr_partial_q_buffer += headdim_alignment;
1071
+ v_cache_block_ptr += v_cache_head_group_stride;
1072
+ }
1073
+
1074
+ // Update
1075
+ remaining_group_num -= curr_group_num_in_block;
1076
+ curr_group_offset = 0;
1077
+ curr_group_num_in_block = token_group_num_per_block;
1078
+ curr_prob_buffer += curr_token_num;
1079
+ curr_partial_q_buffer = partial_q_buffer;
1080
+ accum_c = true;
1081
+ }
1082
+ }
1083
+ // if (debug_info) {
1084
+ // print_logits("output", partial_q_buffer, q_head_num, head_dim,
1085
+ // head_dim);
1086
+ // }
1087
+ }
1088
+
1089
+ void apply_mask(logits_buffer_t* __restrict__ logits_buffer,
1090
+ const int64_t logits_buffer_stride,
1091
+ const int32_t q_tile_start_pos,
1092
+ const int32_t kv_tile_start_pos,
1093
+ const int32_t kv_tile_end_pos, const int32_t q_token_num,
1094
+ const int32_t q_heads_per_kv,
1095
+ const int32_t sliding_window_left,
1096
+ const int32_t sliding_window_right) {
1097
+ // Apply mask
1098
+ constexpr logits_buffer_t neg_inf =
1099
+ -std::numeric_limits<logits_buffer_t>::infinity();
1100
+ logits_buffer_t* __restrict__ curr_logits_buffer = logits_buffer;
1101
+ int32_t curr_token_pos = q_tile_start_pos;
1102
+ for (int32_t token_idx = 0; token_idx < q_token_num; ++token_idx) {
1103
+ int32_t left_kv_pos = [&]() {
1104
+ int32_t pos = kv_tile_start_pos;
1105
+ if (sliding_window_left != -1) {
1106
+ pos = std::max(pos, curr_token_pos - sliding_window_left);
1107
+ }
1108
+ return pos;
1109
+ }();
1110
+
1111
+ int32_t right_kv_pos = [&]() {
1112
+ int32_t pos = kv_tile_end_pos;
1113
+ if (sliding_window_right != -1) {
1114
+ pos = std::min(pos,
1115
+ std::max(kv_tile_start_pos,
1116
+ curr_token_pos + sliding_window_right + 1));
1117
+ }
1118
+ return pos;
1119
+ }();
1120
+
1121
+ int32_t left_invalid_token_num = left_kv_pos - kv_tile_start_pos;
1122
+ int32_t right_invalid_token_num = kv_tile_end_pos - right_kv_pos;
1123
+ for (int32_t head_idx = 0; head_idx < q_heads_per_kv; ++head_idx) {
1124
+ logits_buffer_t* __restrict__ curr_logits_buffer_tail =
1125
+ curr_logits_buffer + right_kv_pos - kv_tile_start_pos;
1126
+ for (int32_t i = 0; i < left_invalid_token_num; ++i) {
1127
+ curr_logits_buffer[i] = neg_inf;
1128
+ }
1129
+ for (int32_t i = 0; i < right_invalid_token_num; ++i) {
1130
+ curr_logits_buffer_tail[i] = neg_inf;
1131
+ }
1132
+
1133
+ curr_logits_buffer += logits_buffer_stride;
1134
+ }
1135
+
1136
+ ++curr_token_pos;
1137
+ }
1138
+ }
1139
+
1140
+ void apply_softmax(logits_buffer_t* __restrict__ logits_buffer,
1141
+ float* __restrict__ partial_q_buffer,
1142
+ float* __restrict__ max_buffer,
1143
+ float* __restrict__ sum_buffer,
1144
+ const int64_t logits_buffer_stride, int32_t q_head_num,
1145
+ int32_t kv_tile_token_num, bool is_first_iter,
1146
+ bool use_sink) {
1147
+#ifdef DEFINE_FAST_EXP
1148
+ DEFINE_FAST_EXP
1149
+#endif
1150
+ using prob_buffer_vec_t = typename VecTypeTrait<prob_buffer_t>::vec_t;
1151
+ static_assert(sizeof(prob_buffer_t) <= sizeof(logits_buffer_t));
1152
+
1153
+ logits_buffer_t* __restrict__ curr_logits_buffer = logits_buffer;
1154
+ float* __restrict__ curr_partial_q_buffer = partial_q_buffer;
1155
+ const int32_t vec_num = kv_tile_token_num / 16;
1156
+ const int32_t head_vec_num = head_dim / 16;
1157
+ for (int32_t i = 0; i < q_head_num; ++i) {
1158
+ float init_max_val = max_buffer[i];
1159
+ float init_sum_val = sum_buffer[i];
1160
+
1161
+ // apply scale and compute max
1162
+ vec_op::FP32Vec16 max_vec(init_max_val);
1163
+ {
1164
+ logits_buffer_t* __restrict__ curr_logits_buffer_iter =
1165
+ curr_logits_buffer;
1166
+ for (int32_t j = 0; j < vec_num; ++j) {
1167
+ vec_op::FP32Vec16 vec(curr_logits_buffer_iter);
1168
+ max_vec = vec.max(max_vec);
1169
+
1170
+ curr_logits_buffer_iter += 16;
1171
+ }
1172
+ }
1173
+ float new_max_val = max_vec.reduce_max();
1174
+ float rescale_factor = init_max_val - new_max_val;
1175
+
1176
+ // use same rescale threshold with FA4.
1177
+ // https://github.com/Dao-AILab/flash-attention/blob/1b8e1e641c6a179be9a0538b7f40fd595050b735/flash_attn/cute/flash_fwd_sm100.py#L1271
1178
+ bool need_rescale = rescale_factor < -8.0;
1179
+ if (!need_rescale) {
1180
+ new_max_val = init_max_val;
1181
+ } else {
1182
+ max_buffer[i] = new_max_val;
1183
+ }
1184
+
1185
+ // sub max, compute exp and sum
1186
+ max_vec = vec_op::FP32Vec16(new_max_val);
1187
+ vec_op::FP32Vec16 sum_vec(0.0);
1188
+ {
1189
+ logits_buffer_t* __restrict__ curr_logits_buffer_iter =
1190
+ curr_logits_buffer;
1191
+ prob_buffer_t* __restrict__ curr_prob_buffer_iter =
1192
+ reinterpret_cast<prob_buffer_t*>(curr_logits_buffer);
1193
+ for (int32_t j = 0; j < vec_num; ++j) {
1194
+ vec_op::FP32Vec16 vec(curr_logits_buffer_iter);
1195
+ vec = vec - max_vec;
1196
+
1197
+ // compute exp
1198
+#ifdef DEFINE_FAST_EXP
1199
+ vec = fast_exp(vec);
1200
+ prob_buffer_vec_t output_vec(vec);
1201
+ output_vec.save(curr_prob_buffer_iter);
1202
+#else
1203
+ vec.save(curr_logits_buffer_iter);
1204
+ for (int32_t k = 0; k < 16; ++k) {
1205
+ curr_logits_buffer_iter[k] = std::exp(curr_logits_buffer_iter[k]);
1206
+ }
1207
+ vec = vec_op::FP32Vec16(curr_logits_buffer_iter);
1208
+#endif
1209
+
1210
+ sum_vec = sum_vec + vec;
1211
+
1212
+ curr_logits_buffer_iter += 16;
1213
+ curr_prob_buffer_iter += 16;
1214
+ }
1215
+ }
1216
+ float new_sum_val = sum_vec.reduce_sum();
1217
+
1218
+ // rescale sum and partial outputs
1219
+ if (need_rescale) {
1220
+ // compute rescale factor
1221
+#ifdef DEFINE_FAST_EXP
1222
+ vec_op::FP32Vec16 rescale_factor_vec(rescale_factor);
1223
+ rescale_factor_vec = fast_exp(rescale_factor_vec);
1224
+ rescale_factor = rescale_factor_vec.get_last_elem();
1225
+#else
1226
+ rescale_factor = std::exp(rescale_factor);
1227
+ vec_op::FP32Vec16 rescale_factor_vec(rescale_factor);
1228
+#endif
1229
+
1230
+ // rescale sum
1231
+ new_sum_val += rescale_factor * init_sum_val;
1232
+
1233
+ // rescale output
1234
+ if (!is_first_iter) {
1235
+ float* __restrict__ curr_partial_q_buffer_iter =
1236
+ curr_partial_q_buffer;
1237
+ for (int32_t j = 0; j < head_vec_num; ++j) {
1238
+ vec_op::FP32Vec16 vec(curr_partial_q_buffer_iter);
1239
+ vec = vec * rescale_factor_vec;
1240
+ vec.save(curr_partial_q_buffer_iter);
1241
+
1242
+ curr_partial_q_buffer_iter += 16;
1243
+ }
1244
+ }
1245
+ } else {
1246
+ new_sum_val += init_sum_val;
1247
+ }
1248
+
1249
+ sum_buffer[i] = new_sum_val;
1250
+
1251
+ curr_logits_buffer += logits_buffer_stride;
1252
+ curr_partial_q_buffer += head_dim;
1253
+ }
1254
+ }
1255
+
1256
+ void apply_softcap(logits_buffer_t* __restrict__ logits_buffer,
1257
+ const int64_t logits_buffer_stride, int32_t q_head_num,
1258
+ int32_t kv_tile_token_num, float softcap_scale) {
1259
+#ifdef DEFINE_FAST_EXP
1260
+ DEFINE_FAST_EXP
1261
+#endif
1262
+ float inv_softcap_scale = 1.0 / softcap_scale;
1263
+ vec_op::FP32Vec16 softcap_scale_vec(softcap_scale);
1264
+ vec_op::FP32Vec16 inv_softcap_scale_vec(inv_softcap_scale);
1265
+ vec_op::FP32Vec16 ones_vec(1.0);
1266
+ logits_buffer_t* __restrict__ curr_logits_buffer = logits_buffer;
1267
+ const int32_t vec_num = kv_tile_token_num / 16;
1268
+ for (int32_t i = 0; i < q_head_num; ++i) {
1269
+ logits_buffer_t* __restrict__ curr_logits_buffer_iter =
1270
+ curr_logits_buffer;
1271
+ for (int32_t j = 0; j < vec_num; ++j) {
1272
+ vec_op::FP32Vec16 vec(curr_logits_buffer_iter);
1273
+ vec = vec * inv_softcap_scale_vec;
1274
+
1275
+#ifdef DEFINE_FAST_EXP
1276
+ vec = fast_exp(vec);
1277
+ vec_op::FP32Vec16 inv_vec = ones_vec / vec;
1278
+ vec = (vec - inv_vec) / (vec + inv_vec);
1279
+#else
1280
+ vec.save(curr_logits_buffer_iter);
1281
+ for (int k = 0; k < 16; ++k) {
1282
+ curr_logits_buffer_iter[k] = std::tanh(curr_logits_buffer_iter[k]);
1283
+ }
1284
+ vec = vec_op::FP32Vec16(curr_logits_buffer_iter);
1285
+#endif
1286
+ vec = vec * softcap_scale_vec;
1287
+ vec.save(curr_logits_buffer_iter);
1288
+
1289
+ curr_logits_buffer_iter += 16;
1290
+ }
1291
+
1292
+ curr_logits_buffer += logits_buffer_stride;
1293
+ }
1294
+ }
1295
+
1296
+ void apply_alibi_slopes(logits_buffer_t* __restrict__ logits_buffer,
1297
+ const float* __restrict__ alibi_slopes,
1298
+ const int64_t logits_buffer_stride,
1299
+ const int32_t q_tile_start_pos,
1300
+ const int32_t kv_tile_start_pos,
1301
+ const int32_t q_token_num,
1302
+ const int32_t kv_tile_token_num,
1303
+ const int32_t q_heads_per_kv) {
1304
+ alignas(64) constexpr float initial_arange_vals[16] = {
1305
+ 0.0f, 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f,
1306
+ 8.0f, 9.0f, 10.0f, 11.0f, 12.0f, 13.0f, 14.0f, 15.0f};
1307
+ const int32_t vec_num = kv_tile_token_num / 16;
1308
+
1309
+ vec_op::FP32Vec16 initial_arange_vals_vec(initial_arange_vals);
1310
+ initial_arange_vals_vec =
1311
+ initial_arange_vals_vec + vec_op::FP32Vec16((float)kv_tile_start_pos);
1312
+ vec_op::FP32Vec16 pos_offset_vec(16.0);
1313
+ logits_buffer_t* __restrict__ curr_logits_buffer = logits_buffer;
1314
+ for (int32_t i = 0; i < q_token_num; ++i) {
1315
+ vec_op::FP32Vec16 curr_q_pos_vec((float)(i + q_tile_start_pos));
1316
+ for (int32_t j = 0; j < q_heads_per_kv; ++j) {
1317
+ vec_op::FP32Vec16 alibi_scale_vec(alibi_slopes[j]);
1318
+ vec_op::FP32Vec16 curr_kv_pos_vec(initial_arange_vals_vec);
1319
+ logits_buffer_t* __restrict__ curr_logits_buffer_iter =
1320
+ curr_logits_buffer;
1321
+ for (int32_t k = 0; k < vec_num; ++k) {
1322
+ vec_op::FP32Vec16 alibi_bias_vec =
1323
+ alibi_scale_vec * (curr_kv_pos_vec - curr_q_pos_vec);
1324
+ vec_op::FP32Vec16 vec(curr_logits_buffer_iter);
1325
+ vec = vec + alibi_bias_vec;
1326
+
1327
+ vec.save(curr_logits_buffer_iter);
1328
+
1329
+ curr_kv_pos_vec = curr_kv_pos_vec + pos_offset_vec;
1330
+ curr_logits_buffer_iter += 16;
1331
+ }
1332
+ curr_logits_buffer += logits_buffer_stride;
1333
+ }
1334
+ }
1335
+ }
1336
+ };
1337
+
1338
+ public:
1339
+ void operator()(const AttentionInput* input) {
1340
+ const int thread_num = omp_get_max_threads();
1341
+ TORCH_CHECK_EQ(input->metadata->thread_num, thread_num);
1342
+ std::atomic<int32_t> guard_counter(0);
1343
+ std::atomic<int32_t>* guard_counter_ptr = &guard_counter;
1344
+
1345
+#pragma omp parallel for schedule(static, 1)
1346
+ for (int thread_id = 0; thread_id < thread_num; ++thread_id) {
1347
+ AttentionMetadata& metadata = *input->metadata;
1348
+ if (metadata.workitem_group_num == 0) {
1349
+ continue;
1350
+ }
1351
+
1352
+ attention_impl_t attn_impl;
1353
+
1354
+ // general information
1355
+ const int32_t q_head_num = input->num_heads;
1356
+ const int32_t kv_head_num = input->num_kv_heads;
1357
+ const int32_t q_heads_per_kv = q_head_num / kv_head_num;
1358
+ const bool use_gqa =
1359
+ (max_q_head_num_per_iter % q_heads_per_kv == 0) ? true : false;
1360
+ const int32_t actual_kv_head_num = use_gqa ? kv_head_num : q_head_num;
1361
+ const int32_t actual_q_heads_per_kv = use_gqa ? q_heads_per_kv : 1;
1362
+ TORCH_CHECK_LE(actual_q_heads_per_kv, max_q_head_num_per_iter);
1363
+ const int32_t max_q_token_num_per_iter =
1364
+ max_q_head_num_per_iter / actual_q_heads_per_kv;
1365
+ const int64_t q_token_num_stride = input->query_num_tokens_stride;
1366
+ const int64_t q_head_num_stride = input->query_num_heads_stride;
1367
+ const int64_t kv_cache_head_num_stride = input->cache_num_kv_heads_stride;
1368
+ const int64_t kv_cache_block_num_stride = input->cache_num_blocks_stride;
1369
+ const int32_t sliding_window_left = input->sliding_window_left;
1370
+ const int32_t sliding_window_right = input->sliding_window_right;
1371
+ const int32_t block_size = input->block_size;
1372
+ const float scale = input->scale;
1373
+ const float softcap_scale = input->softcap;
1374
+ const float* alibi_slopes = input->alibi_slopes;
1375
+ const c10::BFloat16* s_aux = input->s_aux;
1376
+
1377
+ const bool casual = input->causal;
1378
+ int32_t* const block_table = input->block_table;
1379
+ const int64_t block_table_stride = input->blt_num_tokens_stride;
1380
+
1381
+ // init buffers
1382
+ void* scratchpad_ptr =
1383
+ DNNLScratchPadManager::get_dnnl_scratchpad_manager()
1384
+ ->get_data<void>();
1385
+ AttentionScratchPad buffer_manager(thread_id, metadata, scratchpad_ptr);
1386
+
1387
+ const int32_t total_reduction_split_num = metadata.reduction_split_num;
1388
+ if (metadata.reduction_split_num > 0) {
1389
+ // reset split flag
1390
+ for (int32_t head_idx = thread_id; head_idx < actual_kv_head_num;
1391
+ head_idx += thread_num) {
1392
+ buffer_manager.update(head_idx, total_reduction_split_num, head_dim,
1393
+ 0, sizeof(partial_output_buffer_t));
1394
+ volatile bool* __restrict__ curr_flag_ptr =
1395
+ buffer_manager.get_reduce_flag_buffer();
1396
+ for (int32_t split_idx = 0; split_idx < total_reduction_split_num;
1397
+ ++split_idx) {
1398
+ curr_flag_ptr[split_idx] = false;
1399
+ }
1400
+ }
1401
+ }
1402
+
1403
+ const int64_t available_cache_size =
1404
+ AttentionScheduler::get_available_l2_size();
1405
+ const int32_t default_tile_size =
1406
+ AttentionScheduler::calcu_default_tile_size(
1407
+ available_cache_size, head_dim, sizeof(kv_cache_t),
1408
+ sizeof(q_buffer_t), sizeof(logits_buffer_t),
1409
+ sizeof(partial_output_buffer_t), max_q_head_num_per_iter,
1410
+ max_q_head_num_per_iter);
1411
+ const int32_t default_q_tile_token_num =
1412
+ default_tile_size / actual_q_heads_per_kv;
1413
+
1414
+ AttentionWorkItemGroup* const workitem_groups =
1415
+ metadata.workitem_groups_ptr;
1416
+ const int32_t* cu_workitem_num_per_thread =
1417
+ metadata.cu_workitem_num_per_thread;
1418
+ ReductionWorkItemGroup* const reduction_items =
1419
+ metadata.reduction_items_ptr;
1420
+
1421
+ const int32_t effective_thread_num = metadata.effective_thread_num;
1422
+ const int32_t reduction_item_num = metadata.reduction_item_num;
1423
+ const int32_t split_kv_q_token_num_threshold =
1424
+ metadata.split_kv_q_token_num_threshold;
1425
+ const int32_t workitem_groups_counter_num =
1426
+ actual_kv_head_num * effective_thread_num;
1427
+ const int32_t reduction_items_counter_num =
1428
+ actual_kv_head_num * reduction_item_num;
1429
+ const int32_t total_counter_num =
1430
+ workitem_groups_counter_num + reduction_items_counter_num;
1431
+
1432
+ if (metadata.reduction_split_num > 0) {
1433
+ ++(*guard_counter_ptr);
1434
+ while (guard_counter_ptr->load() != thread_num) {
1435
+#ifdef FAST_SPINNING
1436
+ FAST_SPINNING
1437
+#else
1438
+ std::this_thread::yield();
1439
+#endif
1440
+ }
1441
+ }
1442
+
1443
+ // main loop
1444
+ for (;;) {
1445
+ int64_t task_idx = metadata.acquire_counter();
1446
+
1447
+ if (task_idx >= total_counter_num) {
1448
+ // no more tasks, leave loop
1449
+ break;
1450
+ }
1451
+
1452
+ if (task_idx < workitem_groups_counter_num) {
1453
+ // attention task
1454
+ // map task_idx to workitem_groups
1455
+ const int32_t kv_head_idx = task_idx / effective_thread_num;
1456
+ const int32_t thread_offset = task_idx % effective_thread_num;
1457
+ AttentionWorkItemGroup* const curr_workitem_groups =
1458
+ workitem_groups + cu_workitem_num_per_thread[thread_offset];
1459
+ const int32_t curr_workitem_groups_num =
1460
+ cu_workitem_num_per_thread[thread_offset + 1] -
1461
+ cu_workitem_num_per_thread[thread_offset];
1462
+
1463
+ const int32_t q_head_start_idx = kv_head_idx * actual_q_heads_per_kv;
1464
+
1465
+ for (int32_t workitem_group_idx = 0;
1466
+ workitem_group_idx < curr_workitem_groups_num;
1467
+ ++workitem_group_idx) {
1468
+ AttentionWorkItemGroup* const current_workitem_group =
1469
+ &curr_workitem_groups[workitem_group_idx];
1470
+
1471
+ const int32_t current_group_idx = current_workitem_group->req_id;
1472
+ const int32_t kv_start_pos =
1473
+ current_workitem_group->kv_split_pos_start;
1474
+ const int32_t kv_end_pos = current_workitem_group->kv_split_pos_end;
1475
+ const int32_t curr_spilt_id = current_workitem_group->split_id;
1476
+ const int32_t q_token_id_start =
1477
+ current_workitem_group->q_token_id_start;
1478
+ const int32_t q_token_num = current_workitem_group->q_token_num;
1479
+
1480
+ // taskgroup general information
1481
+ const int32_t q_end = input->query_start_loc[current_group_idx + 1];
1482
+ const int32_t q_start = input->query_start_loc[current_group_idx];
1483
+ const int32_t seq_len = input->seq_lens[current_group_idx];
1484
+ const int32_t q_start_pos =
1485
+ (casual ? seq_len - (q_end - q_start) : 0);
1486
+ const int32_t block_num = (seq_len + block_size - 1) / block_size;
1487
+ // Only apply sink for the first KV split
1488
+ bool use_sink = (s_aux != nullptr &&
1489
+ current_workitem_group->local_split_id == 0);
1490
+
1491
+ for (int32_t q_token_offset = 0; q_token_offset < q_token_num;
1492
+ q_token_offset += default_q_tile_token_num) {
1493
+ bool first_iter_flag[AttentionScheduler::MaxQTileIterNum];
1494
+ for (int32_t i = 0; i < AttentionScheduler::MaxQTileIterNum;
1495
+ ++i) {
1496
+ first_iter_flag[i] = true;
1497
+ }
1498
+
1499
+ const int32_t q_token_start_idx =
1500
+ q_start + q_token_offset + q_token_id_start;
1501
+ const int32_t actual_q_token_num = std::min(
1502
+ default_q_tile_token_num, q_token_num - q_token_offset);
1503
+ const int32_t q_head_tile_size =
1504
+ actual_q_token_num * actual_q_heads_per_kv;
1505
+ const int32_t rounded_q_head_tile_size =
1506
+ ((q_head_tile_size + max_q_head_num_per_iter - 1) /
1507
+ max_q_head_num_per_iter) *
1508
+ max_q_head_num_per_iter;
1509
+ const int32_t kv_tile_size =
1510
+ AttentionScheduler::calcu_tile_size_with_constant_q(
1511
+ available_cache_size, head_dim, sizeof(kv_cache_t),
1512
+ sizeof(q_buffer_t), sizeof(logits_buffer_t),
1513
+ sizeof(partial_output_buffer_t), max_q_head_num_per_iter,
1514
+ blocksize_alignment, rounded_q_head_tile_size,
1515
+ rounded_q_head_tile_size <= max_q_head_num_per_iter);
1516
+
1517
+ // update buffers
1518
+ buffer_manager.update(
1519
+ head_dim, sizeof(q_buffer_t), sizeof(logits_buffer_t),
1520
+ sizeof(partial_output_buffer_t), max_q_head_num_per_iter,
1521
+ rounded_q_head_tile_size, kv_tile_size);
1522
+ q_buffer_t* q_buffer = buffer_manager.get_q_buffer<q_buffer_t>();
1523
+ float* logits_buffer = buffer_manager.get_logits_buffer();
1524
+ float* partial_q_buffer = buffer_manager.get_output_buffer();
1525
+ float* max_buffer = buffer_manager.get_max_buffer();
1526
+ float* sum_buffer = buffer_manager.get_sum_buffer();
1527
+
1528
+ const int32_t q_tile_start_pos =
1529
+ q_start_pos + q_token_offset + q_token_id_start;
1530
+ const int32_t q_tile_end_pos =
1531
+ q_tile_start_pos + actual_q_token_num;
1532
+ const auto [kv_tile_start_pos, kv_tile_end_pos] =
1533
+ AttentionScheduler::calcu_kv_tile_pos(
1534
+ kv_start_pos, kv_end_pos, q_tile_start_pos,
1535
+ q_tile_end_pos, sliding_window_left,
1536
+ sliding_window_right);
1537
+ const auto [rounded_kv_tile_start_pos, rounded_kv_tile_end_pos] =
1538
+ AttentionScheduler::align_kv_tile_pos(
1539
+ kv_tile_start_pos, kv_tile_end_pos, blocksize_alignment);
1540
+
1541
+ int32_t curr_kv_head_idx =
1542
+ use_gqa ? kv_head_idx
1543
+ : (kv_head_idx /
1544
+ q_heads_per_kv); // for GQA disabled case
1545
+
1546
+ // std::printf("thread_id: %d, req_id: %d, q_token_start: %d,
1547
+ // q_token_end: %d, q_head_start: %d, q_head_end: %d, kv_head_idx:
1548
+ // %d, kv_pos_start: %d, kv_pos_end: %d\n",
1549
+ // thread_id, current_group_idx,
1550
+ // q_token_start_idx, q_token_start_idx +
1551
+ // actual_q_token_num, q_head_start_idx,
1552
+ // q_head_start_idx + actual_q_heads_per_kv,
1553
+ // curr_kv_head_idx, kv_tile_start_pos,
1554
+ // kv_tile_end_pos);
1555
+
1556
+ // move buffers
1557
+ kv_cache_t* curr_k_cache =
1558
+ reinterpret_cast<kv_cache_t*>(input->key_cache) +
1559
+ curr_kv_head_idx * kv_cache_head_num_stride;
1560
+ kv_cache_t* curr_v_cache =
1561
+ reinterpret_cast<kv_cache_t*>(input->value_cache) +
1562
+ curr_kv_head_idx * kv_cache_head_num_stride;
1563
+ query_t* const q_tile_ptr =
1564
+ reinterpret_cast<query_t*>(input->query) +
1565
+ q_token_start_idx * q_token_num_stride +
1566
+ q_head_start_idx * q_head_num_stride;
1567
+ size_t output_buffer_offset =
1568
+ q_token_start_idx * q_head_num * head_dim +
1569
+ q_head_start_idx * head_dim;
1570
+ int32_t* curr_block_table =
1571
+ block_table + current_group_idx * block_table_stride;
1572
+ const float* curr_alibi_slopes =
1573
+ (alibi_slopes != nullptr ? alibi_slopes + q_head_start_idx
1574
+ : nullptr);
1575
+ const c10::BFloat16* curr_s_aux =
1576
+ (s_aux != nullptr ? s_aux + q_head_start_idx : nullptr);
1577
+
1578
+ // copy the Q tile to q_buffer, the logical layout of q_buffer is
1579
+ // [actual_q_token_num, actual_q_heads_per_kv, head_dim]
1580
+ {
1581
+ attn_impl.copy_q_heads_tile(
1582
+ q_tile_ptr, q_buffer, actual_q_token_num,
1583
+ actual_q_heads_per_kv, q_token_num_stride,
1584
+ q_head_num_stride, scale);
1585
+ }
1586
+
1587
+ if (use_sink) {
1588
+ alignas(64) float s_aux_fp32[16];
1589
+ vec_op::BF16Vec16 vec_bf16(curr_s_aux);
1590
+ vec_op::FP32Vec16 vec_fp32(vec_bf16);
1591
+ vec_fp32.save(s_aux_fp32);
1592
+
1593
+ float* __restrict__ curr_sum_buffer = sum_buffer;
1594
+ float* __restrict__ curr_max_buffer = max_buffer;
1595
+ for (int32_t token_idx = 0; token_idx < actual_q_token_num;
1596
+ ++token_idx) {
1597
+ for (int32_t head_idx = 0; head_idx < actual_q_heads_per_kv;
1598
+ ++head_idx) {
1599
+ curr_sum_buffer[head_idx] = 1.0f;
1600
+ curr_max_buffer[head_idx] = s_aux_fp32[head_idx];
1601
+ }
1602
+
1603
+ curr_sum_buffer += actual_q_heads_per_kv;
1604
+ curr_max_buffer += actual_q_heads_per_kv;
1605
+ }
1606
+ } else {
1607
+ float* __restrict__ curr_sum_buffer = sum_buffer;
1608
+ float* __restrict__ curr_max_buffer = max_buffer;
1609
+ for (int32_t token_idx = 0; token_idx < actual_q_token_num;
1610
+ ++token_idx) {
1611
+ for (int32_t head_idx = 0; head_idx < actual_q_heads_per_kv;
1612
+ ++head_idx) {
1613
+ curr_sum_buffer[head_idx] = 0.0f;
1614
+ curr_max_buffer[head_idx] =
1615
+ std::numeric_limits<float>::lowest();
1616
+ }
1617
+
1618
+ curr_sum_buffer += actual_q_heads_per_kv;
1619
+ curr_max_buffer += actual_q_heads_per_kv;
1620
+ }
1621
+ }
1622
+
1623
+ // compute loop
1624
+ for (int32_t kv_tile_pos = rounded_kv_tile_start_pos;
1625
+ kv_tile_pos < rounded_kv_tile_end_pos;
1626
+ kv_tile_pos += kv_tile_size) {
1627
+ const int32_t kv_tile_pos_left = kv_tile_pos;
1628
+ const int32_t kv_tile_pos_right = std::min(
1629
+ kv_tile_pos_left + kv_tile_size, rounded_kv_tile_end_pos);
1630
+ for (int32_t q_head_tile_token_offset = 0;
1631
+ q_head_tile_token_offset < actual_q_token_num;
1632
+ q_head_tile_token_offset += max_q_token_num_per_iter) {
1633
+ const int32_t q_tile_pos_left =
1634
+ q_tile_start_pos + q_head_tile_token_offset;
1635
+ const int32_t q_tile_token_num =
1636
+ std::min(max_q_token_num_per_iter,
1637
+ actual_q_token_num - q_head_tile_token_offset);
1638
+ const int32_t q_tile_head_offset =
1639
+ q_head_tile_token_offset * actual_q_heads_per_kv;
1640
+ const int32_t q_tile_head_num =
1641
+ q_tile_token_num * actual_q_heads_per_kv;
1642
+ const int32_t q_tile_pos_right =
1643
+ q_tile_pos_left + q_tile_token_num;
1644
+ const auto [actual_kv_tile_pos_left,
1645
+ actual_kv_tile_pos_right] =
1646
+ AttentionScheduler::calcu_kv_tile_pos(
1647
+ kv_tile_pos_left, kv_tile_pos_right, q_tile_pos_left,
1648
+ q_tile_pos_right, sliding_window_left,
1649
+ sliding_window_right);
1650
+ const int32_t q_iter_idx =
1651
+ q_head_tile_token_offset / max_q_token_num_per_iter;
1652
+
1653
+ if (actual_kv_tile_pos_right <= actual_kv_tile_pos_left) {
1654
+ continue;
1655
+ }
1656
+
1657
+ // align kv_pos to blocksize_alignment
1658
+ const auto [aligned_actual_kv_tile_pos_left,
1659
+ aligned_actual_kv_tile_pos_right] =
1660
+ AttentionScheduler::align_kv_tile_pos(
1661
+ actual_kv_tile_pos_left, actual_kv_tile_pos_right,
1662
+ blocksize_alignment);
1663
+ const int32_t actual_kv_token_num =
1664
+ aligned_actual_kv_tile_pos_right -
1665
+ aligned_actual_kv_tile_pos_left;
1666
+
1667
+ // std::printf("\tq_iter_idx: %d, q_token_start: %d,
1668
+ // q_token_end: %d, q_token_num: %d, q_head_num: %d,
1669
+ // q_pos_start: %d, q_pos_end: %d, kv_pos_start: %d,
1670
+ // kv_pos_end: %d\n",
1671
+ // q_iter_idx, q_token_start_idx +
1672
+ // q_head_tile_token_offset, q_token_start_idx +
1673
+ // q_head_tile_token_offset + q_tile_token_num,
1674
+ // q_tile_token_num, q_tile_head_num,
1675
+ // q_tile_pos_left, q_tile_pos_right,
1676
+ // aligned_actual_kv_tile_pos_left,
1677
+ // aligned_actual_kv_tile_pos_right);
1678
+
1679
+ // Move buffers
1680
+ q_buffer_t* curr_q_heads_buffer =
1681
+ q_buffer + q_tile_head_offset * head_dim;
1682
+ float* curr_partial_q_buffer =
1683
+ partial_q_buffer + q_tile_head_offset * head_dim;
1684
+ float* curr_max_buffer = max_buffer + q_tile_head_offset;
1685
+ float* curr_sum_buffer = sum_buffer + q_tile_head_offset;
1686
+
1687
+ bool debug_info = false;
1688
+ // bool debug_info = (
1689
+ // q_head_start_idx == 4 &&
1690
+ // (q_token_start_idx + q_head_tile_token_offset) <=
1691
+ // 4
1692
+ // && (q_token_start_idx + q_head_tile_token_offset +
1693
+ // q_tile_token_num) > 4
1694
+ // );
1695
+ // if (debug_info) {
1696
+ // std::printf("\tq_iter_idx: %d, q_token_start: %d,"
1697
+ // "q_token_end: %d, q_token_num: %d, q_head_num: %d,"
1698
+ // "q_pos_start: %d, q_pos_end: %d, kv_pos_start: %d,"
1699
+ // "kv_pos_end: %d\n",
1700
+ // q_iter_idx, q_token_start_idx +
1701
+ // q_head_tile_token_offset, q_token_start_idx
1702
+ // + q_head_tile_token_offset +
1703
+ // q_tile_token_num, q_tile_token_num,
1704
+ // q_tile_head_num, q_tile_pos_left,
1705
+ // q_tile_pos_right,
1706
+ // aligned_actual_kv_tile_pos_left,
1707
+ // aligned_actual_kv_tile_pos_right);
1708
+ // }
1709
+
1710
+ attn_impl.template execute_attention<Attention>(
1711
+ curr_q_heads_buffer, curr_k_cache, curr_v_cache,
1712
+ logits_buffer, curr_partial_q_buffer, curr_max_buffer,
1713
+ curr_sum_buffer, curr_block_table,
1714
+ aligned_actual_kv_tile_pos_left,
1715
+ aligned_actual_kv_tile_pos_right, actual_kv_token_num,
1716
+ kv_cache_block_num_stride, q_tile_head_num,
1717
+ q_tile_token_num, q_tile_pos_left, actual_q_heads_per_kv,
1718
+ block_size, sliding_window_left, sliding_window_right,
1719
+ scale, softcap_scale, curr_alibi_slopes,
1720
+ first_iter_flag[q_iter_idx], use_sink, debug_info);
1721
+ first_iter_flag[q_iter_idx] = false;
1722
+ }
1723
+ }
1724
+
1725
+ // write back partial results to output buffer or reduction buffer
1726
+ {
1727
+ if (curr_spilt_id == -1) {
1728
+ final_output(partial_q_buffer,
1729
+ reinterpret_cast<query_t*>(input->output) +
1730
+ output_buffer_offset,
1731
+ sum_buffer, actual_q_heads_per_kv,
1732
+ actual_q_token_num, q_head_num);
1733
+ } else {
1734
+ const int32_t stride =
1735
+ actual_q_heads_per_kv * split_kv_q_token_num_threshold;
1736
+ buffer_manager.update(kv_head_idx, total_reduction_split_num,
1737
+ head_dim, stride, sizeof(float));
1738
+ volatile bool* split_flag_buffer =
1739
+ buffer_manager.get_reduce_flag_buffer() + curr_spilt_id;
1740
+ float* split_output_buffer =
1741
+ buffer_manager.get_reduce_output_buffer() +
1742
+ curr_spilt_id * stride * head_dim;
1743
+ float* split_max_buffer =
1744
+ buffer_manager.get_reduce_max_buffer() +
1745
+ curr_spilt_id * stride;
1746
+ float* split_sum_buffer =
1747
+ buffer_manager.get_reduce_sum_buffer() +
1748
+ curr_spilt_id * stride;
1749
+
1750
+ partial_output(partial_q_buffer, max_buffer, sum_buffer,
1751
+ q_head_tile_size, split_output_buffer,
1752
+ split_max_buffer, split_sum_buffer,
1753
+ split_flag_buffer);
1754
+ }
1755
+ }
1756
+ }
1757
+ }
1758
+ } else {
1759
+ task_idx -= workitem_groups_counter_num;
1760
+ const int32_t kv_head_idx = task_idx / reduction_item_num;
1761
+ const int32_t item_offset = task_idx % reduction_item_num;
1762
+ ReductionWorkItemGroup* const curr_workitem_groups =
1763
+ reduction_items + item_offset;
1764
+ const int32_t curr_output_token_idx =
1765
+ curr_workitem_groups->q_token_id_start;
1766
+ const int32_t curr_output_token_num =
1767
+ curr_workitem_groups->q_token_id_num;
1768
+ const int32_t curr_split_id = curr_workitem_groups->split_start_id;
1769
+ const int32_t curr_split_num = curr_workitem_groups->split_num;
1770
+ const int32_t current_group_idx = curr_workitem_groups->req_id;
1771
+ const int32_t curr_output_head_num =
1772
+ curr_output_token_num * actual_q_heads_per_kv;
1773
+
1774
+ const int32_t q_start = input->query_start_loc[current_group_idx];
1775
+ const int32_t q_token_start_idx = q_start + curr_output_token_idx;
1776
+ const int32_t q_head_start_idx = kv_head_idx * actual_q_heads_per_kv;
1777
+ size_t output_buffer_offset =
1778
+ q_token_start_idx * q_head_num * head_dim +
1779
+ q_head_start_idx * head_dim;
1780
+
1781
+ const int32_t stride =
1782
+ actual_q_heads_per_kv * split_kv_q_token_num_threshold;
1783
+ buffer_manager.update(kv_head_idx, total_reduction_split_num,
1784
+ head_dim, stride, sizeof(float));
1785
+ volatile bool* split_flag_buffer =
1786
+ buffer_manager.get_reduce_flag_buffer() + curr_split_id;
1787
+ float* split_output_buffer =
1788
+ buffer_manager.get_reduce_output_buffer() +
1789
+ curr_split_id * stride * head_dim;
1790
+ float* split_max_buffer =
1791
+ buffer_manager.get_reduce_max_buffer() + curr_split_id * stride;
1792
+ float* split_sum_buffer =
1793
+ buffer_manager.get_reduce_sum_buffer() + curr_split_id * stride;
1794
+
1795
+ reduce_splits(split_output_buffer, split_max_buffer, split_sum_buffer,
1796
+ split_flag_buffer, stride, curr_output_head_num,
1797
+ curr_split_num);
1798
+ final_output(
1799
+ split_output_buffer,
1800
+ reinterpret_cast<query_t*>(input->output) + output_buffer_offset,
1801
+ split_sum_buffer, actual_q_heads_per_kv, curr_output_token_num,
1802
+ q_head_num);
1803
+ }
1804
+ }
1805
+ }
1806
+ // Reset counter for next call
1807
+ input->metadata->reset_counter();
1808
+ }
1809
+
1810
+ void reduce_splits(float* __restrict__ split_output_buffer,
1811
+ float* __restrict__ split_max_buffer,
1812
+ float* __restrict__ split_sum_buffer,
1813
+ volatile bool* __restrict__ flags,
1814
+ const int32_t head_num_per_split,
1815
+ const int32_t curr_head_num, const int32_t split_num) {
1816
+#ifdef DEFINE_FAST_EXP
1817
+ DEFINE_FAST_EXP
1818
+#endif
1819
+ // restrict curr_head_num <= 16 in the scheduler
1820
+ // elems in split_max_buffer, split_sum_buffer are not cache alignment, use
1821
+ // local buffers to reduce false-sharing
1822
+ alignas(64) float local_max[16];
1823
+ alignas(64) float local_sum[16];
1824
+
1825
+ float* __restrict__ curr_split_output_buffer = split_output_buffer;
1826
+ float* __restrict__ curr_split_max_buffer = split_max_buffer;
1827
+ float* __restrict__ curr_split_sum_buffer = split_sum_buffer;
1828
+ constexpr int32_t head_dim_group_num = head_dim / 16;
1829
+ for (int32_t split_idx = 0; split_idx < split_num; ++split_idx) {
1830
+ while (!flags[split_idx]) {
1831
+#ifdef FAST_SPINNING
1832
+ FAST_SPINNING
1833
+#else
1834
+ std::this_thread::yield();
1835
+#endif
1836
+ }
1837
+ std::atomic_thread_fence(std::memory_order_acquire);
1838
+
1839
+ if (split_idx > 0) {
1840
+ float* __restrict__ curr_output_buffer = split_output_buffer;
1841
+ float* __restrict__ curr_split_output_buffer_iter =
1842
+ curr_split_output_buffer;
1843
+ for (int32_t head_idx = 0; head_idx < curr_head_num; ++head_idx) {
1844
+ float final_max = local_max[head_idx];
1845
+ float curr_max = curr_split_max_buffer[head_idx];
1846
+ float final_sum = local_sum[head_idx];
1847
+ float curr_sum = curr_split_sum_buffer[head_idx];
1848
+ float* __restrict__ non_scale_output_iter =
1849
+ final_max > curr_max ? curr_output_buffer
1850
+ : curr_split_output_buffer_iter;
1851
+ float* __restrict__ scale_output_iter =
1852
+ final_max > curr_max ? curr_split_output_buffer_iter
1853
+ : curr_output_buffer;
1854
+ float rescale_factor = final_max > curr_max ? curr_max - final_max
1855
+ : final_max - curr_max;
1856
+
1857
+#ifdef DEFINE_FAST_EXP
1858
+ vec_op::FP32Vec16 rescale_factor_vec(rescale_factor);
1859
+ rescale_factor_vec = fast_exp(rescale_factor_vec);
1860
+ rescale_factor = rescale_factor_vec.get_last_elem();
1861
+#else
1862
+ rescale_factor = std::exp(rescale_factor);
1863
+ vec_op::FP32Vec16 rescale_factor_vec(rescale_factor);
1864
+#endif
1865
+
1866
+ local_sum[head_idx] = final_max > curr_max
1867
+ ? final_sum + rescale_factor * curr_sum
1868
+ : rescale_factor * final_sum + curr_sum;
1869
+
1870
+ final_max = std::max(final_max, curr_max);
1871
+ local_max[head_idx] = final_max;
1872
+ for (int32_t i = 0; i < head_dim_group_num; ++i) {
1873
+ vec_op::FP32Vec16 non_scale_vec(non_scale_output_iter);
1874
+ vec_op::FP32Vec16 scale_vec(scale_output_iter);
1875
+ vec_op::FP32Vec16 final_vec =
1876
+ non_scale_vec + scale_vec * rescale_factor_vec;
1877
+ final_vec.save(curr_output_buffer);
1878
+
1879
+ non_scale_output_iter += 16;
1880
+ scale_output_iter += 16;
1881
+ curr_output_buffer += 16;
1882
+ }
1883
+ curr_split_output_buffer_iter += head_dim;
1884
+ }
1885
+ } else {
1886
+ vec_op::FP32Vec16 final_max(split_max_buffer);
1887
+ final_max.save(local_max);
1888
+ vec_op::FP32Vec16 final_sum(split_sum_buffer);
1889
+ final_sum.save(local_sum);
1890
+ }
1891
+
1892
+ curr_split_output_buffer += head_num_per_split * head_dim;
1893
+ curr_split_max_buffer += head_num_per_split;
1894
+ curr_split_sum_buffer += head_num_per_split;
1895
+ }
1896
+ // write back final max and sum
1897
+ for (int32_t i = 0; i < curr_head_num; ++i) {
1898
+ split_max_buffer[i] = local_max[i];
1899
+ split_sum_buffer[i] = local_sum[i];
1900
+ }
1901
+ }
1902
+
1903
+ void partial_output(float* __restrict__ partial_output_buffer,
1904
+ float* __restrict__ partial_max_buffer,
1905
+ float* __restrict__ partial_sum_buffer,
1906
+ int32_t curr_head_num,
1907
+ float* __restrict__ split_output_buffer,
1908
+ float* __restrict__ split_max_buffer,
1909
+ float* __restrict__ split_sum_buffer,
1910
+ volatile bool* __restrict__ flag) {
1911
+ float* __restrict__ curr_partial_output_buffer = partial_output_buffer;
1912
+ float* __restrict__ curr_split_output_buffer = split_output_buffer;
1913
+ constexpr int32_t head_dim_group_num = head_dim / 16;
1914
+ for (int32_t i = 0; i < curr_head_num; ++i) {
1915
+ split_max_buffer[i] = partial_max_buffer[i];
1916
+ split_sum_buffer[i] = partial_sum_buffer[i];
1917
+ for (int32_t j = 0; j < head_dim_group_num; ++j) {
1918
+ vec_op::FP32Vec16 vec(curr_partial_output_buffer);
1919
+ vec.save(curr_split_output_buffer);
1920
+
1921
+ curr_partial_output_buffer += 16;
1922
+ curr_split_output_buffer += 16;
1923
+ }
1924
+ }
1925
+ std::atomic_thread_fence(std::memory_order_release);
1926
+ *flag = true;
1927
+ }
1928
+
1929
+ void final_output(float* __restrict__ partial_q_buffer,
1930
+ query_t* __restrict__ curr_output_buffer,
1931
+ float* __restrict__ sum_buffer,
1932
+ const int32_t q_heads_per_kv,
1933
+ const int32_t actual_q_token_num,
1934
+ const int32_t q_head_num) {
1935
+ // final output
1936
+ using output_vec_t = typename VecTypeTrait<query_t>::vec_t;
1937
+
1938
+ float* __restrict__ curr_partial_output_buffer = partial_q_buffer;
1939
+ float* __restrict__ curr_sum_buffer = sum_buffer;
1940
+ constexpr int32_t group_num_per_head = head_dim / 16;
1941
+ const int32_t partial_q_buffer_stride = q_heads_per_kv * head_dim;
1942
+ const int32_t output_buffer_stride = q_head_num * head_dim;
1943
+ for (int32_t token_idx = 0; token_idx < actual_q_token_num; ++token_idx) {
1944
+ float* __restrict__ curr_partial_output_buffer_iter =
1945
+ curr_partial_output_buffer;
1946
+ query_t* __restrict__ curr_output_buffer_iter = curr_output_buffer;
1947
+ for (int32_t head_idx = 0; head_idx < q_heads_per_kv; ++head_idx) {
1948
+ vec_op::FP32Vec16 inv_sum_scale_vec(1.0 / *curr_sum_buffer);
1949
+
1950
+ for (int32_t i = 0; i < group_num_per_head; ++i) {
1951
+ vec_op::FP32Vec16 vec(curr_partial_output_buffer_iter);
1952
+ // divide the final sum val of softmax here
1953
+ vec = inv_sum_scale_vec * vec;
1954
+
1955
+ // cast to query type
1956
+ output_vec_t output_vec(vec);
1957
+ output_vec.save(curr_output_buffer_iter);
1958
+
1959
+ // update
1960
+ curr_partial_output_buffer_iter += 16;
1961
+ curr_output_buffer_iter += 16;
1962
+ }
1963
+
1964
+ // update
1965
+ curr_sum_buffer += 1;
1966
+ }
1967
+
1968
+ // update
1969
+ curr_partial_output_buffer += partial_q_buffer_stride;
1970
+ curr_output_buffer += output_buffer_stride;
1971
+ }
1972
+ }
1973
+};
1974
+
1975
+} // namespace cpu_attention
1976
+
1977
+#endif
csrc/cpu/cpu_attn_macros.h
new
+63
@@ -0,0 +1,63 @@
1
+#ifndef CPU_ATTN_MACROS_H
2
+#define CPU_ATTN_MACROS_H
3
+
4
+// x86_64
5
+#ifdef __x86_64__
6
+ #define FAST_SPINNING _mm_pause();
7
+
8
+ #ifdef __AVX512F__
9
+ #define DEFINE_FAST_EXP \
10
+ const __m512 vec_factorial_1 = _mm512_set1_ps(0.999999701f); \
11
+ const __m512 vec_factorial_2 = _mm512_set1_ps(0.499991506f); \
12
+ const __m512 vec_factorial_3 = _mm512_set1_ps(0.166676521f); \
13
+ const __m512 vec_factorial_4 = _mm512_set1_ps(0.0418978221f); \
14
+ const __m512 vec_factorial_5 = _mm512_set1_ps(0.00828929059f); \
15
+ const __m512 vec_exp_log2ef = \
16
+ _mm512_castsi512_ps(_mm512_set1_epi32(0x3fb8aa3b)); \
17
+ const __m512 vec_half = _mm512_set1_ps(0.5f); \
18
+ const __m512 vec_one = _mm512_set1_ps(1.f); \
19
+ const __m512 vec_zero = _mm512_set1_ps(0.f); \
20
+ const __m512 vec_two = _mm512_set1_ps(2.f); \
21
+ const __m512 vec_ln2f = \
22
+ _mm512_castsi512_ps(_mm512_set1_epi32(0x3f317218)); \
23
+ const __m512 vec_ln_flt_min = \
24
+ _mm512_castsi512_ps(_mm512_set1_epi32(0xc2aeac50)); \
25
+ const __m512 vec_ln_flt_max = \
26
+ _mm512_castsi512_ps(_mm512_set1_epi32(0x42b17218)); \
27
+ const __m512i vec_127 = _mm512_set1_epi32(0x0000007f); \
28
+ const int n_mantissa_bits = 23; \
29
+ auto fast_exp = [&](vec_op::FP32Vec16& vec) __attribute__(( \
30
+ always_inline)) { \
31
+ __m512 values = vec.reg; \
32
+ auto less_ln_flt_min_mask = \
33
+ _mm512_cmp_ps_mask(values, vec_ln_flt_min, 1 /*_CMP_LT_OS*/); \
34
+ auto vec_src = _mm512_min_ps(values, vec_ln_flt_max); \
35
+ vec_src = _mm512_max_ps(vec_src, vec_ln_flt_min); \
36
+ auto vec_fx = _mm512_fmadd_ps(vec_src, vec_exp_log2ef, vec_half); \
37
+ auto vec_fx_i = _mm512_cvt_roundps_epi32( \
38
+ vec_fx, _MM_FROUND_TO_NEG_INF | _MM_FROUND_NO_EXC); \
39
+ vec_fx = _mm512_cvtepi32_ps(vec_fx_i); \
40
+ auto vec_exp_poly = _mm512_fnmadd_ps(vec_fx, vec_ln2f, vec_src); \
41
+ auto vec_res = \
42
+ _mm512_fmadd_ps(vec_exp_poly, vec_factorial_5, vec_factorial_4); \
43
+ vec_res = _mm512_fmadd_ps(vec_exp_poly, vec_res, vec_factorial_3); \
44
+ vec_res = _mm512_fmadd_ps(vec_exp_poly, vec_res, vec_factorial_2); \
45
+ vec_res = _mm512_fmadd_ps(vec_exp_poly, vec_res, vec_factorial_1); \
46
+ vec_res = _mm512_fmadd_ps(vec_exp_poly, vec_res, vec_one); \
47
+ auto vec_exp_number = _mm512_sub_ps(vec_fx, vec_one); \
48
+ auto vec_exp_number_i = _mm512_cvtps_epi32(vec_exp_number); \
49
+ auto vec_two_pow_n_i = _mm512_add_epi32(vec_exp_number_i, vec_127); \
50
+ vec_two_pow_n_i = _mm512_slli_epi32(vec_two_pow_n_i, n_mantissa_bits); \
51
+ auto vec_two_pow_n = _mm512_castsi512_ps(vec_two_pow_n_i); \
52
+ vec_two_pow_n = _mm512_mask_blend_ps(less_ln_flt_min_mask, \
53
+ vec_two_pow_n, vec_zero); \
54
+ vec_res = _mm512_mul_ps(vec_res, vec_two_pow_n); \
55
+ vec_res = _mm512_mul_ps(vec_res, vec_two); \
56
+ vec_op::FP32Vec16 res(vec_res); \
57
+ return res; \
58
+ };
59
+ #endif
60
+
61
+#endif
62
+
63
+#endif
\ No newline at end of file
csrc/cpu/cpu_attn_vec.hpp
new
+248
@@ -0,0 +1,248 @@
1
+#ifndef CPU_ATTN_VEC_HPP
2
+#define CPU_ATTN_VEC_HPP
3
+
4
+#include "cpu_attn_impl.hpp"
5
+
6
+namespace cpu_attention {
7
+
8
+namespace {
9
+// 8-2-16 pattern, 8 regs for A, 2 regs for B, 16 regs for C, [8, K] @ [k, 32]
10
+template <typename kv_cache_t>
11
+class TileGemm82 {
12
+ public:
13
+ template <AttentionGemmPhase phase, int32_t k_size>
14
+ FORCE_INLINE static void gemm(const int32_t m_size,
15
+ float* __restrict__ a_tile,
16
+ kv_cache_t* __restrict__ b_tile,
17
+ float* __restrict__ c_tile, const int64_t lda,
18
+ const int64_t ldb, const int64_t ldc,
19
+ const int32_t block_size,
20
+ const int32_t dynamic_k_size,
21
+ const bool accum_c) {
22
+ switch (m_size) {
23
+ case 1:
24
+ gemm_micro<1>(a_tile, b_tile, c_tile, lda, ldb, ldc, block_size,
25
+ dynamic_k_size, accum_c);
26
+ break;
27
+ case 2:
28
+ gemm_micro<2>(a_tile, b_tile, c_tile, lda, ldb, ldc, block_size,
29
+ dynamic_k_size, accum_c);
30
+ break;
31
+ case 3:
32
+ case 4:
33
+ gemm_micro<4>(a_tile, b_tile, c_tile, lda, ldb, ldc, block_size,
34
+ dynamic_k_size, accum_c);
35
+ break;
36
+ case 5:
37
+ case 6:
38
+ gemm_micro<6>(a_tile, b_tile, c_tile, lda, ldb, ldc, block_size,
39
+ dynamic_k_size, accum_c);
40
+ break;
41
+ case 7:
42
+ case 8:
43
+ gemm_micro<8>(a_tile, b_tile, c_tile, lda, ldb, ldc, block_size,
44
+ dynamic_k_size, accum_c);
45
+ break;
46
+ }
47
+ }
48
+
49
+ template <int32_t M>
50
+ static void gemm_micro(float* __restrict__ a_tile,
51
+ kv_cache_t* __restrict__ b_tile,
52
+ float* __restrict__ c_tile, const int64_t lda,
53
+ const int64_t ldb, const int64_t ldc,
54
+ const int32_t block_size, const int32_t dynamic_k_size,
55
+ const bool accum_c) {
56
+ static_assert(0 < M <= 8);
57
+ using load_vec_t = typename VecTypeTrait<kv_cache_t>::vec_t;
58
+
59
+ kv_cache_t* __restrict__ curr_b_0 = b_tile;
60
+ kv_cache_t* __restrict__ curr_b_1 = b_tile + 16;
61
+ float* __restrict__ curr_c_0 = c_tile;
62
+ float* __restrict__ curr_c_1 = c_tile + 16;
63
+
64
+ vec_op::FP32Vec16 c_regs[M * 2];
65
+ if (accum_c) {
66
+ float* __restrict__ curr_m_c_0 = curr_c_0;
67
+ float* __restrict__ curr_m_c_1 = curr_c_1;
68
+ vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
69
+ c_regs[i * 2] = vec_op::FP32Vec16(curr_m_c_0);
70
+ c_regs[i * 2 + 1] = vec_op::FP32Vec16(curr_m_c_1);
71
+
72
+ // update
73
+ curr_m_c_0 += ldc;
74
+ curr_m_c_1 += ldc;
75
+ });
76
+ }
77
+
78
+ float* __restrict__ curr_a = a_tile;
79
+ 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);
84
+
85
+ float* __restrict__ curr_m_a = curr_a;
86
+ vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
87
+ float v = *curr_m_a;
88
+ vec_op::FP32Vec16 a_reg(v);
89
+ c_regs[i * 2] = c_regs[i * 2] + a_reg * fp32_b_0_reg;
90
+ c_regs[i * 2 + 1] = c_regs[i * 2 + 1] + a_reg * fp32_b_1_reg;
91
+
92
+ // update
93
+ curr_m_a += lda;
94
+ });
95
+
96
+ // update
97
+ curr_a += 1;
98
+ curr_b_0 += ldb;
99
+ curr_b_1 += ldb;
100
+ }
101
+
102
+ vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
103
+ c_regs[i * 2].save(curr_c_0);
104
+ c_regs[i * 2 + 1].save(curr_c_1);
105
+
106
+ // update
107
+ curr_c_0 += ldc;
108
+ curr_c_1 += ldc;
109
+ });
110
+ }
111
+};
112
+} // namespace
113
+
114
+// 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> {
117
+ public:
118
+ using query_t = scalar_t;
119
+ using q_buffer_t = float;
120
+ using kv_cache_t = scalar_t;
121
+ using logits_buffer_t = float;
122
+ using partial_output_buffer_t = float;
123
+ using prob_buffer_t = float;
124
+
125
+ constexpr static int64_t BlockSizeAlignment =
126
+ 32; // KV token num unit of QK and PV phases
127
+ constexpr static int64_t HeadDimAlignment =
128
+ 32; // headdim num unit of PV phase
129
+ constexpr static int64_t MaxQHeadNumPerIteration = 8;
130
+ constexpr static int64_t HeadDim = head_dim;
131
+ constexpr static ISA ISAType = ISA::VEC;
132
+ constexpr static bool scale_on_logits = false; // apply scale on q_buffer
133
+
134
+ public:
135
+ template <template <typename tile_gemm_t> typename attention>
136
+ FORCE_INLINE void execute_attention(DEFINE_CPU_ATTENTION_PARAMS) {
137
+ attention<TileGemm82<kv_cache_t>> attention_iteration;
138
+ attention_iteration(CPU_ATTENTION_PARAMS);
139
+ }
140
+
141
+ // k_cache_token_group_stride: stride of K cache when move to next
142
+ // BlockSizeAlignment tokens in a block
143
+ constexpr static int64_t k_cache_token_group_stride(
144
+ const int32_t block_size) {
145
+ return BlockSizeAlignment; // layout of k_cache block is [head_dim,
146
+ // block_size], row-major
147
+ }
148
+
149
+ // v_cache_token_group_stride: stride of V cache when move to next
150
+ // BlockSizeAlignment tokens in a block
151
+ constexpr static int64_t v_cache_token_group_stride(
152
+ const int32_t block_size) {
153
+ return head_dim * BlockSizeAlignment; // layout of v_cache is [block_size,
154
+ // head_dim], row-major
155
+ }
156
+
157
+ // v_cache_head_group_stride: stride of V cache when move to next
158
+ // HeadDimAlignment head dims in a block
159
+ constexpr static int64_t v_cache_head_group_stride(const int32_t block_size) {
160
+ return HeadDimAlignment; // layout of v_cache is [block_size, head_dim],
161
+ // row-major
162
+ }
163
+
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) {
170
+ static_assert(head_dim % 16 == 0);
171
+ constexpr int32_t unroll_size = head_dim / 16;
172
+ using load_vec_t = typename VecTypeTrait<scalar_t>::vec_t;
173
+
174
+ vec_op::FP32Vec16 scale_vec(scale);
175
+ for (int32_t q_num_idx = 0; q_num_idx < q_num; ++q_num_idx) {
176
+ for (int32_t q_head_idx = 0; q_head_idx < q_heads_per_kv; ++q_head_idx) {
177
+ scalar_t* __restrict__ curr_q =
178
+ src + q_num_idx * q_num_stride + q_head_idx * q_head_stride;
179
+ float* __restrict__ curr_q_buffer =
180
+ q_buffer + q_num_idx * q_heads_per_kv * head_dim +
181
+ q_head_idx * head_dim;
182
+
183
+ vec_op::unroll_loop<int32_t, unroll_size>([&](int32_t i) {
184
+ load_vec_t vec(curr_q);
185
+ vec_op::FP32Vec16 fp32_vec(vec);
186
+ fp32_vec = fp32_vec * scale_vec;
187
+ fp32_vec.save(curr_q_buffer);
188
+
189
+ curr_q += 16;
190
+ curr_q_buffer += 16;
191
+ });
192
+ }
193
+ }
194
+ }
195
+
196
+ // reshape K as column-major and V as row-major
197
+ static void reshape_and_cache(
198
+ const scalar_t* __restrict__ key, const scalar_t* __restrict__ value,
199
+ scalar_t* __restrict__ key_cache, scalar_t* __restrict__ value_cache,
200
+ const int64_t* __restrict__ slot_mapping, const int64_t token_num,
201
+ const int64_t key_token_num_stride, const int64_t value_token_num_stride,
202
+ const int64_t head_num, const int64_t key_head_num_stride,
203
+ const int64_t value_head_num_stride, const int64_t num_blocks,
204
+ 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) {
206
+#pragma omp parallel for collapse(2)
207
+ for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
208
+ for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
209
+ const int64_t pos = slot_mapping[token_idx];
210
+ if (pos < 0) {
211
+ // skip
212
+ continue;
213
+ }
214
+
215
+ const int64_t block_idx = pos / block_size;
216
+ const int64_t block_offset = pos % block_size;
217
+ {
218
+ // Write Key as column-major
219
+ const scalar_t* key_start_ptr = key +
220
+ token_idx * key_token_num_stride +
221
+ head_idx * key_head_num_stride;
222
+ scalar_t* key_cache_start_ptr =
223
+ key_cache + block_idx * num_blocks_stride +
224
+ head_idx * cache_head_num_stride + block_offset;
225
+
226
+#pragma GCC unroll 8
227
+ for (int64_t i = 0, j = 0; i < head_dim; ++i, j += block_size) {
228
+ key_cache_start_ptr[j] = key_start_ptr[i];
229
+ }
230
+ }
231
+ {
232
+ // Write Value as row-major
233
+ const scalar_t* value_start_ptr = value +
234
+ token_idx * value_token_num_stride +
235
+ head_idx * value_head_num_stride;
236
+ 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;
239
+ std::memcpy(value_cache_start_ptr, value_start_ptr,
240
+ sizeof(scalar_t) * head_dim);
241
+ }
242
+ }
243
+ }
244
+ }
245
+};
246
+} // namespace cpu_attention
247
+
248
+#endif
csrc/cpu/cpu_attn_vec16.hpp
new
+171
@@ -0,0 +1,171 @@
1
+#ifndef CPU_ATTN_VEC16_HPP
2
+#define CPU_ATTN_VEC16_HPP
3
+
4
+#include "cpu_attn_vec.hpp"
5
+
6
+namespace cpu_attention {
7
+
8
+namespace {
9
+// 16-1-16 pattern, 16 regs for A, 1 regs for B, 16 regs for C, [16, K] @ [k,
10
+// 16]
11
+template <typename kv_cache_t>
12
+class TileGemm161 {
13
+ public:
14
+ template <AttentionGemmPhase phase, int32_t k_size>
15
+ FORCE_INLINE static void gemm(const int32_t m_size,
16
+ float* __restrict__ a_tile,
17
+ kv_cache_t* __restrict__ b_tile,
18
+ float* __restrict__ c_tile, const int64_t lda,
19
+ const int64_t ldb, const int64_t ldc,
20
+ const int32_t block_size,
21
+ const int32_t dynamic_k_size,
22
+ const bool accum_c) {
23
+ switch (m_size) {
24
+ case 1:
25
+ gemm_micro<1>(a_tile, b_tile, c_tile, lda, ldb, ldc, block_size,
26
+ dynamic_k_size, accum_c);
27
+ break;
28
+ case 2:
29
+ gemm_micro<2>(a_tile, b_tile, c_tile, lda, ldb, ldc, block_size,
30
+ dynamic_k_size, accum_c);
31
+ break;
32
+ case 3:
33
+ case 4:
34
+ gemm_micro<4>(a_tile, b_tile, c_tile, lda, ldb, ldc, block_size,
35
+ dynamic_k_size, accum_c);
36
+ break;
37
+ case 5:
38
+ case 6:
39
+ gemm_micro<6>(a_tile, b_tile, c_tile, lda, ldb, ldc, block_size,
40
+ dynamic_k_size, accum_c);
41
+ break;
42
+ case 7:
43
+ case 8:
44
+ gemm_micro<8>(a_tile, b_tile, c_tile, lda, ldb, ldc, block_size,
45
+ dynamic_k_size, accum_c);
46
+ break;
47
+ case 9:
48
+ case 10:
49
+ case 11:
50
+ case 12:
51
+ gemm_micro<12>(a_tile, b_tile, c_tile, lda, ldb, ldc, block_size,
52
+ dynamic_k_size, accum_c);
53
+ break;
54
+ case 13:
55
+ case 14:
56
+ case 15:
57
+ case 16:
58
+ gemm_micro<16>(a_tile, b_tile, c_tile, lda, ldb, ldc, block_size,
59
+ dynamic_k_size, accum_c);
60
+ break;
61
+ }
62
+ }
63
+
64
+ template <int32_t M>
65
+ static void gemm_micro(float* __restrict__ a_tile,
66
+ kv_cache_t* __restrict__ b_tile,
67
+ float* __restrict__ c_tile, const int64_t lda,
68
+ const int64_t ldb, const int64_t ldc,
69
+ const int32_t block_size, const int32_t dynamic_k_size,
70
+ const bool accum_c) {
71
+ static_assert(0 < M <= 16);
72
+ using load_vec_t = typename VecTypeTrait<kv_cache_t>::vec_t;
73
+
74
+ kv_cache_t* __restrict__ curr_b_0 = b_tile;
75
+ float* __restrict__ curr_c_0 = c_tile;
76
+
77
+ vec_op::FP32Vec16 c_regs[M];
78
+ if (accum_c) {
79
+ float* __restrict__ curr_m_c_0 = curr_c_0;
80
+ vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
81
+ c_regs[i] = vec_op::FP32Vec16(curr_m_c_0);
82
+
83
+ // update
84
+ curr_m_c_0 += ldc;
85
+ });
86
+ }
87
+
88
+ float* __restrict__ curr_a = a_tile;
89
+ for (int32_t k = 0; k < dynamic_k_size; ++k) {
90
+ load_vec_t b_0_reg(curr_b_0);
91
+ vec_op::FP32Vec16 fp32_b_0_reg(b_0_reg);
92
+
93
+ float* __restrict__ curr_m_a = curr_a;
94
+ vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
95
+ float v = *curr_m_a;
96
+ vec_op::FP32Vec16 a_reg(v);
97
+ c_regs[i] = c_regs[i] + a_reg * fp32_b_0_reg;
98
+
99
+ // update
100
+ curr_m_a += lda;
101
+ });
102
+
103
+ // update
104
+ curr_a += 1;
105
+ curr_b_0 += ldb;
106
+ }
107
+
108
+ vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
109
+ c_regs[i].save(curr_c_0);
110
+
111
+ // update
112
+ curr_c_0 += ldc;
113
+ });
114
+ }
115
+};
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> {
122
+ public:
123
+ using query_t = scalar_t;
124
+ using q_buffer_t = float;
125
+ using kv_cache_t = scalar_t;
126
+ using logits_buffer_t = float;
127
+ using partial_output_buffer_t = float;
128
+ using prob_buffer_t = float;
129
+
130
+ constexpr static int64_t BlockSizeAlignment =
131
+ 16; // KV token num unit of QK and PV phases
132
+ constexpr static int64_t HeadDimAlignment =
133
+ 16; // headdim num unit of PV phase
134
+ constexpr static int64_t MaxQHeadNumPerIteration = 16;
135
+ constexpr static int64_t HeadDim = head_dim;
136
+ constexpr static ISA ISAType = ISA::VEC16;
137
+ constexpr static bool scale_on_logits = false; // apply scale on q_buffer
138
+
139
+ public:
140
+ template <template <typename tile_gemm_t> typename attention>
141
+ FORCE_INLINE void execute_attention(DEFINE_CPU_ATTENTION_PARAMS) {
142
+ attention<TileGemm161<kv_cache_t>> attention_iteration;
143
+ attention_iteration(CPU_ATTENTION_PARAMS);
144
+ }
145
+
146
+ // k_cache_token_group_stride: stride of K cache when move to next
147
+ // BlockSizeAlignment tokens in a block
148
+ constexpr static int64_t k_cache_token_group_stride(
149
+ const int32_t block_size) {
150
+ return BlockSizeAlignment; // layout of k_cache block is [head_dim,
151
+ // block_size], row-major
152
+ }
153
+
154
+ // v_cache_token_group_stride: stride of V cache when move to next
155
+ // BlockSizeAlignment tokens in a block
156
+ constexpr static int64_t v_cache_token_group_stride(
157
+ const int32_t block_size) {
158
+ return head_dim * BlockSizeAlignment; // layout of v_cache is [block_size,
159
+ // head_dim], row-major
160
+ }
161
+
162
+ // v_cache_head_group_stride: stride of V cache when move to next
163
+ // HeadDimAlignment head dims in a block
164
+ constexpr static int64_t v_cache_head_group_stride(const int32_t block_size) {
165
+ return HeadDimAlignment; // layout of v_cache is [block_size, head_dim],
166
+ // row-major
167
+ }
168
+};
169
+} // namespace cpu_attention
170
+
171
+#endif
csrc/cpu/cpu_types_x86.hpp
+46
-4
@@ -40,6 +40,23 @@ namespace vec_op {
40
41
#define FORCE_INLINE __attribute__((always_inline)) inline
42
43
+// Function to get the timestamp using RDTSCP
44
+FORCE_INLINE uint64_t bench_timestamp() {
45
+ unsigned int cycles_low, cycles_high;
46
+ asm volatile(
47
+ ".intel_syntax noprefix\n\t"
48
+ "CPUID\n\t" // Serialize instruction stream to ensure previous
49
+ // instructions complete
50
+ "RDTSCP\n\t" // Read TSC and core ID
51
+ "mov %0, edx\n\t" // Store high 32 bits of TSC
52
+ "mov %1, eax\n\t" // Store low 32 bits of TSC
53
+ ".att_syntax"
54
+ : "=r"(cycles_high), "=r"(cycles_low)::"rax", "rbx", "rcx",
55
+ "rdx" // Clobbered registers
56
+ );
57
+ return (uint64_t)cycles_high << 32 | cycles_low;
58
+}
59
+
60
namespace {
61
template <typename T, T... indexes, typename F>
62
constexpr void unroll_loop_item(std::integer_sequence<T, indexes...>, F&& f) {
@@ -407,6 +424,8 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
424
425
float reduce_min() const { return _mm512_reduce_min_ps(reg); }
426
427
+ float get_last_elem() const { return _mm512_cvtss_f32(reg); }
428
+
429
template <int group_size>
430
float reduce_sub_sum(int idx) {
431
static_assert(VEC_ELEM_NUM % group_size == 0);
@@ -446,9 +465,6 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
465
466
explicit FP32Vec16(__m256 low, __m256 high) : reg_low(low), reg_high(high) {}
467
449
- explicit FP32Vec16(const FP32Vec16& data)
450
- : reg_low(data.reg_low), reg_high(data.reg_high) {}
451
-
468
explicit FP32Vec16(const FP32Vec4& data)
469
: reg_low((__m256)_mm256_inserti128_si256(
470
_mm256_castsi128_si256((__m128i)data.reg), (__m128i)data.reg, 1)),
@@ -504,6 +520,32 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
520
_mm256_div_ps(reg_high, b.reg_high));
521
}
522
523
+ FP32Vec16 max(const FP32Vec16& b) const {
524
+ return FP32Vec16(_mm256_max_ps(reg_low, b.reg_low),
525
+ _mm256_max_ps(reg_high, b.reg_high));
526
+ }
527
+
528
+ float reduce_max() const {
529
+ __m256 v = _mm256_max_ps(reg_low, reg_high);
530
+ // Permute to compare elements within 128-bit lanes
531
+ __m256 v_shuffled = _mm256_permute_ps(
532
+ v, 0b00001011); // Swap halves within each 128-bit lane
533
+ __m256 v_max = _mm256_max_ps(v, v_shuffled);
534
+
535
+ v_shuffled = _mm256_permute_ps(
536
+ v_max, 0b00000001); // Shuffle elements within each 128-bit lane
537
+ v_max = _mm256_max_ps(v_max, v_shuffled);
538
+
539
+ // Permute to compare elements between 128-bit lanes
540
+ v_shuffled =
541
+ _mm256_permute2f128_ps(v_max, v_max, 0b00000001); // Swap 128-bit lanes
542
+ v_max = _mm256_max_ps(v_max, v_shuffled);
543
+
544
+ // At this point, the maximum value is present in all elements of v_max.
545
+ // Extract the first element for the scalar result.
546
+ return _mm256_cvtss_f32(v_max); // Extract the lowest 32-bit float
547
+ }
548
+
549
float reduce_sum() const {
550
FP32Vec8 low = FP32Vec8(reg_low);
551
FP32Vec8 high = FP32Vec8(reg_high);
@@ -642,7 +684,7 @@ inline FP16Vec16::FP16Vec16(const FP32Vec16& v)
684
inline FP16Vec16::FP16Vec16(const FP32Vec16& v)
685
: reg(_mm256_insertf128_si256(
686
_mm256_castsi128_si256(FP16Vec8(FP32Vec8(v.reg_low)).reg),
645
- FP16Vec8(FP32Vec8(v.reg_low)).reg, 1)) {}
687
+ FP16Vec8(FP32Vec8(v.reg_high)).reg, 1)) {}
688
#endif
689
690
#ifdef __AVX512BF16__
csrc/cpu/dnnl_helper.cpp
+1
-17
@@ -5,6 +5,7 @@
5
#include "common/memory.hpp"
6
7
#include "dnnl_helper.h"
8
+#include "scratchpad_manager.h"
9
10
static dnnl::engine& default_engine() {
11
static dnnl::engine engine(dnnl::engine::kind::cpu, 0);
@@ -22,23 +23,6 @@ void release_dnnl_matmul_handler(int64_t handler) {
23
delete ptr;
24
}
25
25
-DNNLScratchPadManager::DNNLScratchPadManager() : size_(0), ptr_(nullptr) {
26
- this->realloc(allocation_unit * 128);
27
-}
28
-
29
-void DNNLScratchPadManager::realloc(size_t new_size) {
30
- new_size = round(new_size);
31
- if (new_size > size_) {
32
- ptr_ = std::aligned_alloc(64, new_size);
33
- size_ = new_size;
34
- }
35
-}
36
-
37
-DNNLScratchPadManager* DNNLScratchPadManager::get_dnnl_scratchpad_manager() {
38
- static DNNLScratchPadManager manager;
39
- return &manager;
40
-}
41
-
26
template <typename KT, typename VT>
27
class DNNLPrimitiveCache {
28
public:
csrc/cpu/dnnl_helper.h
-24
@@ -59,30 +59,6 @@ constexpr inline dnnl::memory::data_type get_dnnl_type() {
59
return DNNLType<std::decay_t<T>>::type;
60
}
61
62
-class DNNLScratchPadManager {
63
- public:
64
- static constexpr size_t allocation_unit = 4 * 1024 * 1024; // 4KB
65
-
66
- static DNNLScratchPadManager* get_dnnl_scratchpad_manager();
67
-
68
- DNNLScratchPadManager();
69
-
70
- template <typename T>
71
- T* get_data() {
72
- return reinterpret_cast<T*>(ptr_);
73
- }
74
-
75
- static size_t round(size_t size) {
76
- return ((size + allocation_unit - 1) / allocation_unit) * allocation_unit;
77
- }
78
-
79
- void realloc(size_t new_size);
80
-
81
- private:
82
- size_t size_;
83
- void* ptr_;
84
-};
85
-
62
class DNNLMatMulPrimitiveHandler {
63
public:
64
virtual ~DNNLMatMulPrimitiveHandler() = default;
csrc/cpu/scratchpad_manager.cpp
new
+23
@@ -0,0 +1,23 @@
1
+#include <cstdlib>
2
+
3
+#include "scratchpad_manager.h"
4
+
5
+DNNLScratchPadManager::DNNLScratchPadManager() : size_(0), ptr_(nullptr) {
6
+ this->realloc(allocation_unit * 128);
7
+}
8
+
9
+void DNNLScratchPadManager::realloc(size_t new_size) {
10
+ new_size = round(new_size);
11
+ if (new_size > size_) {
12
+ if (ptr_ != nullptr) {
13
+ std::free(ptr_);
14
+ }
15
+ ptr_ = std::aligned_alloc(64, new_size);
16
+ size_ = new_size;
17
+ }
18
+}
19
+
20
+DNNLScratchPadManager* DNNLScratchPadManager::get_dnnl_scratchpad_manager() {
21
+ static DNNLScratchPadManager manager;
22
+ return &manager;
23
+}
csrc/cpu/scratchpad_manager.h
new
+31
@@ -0,0 +1,31 @@
1
+#ifndef SCRATCHPAD_MANAGER_H
2
+#define SCRATCHPAD_MANAGER_H
3
+
4
+#include <cstddef>
5
+#include <cstdio>
6
+
7
+class DNNLScratchPadManager {
8
+ public:
9
+ static constexpr size_t allocation_unit = 4 * 1024; // 4KB
10
+
11
+ static DNNLScratchPadManager* get_dnnl_scratchpad_manager();
12
+
13
+ DNNLScratchPadManager();
14
+
15
+ template <typename T>
16
+ T* get_data() {
17
+ return reinterpret_cast<T*>(ptr_);
18
+ }
19
+
20
+ static size_t round(size_t size) {
21
+ return ((size + allocation_unit - 1) / allocation_unit) * allocation_unit;
22
+ }
23
+
24
+ void realloc(size_t new_size);
25
+
26
+ private:
27
+ size_t size_;
28
+ void* ptr_;
29
+};
30
+
31
+#endif
csrc/cpu/shm.cpp
+1
-1
@@ -192,7 +192,7 @@ class SHMManager {
192
const int group_size)
193
: _rank(rank),
194
_group_size(group_size),
195
- _thread_num(torch::get_num_threads()),
195
+ _thread_num(omp_get_max_threads()),
196
_shm_names({""}),
197
_shared_mem_ptrs({nullptr}),
198
_shm_ctx(nullptr) {
csrc/cpu/torch_bindings.cpp
+45
-60
@@ -74,25 +74,35 @@ at::Tensor int8_scaled_mm_with_quant(at::Tensor& mat1, at::Tensor& mat2,
74
const std::optional<at::Tensor>& bias,
75
at::ScalarType out_dtype, bool is_vnni);
76
77
+torch::Tensor get_scheduler_metadata(
78
+ const int64_t num_req, const int64_t num_heads_q,
79
+ const int64_t num_heads_kv, const int64_t head_dim,
80
+ const torch::Tensor& seq_lens, at::ScalarType dtype,
81
+ const torch::Tensor& query_start_loc, const bool casual,
82
+ const int64_t window_size, const std::string& isa_hint,
83
+ const bool enable_kv_split);
84
+
85
+void cpu_attn_reshape_and_cache(const torch::Tensor& key,
86
+ const torch::Tensor& value,
87
+ torch::Tensor& key_cache,
88
+ torch::Tensor& value_cache,
89
+ const torch::Tensor& slot_mapping,
90
+ const std::string& isa);
91
+
92
+void cpu_attention_with_kv_cache(
93
+ const torch::Tensor& query, const torch::Tensor& key_cache,
94
+ const torch::Tensor& value_cache, torch::Tensor& output,
95
+ const torch::Tensor& query_start_loc, const torch::Tensor& seq_lens,
96
+ const double scale, const bool causal,
97
+ const std::optional<torch::Tensor>& alibi_slopes,
98
+ const int64_t sliding_window_left, const int64_t sliding_window_right,
99
+ const torch::Tensor& block_table, const double softcap,
100
+ const torch::Tensor& scheduler_metadata,
101
+ const std::optional<torch::Tensor>& s_aux);
102
+
103
TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
104
// vLLM custom ops
105
80
- // Attention ops
81
- // Compute the attention between an input query and the cached keys/values
82
- // using PagedAttention.
83
- ops.def(
84
- "paged_attention_v1("
85
- " Tensor! out, Tensor query, Tensor key_cache,"
86
- " Tensor value_cache, int num_kv_heads, float scale,"
87
- " Tensor block_tables, Tensor seq_lens, int block_size,"
88
- " int max_seq_len, Tensor? alibi_slopes,"
89
- " str kv_cache_dtype, Tensor k_scale, Tensor v_scale,"
90
- " int tp_rank, int blocksparse_local_blocks,"
91
- " int blocksparse_vert_stride, int blocksparse_block_size,"
92
- " int blocksparse_head_sliding_step) -> ()");
93
-
94
- ops.impl("paged_attention_v1", torch::kCPU, &paged_attention_v1);
95
-
106
ops.def(
107
"dynamic_4bit_int_moe("
108
"Tensor x, Tensor topk_ids, Tensor topk_weights,"
@@ -102,20 +112,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
112
113
ops.impl("dynamic_4bit_int_moe", torch::kCPU, &dynamic_4bit_int_moe_cpu);
114
105
- // PagedAttention V2.
106
- ops.def(
107
- "paged_attention_v2("
108
- " Tensor! out, Tensor! exp_sums, Tensor! max_logits,"
109
- " Tensor! tmp_out, Tensor query, Tensor key_cache,"
110
- " Tensor value_cache, int num_kv_heads, float scale,"
111
- " Tensor block_tables, Tensor seq_lens, int block_size,"
112
- " int max_seq_len, Tensor? alibi_slopes,"
113
- " str kv_cache_dtype, Tensor k_scale, Tensor v_scale,"
114
- " int tp_rank, int blocksparse_local_blocks,"
115
- " int blocksparse_vert_stride, int blocksparse_block_size,"
116
- " int blocksparse_head_sliding_step) -> ()");
117
- ops.impl("paged_attention_v2", torch::kCPU, &paged_attention_v2);
118
-
115
// Activation ops
116
117
// Activation function used in SwiGLU.
@@ -259,37 +255,26 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
255
ops.impl("int8_scaled_mm_with_quant", torch::kCPU,
256
&int8_scaled_mm_with_quant);
257
#endif
262
-}
258
264
-TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cache_ops), cache_ops) {
265
- // Cache ops
266
- // Swap in (out) the cache blocks from src to dst.
267
- cache_ops.def(
268
- "swap_blocks(Tensor src, Tensor! dst, Tensor block_mapping) -> ()");
269
- cache_ops.impl("swap_blocks", torch::kCPU, &swap_blocks);
270
-
271
- // Copy the cache blocks from src to dst.
272
- cache_ops.def(
273
- "copy_blocks(Tensor(a!)[] key_caches, Tensor[](b!) value_caches, "
274
- "Tensor block_mapping) -> ()");
275
- cache_ops.impl("copy_blocks", torch::kCPU, ©_blocks);
276
-
277
- // Reshape the key and value tensors and cache them.
278
- cache_ops.def(
279
- "reshape_and_cache(Tensor key, Tensor value,"
280
- " Tensor! key_cache, Tensor! value_cache,"
281
- " Tensor slot_mapping,"
282
- " str kv_cache_dtype,"
283
- " Tensor k_scale, Tensor v_scale) -> ()");
284
- cache_ops.impl("reshape_and_cache", torch::kCPU, &reshape_and_cache);
285
-
286
- cache_ops.def(
287
- "concat_and_cache_mla(Tensor kv_c, Tensor k_pe,"
288
- " Tensor! kv_cache,"
289
- " Tensor slot_mapping,"
290
- " str kv_cache_dtype,"
291
- " Tensor scale) -> ()");
292
- cache_ops.impl("concat_and_cache_mla", torch::kCPU, &concat_and_cache_mla);
259
+ // CPU attention kernels
260
+ ops.def(
261
+ "get_scheduler_metadata(int num_req, int num_heads_q, int num_heads_kv, "
262
+ "int head_dim, Tensor seq_lens, ScalarType dtype, Tensor "
263
+ "query_start_loc, bool casual, int window_size, str isa_hint, bool "
264
+ "enable_kv_split) -> Tensor",
265
+ &get_scheduler_metadata);
266
+ ops.def(
267
+ "cpu_attn_reshape_and_cache(Tensor key, Tensor value, Tensor(a2!) "
268
+ "key_cache, Tensor(a3!) value_cache, Tensor slot_mapping, str "
269
+ "isa) -> ()",
270
+ &cpu_attn_reshape_and_cache);
271
+ ops.def(
272
+ "cpu_attention_with_kv_cache(Tensor query, Tensor key_cache, Tensor "
273
+ "value_cache, Tensor(a3!) output, Tensor query_start_loc, Tensor "
274
+ "seq_lens, float scale, bool causal, Tensor? alibi_slopes, SymInt "
275
+ "sliding_window_left, SymInt sliding_window_right, Tensor block_table, "
276
+ "float softcap, Tensor sheduler_metadata, Tensor? s_aux) -> ()",
277
+ &cpu_attention_with_kv_cache);
278
}
279
280
TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _utils), utils) {
docker/Dockerfile.cpu
+4
@@ -17,6 +17,7 @@
17
# VLLM_CPU_DISABLE_AVX512=false (default)|true
18
# VLLM_CPU_AVX512BF16=false (default)|true
19
# VLLM_CPU_AVX512VNNI=false (default)|true
20
+# VLLM_CPU_AMXBF16=false (default)|true
21
#
22
23
######################### COMMON BASE IMAGE #########################
@@ -92,6 +93,9 @@ ENV VLLM_CPU_AVX512BF16=${VLLM_CPU_AVX512BF16}
93
# Support for building with AVX512VNNI ISA: docker build --build-arg VLLM_CPU_AVX512VNNI="true" ...
94
ARG VLLM_CPU_AVX512VNNI=0
95
ENV VLLM_CPU_AVX512VNNI=${VLLM_CPU_AVX512VNNI}
96
+# Support for building with AMXBF16 ISA: docker build --build-arg VLLM_CPU_AMXBF16="true" ...
97
+ARG VLLM_CPU_AMXBF16=0
98
+ENV VLLM_CPU_AMXBF16=${VLLM_CPU_AMXBF16}
99
100
WORKDIR /workspace/vllm
101
docs/getting_started/installation/cpu.md
+2
@@ -171,6 +171,8 @@ This value is 4GB by default. Larger space can support more concurrent requests,
171
172
First of all, please make sure the thread-binding and KV cache space are properly set and take effect. You can check the thread-binding by running a vLLM benchmark and observing CPU cores usage via `htop`.
173
174
+Use multiples of 32 as `--block-size`, which is 128 by default.
175
+
176
Inference batch size is an important parameter for the performance. A larger batch usually provides higher throughput, a smaller batch provides lower latency. Tuning the max batch size starting from the default value to balance throughput and latency is an effective way to improve vLLM CPU performance on specific platforms. There are two important related parameters in vLLM:
177
178
- `--max-num-batched-tokens`, defines the limit of token numbers in a single batch, has more impacts on the first token performance. The default value is set as:
tests/kernels/attention/test_attention_selector.py
+3
-3
@@ -35,7 +35,7 @@ DEVICE_MLA_BACKENDS = {
35
DEVICE_REGULAR_ATTN_BACKENDS = {
36
"cuda": ["XFORMERS", "FLASHINFER", "FLASH_ATTN"],
37
"hip": ["ROCM_ATTN"],
38
- "cpu": ["TORCH_SDPA"],
38
+ "cpu": ["CPU_ATTN"],
39
}
40
41
DEVICE_MLA_BLOCK_SIZES = {
@@ -86,7 +86,7 @@ def test_env(
86
if device == "cpu":
87
with patch("vllm.platforms.current_platform", CpuPlatform()):
88
backend = get_attn_backend(16, torch.float16, None, block_size)
89
- assert backend.get_name() == "TORCH_SDPA"
89
+ assert backend.get_name() == "CPU_ATTN"
90
91
elif device == "hip":
92
with patch("vllm.platforms.current_platform", RocmPlatform()):
@@ -224,7 +224,7 @@ def test_fp32_fallback(device: str):
224
if device == "cpu":
225
with patch("vllm.platforms.current_platform", CpuPlatform()):
226
backend = get_attn_backend(16, torch.float32, None, 16)
227
- assert backend.get_name() == "TORCH_SDPA"
227
+ assert backend.get_name() == "CPU_ATTN"
228
229
elif device == "cuda":
230
with patch("vllm.platforms.current_platform", CudaPlatform()):
tests/kernels/attention/test_cpu_attn.py
new
+575
@@ -0,0 +1,575 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import functools
5
+import math
6
+
7
+import pytest
8
+import torch
9
+
10
+from vllm.platforms import current_platform
11
+
12
+if not current_platform.is_cpu():
13
+ pytest.skip("skipping CPU-only tests", allow_module_level=True)
14
+
15
+from vllm._custom_ops import (
16
+ cpu_attention_with_kv_cache,
17
+ cpu_attn_get_scheduler_metadata,
18
+ cpu_attn_reshape_and_cache,
19
+)
20
+
21
+NUM_HEADS = [
22
+ (4, 4),
23
+ (8, 2),
24
+ (9, 3),
25
+]
26
+HEAD_SIZES = [96, 128]
27
+QTYPES = [torch.bfloat16, torch.half, torch.float32]
28
+SLIDING_WINDOWS = [None, 256]
29
+NUM_BLOCKS = [
30
+ 1024,
31
+]
32
+SEQ_LENS = [ # (q_len, kv_len)
33
+ [(1, 213), (1, 1), (1, 312), (1, 7), (1, 7812)], # decode batch
34
+ [(2345, 2345), (5, 5), (3, 16), (134, 5131)], # prefill batch
35
+ [(992, 2456), (1, 1234), (98, 1145), (1, 4162), (2345, 2345)], # mixed batch
36
+]
37
+
38
+
39
+# rand number generation takes too much time, cache rand tensors
40
+@functools.lru_cache(maxsize=128, typed=False)
41
+def tensor_cache(
42
+ elem_num: int,
43
+ dtype: torch.dtype,
44
+) -> torch.Tensor:
45
+ tensor = torch.randn(elem_num, dtype=dtype)
46
+
47
+ return tensor
48
+
49
+
50
+def _get_alibi_slopes(total_num_heads: int) -> torch.Tensor:
51
+ closest_power_of_2 = 2 ** math.floor(math.log2(total_num_heads))
52
+ base = torch.tensor(
53
+ 2 ** (-(2 ** -(math.log2(closest_power_of_2) - 3))),
54
+ dtype=torch.float32,
55
+ )
56
+ powers = torch.arange(1, 1 + closest_power_of_2, dtype=torch.int32)
57
+ slopes = torch.pow(base, powers)
58
+
59
+ if closest_power_of_2 != total_num_heads:
60
+ extra_base = torch.tensor(
61
+ 2 ** (-(2 ** -(math.log2(2 * closest_power_of_2) - 3))),
62
+ dtype=torch.float32,
63
+ )
64
+ num_remaining_heads = min(
65
+ closest_power_of_2, total_num_heads - closest_power_of_2
66
+ )
67
+ extra_powers = torch.arange(
68
+ start=1, end=1 + 2 * num_remaining_heads, step=2, dtype=torch.int32
69
+ )
70
+ slopes = torch.cat([slopes, torch.pow(extra_base, extra_powers)], dim=0)
71
+ return slopes.float()
72
+
73
+
74
+def ref_paged_attn(
75
+ query: torch.Tensor,
76
+ key_cache: torch.Tensor,
77
+ value_cache: torch.Tensor,
78
+ query_lens: list[int],
79
+ kv_lens: list[int],
80
+ block_tables: torch.Tensor,
81
+ scale: float,
82
+ sliding_window: int | None = None,
83
+ soft_cap: float | None = None,
84
+ alibi_slopes: torch.Tensor | None = None,
85
+ s_aux: torch.Tensor | None = None,
86
+) -> torch.Tensor:
87
+ num_seqs = len(query_lens)
88
+ block_tables = block_tables.cpu().numpy()
89
+ _, block_size, num_kv_heads, head_size = key_cache.shape
90
+ dtype = query.dtype
91
+
92
+ outputs: list[torch.Tensor] = []
93
+ start_idx = 0
94
+
95
+ if alibi_slopes is not None:
96
+ alibi_slopes = alibi_slopes[:, None, None]
97
+
98
+ if s_aux is not None:
99
+ s_aux = s_aux.float()
100
+ s_aux = s_aux[:, None, None]
101
+
102
+ for i in range(num_seqs):
103
+ query_len = query_lens[i]
104
+ kv_len = kv_lens[i]
105
+ q = query[start_idx : start_idx + query_len].float()
106
+ q *= scale
107
+
108
+ num_kv_blocks = (kv_len + block_size - 1) // block_size
109
+ block_indices = block_tables[i, :num_kv_blocks]
110
+
111
+ k = key_cache[block_indices].view(-1, num_kv_heads, head_size)
112
+ k = k[:kv_len].float()
113
+ v = value_cache[block_indices].view(-1, num_kv_heads, head_size)
114
+ v = v[:kv_len].float()
115
+
116
+ if q.shape[1] != k.shape[1]:
117
+ k = torch.repeat_interleave(k, q.shape[1] // k.shape[1], dim=1)
118
+ v = torch.repeat_interleave(v, q.shape[1] // v.shape[1], dim=1)
119
+ attn = torch.einsum("qhd,khd->hqk", q, k).float()
120
+ empty_mask = torch.ones(query_len, kv_len)
121
+ mask = torch.triu(empty_mask, diagonal=kv_len - query_len + 1).bool()
122
+
123
+ if sliding_window is not None:
124
+ sliding_window_mask = (
125
+ torch.triu(
126
+ empty_mask, diagonal=kv_len - (query_len + sliding_window) + 1
127
+ )
128
+ .bool()
129
+ .logical_not()
130
+ )
131
+ mask |= sliding_window_mask
132
+
133
+ if soft_cap is not None:
134
+ attn = soft_cap * torch.tanh(attn / soft_cap)
135
+
136
+ if alibi_slopes is not None:
137
+ q_start_pos = kv_len - query_len
138
+ q_pos = q_start_pos + torch.arange(0, query_len)[None, :, None]
139
+ kv_pos = torch.arange(0, kv_len)[None, None, :]
140
+ dist = q_pos - kv_pos
141
+ alibi_bias = -alibi_slopes * dist
142
+ attn += alibi_bias
143
+
144
+ attn.masked_fill_(mask, float("-inf"))
145
+
146
+ if s_aux is not None:
147
+ s_aux_ext = s_aux.repeat(1, query_len, 1)
148
+ attn = torch.cat((s_aux_ext, attn), dim=-1)
149
+
150
+ attn = torch.softmax(attn, dim=-1)
151
+
152
+ if s_aux is not None:
153
+ attn = attn[:, :, 1:]
154
+
155
+ out = torch.einsum("hqk,khd->qhd", attn, v).to(dtype=dtype)
156
+
157
+ outputs.append(out)
158
+ start_idx += query_len
159
+
160
+ return torch.cat(outputs, dim=0)
161
+
162
+
163
+@torch.inference_mode()
164
+def varlen_with_paged_kv(
165
+ seq_lens: list[tuple[int, int]],
166
+ num_heads: tuple[int, int],
167
+ head_size: int,
168
+ sliding_window: int | None,
169
+ dtype: torch.dtype,
170
+ block_size: int,
171
+ soft_cap: float | None,
172
+ num_blocks: int,
173
+ use_alibi: bool,
174
+ use_sink: bool,
175
+ isa: str,
176
+) -> None:
177
+ current_platform.seed_everything(0)
178
+ num_seqs = len(seq_lens)
179
+ query_lens = [x[0] for x in seq_lens]
180
+ kv_lens = [x[1] for x in seq_lens]
181
+ num_query_heads = num_heads[0]
182
+ num_kv_heads = num_heads[1]
183
+ assert num_query_heads % num_kv_heads == 0
184
+ max_kv_len = max(kv_lens)
185
+ window_size = (sliding_window - 1, 0) if sliding_window is not None else (-1, -1)
186
+ scale = head_size**-0.5
187
+ token_num = sum(query_lens)
188
+
189
+ # for n heads the set of slopes is the geometric sequence that starts
190
+ # 2^(-8/n)
191
+ alibi_slopes = _get_alibi_slopes(num_query_heads) if use_alibi else None
192
+
193
+ s_aux = (
194
+ 15 * torch.rand((num_query_heads,), dtype=torch.bfloat16) if use_sink else None
195
+ )
196
+
197
+ query = tensor_cache(
198
+ elem_num=token_num * num_query_heads * head_size,
199
+ dtype=dtype,
200
+ )
201
+ query = query.view(
202
+ token_num,
203
+ num_query_heads,
204
+ head_size,
205
+ )
206
+
207
+ key_value = tensor_cache(
208
+ elem_num=2 * num_blocks * num_kv_heads * block_size * head_size,
209
+ dtype=dtype,
210
+ )
211
+ key_value = key_value.view(
212
+ 2,
213
+ num_blocks,
214
+ block_size,
215
+ num_kv_heads,
216
+ head_size,
217
+ )
218
+ key_cache, value_cache = key_value.unbind(0)
219
+
220
+ # KV cache for CPU attention
221
+ packed_key_cache = torch.empty(
222
+ num_blocks, num_kv_heads, block_size, head_size, dtype=dtype
223
+ )
224
+ packed_value_cache = torch.empty_like(packed_key_cache)
225
+
226
+ cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum(
227
+ dim=0, dtype=torch.int32
228
+ )
229
+ kv_lens_tensor = torch.tensor(kv_lens, dtype=torch.int32)
230
+ max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size
231
+ block_tables = torch.randint(
232
+ 0, num_blocks, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32
233
+ )
234
+
235
+ # use reshape_and_cache to pack key_cache and value_cache
236
+ slot_mapping = torch.arange(0, num_blocks * block_size, dtype=torch.int64)
237
+ cpu_attn_reshape_and_cache(
238
+ key=key_cache.view(-1, num_kv_heads, head_size),
239
+ value=value_cache.view(-1, num_kv_heads, head_size),
240
+ key_cache=packed_key_cache,
241
+ value_cache=packed_value_cache,
242
+ slot_mapping=slot_mapping,
243
+ isa=isa,
244
+ )
245
+
246
+ metadata = cpu_attn_get_scheduler_metadata(
247
+ num_reqs=num_seqs,
248
+ num_heads=num_query_heads,
249
+ num_kv_heads=num_kv_heads,
250
+ head_dim=head_size,
251
+ seq_lens=kv_lens_tensor,
252
+ dtype=dtype,
253
+ query_start_loc=cu_query_lens,
254
+ causal=True,
255
+ sliding_window_size=sliding_window if sliding_window is not None else -1,
256
+ isa=isa,
257
+ enable_kv_split=False,
258
+ )
259
+
260
+ out_without_split = torch.empty_like(query)
261
+ cpu_attention_with_kv_cache(
262
+ query=query,
263
+ key_cache=packed_key_cache,
264
+ value_cache=packed_value_cache,
265
+ output=out_without_split,
266
+ query_start_loc=cu_query_lens,
267
+ seq_lens=kv_lens_tensor,
268
+ scale=scale,
269
+ causal=True,
270
+ alibi_slopes=alibi_slopes,
271
+ sliding_window=window_size,
272
+ block_table=block_tables,
273
+ softcap=soft_cap if soft_cap is not None else 0,
274
+ scheduler_metadata=metadata,
275
+ s_aux=s_aux,
276
+ )
277
+
278
+ metadata = cpu_attn_get_scheduler_metadata(
279
+ num_reqs=num_seqs,
280
+ num_heads=num_query_heads,
281
+ num_kv_heads=num_kv_heads,
282
+ head_dim=head_size,
283
+ seq_lens=kv_lens_tensor,
284
+ dtype=dtype,
285
+ query_start_loc=cu_query_lens,
286
+ causal=True,
287
+ sliding_window_size=sliding_window if sliding_window is not None else -1,
288
+ isa=isa,
289
+ enable_kv_split=True,
290
+ )
291
+
292
+ out_with_split = torch.empty_like(query)
293
+ cpu_attention_with_kv_cache(
294
+ query=query,
295
+ key_cache=packed_key_cache,
296
+ value_cache=packed_value_cache,
297
+ output=out_with_split,
298
+ query_start_loc=cu_query_lens,
299
+ seq_lens=kv_lens_tensor,
300
+ scale=scale,
301
+ causal=True,
302
+ alibi_slopes=alibi_slopes,
303
+ sliding_window=window_size,
304
+ block_table=block_tables,
305
+ softcap=soft_cap if soft_cap is not None else 0,
306
+ scheduler_metadata=metadata,
307
+ s_aux=s_aux,
308
+ )
309
+
310
+ ref_output = ref_paged_attn(
311
+ query=query,
312
+ key_cache=key_cache,
313
+ value_cache=value_cache,
314
+ query_lens=query_lens,
315
+ kv_lens=kv_lens,
316
+ block_tables=block_tables,
317
+ scale=scale,
318
+ sliding_window=sliding_window,
319
+ soft_cap=soft_cap,
320
+ alibi_slopes=alibi_slopes,
321
+ s_aux=s_aux,
322
+ )
323
+
324
+ atol, rtol = 1.5e-2, 1e-2
325
+ (
326
+ torch.testing.assert_close(out_with_split, ref_output, atol=atol, rtol=rtol),
327
+ f"{torch.max(torch.abs(out_with_split - ref_output))}",
328
+ )
329
+ (
330
+ torch.testing.assert_close(out_without_split, ref_output, atol=atol, rtol=rtol),
331
+ f"{torch.max(torch.abs(out_without_split - ref_output))}",
332
+ )
333
+
334
+
335
+@pytest.mark.parametrize("seq_lens", SEQ_LENS)
336
+@pytest.mark.parametrize("num_heads", NUM_HEADS)
337
+@pytest.mark.parametrize("head_size", HEAD_SIZES)
338
+@pytest.mark.parametrize("block_size", [96, 128])
339
+@pytest.mark.parametrize("sliding_window", SLIDING_WINDOWS)
340
+@pytest.mark.parametrize("dtype", QTYPES)
341
+@pytest.mark.parametrize("soft_cap", [None])
342
+@pytest.mark.parametrize("num_blocks", NUM_BLOCKS)
343
+@pytest.mark.parametrize("use_alibi", [False])
344
+@pytest.mark.parametrize("use_sink", [False])
345
+@pytest.mark.parametrize("isa", ["vec"])
346
+def test_varlen_with_paged_kv_normal_vec(
347
+ seq_lens: list[tuple[int, int]],
348
+ num_heads: tuple[int, int],
349
+ head_size: int,
350
+ sliding_window: int | None,
351
+ dtype: torch.dtype,
352
+ block_size: int,
353
+ soft_cap: float | None,
354
+ num_blocks: int,
355
+ use_alibi: bool,
356
+ use_sink: bool,
357
+ isa: str,
358
+) -> None:
359
+ varlen_with_paged_kv(
360
+ seq_lens=seq_lens,
361
+ num_heads=num_heads,
362
+ head_size=head_size,
363
+ sliding_window=sliding_window,
364
+ dtype=dtype,
365
+ block_size=block_size,
366
+ soft_cap=soft_cap,
367
+ num_blocks=num_blocks,
368
+ use_alibi=use_alibi,
369
+ use_sink=use_sink,
370
+ isa=isa,
371
+ )
372
+
373
+
374
+@pytest.mark.parametrize("seq_lens", SEQ_LENS)
375
+@pytest.mark.parametrize("num_heads", NUM_HEADS)
376
+@pytest.mark.parametrize("head_size", HEAD_SIZES)
377
+@pytest.mark.parametrize("block_size", [96, 128])
378
+@pytest.mark.parametrize("sliding_window", SLIDING_WINDOWS)
379
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
380
+@pytest.mark.parametrize("soft_cap", [None])
381
+@pytest.mark.parametrize("num_blocks", NUM_BLOCKS)
382
+@pytest.mark.parametrize("use_alibi", [False])
383
+@pytest.mark.parametrize("use_sink", [False])
384
+@pytest.mark.parametrize("isa", ["amx"])
385
+@pytest.mark.skipif(
386
+ not torch._C._cpu._is_amx_tile_supported(), reason="no AMX support."
387
+)
388
+def test_varlen_with_paged_kv_normal_amx(
389
+ seq_lens: list[tuple[int, int]],
390
+ num_heads: tuple[int, int],
391
+ head_size: int,
392
+ sliding_window: int | None,
393
+ dtype: torch.dtype,
394
+ block_size: int,
395
+ soft_cap: float | None,
396
+ num_blocks: int,
397
+ use_alibi: bool,
398
+ use_sink: bool,
399
+ isa: str,
400
+) -> None:
401
+ varlen_with_paged_kv(
402
+ seq_lens=seq_lens,
403
+ num_heads=num_heads,
404
+ head_size=head_size,
405
+ sliding_window=sliding_window,
406
+ dtype=dtype,
407
+ block_size=block_size,
408
+ soft_cap=soft_cap,
409
+ num_blocks=num_blocks,
410
+ use_alibi=use_alibi,
411
+ use_sink=use_sink,
412
+ isa=isa,
413
+ )
414
+
415
+
416
+@pytest.mark.parametrize("seq_lens", SEQ_LENS)
417
+@pytest.mark.parametrize("num_heads", NUM_HEADS)
418
+@pytest.mark.parametrize("head_size", HEAD_SIZES)
419
+@pytest.mark.parametrize("block_size", [48])
420
+@pytest.mark.parametrize("sliding_window", SLIDING_WINDOWS)
421
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
422
+@pytest.mark.parametrize("soft_cap", [None])
423
+@pytest.mark.parametrize("num_blocks", NUM_BLOCKS)
424
+@pytest.mark.parametrize("use_alibi", [False])
425
+@pytest.mark.parametrize("use_sink", [False])
426
+@pytest.mark.parametrize("isa", ["vec16"])
427
+def test_varlen_with_paged_kv_normal_vec16(
428
+ seq_lens: list[tuple[int, int]],
429
+ num_heads: tuple[int, int],
430
+ head_size: int,
431
+ sliding_window: int | None,
432
+ dtype: torch.dtype,
433
+ block_size: int,
434
+ soft_cap: float | None,
435
+ num_blocks: int,
436
+ use_alibi: bool,
437
+ use_sink: bool,
438
+ isa: str,
439
+) -> None:
440
+ varlen_with_paged_kv(
441
+ seq_lens=seq_lens,
442
+ num_heads=num_heads,
443
+ head_size=head_size,
444
+ sliding_window=sliding_window,
445
+ dtype=dtype,
446
+ block_size=block_size,
447
+ soft_cap=soft_cap,
448
+ num_blocks=num_blocks,
449
+ use_alibi=use_alibi,
450
+ use_sink=use_sink,
451
+ isa=isa,
452
+ )
453
+
454
+
455
+@pytest.mark.parametrize("seq_lens", SEQ_LENS)
456
+@pytest.mark.parametrize("num_heads", NUM_HEADS)
457
+@pytest.mark.parametrize("head_size", [96])
458
+@pytest.mark.parametrize("block_size", [128])
459
+@pytest.mark.parametrize("sliding_window", SLIDING_WINDOWS)
460
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
461
+@pytest.mark.parametrize("soft_cap", [50])
462
+@pytest.mark.parametrize("num_blocks", NUM_BLOCKS)
463
+@pytest.mark.parametrize("use_alibi", [False])
464
+@pytest.mark.parametrize("use_sink", [False])
465
+@pytest.mark.parametrize(
466
+ "isa", ["amx"] if torch._C._cpu._is_amx_tile_supported() else ["vec"]
467
+)
468
+def test_varlen_with_paged_kv_softcap(
469
+ seq_lens: list[tuple[int, int]],
470
+ num_heads: tuple[int, int],
471
+ head_size: int,
472
+ sliding_window: int | None,
473
+ dtype: torch.dtype,
474
+ block_size: int,
475
+ soft_cap: float | None,
476
+ num_blocks: int,
477
+ use_alibi: bool,
478
+ use_sink: bool,
479
+ isa: str,
480
+) -> None:
481
+ varlen_with_paged_kv(
482
+ seq_lens=seq_lens,
483
+ num_heads=num_heads,
484
+ head_size=head_size,
485
+ sliding_window=sliding_window,
486
+ dtype=dtype,
487
+ block_size=block_size,
488
+ soft_cap=soft_cap,
489
+ num_blocks=num_blocks,
490
+ use_alibi=use_alibi,
491
+ use_sink=use_sink,
492
+ isa=isa,
493
+ )
494
+
495
+
496
+@pytest.mark.parametrize("seq_lens", SEQ_LENS)
497
+@pytest.mark.parametrize("num_heads", NUM_HEADS)
498
+@pytest.mark.parametrize("head_size", [96])
499
+@pytest.mark.parametrize("block_size", [128])
500
+@pytest.mark.parametrize("sliding_window", SLIDING_WINDOWS)
501
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
502
+@pytest.mark.parametrize("soft_cap", [None])
503
+@pytest.mark.parametrize("num_blocks", NUM_BLOCKS)
504
+@pytest.mark.parametrize("use_alibi", [True])
505
+@pytest.mark.parametrize("use_sink", [False])
506
+@pytest.mark.parametrize(
507
+ "isa", ["amx"] if torch._C._cpu._is_amx_tile_supported() else ["vec"]
508
+)
509
+def test_varlen_with_paged_kv_alibi(
510
+ seq_lens: list[tuple[int, int]],
511
+ num_heads: tuple[int, int],
512
+ head_size: int,
513
+ sliding_window: int | None,
514
+ dtype: torch.dtype,
515
+ block_size: int,
516
+ soft_cap: float | None,
517
+ num_blocks: int,
518
+ use_alibi: bool,
519
+ use_sink: bool,
520
+ isa: str,
521
+) -> None:
522
+ varlen_with_paged_kv(
523
+ seq_lens=seq_lens,
524
+ num_heads=num_heads,
525
+ head_size=head_size,
526
+ sliding_window=sliding_window,
527
+ dtype=dtype,
528
+ block_size=block_size,
529
+ soft_cap=soft_cap,
530
+ num_blocks=num_blocks,
531
+ use_alibi=use_alibi,
532
+ use_sink=use_sink,
533
+ isa=isa,
534
+ )
535
+
536
+
537
+@pytest.mark.parametrize("seq_lens", SEQ_LENS)
538
+@pytest.mark.parametrize("num_heads", NUM_HEADS)
539
+@pytest.mark.parametrize("head_size", [96])
540
+@pytest.mark.parametrize("block_size", [128])
541
+@pytest.mark.parametrize("sliding_window", SLIDING_WINDOWS)
542
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
543
+@pytest.mark.parametrize("soft_cap", [None])
544
+@pytest.mark.parametrize("num_blocks", NUM_BLOCKS)
545
+@pytest.mark.parametrize("use_alibi", [False])
546
+@pytest.mark.parametrize("use_sink", [True])
547
+@pytest.mark.parametrize(
548
+ "isa", ["amx"] if torch._C._cpu._is_amx_tile_supported() else ["vec"]
549
+)
550
+def test_varlen_with_paged_kv_sink(
551
+ seq_lens: list[tuple[int, int]],
552
+ num_heads: tuple[int, int],
553
+ head_size: int,
554
+ sliding_window: int | None,
555
+ dtype: torch.dtype,
556
+ block_size: int,
557
+ soft_cap: float | None,
558
+ num_blocks: int,
559
+ use_alibi: bool,
560
+ use_sink: bool,
561
+ isa: str,
562
+) -> None:
563
+ varlen_with_paged_kv(
564
+ seq_lens=seq_lens,
565
+ num_heads=num_heads,
566
+ head_size=head_size,
567
+ sliding_window=sliding_window,
568
+ dtype=dtype,
569
+ block_size=block_size,
570
+ soft_cap=soft_cap,
571
+ num_blocks=num_blocks,
572
+ use_alibi=use_alibi,
573
+ use_sink=use_sink,
574
+ isa=isa,
575
+ )
tests/kernels/test_onednn.py
-1
@@ -1,6 +1,5 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Integration tests for FlexAttention backend vs default backend"""
3
4
import pytest
5
import torch
tests/models/language/generation/test_common.py
+10
-7
@@ -38,7 +38,11 @@ AITER_MODEL_LIST = [
38
[
39
pytest.param(
40
"bigscience/bloom-560m", # bloom - testing alibi slopes
41
- marks=[pytest.mark.core_model, pytest.mark.slow_test],
41
+ marks=[
42
+ pytest.mark.core_model,
43
+ pytest.mark.slow_test,
44
+ pytest.mark.cpu_model,
45
+ ],
46
),
47
pytest.param(
48
"openai-community/gpt2", # gpt2
@@ -55,6 +59,10 @@ AITER_MODEL_LIST = [
59
pytest.mark.slow_test,
60
],
61
),
62
+ pytest.param(
63
+ "google/gemma-2-2b-it", # test hybrid attention
64
+ marks=[pytest.mark.cpu_model],
65
+ ),
66
pytest.param(
67
"zai-org/chatglm3-6b", # chatglm (text-only)
68
),
@@ -64,7 +72,6 @@ AITER_MODEL_LIST = [
72
),
73
pytest.param(
74
"openbmb/MiniCPM3-4B",
67
- # fused_moe not supported on CPU
75
marks=[pytest.mark.core_model, large_gpu_mark(min_gb=32)],
76
),
77
pytest.param(
@@ -93,11 +100,7 @@ AITER_MODEL_LIST = [
100
pytest.param("bigcode/starcoder2-3b"), # starcoder2
101
pytest.param(
102
"TitanML/tiny-mixtral", # mixtral
96
- marks=[pytest.mark.core_model],
97
- ),
98
- pytest.param(
99
- "allenai/OLMoE-1B-7B-0924-Instruct",
100
- marks=[pytest.mark.cpu_model],
103
+ marks=[pytest.mark.core_model, pytest.mark.cpu_model],
104
),
105
pytest.param("swiss-ai/Apertus-8B-Instruct-2509"), # apertus
106
],
tests/models/language/pooling/test_embedding.py
+1
-2
@@ -23,8 +23,7 @@ from ...utils import check_embeddings_close
23
),
24
pytest.param(
25
"intfloat/e5-mistral-7b-instruct",
26
- # CPU v1 doesn't support sliding window
27
- marks=[pytest.mark.core_model],
26
+ marks=[pytest.mark.core_model, pytest.mark.cpu_model],
27
),
28
pytest.param(
29
"ssmits/Qwen2-7B-Instruct-embed-base", marks=[pytest.mark.cpu_model]
tests/models/registry.py
+3
-1
@@ -243,7 +243,9 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
243
"FalconH1ForCausalLM": _HfExamplesInfo("tiiuae/Falcon-H1-0.5B-Base"),
244
"FlexOlmoForCausalLM": _HfExamplesInfo("allenai/Flex-reddit-2x7B-1T"),
245
"GemmaForCausalLM": _HfExamplesInfo("google/gemma-1.1-2b-it"),
246
- "Gemma2ForCausalLM": _HfExamplesInfo("google/gemma-2-9b"),
246
+ "Gemma2ForCausalLM": _HfExamplesInfo(
247
+ "google/gemma-2-9b", extras={"tiny": "google/gemma-2-2b-it"}
248
+ ),
249
"Gemma3ForCausalLM": _HfExamplesInfo("google/gemma-3-1b-it"),
250
"Gemma3nForCausalLM": _HfExamplesInfo("google/gemma-3n-E2B-it"),
251
"GlmForCausalLM": _HfExamplesInfo("zai-org/glm-4-9b-chat-hf"),
vllm/_custom_ops.py
+82
@@ -2583,6 +2583,88 @@ def onednn_scaled_mm(
2583
return output
2584
2585
2586
+def cpu_attn_get_scheduler_metadata(
2587
+ num_reqs: int,
2588
+ num_heads: int,
2589
+ num_kv_heads: int,
2590
+ head_dim: int,
2591
+ seq_lens: torch.Tensor,
2592
+ dtype: torch.dtype,
2593
+ query_start_loc: torch.Tensor,
2594
+ causal: bool,
2595
+ sliding_window_size: int,
2596
+ isa: str,
2597
+ enable_kv_split: bool,
2598
+) -> torch.Tensor:
2599
+ sheduler_metadata = torch.ops._C.get_scheduler_metadata(
2600
+ num_reqs,
2601
+ num_heads,
2602
+ num_kv_heads,
2603
+ head_dim,
2604
+ seq_lens,
2605
+ dtype,
2606
+ query_start_loc,
2607
+ causal,
2608
+ sliding_window_size,
2609
+ isa,
2610
+ enable_kv_split,
2611
+ )
2612
+ return sheduler_metadata
2613
+
2614
+
2615
+def cpu_attn_reshape_and_cache(
2616
+ key: torch.Tensor,
2617
+ value: torch.Tensor,
2618
+ key_cache: torch.Tensor,
2619
+ value_cache: torch.Tensor,
2620
+ slot_mapping: torch.Tensor,
2621
+ isa: str,
2622
+) -> None:
2623
+ torch.ops._C.cpu_attn_reshape_and_cache(
2624
+ key,
2625
+ value,
2626
+ key_cache,
2627
+ value_cache,
2628
+ slot_mapping,
2629
+ isa,
2630
+ )
2631
+
2632
+
2633
+def cpu_attention_with_kv_cache(
2634
+ query: torch.Tensor,
2635
+ key_cache: torch.Tensor,
2636
+ value_cache: torch.Tensor,
2637
+ output: torch.Tensor,
2638
+ query_start_loc: torch.Tensor,
2639
+ seq_lens: torch.Tensor,
2640
+ scale: float,
2641
+ causal: bool,
2642
+ alibi_slopes: torch.Tensor | None,
2643
+ sliding_window: tuple[int, int],
2644
+ block_table: torch.Tensor,
2645
+ softcap: float,
2646
+ scheduler_metadata: torch.Tensor,
2647
+ s_aux: torch.Tensor | None,
2648
+) -> None:
2649
+ torch.ops._C.cpu_attention_with_kv_cache(
2650
+ query,
2651
+ key_cache,
2652
+ value_cache,
2653
+ output,
2654
+ query_start_loc,
2655
+ seq_lens,
2656
+ scale,
2657
+ causal,
2658
+ alibi_slopes,
2659
+ sliding_window[0],
2660
+ sliding_window[1],
2661
+ block_table,
2662
+ softcap,
2663
+ scheduler_metadata,
2664
+ s_aux,
2665
+ )
2666
+
2667
+
2668
if hasattr(torch.ops._qutlass_C, "matmul_mxf4_bf16_tn"):
2669
2670
@register_fake("_qutlass_C::matmul_mxf4_bf16_tn")
vllm/attention/backends/registry.py
+2
-1
@@ -49,7 +49,7 @@ class AttentionBackendEnum(enum.Enum, metaclass=_AttentionBackendEnumMeta):
49
ROCM_AITER_FA = (
50
"vllm.v1.attention.backends.rocm_aiter_fa.AiterFlashAttentionBackend"
51
)
52
- TORCH_SDPA = "vllm.v1.attention.backends.cpu_attn.TorchSDPABackend"
52
+ TORCH_SDPA = "" # this tag is only used for ViT
53
FLASHINFER = "vllm.v1.attention.backends.flashinfer.FlashInferBackend"
54
FLASHINFER_MLA = (
55
"vllm.v1.attention.backends.mla.flashinfer_mla.FlashInferMLABackend"
@@ -70,6 +70,7 @@ class AttentionBackendEnum(enum.Enum, metaclass=_AttentionBackendEnumMeta):
70
"vllm.v1.attention.backends.rocm_aiter_unified_attn."
71
"RocmAiterUnifiedAttentionBackend"
72
)
73
+ CPU_ATTN = "vllm.v1.attention.backends.cpu_attn.CPUAttentionBackend"
74
# Placeholder for third-party/custom backends - must be registered before use
75
CUSTOM = ""
76
vllm/engine/arg_utils.py
-3
@@ -1726,9 +1726,6 @@ class EngineArgs:
1726
)
1727
_raise_unsupported_error(feature_name=name)
1728
1729
- if current_platform.is_cpu() and model_config.get_sliding_window() is not None:
1730
- _raise_unsupported_error(feature_name="sliding window (CPU backend)")
1731
-
1729
def _set_default_args(
1730
self, usage_context: UsageContext, model_config: ModelConfig
1731
) -> None:
vllm/platforms/cpu.py
+11
-26
@@ -8,7 +8,6 @@ import platform
8
import subprocess
9
import sys
10
from dataclasses import dataclass
11
-from importlib.util import find_spec
11
from typing import TYPE_CHECKING
12
13
import regex as re
@@ -139,16 +138,15 @@ class CpuPlatform(Platform):
138
) -> str:
139
from vllm.attention.backends.registry import AttentionBackendEnum
140
142
- if selected_backend and selected_backend != AttentionBackendEnum.TORCH_SDPA:
141
+ if selected_backend and selected_backend != AttentionBackendEnum.CPU_ATTN:
142
logger.info("Cannot use %s backend on CPU.", selected_backend)
143
if use_mla:
144
raise NotImplementedError("MLA is not supported on CPU.")
145
if use_sparse:
146
raise NotImplementedError("Sparse Attention is not supported on CPU.")
148
- logger.info("Using Torch SDPA backend.")
147
if not use_v1:
148
raise ValueError("CPU backend only supports V1.")
151
- return AttentionBackendEnum.TORCH_SDPA.get_path()
149
+ return AttentionBackendEnum.CPU_ATTN.get_path()
150
151
@classmethod
152
def get_device_total_memory(cls, device_id: int = 0) -> int:
@@ -186,15 +184,13 @@ class CpuPlatform(Platform):
184
185
cache_config = vllm_config.cache_config
186
189
- ipex_available = find_spec("intel_extension_for_pytorch") is not None
187
+ if cache_config.block_size is None:
188
+ cache_config.block_size = 128
189
191
- if cache_config and cache_config.block_size is None:
192
- cache_config.block_size = 128 if ipex_available else 16
193
-
194
- if not ipex_available and cache_config.block_size != 16:
195
- raise RuntimeError(
196
- f"--block-size={cache_config.block_size} requires"
197
- " intel_extension_for_pytorch"
190
+ if cache_config.block_size % 32 != 0:
191
+ logger.warning(
192
+ "CPU backend prefers block_size is multiples of 32, "
193
+ "otherwise the performance is not optimized."
194
)
195
196
scheduler_config = vllm_config.scheduler_config
@@ -207,22 +203,11 @@ class CpuPlatform(Platform):
203
"backend is not compatible with FP8 KV cache."
204
)
205
210
- if cache_config.cache_dtype == "fp8_e4m3":
211
- cache_config.cache_dtype = "fp8_e5m2"
212
- logger.warning(
213
- "CPU backend doesn't support fp8_e4m3 KV cache type, cast to fp8_e5m2."
214
- )
215
-
216
- if (
217
- cache_config.cache_dtype != "auto"
218
- and model_config is not None
219
- and model_config.dtype == torch.half
220
- ):
206
+ if cache_config.cache_dtype != "auto":
207
logger.warning(
222
- "FP8 KV cache on the CPU backend only does not"
223
- " support fp16 for now, cast to bf16."
208
+ "CPU backend doesn't support KV cache quantization fallback to auto."
209
)
225
- model_config.dtype = torch.bfloat16
210
+ cache_config.cache_dtype = "auto"
211
212
cache_config.cpu_kvcache_space_bytes = CpuPlatform.get_device_total_memory()
213
vllm/utils/__init__.py
-1
@@ -57,7 +57,6 @@ STR_BACKEND_ENV_VAR: str = "VLLM_ATTENTION_BACKEND"
57
# Possible string values of STR_BACKEND_ENV_VAR
58
# register, corresponding to possible backends
59
STR_FLASHINFER_ATTN_VAL: str = "FLASHINFER"
60
-STR_TORCH_SDPA_ATTN_VAL: str = "TORCH_SDPA"
60
STR_XFORMERS_ATTN_VAL: str = "XFORMERS"
61
STR_FLASH_ATTN_VAL: str = "FLASH_ATTN"
62
STR_INVALID_VAL: str = "INVALID"
vllm/v1/attention/backends/cpu_attn.py
+264
-717
@@ -1,22 +1,21 @@
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, Optional
4
+from typing import ClassVar
5
6
-import numpy as np
6
import torch
8
-from torch.nn.functional import scaled_dot_product_attention
7
8
+from vllm import _custom_ops as ops
9
from vllm.attention.backends.abstract import (
10
AttentionBackend,
11
AttentionImpl,
12
AttentionLayer,
14
- AttentionMetadata,
13
AttentionType,
14
is_quantized_kv_cache,
15
)
16
from vllm.config import VllmConfig
17
from vllm.logger import init_logger
18
+from vllm.platforms import CpuArchEnum, current_platform
19
from vllm.v1.attention.backends.utils import (
20
AttentionMetadataBuilder,
21
CommonAttentionMetadata,
@@ -24,44 +23,38 @@ from vllm.v1.attention.backends.utils import (
23
)
24
from vllm.v1.kv_cache_interface import AttentionSpec
25
27
-try:
28
- import intel_extension_for_pytorch.llm.modules as ipex_modules
29
-
30
- _use_ipex = True
31
-# AttributeError is to handle a bug in ipex
32
-# https://github.com/intel/intel-extension-for-pytorch/pull/813
33
-except (ImportError, AttributeError):
34
- _use_ipex = False
35
-
36
-from vllm import _custom_ops as ops
37
-
26
logger = init_logger(__name__)
27
28
+_CPU_ARCH_PREFER_MIXED_BATCH = (CpuArchEnum.X86,)
29
+
30
41
-class TorchSDPABackend(AttentionBackend):
42
- accept_output_buffer: bool = False
31
+class CPUAttentionBackend(AttentionBackend):
32
+ accept_output_buffer: bool = True
33
supported_dtypes: ClassVar[list[torch.dtype]] = [
34
torch.float16,
35
torch.bfloat16,
36
torch.float32,
37
]
38
39
+ @classmethod
40
+ def get_supported_dtypes(cls) -> list[torch.dtype]:
41
+ return [torch.float16, torch.bfloat16, torch.float32]
42
+
43
@classmethod
44
def get_supported_head_sizes(cls) -> list[int]:
51
- attn_impl = _get_paged_attn_impl()
52
- return attn_impl.get_supported_head_sizes()
45
+ return [32, 64, 96, 128, 160, 192, 224, 256]
46
47
@staticmethod
48
def get_name() -> str:
56
- return "TORCH_SDPA"
49
+ return "CPU_ATTN"
50
51
@staticmethod
59
- def get_impl_cls() -> type["TorchSDPABackendImpl"]:
60
- return TorchSDPABackendImpl
52
+ def get_impl_cls() -> type["CPUAttentionBackendImpl"]:
53
+ return CPUAttentionBackendImpl
54
55
@staticmethod
63
- def get_builder_cls() -> type["TorchSDPAMetadataBuilderV1"]:
64
- return TorchSDPAMetadataBuilderV1
56
+ def get_builder_cls() -> type["CPUAttentionMetadataBuilder"]:
57
+ return CPUAttentionMetadataBuilder
58
59
@staticmethod
60
def get_kv_cache_shape(
@@ -71,9 +64,7 @@ class TorchSDPABackend(AttentionBackend):
64
head_size: int,
65
cache_dtype_str: str = "auto",
66
) -> tuple[int, ...]:
74
- return _get_paged_attn_impl().get_kv_cache_shape(
75
- num_blocks, block_size, num_kv_heads, head_size
76
- )
67
+ return 2, num_blocks, num_kv_heads, block_size, head_size
68
69
@staticmethod
70
def use_cascade_attention(*args, **kwargs) -> bool:
@@ -81,264 +72,26 @@ class TorchSDPABackend(AttentionBackend):
72
73
74
@dataclass
84
-class TorchSDPAMetadata(AttentionMetadata):
85
- """Attention metadata for prefill and decode batched together."""
86
-
87
- # Total number of prefill requests.
88
- num_prefills: int
89
- # Number of prefill tokens.
90
- num_prefill_tokens: int
91
- # Number of decode tokens. Note that it is equivalent to the number of
92
- # decode requests.
93
- num_decode_tokens: int
94
- # (num_tokens,). The indices of the token slots that input tokens will be
95
- # stored into. E.g., if `slot_mapping` is [35, 2, 17] and the block size
96
- # is 16, the three tokens are stored in the 3rd slot in block 2, 2nd slot
97
- # in block 0, and 1st slot in block 1, respectively.
75
+class CPUAttentionMetadata:
76
+ isa: str
77
+ num_actual_tokens: int # Number of tokens excluding padding.
78
+ max_query_len: int
79
+ query_start_loc: torch.Tensor
80
+ max_seq_len: int
81
+ seq_lens: torch.Tensor
82
+ block_table: torch.Tensor
83
slot_mapping: torch.Tensor
99
- """Metadata for PagedAttention."""
100
- # (batch_size,). The length of sequences (entire tokens seen so far) per
101
- # sequence.
102
- decode_seq_lens_tensor: torch.Tensor | None
103
- # Maximum sequence length in the batch. 0 if it is prefill-only batch.
104
- decode_max_seq_len: int
105
- # (batch_size, max_blocks_per_seq).
106
- # Block addresses per sequence. (Seq id -> list of physical block)
107
- # E.g., [0, 1, 2] means tokens are stored in 0th, 1st, and 2nd blocks
108
- # in the kv cache. Each block can contain up to block_size tokens.
109
- # 2nd dimensions are padded up to max_blocks_per_seq if it is cuda-graph
110
- # captured.
111
- decode_block_tables: torch.Tensor | None
112
- """Metadata for TorchSDPABackend.
113
- """
114
- # Currently, input sequences can only contain all prompts
115
- # or all decoding. True if all sequences are prompts.
116
- chunked_prefill: bool
117
- seq_lens: list[int] | None = None # For non-chunked prefill
118
-
119
- # For chunked prefill only
120
- max_query_len: int | None = None
121
- prefill_max_seq_len: int | None = None
122
- prefill_query_start_loc: torch.Tensor | None = None
123
- prefill_seq_start_loc: torch.Tensor | None = None
124
- prefill_block_tables: torch.Tensor | None = None
125
-
126
- # For V1 logits index only
127
- query_start_loc: torch.Tensor | None = None
128
-
129
- # Begin encoder attn & enc/dec cross-attn fields...
130
- # Encoder sequence lengths representation
131
- encoder_seq_lens: list[int] | None = None
132
- encoder_seq_lens_tensor: torch.Tensor | None = None
133
-
134
- # Maximum sequence length among encoder sequences
135
- max_encoder_seq_len: int | None = None
136
-
137
- # Number of tokens input to encoder
138
- num_encoder_tokens: int | None = None
139
-
140
- # Cross-attention memory-mapping data structures: slot mapping
141
- # and block tables
142
- cross_slot_mapping: torch.Tensor | None = None
143
- cross_block_tables: torch.Tensor | None = None
144
-
145
- def __post_init__(self):
146
- # Set during the execution of the first attention op.
147
- # It is a list because it is needed to set per prompt
148
- # when alibi slopes is used. It is because of the limitation
149
- # from xformer API.
150
- # will not appear in the __repr__ and __init__
151
- self.attn_bias: list[torch.Tensor] | None = None
152
- self.encoder_attn_bias: list[torch.Tensor] | None = None
153
- self.cross_attn_bias: list[torch.Tensor] | None = None
154
-
155
- @property
156
- def is_all_encoder_attn_metadata_set(self):
157
- """
158
- All attention metadata required for encoder attention is set.
159
- """
160
- return (
161
- (self.encoder_seq_lens is not None)
162
- and (self.encoder_seq_lens_tensor is not None)
163
- and (self.max_encoder_seq_len is not None)
164
- )
165
-
166
- @property
167
- def is_all_cross_attn_metadata_set(self):
168
- """
169
- All attention metadata required for enc/dec cross-attention is set.
170
-
171
- Superset of encoder attention required metadata.
172
- """
173
- return (
174
- self.is_all_encoder_attn_metadata_set
175
- and (self.cross_slot_mapping is not None)
176
- and (self.cross_block_tables is not None)
177
- )
178
-
179
- @property
180
- def prefill_metadata(self) -> Optional["TorchSDPAMetadata"]:
181
- if self.num_prefill_tokens == 0:
182
- return None
183
- return self
184
-
185
- @property
186
- def decode_metadata(self) -> Optional["TorchSDPAMetadata"]:
187
- if self.num_decode_tokens == 0:
188
- return None
189
- return self
190
-
191
- def get_seq_lens(
192
- self,
193
- attn_type: str,
194
- ):
195
- """
196
- Extract appropriate sequence lengths from attention metadata
197
- according to attention type.
198
-
199
- Arguments:
200
-
201
- * attn_metadata: Attention metadata structure associated with attention
202
- * attn_type: encoder attention, decoder self-attention,
203
- encoder/decoder cross-attention
204
-
205
- Returns:
206
- * Appropriate sequence lengths tensor for query
207
- * Appropriate sequence lengths tensor for key & value
208
- """
209
-
210
- if (
211
- attn_type == AttentionType.DECODER
212
- or attn_type == AttentionType.ENCODER_ONLY
213
- ):
214
- seq_lens_q = self.seq_lens
215
- seq_lens_kv = self.seq_lens
216
- elif attn_type == AttentionType.ENCODER:
217
- seq_lens_q = self.encoder_seq_lens
218
- seq_lens_kv = self.encoder_seq_lens
219
- elif attn_type == AttentionType.ENCODER_DECODER:
220
- seq_lens_q = self.seq_lens
221
- seq_lens_kv = self.encoder_seq_lens
222
- else:
223
- raise AttributeError(f"Invalid attention type {str(attn_type)}")
224
- return seq_lens_q, seq_lens_kv
225
-
226
- def get_attn_bias(
227
- self,
228
- attn_type: str,
229
- ) -> list[torch.Tensor] | None:
230
- """
231
- Extract appropriate attention bias from attention metadata
232
- according to attention type.
233
-
234
- Arguments:
235
-
236
- * attn_metadata: Attention metadata structure associated with attention
237
- * attn_type: encoder attention, decoder self-attention,
238
- encoder/decoder cross-attention
239
-
240
- Returns:
241
- * Appropriate attention bias value given the attention type
242
- """
243
-
244
- if (
245
- attn_type == AttentionType.DECODER
246
- or attn_type == AttentionType.ENCODER_ONLY
247
- ):
248
- return self.attn_bias
249
- elif attn_type == AttentionType.ENCODER:
250
- return self.encoder_attn_bias
251
- elif attn_type == AttentionType.ENCODER_DECODER:
252
- return self.cross_attn_bias
253
- else:
254
- raise AttributeError(f"Invalid attention type {str(attn_type)}")
255
-
256
- def set_attn_bias(
257
- self,
258
- attn_bias: list[torch.Tensor],
259
- attn_type: str,
260
- ) -> None:
261
- """
262
- Update appropriate attention bias field of attention metadata,
263
- according to attention type.
264
-
265
- Arguments:
266
-
267
- * attn_metadata: Attention metadata structure associated with attention
268
- * attn_bias: The desired attention bias value
269
- * attn_type: encoder attention, decoder self-attention,
270
- encoder/decoder cross-attention
271
- """
272
-
273
- if (
274
- attn_type == AttentionType.DECODER
275
- or attn_type == AttentionType.ENCODER_ONLY
276
- ):
277
- self.attn_bias = attn_bias
278
- elif attn_type == AttentionType.ENCODER:
279
- self.encoder_attn_bias = attn_bias
280
- elif attn_type == AttentionType.ENCODER_DECODER:
281
- self.cross_attn_bias = attn_bias
282
- else:
283
- raise AttributeError(f"Invalid attention type {str(attn_type)}")
284
-
285
- def get_seq_len_block_table_args(
286
- self,
287
- attn_type: str,
288
- ) -> tuple:
289
- """
290
- The particular choice of sequence-length- and block-table-related
291
- attributes which should be extracted from attn_metadata is dependent
292
- on the type of attention operation.
293
-
294
- Decoder attn -> select entirely decoder self-attention-related fields
295
- Encoder/decoder cross-attn -> select encoder sequence lengths &
296
- cross-attn block-tables fields
297
- Encoder attn -> select encoder sequence lengths fields & no block tables
298
-
299
- Arguments:
300
-
301
- * attn_metadata: Attention metadata structure associated with attention
302
- * is_prompt: True if prefill, False otherwise
303
- * attn_type: encoder attention, decoder self-attention,
304
- encoder/decoder cross-attention
305
-
306
- Returns:
307
-
308
- * Appropriate sequence-lengths tensor
309
- * Appropriate max sequence-length scalar
310
- * Appropriate block tables (or None)
311
- """
312
-
313
- if (
314
- attn_type == AttentionType.DECODER
315
- or attn_type == AttentionType.ENCODER_ONLY
316
- ):
317
- # Decoder self-attention
318
- # Choose max_seq_len based on whether we are in prompt_run
319
- return (
320
- self.decode_seq_lens_tensor,
321
- self.decode_max_seq_len,
322
- self.decode_block_tables,
323
- )
324
- elif attn_type == AttentionType.ENCODER_DECODER:
325
- # Enc/dec cross-attention KVs match encoder sequence length;
326
- # cross-attention utilizes special "cross" block tables
327
- return (
328
- self.encoder_seq_lens_tensor,
329
- self.max_encoder_seq_len,
330
- self.cross_block_tables,
331
- )
332
- elif attn_type == AttentionType.ENCODER:
333
- # No block tables associated with encoder attention
334
- return (self.encoder_seq_lens_tensor, self.max_encoder_seq_len, None)
335
- else:
336
- raise AttributeError(f"Invalid attention type {str(attn_type)}")
84
+ scheduler_metadata: torch.Tensor | None
85
+ causal: bool = True
86
87
+ # can be removed after deprecate sdpa
88
+ use_sdpa_prefill: bool = False
89
+ num_decode_tokens: int = 0
90
+ sdpa_attn_masks: list[torch.Tensor | None] | None = None
91
+ sdpa_start_loc: torch.Tensor | None = None
92
339
-class TorchSDPAMetadataBuilderV1(AttentionMetadataBuilder[TorchSDPAMetadata]):
340
- reorder_batch_threshold: int = 1
93
94
+class CPUAttentionMetadataBuilder(AttentionMetadataBuilder[CPUAttentionMetadata]):
95
def __init__(
96
self,
97
kv_cache_spec: AttentionSpec,
@@ -348,80 +101,104 @@ class TorchSDPAMetadataBuilderV1(AttentionMetadataBuilder[TorchSDPAMetadata]):
101
) -> None:
102
super().__init__(kv_cache_spec, layer_names, vllm_config, device)
103
351
- self.scheduler_config = vllm_config.scheduler_config
352
- self._init_reorder_batch_threshold(1, False)
104
+ self.use_sdpa_prefill = False
105
+ reorder_batch_threshold = None
106
+ if current_platform.get_cpu_architecture() not in _CPU_ARCH_PREFER_MIXED_BATCH:
107
+ # in this case, decode seqs are reordered to the front of prefill seqs
108
+ # to split decode and prefill. Then use SDPA for prefill and
109
+ # cpu_attention_with_kv_cache for decode
110
+ reorder_batch_threshold = 1
111
+ self.use_sdpa_prefill = True
112
354
- self.seq_start_loc_cpu = torch.zeros(
355
- vllm_config.scheduler_config.max_num_seqs + 1,
356
- dtype=torch.int32,
357
- device="cpu",
113
+ self._init_reorder_batch_threshold(reorder_batch_threshold, False)
114
+
115
+ self.kv_cache_spec = kv_cache_spec
116
+ self.vllm_config = vllm_config
117
+
118
+ parallel_config = vllm_config.parallel_config
119
+ self.num_kv_heads = vllm_config.model_config.get_num_kv_heads(parallel_config)
120
+ self.num_heads = vllm_config.model_config.get_num_attention_heads(
121
+ parallel_config
122
)
359
- self.seq_start_loc_np = self.seq_start_loc_cpu.numpy()
123
+ self.head_dim = kv_cache_spec.head_size
124
+ self.dtype = vllm_config.model_config.dtype
125
+ self.window_size = getattr(kv_cache_spec, "sliding_window", -1)
126
+ if self.window_size is None:
127
+ self.window_size = -1
128
+ self.block_size = vllm_config.cache_config.block_size
129
+ self.isa = _get_attn_isa(self.dtype, self.block_size)
130
131
def build(
132
self,
133
common_prefix_len: int,
134
common_attn_metadata: CommonAttentionMetadata,
135
fast_build: bool = False,
366
- ) -> TorchSDPAMetadata:
136
+ ) -> CPUAttentionMetadata:
137
num_reqs = common_attn_metadata.num_reqs
138
+ num_actual_tokens = common_attn_metadata.num_actual_tokens
139
max_query_len = common_attn_metadata.max_query_len
369
-
370
- seq_lens_cpu = common_attn_metadata.seq_lens_cpu
371
- seq_lens_np = seq_lens_cpu.numpy()
372
-
373
- query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu
374
- query_start_loc_np = query_start_loc_cpu.numpy()
375
-
376
- num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
377
- split_decodes_and_prefills(
378
- common_attn_metadata,
379
- decode_threshold=self.reorder_batch_threshold,
380
- require_uniform=True,
381
- )
382
- )
383
-
384
- max_prefill_seq_len = (
385
- seq_lens_np[num_decodes:num_reqs].max().item() if num_prefills > 0 else 0
386
- )
387
- max_decode_seq_len = (
388
- seq_lens_np[:num_decodes].max().item() if num_prefills < num_reqs else 0
389
- )
390
- self.seq_start_loc_np[0] = 0
391
- np.cumsum(seq_lens_np, out=self.seq_start_loc_np[1 : num_reqs + 1])
392
-
393
- slot_mapping = common_attn_metadata.slot_mapping.long()
140
+ max_seq_len = common_attn_metadata.max_seq_len
141
+ query_start_loc = common_attn_metadata.query_start_loc
142
+ seq_lens = common_attn_metadata.seq_lens
143
block_table_tensor = common_attn_metadata.block_table_tensor
395
- query_start_loc_np = query_start_loc_cpu.numpy()
396
- query_start_loc_np[num_decodes : num_reqs + 1] -= num_decode_tokens
144
+ slot_mapping = common_attn_metadata.slot_mapping
145
+ causal = common_attn_metadata.causal
146
+
147
+ sdpa_start_loc = query_start_loc
148
+ num_decode_tokens = 0
149
+ if self.use_sdpa_prefill and causal:
150
+ # Decoder, need reorder and truncate
151
+ assert self.reorder_batch_threshold
152
+ (num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens) = (
153
+ split_decodes_and_prefills(
154
+ common_attn_metadata,
155
+ decode_threshold=self.reorder_batch_threshold,
156
+ require_uniform=True,
157
+ )
158
+ )
159
+ num_reqs = num_decodes
160
+ sdpa_start_loc = sdpa_start_loc[num_decodes:] - num_decode_tokens
161
+ seq_lens = seq_lens[:num_decodes]
162
+ query_start_loc = query_start_loc[: num_decodes + 1]
163
+ block_table_tensor = block_table_tensor[:num_decodes]
164
+
165
+ sheduler_metadata = None
166
+ if causal:
167
+ # for decode batch, use the custom kernel
168
+ sheduler_metadata = ops.cpu_attn_get_scheduler_metadata(
169
+ num_reqs=num_reqs,
170
+ num_heads=self.num_heads,
171
+ num_kv_heads=self.num_kv_heads,
172
+ head_dim=self.head_dim,
173
+ seq_lens=seq_lens,
174
+ dtype=self.dtype,
175
+ query_start_loc=query_start_loc,
176
+ causal=causal,
177
+ sliding_window_size=self.window_size,
178
+ isa=self.isa,
179
+ enable_kv_split=True,
180
+ )
181
398
- attn_metadata = TorchSDPAMetadata(
399
- num_prefills=num_prefills,
400
- num_prefill_tokens=num_prefill_tokens,
401
- num_decode_tokens=num_decode_tokens,
402
- slot_mapping=slot_mapping,
403
- # to ensure inference when chunked_prefill is disabled
404
- seq_lens=seq_lens_cpu.tolist()[num_decodes:], # prefill
405
- decode_seq_lens_tensor=seq_lens_cpu[:num_decodes], # decode
406
- decode_max_seq_len=max_decode_seq_len, # decode
407
- decode_block_tables=block_table_tensor[:num_decodes], # decode
408
- chunked_prefill=self.scheduler_config.chunked_prefill_enabled,
182
+ attn_metadata = CPUAttentionMetadata(
183
+ isa=self.isa,
184
+ num_actual_tokens=num_actual_tokens,
185
max_query_len=max_query_len,
410
- prefill_max_seq_len=max_prefill_seq_len,
411
- prefill_query_start_loc=query_start_loc_cpu[
412
- num_decodes : num_reqs + 1
413
- ], # prefill
414
- prefill_seq_start_loc=self.seq_start_loc_cpu[
415
- num_decodes : num_reqs + 1
416
- ], # prefill
417
- prefill_block_tables=block_table_tensor[num_decodes:num_reqs], # prefill
418
- query_start_loc=query_start_loc_cpu[: num_reqs + 1], # for logits index
186
+ query_start_loc=query_start_loc,
187
+ max_seq_len=max_seq_len,
188
+ seq_lens=seq_lens,
189
+ block_table=block_table_tensor,
190
+ slot_mapping=slot_mapping,
191
+ scheduler_metadata=sheduler_metadata,
192
+ causal=causal,
193
+ use_sdpa_prefill=self.use_sdpa_prefill,
194
+ num_decode_tokens=num_decode_tokens,
195
+ sdpa_start_loc=sdpa_start_loc,
196
)
197
198
return attn_metadata
199
200
424
-class TorchSDPABackendImpl(AttentionImpl[TorchSDPAMetadata]):
201
+class CPUAttentionBackendImpl(AttentionImpl):
202
def __init__(
203
self,
204
num_heads: int,
@@ -434,37 +211,48 @@ class TorchSDPABackendImpl(AttentionImpl[TorchSDPAMetadata]):
211
logits_soft_cap: float | None = None,
212
attn_type: str = AttentionType.DECODER,
213
kv_sharing_target_layer_name: str | None = None,
214
+ sinks: torch.Tensor | None = None,
215
) -> None:
438
- if kv_sharing_target_layer_name is not None:
439
- raise NotImplementedError("KV sharing is not supported in V0.")
440
- if logits_soft_cap is not None:
441
- logger.warning_once(
442
- "Torch SPDA does not support logits soft cap. "
443
- "Outputs may be slightly off."
444
- )
445
- self.paged_attn_impl = _get_paged_attn_impl()
216
+ self.kv_sharing_target_layer_name = kv_sharing_target_layer_name
217
self.num_heads = num_heads
218
self.head_size = head_size
219
self.scale = float(scale)
220
+ if logits_soft_cap is not None and attn_type in (
221
+ AttentionType.ENCODER,
222
+ AttentionType.ENCODER_ONLY,
223
+ ):
224
+ logger.warning_once(
225
+ "CPU_ATTN does not support logits softcap for"
226
+ " ENCODER and ENCODER_ONLY, outputs may be slightly off"
227
+ )
228
+ if logits_soft_cap is None:
229
+ logits_soft_cap = 0
230
+ self.logits_soft_cap = logits_soft_cap
231
+
232
self.num_kv_heads = num_kv_heads
233
if alibi_slopes is not None:
234
alibi_slopes = torch.tensor(alibi_slopes, dtype=torch.float32)
235
self.alibi_slopes = alibi_slopes
453
- self.sliding_window = sliding_window
236
+ if sliding_window is None:
237
+ self.sliding_window = (-1, -1)
238
+ elif attn_type == AttentionType.ENCODER_ONLY:
239
+ self.sliding_window = (sliding_window - 1, sliding_window - 1)
240
+ else:
241
+ self.sliding_window = (sliding_window - 1, 0)
242
self.kv_cache_dtype = kv_cache_dtype
455
-
243
self.num_queries_per_kv = self.num_heads // self.num_kv_heads
457
- self.need_mask = (
458
- self.alibi_slopes is not None or self.sliding_window is not None
459
- )
244
461
- if is_quantized_kv_cache(kv_cache_dtype) and not _use_ipex:
462
- raise NotImplementedError(
463
- "Torch SDPA backend FP8 KV cache requires "
464
- "intel_extension_for_pytorch support."
465
- )
245
+ if is_quantized_kv_cache(kv_cache_dtype):
246
+ raise NotImplementedError("FP8 KV cache is unsupported in CPU_ATTN")
247
self.attn_type = attn_type
248
249
+ self.sinks = sinks
250
+ if self.sinks is not None:
251
+ assert self.sinks.shape[0] == num_heads, (
252
+ "Sinks must have the same number of heads as the number of "
253
+ "heads in the layer"
254
+ )
255
+
256
def forward(
257
self,
258
layer: AttentionLayer,
@@ -472,196 +260,130 @@ class TorchSDPABackendImpl(AttentionImpl[TorchSDPAMetadata]):
260
key: torch.Tensor,
261
value: torch.Tensor,
262
kv_cache: torch.Tensor,
475
- attn_metadata: TorchSDPAMetadata, # type: ignore
263
+ attn_metadata: CPUAttentionMetadata | None,
264
output: torch.Tensor | None = None,
265
output_scale: torch.Tensor | None = None,
266
output_block_scale: torch.Tensor | None = None,
267
) -> torch.Tensor:
480
- """Forward pass with torch SDPA and PagedAttention.
268
+ """Forward pass for CPU attention backend.
269
270
Args:
483
- query: shape = [num_tokens, num_heads * head_size]
484
- key: shape = [num_tokens, num_kv_heads * head_size]
485
- value: shape = [num_tokens, num_kv_heads * head_size]
271
+ query: shape = [num_tokens, num_heads, head_size]
272
+ key: shape = [num_tokens, num_kv_heads, head_size]
273
+ value: shape = [num_tokens, num_kv_heads, head_size]
274
kv_cache: shape =
487
- [2, num_blocks, block_size * num_kv_heads * head_size]
488
- NOTE: kv_cache will be an empty tensor with shape [0]
489
- for profiling run.
275
+ [2, num_blocks, num_kv_heads, block_size, head_size]
276
attn_metadata: Metadata for attention.
277
Returns:
278
shape = [num_tokens, num_heads * head_size]
279
"""
280
+ assert output is not None, "Output tensor must be provided."
281
if output_scale is not None or output_block_scale is not None:
282
raise NotImplementedError(
283
"fused output quantization is not yet supported"
497
- " for TorchSDPABackendImpl"
284
+ " for CPUAttentionBackendImpl"
285
)
286
287
# For warming-up
288
if attn_metadata is None:
502
- return query
503
-
504
- attn_type = self.attn_type
505
- if attn_type == AttentionType.ENCODER and (
506
- not attn_metadata.is_all_encoder_attn_metadata_set
507
- ):
508
- raise AttributeError(
509
- "Encoder attention requires setting encoder metadata attributes."
510
- )
511
- elif attn_type == AttentionType.ENCODER_DECODER and (
512
- not attn_metadata.is_all_cross_attn_metadata_set
513
- ):
514
- raise AttributeError(
515
- "Encoder/decoder cross-attention "
516
- "requires setting cross-attention "
517
- "metadata attributes."
289
+ return output
290
+
291
+ num_actual_tokens = attn_metadata.num_actual_tokens
292
+
293
+ # Handle encoder attention differently - no KV cache needed
294
+ if self.attn_type in (AttentionType.ENCODER_ONLY, AttentionType.ENCODER):
295
+ # For encoder attention,
296
+ return self._run_sdpa_forward(
297
+ query[:num_actual_tokens],
298
+ key[:num_actual_tokens],
299
+ value[:num_actual_tokens],
300
+ output[:num_actual_tokens],
301
+ attn_metadata,
302
+ self.attn_type,
303
)
304
520
- # Reshape the query, key, and value tensors.
521
- query = query.view(-1, self.num_heads, self.head_size)
522
- if key is not None:
523
- assert value is not None
524
- key = key.view(-1, self.num_kv_heads, self.head_size)
525
- value = value.view(-1, self.num_kv_heads, self.head_size)
526
- else:
527
- assert value is None
528
-
529
- if attn_type != AttentionType.ENCODER and kv_cache.numel() > 0:
530
- # KV-cache during decoder-self- or
531
- # encoder-decoder-cross-attention, but not
532
- # during encoder attention.
533
- #
534
- # Even if there are no new key/value pairs to cache,
535
- # we still need to break out key_cache and value_cache
536
- # i.e. for later use by paged attention
537
- key_cache, value_cache = self.paged_attn_impl.split_kv_cache(
538
- kv_cache, self.num_kv_heads, self.head_size
539
- )
305
+ # For decoder and cross-attention, use KV cache, size are
306
+ # [num_blocks, num_kv_heads, block_size, head_size]
307
+ key_cache, value_cache = kv_cache.unbind(0)
308
541
- if (key is not None) and (value is not None):
542
- if attn_type == AttentionType.ENCODER_DECODER:
543
- # Update cross-attention KV cache (prefill-only)
544
- # During cross-attention decode, key & value will be None,
545
- # preventing this IF-statement branch from running
546
- updated_slot_mapping = attn_metadata.cross_slot_mapping
547
- else:
548
- # Update self-attention KV cache (prefill/decode)
549
- updated_slot_mapping = attn_metadata.slot_mapping
550
-
551
- self.paged_attn_impl.write_to_paged_cache(
552
- key,
553
- value,
554
- key_cache,
555
- value_cache,
556
- updated_slot_mapping,
557
- self.kv_cache_dtype,
558
- layer._k_scale,
559
- layer._v_scale,
560
- )
309
+ # key and value may be None in the case of cross attention. They are
310
+ # calculated once based on the output from the encoder and then cached
311
+ # in KV cache.
312
+ if (
313
+ self.kv_sharing_target_layer_name is None
314
+ and key is not None
315
+ and value is not None
316
+ ):
317
+ ops.cpu_attn_reshape_and_cache(
318
+ key,
319
+ value,
320
+ key_cache,
321
+ value_cache,
322
+ attn_metadata.slot_mapping,
323
+ attn_metadata.isa,
324
+ )
325
562
- if attn_type != AttentionType.ENCODER:
563
- # Decoder self-attention supports chunked prefill.
564
- # Encoder/decoder cross-attention requires no chunked
565
- # prefill (100% prefill or 100% decode tokens, no mix)
566
- num_prefill_tokens = attn_metadata.num_prefill_tokens
326
+ if attn_metadata.use_sdpa_prefill:
327
+ assert self.sinks is None, "Attention sink is unsupported in SDPA prefill"
328
num_decode_tokens = attn_metadata.num_decode_tokens
568
- else:
569
- # Encoder attention - chunked prefill is not applicable;
570
- # derive token-count from query shape & and treat them
571
- # as 100% prefill tokens
572
- assert attn_metadata.num_encoder_tokens is not None
573
- num_prefill_tokens = attn_metadata.num_encoder_tokens
574
- num_decode_tokens = 0
575
-
576
- if attn_type == AttentionType.DECODER:
577
- # Only enforce this shape-constraint for decoder
578
- # self-attention
579
- assert key.shape[0] == num_prefill_tokens + num_decode_tokens
580
- assert value.shape[0] == num_prefill_tokens + num_decode_tokens
581
-
582
- output = torch.empty_like(query)
583
- if prefill_meta := attn_metadata.prefill_metadata:
584
- if not prefill_meta.prefill_metadata.chunked_prefill: # type: ignore
585
- assert attn_metadata.seq_lens is not None
586
- self._run_sdpa_forward(
587
- output, query, key, value, prefill_meta, attn_type=attn_type
588
- )
589
- else:
590
- # prefix-enabled attention
591
- assert not self.need_mask
592
- import intel_extension_for_pytorch.llm.modules as ipex_modules
593
-
594
- output = torch.empty_like(query)
595
- ipex_modules.PagedAttention.flash_attn_varlen_func(
596
- output[prefill_meta.num_decode_tokens :, :, :],
597
- query[prefill_meta.num_decode_tokens :, :, :],
598
- key_cache,
599
- value_cache,
600
- prefill_meta.prefill_query_start_loc,
601
- prefill_meta.prefill_seq_start_loc,
602
- prefill_meta.max_query_len,
603
- prefill_meta.prefill_max_seq_len,
604
- self.scale,
605
- True,
606
- prefill_meta.prefill_block_tables,
607
- self.alibi_slopes,
608
- )
609
- if decode_meta := attn_metadata.decode_metadata:
610
- assert attn_type != AttentionType.ENCODER_ONLY, (
611
- "Encoder-only models should not have decode metadata."
329
+ self._run_sdpa_forward(
330
+ query[num_decode_tokens:num_actual_tokens],
331
+ key[num_decode_tokens:num_actual_tokens],
332
+ value[num_decode_tokens:num_actual_tokens],
333
+ output[num_decode_tokens:num_actual_tokens],
334
+ attn_metadata,
335
+ self.attn_type,
336
)
613
- # Decoding run.
614
- (
615
- seq_lens_arg,
616
- max_seq_len_arg,
617
- block_tables_arg,
618
- ) = decode_meta.get_seq_len_block_table_args(attn_type)
619
-
620
- self.paged_attn_impl.forward_decode(
621
- output[: attn_metadata.num_decode_tokens, :, :],
622
- query[: attn_metadata.num_decode_tokens, :, :],
623
- key_cache,
624
- value_cache,
625
- block_tables_arg,
626
- seq_lens_arg,
627
- max_seq_len_arg,
628
- self.kv_cache_dtype,
629
- self.num_kv_heads,
630
- self.scale,
631
- self.alibi_slopes,
632
- layer._k_scale,
633
- layer._v_scale,
337
+ num_actual_tokens = num_decode_tokens
338
+
339
+ if num_actual_tokens > 0:
340
+ ops.cpu_attention_with_kv_cache(
341
+ query=query[:num_actual_tokens],
342
+ key_cache=key_cache,
343
+ value_cache=value_cache,
344
+ output=output[:num_actual_tokens], # type: ignore
345
+ query_start_loc=attn_metadata.query_start_loc,
346
+ seq_lens=attn_metadata.seq_lens,
347
+ scale=self.scale,
348
+ causal=attn_metadata.causal,
349
+ alibi_slopes=self.alibi_slopes, # type: ignore
350
+ sliding_window=self.sliding_window,
351
+ block_table=attn_metadata.block_table,
352
+ softcap=self.logits_soft_cap,
353
+ scheduler_metadata=attn_metadata.scheduler_metadata,
354
+ s_aux=self.sinks,
355
)
356
636
- # Reshape the output tensor.
637
- return output.view(-1, self.num_heads * self.head_size)
357
+ return output
358
359
def _run_sdpa_forward(
360
self,
641
- output: torch.Tensor,
361
query: torch.Tensor,
362
key: torch.Tensor,
363
value: torch.Tensor,
645
- attn_metadata: TorchSDPAMetadata,
646
- attn_type: str = AttentionType.DECODER,
647
- ) -> None:
648
- attn_masks = attn_metadata.get_attn_bias(attn_type)
364
+ output: torch.Tensor,
365
+ attn_metadata: CPUAttentionMetadata,
366
+ attn_type: str,
367
+ ) -> torch.Tensor:
368
+ attn_masks = attn_metadata.sdpa_attn_masks
369
if attn_masks is None:
370
if self.alibi_slopes is not None:
371
attn_masks = _make_alibi_bias(
372
self.alibi_slopes,
373
query.dtype,
654
- attn_metadata.seq_lens, # type: ignore
374
+ attn_metadata.sdpa_start_loc,
375
)
656
- elif self.sliding_window is not None:
376
+ elif self.sliding_window[0] != -1 or self.sliding_window[1] != -1:
377
assert attn_metadata.seq_lens is not None
378
attn_masks = _make_sliding_window_bias(
659
- attn_metadata.seq_lens, self.sliding_window, query.dtype
379
+ attn_metadata.sdpa_start_loc,
380
+ self.sliding_window[0],
381
+ self.sliding_window[1],
382
+ query.dtype,
383
)
384
else:
662
- seq_lens, _ = attn_metadata.get_seq_lens(attn_type)
663
- attn_masks = [None] * len(seq_lens)
664
- attn_metadata.set_attn_bias(attn_masks, attn_type)
385
+ attn_masks = [None] * (attn_metadata.sdpa_start_loc.size(0) - 1) # type: ignore
386
+ attn_metadata.sdpa_attn_masks = attn_masks
387
388
query = query.movedim(0, query.dim() - 2)
389
key = key.movedim(0, key.dim() - 2)
@@ -673,21 +395,16 @@ class TorchSDPABackendImpl(AttentionImpl[TorchSDPAMetadata]):
395
396
causal_attn = attn_type == AttentionType.DECODER
397
676
- seq_lens_q, seq_lens_kv = attn_metadata.get_seq_lens(attn_type)
677
- # Incoming Q and KV contain decoded tokens as well, hence start at an offset
678
- # equal to num_decode_tokens since decode requests appear first
679
- start_q, start_kv = (
680
- attn_metadata.num_decode_tokens,
681
- attn_metadata.num_decode_tokens,
682
- )
683
- for seq_len_q, seq_len_kv, mask in zip(seq_lens_q, seq_lens_kv, attn_masks):
684
- end_q = start_q + seq_len_q
685
- end_kv = start_kv + seq_len_kv
398
+ sdpa_start_loc = attn_metadata.sdpa_start_loc.numpy() # type: ignore
399
+ for i in range(len(attn_masks)):
400
+ mask = attn_masks[i]
401
+ start_q = sdpa_start_loc[i]
402
+ end_q = sdpa_start_loc[i + 1]
403
sub_out = (
687
- scaled_dot_product_attention(
404
+ torch.nn.functional.scaled_dot_product_attention(
405
query[None, :, start_q:end_q, :],
689
- key[None, :, start_kv:end_kv, :],
690
- value[None, :, start_kv:end_kv, :],
406
+ key[None, :, start_q:end_q, :],
407
+ value[None, :, start_q:end_q, :],
408
attn_mask=mask,
409
dropout_p=0.0,
410
is_causal=causal_attn and mask is None,
@@ -697,17 +414,20 @@ class TorchSDPABackendImpl(AttentionImpl[TorchSDPAMetadata]):
414
.movedim(query.dim() - 2, 0)
415
)
416
output[start_q:end_q, :, :] = sub_out
700
- start_q, start_kv = end_q, end_kv
417
+ return output
418
419
420
def _make_alibi_bias(
421
alibi_slopes: torch.Tensor,
422
dtype: torch.dtype,
706
- seq_lens: list[int],
423
+ sdpa_start_loc: torch.Tensor,
424
) -> list[torch.Tensor]:
425
attn_biases: list[torch.Tensor] = []
709
- for seq_len in seq_lens:
710
- bias = torch.arange(seq_len, dtype=dtype)
426
+ seq_num = sdpa_start_loc.size(0) - 1
427
+ sdpa_start_loc = sdpa_start_loc.numpy() # type: ignore
428
+ for i in range(seq_num):
429
+ seq_len = sdpa_start_loc[i + 1] - sdpa_start_loc[i]
430
+ bias = torch.arange(seq_len, dtype=dtype) # type: ignore
431
# NOTE(zhuohan): HF uses
432
# `bias = bias[None, :].repeat(seq_len, 1)`
433
# here. We find that both biases give the same results, but
@@ -719,7 +439,7 @@ def _make_alibi_bias(
439
bias = bias[None, :].repeat((num_heads, 1, 1))
440
bias.mul_(alibi_slopes[:, None, None]).unsqueeze_(0)
441
inf_mask = (
722
- torch.empty((1, seq_len, seq_len), dtype=bias.dtype)
442
+ torch.empty((1, seq_len, seq_len), dtype=bias.dtype) # type: ignore
443
.fill_(-torch.inf)
444
.triu_(diagonal=1)
445
)
@@ -729,210 +449,37 @@ def _make_alibi_bias(
449
450
451
def _make_sliding_window_bias(
732
- seq_lens: list[int],
733
- window_size: int | None,
452
+ sdpa_start_loc: torch.Tensor,
453
+ left_window_size: int,
454
+ right_window_size: int,
455
dtype: torch.dtype,
456
) -> list[torch.Tensor]:
457
attn_biases: list[torch.Tensor] = []
737
- for seq_len in seq_lens:
738
- tensor = torch.full(
739
- (1, seq_len, seq_len),
740
- dtype=dtype,
458
+ seq_num = sdpa_start_loc.size(0) - 1
459
+ sdpa_start_loc = sdpa_start_loc.numpy() # type: ignore
460
+ for i in range(seq_num):
461
+ seq_len = sdpa_start_loc[i + 1] - sdpa_start_loc[i]
462
+ mask = torch.full( # type: ignore
463
+ (1, seq_len, seq_len), # type: ignore
464
fill_value=1,
465
+ dtype=dtype,
466
)
743
- shift = 0
744
- mask = torch.tril(tensor, diagonal=shift).to(dtype) # type: ignore
745
- if window_size is not None:
746
- mask = torch.triu(mask, diagonal=shift - window_size + 1)
467
+
468
+ if right_window_size != -1:
469
+ mask = torch.tril(mask, diagonal=right_window_size)
470
+ if left_window_size != -1:
471
+ mask = torch.triu(mask, diagonal=-left_window_size)
472
mask = torch.log(mask)
748
- attn_biases.append(mask.to(dtype))
473
+ attn_biases.append(mask)
474
475
return attn_biases
476
477
753
-class _PagedAttention:
754
- @staticmethod
755
- def get_supported_head_sizes() -> list[int]:
756
- return [32, 64, 80, 96, 112, 128, 192, 256]
757
-
758
- @staticmethod
759
- def get_kv_cache_shape(
760
- num_blocks: int,
761
- block_size: int,
762
- num_kv_heads: int,
763
- head_size: int,
764
- *args,
765
- ) -> tuple[int, ...]:
766
- return 2, num_blocks, block_size * num_kv_heads * head_size
767
-
768
- @staticmethod
769
- def split_kv_cache(
770
- kv_cache: torch.Tensor,
771
- num_kv_heads: int,
772
- head_size: int,
773
- *args,
774
- ) -> tuple[torch.Tensor, torch.Tensor]:
775
- x = 16 // kv_cache.element_size()
776
- num_blocks = kv_cache.shape[1]
777
-
778
- key_cache = kv_cache[0]
779
- key_cache = key_cache.view(num_blocks, num_kv_heads, head_size // x, -1, x)
780
- value_cache = kv_cache[1]
781
- value_cache = value_cache.view(num_blocks, num_kv_heads, head_size, -1)
782
- return key_cache, value_cache
783
-
784
- @staticmethod
785
- def write_to_paged_cache(
786
- key: torch.Tensor,
787
- value: torch.Tensor,
788
- key_cache: torch.Tensor,
789
- value_cache: torch.Tensor,
790
- slot_mapping: torch.Tensor,
791
- kv_cache_dtype: str,
792
- k_scale: torch.Tensor,
793
- v_scale: torch.Tensor,
794
- *args,
795
- ) -> None:
796
- ops.reshape_and_cache(
797
- key,
798
- value,
799
- key_cache,
800
- value_cache,
801
- slot_mapping.flatten(),
802
- kv_cache_dtype,
803
- k_scale,
804
- v_scale,
805
- )
806
-
807
- @staticmethod
808
- def forward_decode(
809
- output: torch.Tensor,
810
- query: torch.Tensor,
811
- key_cache: torch.Tensor,
812
- value_cache: torch.Tensor,
813
- block_tables: torch.Tensor,
814
- context_lens: torch.Tensor,
815
- max_context_len: int,
816
- kv_cache_dtype: str,
817
- num_kv_heads: int,
818
- scale: float,
819
- alibi_slopes: torch.Tensor | None,
820
- k_scale: torch.Tensor,
821
- v_scale: torch.Tensor,
822
- *args,
823
- ) -> None:
824
- tp_rank: int = 0
825
- blocksparse_local_blocks: int = 0
826
- blocksparse_vert_stride: int = 0
827
- blocksparse_block_size: int = 64
828
- blocksparse_head_sliding_step: int = 0
829
- block_size = value_cache.shape[3]
830
-
831
- ops.paged_attention_v1(
832
- output,
833
- query,
834
- key_cache,
835
- value_cache,
836
- num_kv_heads,
837
- scale,
838
- block_tables,
839
- context_lens,
840
- block_size,
841
- max_context_len,
842
- alibi_slopes,
843
- kv_cache_dtype,
844
- k_scale,
845
- v_scale,
846
- tp_rank,
847
- blocksparse_local_blocks,
848
- blocksparse_vert_stride,
849
- blocksparse_block_size,
850
- blocksparse_head_sliding_step,
851
- )
852
-
853
-
854
-class _IPEXPagedAttention(_PagedAttention):
855
- @staticmethod
856
- def get_supported_head_sizes() -> list[int]:
857
- return []
858
-
859
- @staticmethod
860
- def split_kv_cache(
861
- kv_cache: torch.Tensor,
862
- num_kv_heads: int,
863
- head_size: int,
864
- *args,
865
- ) -> tuple[torch.Tensor, torch.Tensor]:
866
- num_blocks = kv_cache.shape[1]
867
-
868
- key_cache = kv_cache[0]
869
- key_cache = key_cache.view(num_blocks, num_kv_heads, -1, head_size)
870
- value_cache = kv_cache[1]
871
- value_cache = value_cache.view(num_blocks, num_kv_heads, -1, head_size)
872
- return key_cache, value_cache
873
-
874
- @staticmethod
875
- def write_to_paged_cache(
876
- key: torch.Tensor,
877
- value: torch.Tensor,
878
- key_cache: torch.Tensor,
879
- value_cache: torch.Tensor,
880
- slot_mapping: torch.Tensor,
881
- kv_cache_dtype: str,
882
- k_scale: torch.Tensor,
883
- v_scale: torch.Tensor,
884
- *args,
885
- ) -> None:
886
- ipex_modules.PagedAttention.reshape_and_cache(
887
- key, value, key_cache, value_cache, slot_mapping.flatten().int()
888
- )
889
-
890
- @staticmethod
891
- def forward_decode(
892
- output: torch.Tensor,
893
- query: torch.Tensor,
894
- key_cache: torch.Tensor,
895
- value_cache: torch.Tensor,
896
- block_tables: torch.Tensor,
897
- context_lens: torch.Tensor,
898
- max_context_len: int,
899
- kv_cache_dtype: str,
900
- num_kv_heads: int,
901
- scale: float,
902
- alibi_slopes: torch.Tensor | None,
903
- k_scale: torch.Tensor,
904
- v_scale: torch.Tensor,
905
- *args,
906
- ) -> None:
907
- block_size = value_cache.shape[2]
908
- head_mapping = (
909
- torch.arange(
910
- 0,
911
- num_kv_heads,
912
- device="cpu",
913
- dtype=torch.int32,
914
- )
915
- .view(num_kv_heads, 1)
916
- .repeat_interleave(query.size(1) // num_kv_heads)
917
- .flatten()
918
- )
919
- ipex_modules.PagedAttention.single_query_cached_kv_attention(
920
- output,
921
- query.contiguous(),
922
- key_cache,
923
- value_cache,
924
- head_mapping,
925
- scale,
926
- block_tables,
927
- context_lens,
928
- block_size,
929
- max_context_len,
930
- alibi_slopes,
931
- )
932
-
933
-
934
-def _get_paged_attn_impl():
935
- if _use_ipex:
936
- return _IPEXPagedAttention
478
+def _get_attn_isa(dtype: torch.dtype, block_size: int) -> str:
479
+ supports_amx = torch._C._cpu._is_amx_tile_supported()
480
+ if supports_amx and dtype in (torch.bfloat16,) and block_size % 32 == 0:
481
+ return "amx"
482
+ elif block_size % 32 == 0:
483
+ return "vec"
484
else:
938
- return _PagedAttention
485
+ return "vec16"
vllm/v1/attention/backends/utils.py
+1
-1
@@ -265,7 +265,7 @@ class AttentionMetadataBuilder(abc.ABC, Generic[M]):
265
266
def _init_reorder_batch_threshold(
267
self,
268
- reorder_batch_threshold: int = 1,
268
+ reorder_batch_threshold: int | None = 1,
269
supports_spec_as_decode: bool = False,
270
supports_dcp_with_varlen: bool = False,
271
) -> None:
vllm/v1/worker/cpu_model_runner.py
+1
-13
@@ -1,7 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from contextlib import contextmanager
4
-from typing import TYPE_CHECKING, Any
4
+from typing import Any
5
6
import torch
7
import torch.nn as nn
@@ -12,9 +12,6 @@ from vllm.model_executor.model_loader import get_model
12
from vllm.v1.utils import CpuGpuBuffer
13
from vllm.v1.worker.gpu_model_runner import GPUModelRunner
14
15
-if TYPE_CHECKING:
16
- from vllm.v1.core.sched.output import SchedulerOutput
17
-
15
logger = init_logger(__name__)
16
17
@@ -31,15 +28,6 @@ class CPUModelRunner(GPUModelRunner):
28
29
self._postprocess_tensors()
30
34
- # Note: Remove the override after new attention backend finished
35
- def _may_reorder_batch(self, scheduler_output: "SchedulerOutput") -> None:
36
- if len(self.kv_cache_config.kv_cache_groups) > 1:
37
- raise ValueError(
38
- "Multiple KVCacheGroups is not"
39
- "currently supported with CPU model runner."
40
- )
41
- super()._may_reorder_batch(scheduler_output)
42
-
31
def _postprocess_tensors(self) -> None:
32
# Note: replace device tensors with cpu tensors
33
def replace_tensor(obj: Any, cpu_attr_name: str, device_attr_name) -> None: