@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 4f044b1d67
Fork 0

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

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

Synced Nov 15, 17:15 UTC

[Kernel][CPU] CPU MLA (#14744)

Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>

Thien Tran committed Mar 25, 2025 at 17:34 UTC 4f044b1d67964e53587a4d0c7f00233a04b7be4e
15 files changed +1010 -17
.buildkite/run-cpu-test.sh
+2
@@ -38,6 +38,8 @@ function cpu_tests() {
38 set -e
39 pip install -r vllm/requirements/test.txt
40 pip install -r vllm/requirements/cpu.txt
41 + pytest -v -s tests/kernels/test_cache.py -m cpu_model
42 + pytest -v -s tests/kernels/test_mla_decode_cpu.py -m cpu_model
43 pytest -v -s tests/models/decoder_only/language -m cpu_model
44 pytest -v -s tests/models/embedding/language -m cpu_model
45 pytest -v -s tests/models/encoder_decoder/language -m cpu_model
cmake/cpu_extension.cmake
+1
@@ -190,6 +190,7 @@ set(VLLM_EXT_SRC
190 "csrc/cpu/cache.cpp"
191 "csrc/cpu/utils.cpp"
192 "csrc/cpu/layernorm.cpp"
193 + "csrc/cpu/mla_decode.cpp"
194 "csrc/cpu/pos_encoding.cpp"
195 "csrc/cpu/torch_bindings.cpp")
196
csrc/cpu/cache.cpp
+74
@@ -88,6 +88,48 @@ void reshape_and_cache_cpu_impl(
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.
@@ -134,6 +176,38 @@ void reshape_and_cache(torch::Tensor& key, torch::Tensor& value,
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.")
csrc/cpu/cpu_types_x86.hpp
+2
@@ -130,6 +130,8 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
130
131 __m512i reg;
132
133 + explicit BF16Vec32() : reg(_mm512_setzero_si512()) {}
134 +
135 explicit BF16Vec32(const void* ptr) : reg((__m512i)_mm512_loadu_si512(ptr)) {}
136
137 explicit BF16Vec32(__m512i data) : reg(data) {}
csrc/cpu/mla_decode.cpp new
+393
@@ -0,0 +1,393 @@
1 +#include "cpu_types.hpp"
2 +#include <float.h>
3 +
4 +namespace {
5 +template <typename scalar_t>
6 +struct KernelVecType {
7 + using qk_load_vec_type = void;
8 + using qk_vec_type = void;
9 + using v_load_vec_type = void;
10 +};
11 +
12 +template <>
13 +struct KernelVecType<float> {
14 + using qk_load_vec_type = vec_op::FP32Vec16;
15 + using qk_vec_type = vec_op::FP32Vec16;
16 + using v_load_vec_type = vec_op::FP32Vec16;
17 +};
18 +
19 +template <>
20 +struct KernelVecType<c10::Half> {
21 +#if defined(__powerpc64__) || defined(__s390x__)
22 + // Power and s390x architecture-specific vector types
23 + using qk_load_vec_type = vec_op::FP32Vec16;
24 + using qk_vec_type = vec_op::FP32Vec16;
25 + using v_load_vec_type = vec_op::FP32Vec16;
26 +#else
27 + // Fallback for other architectures, including x86
28 + using qk_load_vec_type = vec_op::FP16Vec16;
29 + using qk_vec_type = vec_op::FP32Vec16;
30 + using v_load_vec_type = vec_op::FP16Vec16;
31 +#endif
32 +};
33 +
34 +#ifdef __AVX512BF16__
35 +template <>
36 +struct KernelVecType<c10::BFloat16> {
37 + using qk_load_vec_type = vec_op::BF16Vec32;
38 + using qk_vec_type = vec_op::BF16Vec32;
39 + using v_load_vec_type = vec_op::BF16Vec16;
40 +};
41 +#elif defined(__aarch64__) && !defined(ARM_BF16_SUPPORT)
42 +// pass
43 +#else
44 +template <>
45 +struct KernelVecType<c10::BFloat16> {
46 + using qk_load_vec_type = vec_op::BF16Vec16;
47 + using qk_vec_type = vec_op::FP32Vec16;
48 + using v_load_vec_type = vec_op::BF16Vec16;
49 +};
50 +#endif
51 +
52 +template <int HEAD_DIM, int V_HEAD_DIM, int BLOCK_SIZE, int HEAD_UNROLL,
53 + typename qk_vec_type>
54 +void mla_decode_block_head(
55 + const qk_vec_type* __restrict__ q_vecs, // [HEAD_UNROLL, head_dim]
56 + const qk_vec_type* __restrict__ k_vecs, // [block_size, head_dim]
57 + const vec_op::FP32Vec16* __restrict v_vecs_f32, // [block_size, v_head_dim]
58 + float* __restrict__ acc_out, // [HEAD_UNROLL, v_head_dim]
59 + float* __restrict__ acc_lse, // [HEAD_UNROLL]
60 + const float scale, const int num_tokens) {
61 + using f32_vec_type = vec_op::FP32Vec16;
62 + constexpr int QK_NUM_ELEM = qk_vec_type::VEC_ELEM_NUM;
63 + constexpr int V_NUM_ELEM = f32_vec_type::VEC_ELEM_NUM;
64 +
65 + float logits[BLOCK_SIZE][HEAD_UNROLL] = {}; // initialize to zeros
66 + float max_val[HEAD_UNROLL];
67 + std::fill(max_val, max_val + HEAD_UNROLL, -FLT_MAX);
68 +
69 + f32_vec_type acc_vec[BLOCK_SIZE][HEAD_UNROLL];
70 + for (int i = 0; i < HEAD_DIM; i += QK_NUM_ELEM) {
71 + // load to registers
72 + qk_vec_type q_vec[HEAD_UNROLL];
73 +
74 +#pragma unroll
75 + for (int unroll = 0; unroll < HEAD_UNROLL; ++unroll)
76 + q_vec[unroll] =
77 + qk_vec_type{q_vecs[(i + unroll * HEAD_DIM) / QK_NUM_ELEM]};
78 +
79 + for (int block_offset = 0; block_offset < num_tokens; ++block_offset) {
80 + qk_vec_type k_vec(k_vecs[(block_offset * HEAD_DIM + i) / QK_NUM_ELEM]);
81 +
82 +#pragma unroll
83 + for (int unroll = 0; unroll < HEAD_UNROLL; ++unroll)
84 + vec_op::fma(acc_vec[block_offset][unroll], q_vec[unroll], k_vec);
85 + }
86 + }
87 +
88 + for (int block_offset = 0; block_offset < num_tokens; ++block_offset) {
89 +#pragma unroll
90 + for (int unroll = 0; unroll < HEAD_UNROLL; ++unroll) {
91 + const float acc = acc_vec[block_offset][unroll].reduce_sum() * scale;
92 + logits[block_offset][unroll] = acc;
93 + max_val[unroll] = std::max(max_val[unroll], acc);
94 + }
95 + }
96 +
97 + float sum_exp[HEAD_UNROLL] = {};
98 + for (int block_offset = 0; block_offset < num_tokens; ++block_offset) {
99 +#pragma unroll
100 + for (int unroll = 0; unroll < HEAD_UNROLL; ++unroll) {
101 + const float val =
102 + std::exp(logits[block_offset][unroll] - max_val[unroll]);
103 + logits[block_offset][unroll] = val;
104 + sum_exp[unroll] += val;
105 + }
106 + }
107 +
108 + f32_vec_type this_out[V_HEAD_DIM / V_NUM_ELEM][HEAD_UNROLL];
109 +
110 + for (int block_offset = 0; block_offset < num_tokens; ++block_offset) {
111 + // load to registers
112 + f32_vec_type scale_[HEAD_UNROLL];
113 +
114 +#pragma unroll
115 + for (int unroll = 0; unroll < HEAD_UNROLL; ++unroll)
116 + scale_[unroll] =
117 + f32_vec_type{logits[block_offset][unroll] / sum_exp[unroll]};
118 +
119 + for (int i = 0; i < V_HEAD_DIM; i += V_NUM_ELEM) {
120 + f32_vec_type v_vec(
121 + v_vecs_f32[(block_offset * HEAD_DIM + i) / V_NUM_ELEM]);
122 +
123 +#pragma unroll
124 + for (int unroll = 0; unroll < HEAD_UNROLL; ++unroll)
125 + vec_op::fma(this_out[i / V_NUM_ELEM][unroll], v_vec, scale_[unroll]);
126 + }
127 + }
128 +
129 + // merge attention state
130 + // section 2.2 in https://arxiv.org/pdf/2501.01005
131 + f32_vec_type prev_scale[HEAD_UNROLL];
132 + f32_vec_type curr_scale[HEAD_UNROLL];
133 +
134 +#pragma unroll
135 + for (int unroll = 0; unroll < HEAD_UNROLL; ++unroll) {
136 + const float prev_lse = acc_lse[unroll];
137 + const float curr_lse = std::log(sum_exp[unroll]) +
138 + max_val[unroll]; // add back max_val to get true lse
139 + // softmax trick
140 + const float max_lse = std::max(prev_lse, curr_lse);
141 + const float prev_sum_exp = std::exp(prev_lse - max_lse);
142 + const float curr_sum_exp = std::exp(curr_lse - max_lse);
143 +
144 + const float new_sum_exp = prev_sum_exp + curr_sum_exp;
145 + acc_lse[unroll] = std::log(new_sum_exp) + max_lse;
146 +
147 + prev_scale[unroll] = f32_vec_type{prev_sum_exp / new_sum_exp};
148 + curr_scale[unroll] = f32_vec_type{curr_sum_exp / new_sum_exp};
149 + }
150 +
151 + for (int i = 0; i < V_HEAD_DIM; i += V_NUM_ELEM) {
152 +#pragma unroll
153 + for (int unroll = 0; unroll < HEAD_UNROLL; ++unroll) {
154 + f32_vec_type o_vec(acc_out + i + V_HEAD_DIM * unroll);
155 + o_vec = o_vec * prev_scale[unroll] +
156 + this_out[i / V_NUM_ELEM][unroll] * curr_scale[unroll];
157 + o_vec.save(acc_out + i + V_HEAD_DIM * unroll);
158 + }
159 + }
160 +
161 + q_vecs += HEAD_DIM / QK_NUM_ELEM * HEAD_UNROLL;
162 + acc_out += V_HEAD_DIM * HEAD_UNROLL;
163 +}
164 +
165 +template <typename scalar_t, int HEAD_DIM, int V_HEAD_DIM, int BLOCK_SIZE,
166 + typename qk_vec_type>
167 +void mla_decode_block(
168 + const qk_vec_type* __restrict__ q_vecs, // [num_heads, head_dim]
169 + const scalar_t* __restrict__ kv_cache, // [block_size, head_dim]
170 + float* __restrict__ acc_out, // [num_heads, v_head_dim]
171 + float* __restrict__ acc_lse, // [num_heads]
172 + const int num_heads, const float scale, const int num_tokens) {
173 + using qk_load_vec_type = typename KernelVecType<scalar_t>::qk_load_vec_type;
174 + static_assert(
175 + std::is_same<qk_vec_type,
176 + typename KernelVecType<scalar_t>::qk_vec_type>::value);
177 + using v_load_vec_type = typename KernelVecType<scalar_t>::v_load_vec_type;
178 + using f32_vec_type = vec_op::FP32Vec16;
179 + static_assert(qk_load_vec_type::VEC_ELEM_NUM == qk_vec_type::VEC_ELEM_NUM);
180 + static_assert(v_load_vec_type::VEC_ELEM_NUM == f32_vec_type::VEC_ELEM_NUM);
181 + constexpr int QK_NUM_ELEM = qk_vec_type::VEC_ELEM_NUM;
182 + constexpr int V_NUM_ELEM = v_load_vec_type::VEC_ELEM_NUM;
183 +
184 + const qk_vec_type* k_vecs;
185 + const f32_vec_type* v_vecs_f32;
186 + float* kv_cache_f32 = nullptr;
187 +
188 + if constexpr (!std::is_same<scalar_t, float>::value) {
189 + // convert KV cache block to FP32 to reuse it across query heads and
190 + // attn @ V computation, since FP16/BF16->FP32 is expensive.
191 + // TODO: move malloc outside of this fn to reuse across iterations.
192 + const int nbytes = BLOCK_SIZE * HEAD_DIM * sizeof(float);
193 + kv_cache_f32 = static_cast<float*>(std::aligned_alloc(64, nbytes));
194 +
195 + for (int block_offset = 0; block_offset < num_tokens; ++block_offset)
196 + for (int i = 0; i < HEAD_DIM; i += V_NUM_ELEM) {
197 + v_load_vec_type kv_load_vec(kv_cache + block_offset * HEAD_DIM + i);
198 + f32_vec_type kv_vec_f32(kv_load_vec);
199 + kv_vec_f32.save(kv_cache_f32 + block_offset * HEAD_DIM + i);
200 + }
201 +
202 + if constexpr (std::is_same<qk_load_vec_type, qk_vec_type>::value) {
203 + // for AVX512_BF16, Q @ K.T uses BF16 for K (no conversion)
204 + // NOTE: in this case, we only need to convert the V section to FP32.
205 + // But for simplicity, we will convert the whole KV block to FP32.
206 + k_vecs = reinterpret_cast<const qk_vec_type*>(kv_cache);
207 + } else {
208 + k_vecs = reinterpret_cast<const qk_vec_type*>(kv_cache_f32);
209 + }
210 +
211 + // attn @ V always use FP32 for V, since attn is FP32.
212 + v_vecs_f32 = reinterpret_cast<const f32_vec_type*>(kv_cache_f32);
213 +
214 + } else {
215 + // KV cache is FP32. don't need to do anything.
216 + k_vecs = reinterpret_cast<const qk_vec_type*>(kv_cache);
217 + v_vecs_f32 = reinterpret_cast<const f32_vec_type*>(kv_cache);
218 + }
219 +
220 + // compute 2 heads at the same time to improve ILP and
221 + // take advantage of register cache for K and V.
222 + constexpr int HEAD_UNROLL = 2;
223 + for (int iter = 0; iter < num_heads / HEAD_UNROLL; ++iter) {
224 + mla_decode_block_head<HEAD_DIM, V_HEAD_DIM, BLOCK_SIZE, HEAD_UNROLL>(
225 + q_vecs, k_vecs, v_vecs_f32, acc_out, acc_lse, scale, num_tokens);
226 +
227 + q_vecs += HEAD_UNROLL * HEAD_DIM / QK_NUM_ELEM;
228 + acc_out += HEAD_UNROLL * V_HEAD_DIM;
229 + acc_lse += HEAD_UNROLL;
230 + }
231 +
232 + // take care of the remaining heads
233 + for (int iter = 0; iter < num_heads % HEAD_UNROLL; ++iter) {
234 + mla_decode_block_head<HEAD_DIM, V_HEAD_DIM, BLOCK_SIZE, 1>(
235 + q_vecs, k_vecs, v_vecs_f32, acc_out, acc_lse, scale, num_tokens);
236 +
237 + q_vecs += HEAD_DIM / QK_NUM_ELEM;
238 + acc_out += V_HEAD_DIM;
239 + acc_lse += 1;
240 + }
241 +
242 + if (kv_cache_f32 != nullptr) {
243 + std::free(kv_cache_f32);
244 + }
245 +}
246 +} // namespace
247 +
248 +template <typename scalar_t, int HEAD_DIM, int V_HEAD_DIM, int BLOCK_SIZE>
249 +void mla_decode_kvcache_cpu_impl(
250 + scalar_t* __restrict__ out, // [num_seqs, num_heads, v_head_dim]
251 + const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_dim]
252 + const scalar_t* __restrict__ kv_cache, // [num_blocks, block_size,
253 + // head_dim]
254 + const int num_heads, const float scale,
255 + const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq]
256 + const int* __restrict__ seq_lens, // [num_seqs]
257 + const int max_num_blocks_per_seq, const int o_stride, const int q_stride,
258 + const int kv_stride, const int num_seqs) {
259 + using qk_load_vec_type = typename KernelVecType<scalar_t>::qk_load_vec_type;
260 + using qk_vec_type = typename KernelVecType<scalar_t>::qk_vec_type;
261 + constexpr int QK_NUM_ELEM = qk_vec_type::VEC_ELEM_NUM;
262 +
263 + // shared across threads
264 + const int max_threads = omp_get_max_threads();
265 + const int acc_out_nbytes =
266 + max_threads * num_heads * V_HEAD_DIM * sizeof(float);
267 + float* acc_out = static_cast<float*>(std::aligned_alloc(64, acc_out_nbytes));
268 + std::vector<float> acc_lse(max_threads * num_heads);
269 +
270 + // allocate memory to pre-convert query to FP32 later
271 + float* q_f32;
272 + constexpr bool PRE_CONVERT_QUERY =
273 + !std::is_same<scalar_t, float>::value &&
274 + std::is_same<qk_vec_type, vec_op::FP32Vec16>::value;
275 + if constexpr (PRE_CONVERT_QUERY) {
276 + const int q_f32_nbytes = num_heads * HEAD_DIM * sizeof(float);
277 + q_f32 = static_cast<float*>(std::aligned_alloc(64, q_f32_nbytes));
278 + }
279 +
280 +#pragma omp parallel
281 + {
282 + const int num_threads = omp_get_num_threads();
283 + const int thread_id = omp_get_thread_num();
284 + float* __restrict__ acc_out_thread =
285 + acc_out + thread_id * num_heads * V_HEAD_DIM;
286 + float* __restrict__ acc_lse_thread = acc_lse.data() + thread_id * num_heads;
287 +
288 + for (int seq_idx = 0; seq_idx < num_seqs; ++seq_idx) {
289 + // reset accumulator
290 + std::fill(acc_out_thread, acc_out_thread + num_heads * V_HEAD_DIM, 0.0f);
291 + std::fill(acc_lse_thread, acc_lse_thread + num_heads, -FLT_MAX);
292 +
293 + const int seq_len = seq_lens[seq_idx];
294 + const int block_num = (seq_len + BLOCK_SIZE - 1) / BLOCK_SIZE;
295 + const int last_block_size = seq_len - (block_num - 1) * BLOCK_SIZE;
296 +
297 + const qk_vec_type* q_vecs;
298 + if constexpr (PRE_CONVERT_QUERY) {
299 +// pre-convert query to FP32 since FP16/BF16->FP32 is slow.
300 +#pragma omp for
301 + for (int i = 0; i < num_heads * HEAD_DIM; i += QK_NUM_ELEM) {
302 + qk_load_vec_type q_load_vec(q + seq_idx * q_stride + i);
303 + qk_vec_type q_vec(q_load_vec);
304 + q_vec.save(q_f32 + i);
305 + }
306 + q_vecs = reinterpret_cast<const qk_vec_type*>(q_f32);
307 + } else {
308 + q_vecs = reinterpret_cast<const qk_vec_type*>(q + seq_idx * q_stride);
309 + }
310 +
311 +#pragma omp for
312 + for (int block_idx = 0; block_idx < block_num; ++block_idx) {
313 + const int physical_block_idx =
314 + block_tables[seq_idx * max_num_blocks_per_seq + block_idx];
315 + const int num_tokens =
316 + block_idx < block_num - 1 ? BLOCK_SIZE : last_block_size;
317 +
318 + mla_decode_block<scalar_t, HEAD_DIM, V_HEAD_DIM, BLOCK_SIZE>(
319 + q_vecs, kv_cache + physical_block_idx * kv_stride, acc_out_thread,
320 + acc_lse_thread, num_heads, scale, num_tokens);
321 + }
322 +
323 +// merge attention states across threads
324 +// section 2.2 in https://arxiv.org/pdf/2501.01005
325 +// each thread is responsible for 1 head
326 +#pragma omp for
327 + for (int head_idx = 0; head_idx < num_heads; ++head_idx) {
328 + float* acc_lse_head = acc_lse.data() + head_idx;
329 + float* acc_out_head = acc_out + head_idx * V_HEAD_DIM;
330 +
331 + float max_val = -FLT_MAX;
332 + for (int thread_id_ = 0; thread_id_ < num_threads; ++thread_id_) {
333 + max_val = std::max(max_val, acc_lse_head[thread_id_ * num_heads]);
334 + }
335 +
336 + float sum_exp = 0.0f;
337 + for (int thread_id_ = 0; thread_id_ < num_threads; ++thread_id_) {
338 + float val = std::exp(acc_lse_head[thread_id_ * num_heads] - max_val);
339 + acc_lse_head[thread_id_ * num_heads] = val;
340 + sum_exp += val;
341 + }
342 +
343 + float inv_sum = 1.0f / sum_exp;
344 + float out_head[V_HEAD_DIM] = {};
345 + for (int thread_id_ = 0; thread_id_ < num_threads; ++thread_id_) {
346 + float scale_ = acc_lse_head[thread_id_ * num_heads] * inv_sum;
347 + for (int i = 0; i < V_HEAD_DIM; ++i) {
348 + out_head[i] +=
349 + acc_out_head[thread_id_ * num_heads * V_HEAD_DIM + i] * scale_;
350 + }
351 + }
352 +
353 + for (int i = 0; i < V_HEAD_DIM; ++i) {
354 + vec_op::storeFP32(out_head[i], out + seq_idx * o_stride +
355 + head_idx * V_HEAD_DIM + i);
356 + }
357 + }
358 + }
359 + }
360 + if (PRE_CONVERT_QUERY) {
361 + std::free(q_f32);
362 + }
363 + std::free(acc_out);
364 +}
365 +
366 +void mla_decode_kvcache(torch::Tensor& out, torch::Tensor& query,
367 + torch::Tensor& kv_cache, double scale,
368 + torch::Tensor& block_tables, torch::Tensor& seq_lens) {
369 + const int num_seqs = query.size(0);
370 + const int num_heads = query.size(1);
371 + const int head_dim = query.size(2);
372 + const int block_size = kv_cache.size(1);
373 + const int v_head_dim = out.size(2);
374 +
375 + const int max_num_blocks_per_seq = block_tables.size(1);
376 + const int o_stride = out.stride(0);
377 + const int q_stride = query.stride(0);
378 + const int kv_stride = kv_cache.stride(0);
379 +
380 + VLLM_DISPATCH_FLOATING_TYPES(
381 + query.scalar_type(), "mla_decode_kvcache_cpu_impl", [&] {
382 + CPU_KERNEL_GUARD_IN(mla_decode_kvcache_cpu_impl)
383 + if (head_dim == 576 && v_head_dim == 512 && block_size == 16)
384 + mla_decode_kvcache_cpu_impl<scalar_t, 576, 512, 16>(
385 + out.data_ptr<scalar_t>(), query.data_ptr<scalar_t>(),
386 + kv_cache.data_ptr<scalar_t>(), num_heads, scale,
387 + block_tables.data_ptr<int>(), seq_lens.data_ptr<int>(),
388 + max_num_blocks_per_seq, o_stride, q_stride, kv_stride, num_seqs);
389 + else
390 + TORCH_CHECK(false, "Unsupported block size: ", block_size);
391 + CPU_KERNEL_GUARD_OUT(mla_decode_kvcache_cpu_impl)
392 + });
393 +}
\ No newline at end of file
csrc/cpu/torch_bindings.cpp
+20
@@ -18,6 +18,10 @@ void int8_scaled_mm_azp(torch::Tensor& c, const torch::Tensor& a,
18 const std::optional<torch::Tensor>& azp,
19 const std::optional<torch::Tensor>& bias);
20
21 +void mla_decode_kvcache(torch::Tensor& out, torch::Tensor& query,
22 + torch::Tensor& kv_cache, double scale,
23 + torch::Tensor& block_tables, torch::Tensor& seq_lens);
24 +
25 TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
26 // vLLM custom ops
27
@@ -150,6 +154,14 @@ TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cache_ops), cache_ops) {
154 " str kv_cache_dtype,"
155 " Tensor k_scale, Tensor v_scale) -> ()");
156 cache_ops.impl("reshape_and_cache", torch::kCPU, &reshape_and_cache);
157 +
158 + cache_ops.def(
159 + "concat_and_cache_mla(Tensor kv_c, Tensor k_pe,"
160 + " Tensor! kv_cache,"
161 + " Tensor slot_mapping,"
162 + " str kv_cache_dtype,"
163 + " Tensor scale) -> ()");
164 + cache_ops.impl("concat_and_cache_mla", torch::kCPU, &concat_and_cache_mla);
165 }
166
167 TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _utils), utils) {
@@ -157,4 +169,12 @@ TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _utils), utils) {
169 utils.def("init_cpu_threads_env(str cpu_ids) -> str", &init_cpu_threads_env);
170 }
171
172 +TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cpu), cpu_ops) {
173 + cpu_ops.def(
174 + "mla_decode_kvcache("
175 + " Tensor! out, Tensor query, Tensor kv_cache,"
176 + " float scale, Tensor block_tables, Tensor seq_lens) -> ()");
177 + cpu_ops.impl("mla_decode_kvcache", torch::kCPU, &mla_decode_kvcache);
178 +}
179 +
180 REGISTER_EXTENSION(TORCH_EXTENSION_NAME)
tests/kernels/test_cache.py
+69
@@ -749,3 +749,72 @@ def test_gather_cache_mla(kv_lora_rank, qk_rope_head_dim, block_size,
749
750 ops.gather_cache(src_cache, dst, block_table, cu_seq_lens, batch_size)
751 torch.testing.assert_close(dst, expected)
752 +
753 +
754 +@pytest.mark.parametrize("kv_lora_rank", KV_LORA_RANKS)
755 +@pytest.mark.parametrize("qk_rope_head_dim", QK_ROPE_HEAD_DIMS)
756 +@pytest.mark.parametrize("num_tokens", NUM_TOKENS_MLA)
757 +@pytest.mark.parametrize("block_size", BLOCK_SIZES_MLA)
758 +@pytest.mark.parametrize("num_blocks", NUM_BLOCKS_MLA)
759 +@pytest.mark.parametrize("dtype", DTYPES)
760 +@pytest.mark.parametrize("seed", SEEDS)
761 +@pytest.mark.cpu_model
762 +@pytest.mark.skipif(not current_platform.is_cpu(), reason="CPU only")
763 +@torch.inference_mode()
764 +def test_concat_and_cache_mla_cpu(
765 + kv_lora_rank: int,
766 + qk_rope_head_dim: int,
767 + num_tokens: int,
768 + block_size: int,
769 + num_blocks: int,
770 + dtype: torch.dtype,
771 + seed: int,
772 +) -> None:
773 + device = "cpu"
774 + kv_cache_dtype = "auto"
775 + current_platform.seed_everything(seed)
776 + torch.set_default_device(device)
777 +
778 + total_slots = num_blocks * block_size
779 + slot_mapping_lst = random.sample(range(total_slots), num_tokens)
780 + slot_mapping = torch.tensor(slot_mapping_lst,
781 + dtype=torch.long,
782 + device=device)
783 +
784 + kv_c = torch.randn(num_tokens, kv_lora_rank, dtype=dtype, device=device)
785 + k_pe = torch.randn(num_tokens,
786 + qk_rope_head_dim,
787 + dtype=dtype,
788 + device=device)
789 + entry_size = kv_lora_rank + qk_rope_head_dim
790 +
791 + scale = torch.tensor(0.1, dtype=torch.float32, device=device)
792 + kv_cache = _create_mla_cache(num_blocks, block_size, entry_size, dtype,
793 + kv_cache_dtype, device)
794 + ref_temp = torch.zeros(*kv_cache.shape, dtype=dtype, device=device)
795 +
796 + for i in range(num_tokens):
797 + slot = slot_mapping[i].item()
798 + block_idx = slot // block_size
799 + block_offset = slot % block_size
800 + ref_temp[block_idx, block_offset, :kv_lora_rank] = kv_c[i]
801 + ref_temp[block_idx, block_offset, kv_lora_rank:] = k_pe[i]
802 +
803 + if kv_cache_dtype == "fp8":
804 + ref_kv_cache = torch.empty_like(ref_temp, dtype=kv_cache.dtype)
805 + ops.convert_fp8(ref_kv_cache,
806 + ref_temp,
807 + scale.item(),
808 + kv_dtype=kv_cache_dtype)
809 + else:
810 + ref_kv_cache = ref_temp
811 +
812 + opcheck(
813 + torch.ops._C_cache_ops.concat_and_cache_mla,
814 + (kv_c, k_pe, kv_cache, slot_mapping, kv_cache_dtype, scale),
815 + test_utils=DEFAULT_OPCHECK_TEST_UTILS,
816 + )
817 +
818 + ops.concat_and_cache_mla(kv_c, k_pe, kv_cache, slot_mapping,
819 + kv_cache_dtype, scale)
820 + torch.testing.assert_close(kv_cache, ref_kv_cache)
tests/kernels/test_mla_decode_cpu.py new
+94
@@ -0,0 +1,94 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +import pytest
3 +import torch
4 +import torch.nn.functional as F
5 +from torch import Tensor
6 +
7 +import vllm._custom_ops as ops
8 +from vllm.platforms import current_platform
9 +
10 +
11 +def cdiv(a, b):
12 + return (a + b - 1) // b
13 +
14 +
15 +def ref_mla(
16 + out: Tensor, # (bs, num_heads, v_head_dim)
17 + query: Tensor, # (bs, num_heads, head_dim)
18 + kv_cache: Tensor, # (num_blocks, block_size, head_dim)
19 + scale: float,
20 + block_tables: Tensor, # (bs, max_num_blocks)
21 + seq_lens: Tensor, # (bs,)
22 +):
23 + bs, num_heads, v_head_dim = out.shape
24 + head_dim = query.shape[2]
25 +
26 + for i in range(bs):
27 + # gather and flatten KV-cache
28 + kv = kv_cache[
29 + block_tables[i]] # (max_num_blocks, block_size, head_dim)
30 + kv = kv.view(1, -1,
31 + head_dim)[:, :seq_lens[i]] # (1, seq_len, head_dim)
32 + v = kv[:, :, :v_head_dim]
33 +
34 + q = query[i].view(num_heads, 1, head_dim)
35 + o = F.scaled_dot_product_attention(q,
36 + kv,
37 + v,
38 + scale=scale,
39 + enable_gqa=True)
40 + out[i] = o.view(num_heads, v_head_dim)
41 +
42 + return out
43 +
44 +
45 +@pytest.mark.parametrize("bs", [4])
46 +@pytest.mark.parametrize("mean_seq_len", [256])
47 +@pytest.mark.parametrize("h_q", [16])
48 +@pytest.mark.parametrize("d", [576])
49 +@pytest.mark.parametrize("dv", [512])
50 +@pytest.mark.parametrize("block_size", [16])
51 +@pytest.mark.parametrize("dtype", [torch.float, torch.half, torch.bfloat16])
52 +@pytest.mark.parametrize("varlen", [False, True])
53 +@pytest.mark.cpu_model
54 +@pytest.mark.skipif(not current_platform.is_cpu(), reason="CPU only")
55 +def test_mla_decode_cpu(
56 + bs: int,
57 + mean_seq_len: int,
58 + h_q: int,
59 + d: int,
60 + dv: int,
61 + block_size: int,
62 + dtype: torch.dtype,
63 + varlen: bool,
64 +):
65 + torch.set_default_dtype(dtype)
66 + torch.manual_seed(0)
67 +
68 + scale = d**(-0.5)
69 + if varlen:
70 + seq_lens = torch.empty(bs).normal_(mean_seq_len, mean_seq_len / 2)
71 + seq_lens = seq_lens.clip(2).to(torch.int32)
72 + else:
73 + seq_lens = torch.full((bs, ), mean_seq_len, dtype=torch.int32)
74 + max_seq_len = seq_lens.max().item()
75 + seqlen_pad = cdiv(max_seq_len, 256) * 256 # is this necessary?
76 +
77 + q = torch.randn(bs, h_q, d)
78 + block_table = torch.arange(bs * seqlen_pad // block_size,
79 + dtype=torch.int32)
80 + block_table = block_table.view(bs, seqlen_pad // block_size)
81 +
82 + kv_cache = torch.randn(block_table.numel(), block_size, d)
83 + for i, seq_len in enumerate(seq_lens.tolist()):
84 + kv_cache.view(bs, seqlen_pad, d)[i, seq_len:] = float("nan")
85 +
86 + out_mla = q.new_zeros(bs, h_q, dv)
87 + ops.mla_decode_kvcache_cpu(out_mla, q, kv_cache, scale, block_table,
88 + seq_lens)
89 +
90 + out_ref = q.new_zeros(bs, h_q, dv)
91 + ref_mla(out_ref, q, kv_cache, scale, block_table, seq_lens)
92 +
93 + assert not out_mla.isnan().any(), "Likely read out of bounds"
94 + torch.testing.assert_close(out_mla, out_ref)
vllm/_custom_ops.py
+12
@@ -124,6 +124,18 @@ def paged_attention_rocm(
124 kv_cache_dtype, k_scale, v_scale)
125
126
127 +def mla_decode_kvcache_cpu(
128 + out: torch.Tensor,
129 + query: torch.Tensor,
130 + kv_cache: torch.Tensor,
131 + scale: float,
132 + block_tables: torch.Tensor,
133 + seq_lens: torch.Tensor,
134 +) -> None:
135 + torch.ops._C_cpu.mla_decode_kvcache(out, query, kv_cache, scale,
136 + block_tables, seq_lens)
137 +
138 +
139 # pos encoding ops
140 def rotary_embedding(
141 positions: torch.Tensor,
vllm/_ipex_ops.py
+22 -9
@@ -187,15 +187,28 @@ class ipex_ops:
187 gen_: torch.Generator,
188 logits_soft_cap: float,
189 ) -> None:
190 - ipex.llm.functional.varlen_attention(query.contiguous(),
191 - key.contiguous(),
192 - value.contiguous(), out,
193 - seqlen_q.int(), seqlen_k.int(),
194 - max_seqlen_q, max_seqlen_k,
195 - pdropout, softmax_scale,
196 - zero_tensors, is_causal,
197 - return_softmax, gen_,
198 - logits_soft_cap)
190 + if ipex.__version__.endswith("cpu"):
191 + if logits_soft_cap != 0.0:
192 + raise ValueError("IPEX CPU does not support logits_soft_cap")
193 + ipex.llm.functional.varlen_attention(query.contiguous(),
194 + key.contiguous(),
195 + value.contiguous(), out,
196 + seqlen_q.int(),
197 + seqlen_k.int(), max_seqlen_q,
198 + max_seqlen_k, pdropout,
199 + softmax_scale, zero_tensors,
200 + is_causal, return_softmax,
201 + gen_)
202 + else: # XPU build
203 + ipex.llm.functional.varlen_attention(query.contiguous(),
204 + key.contiguous(),
205 + value.contiguous(), out,
206 + seqlen_q.int(),
207 + seqlen_k.int(), max_seqlen_q,
208 + max_seqlen_k, pdropout,
209 + softmax_scale, zero_tensors,
210 + is_causal, return_softmax,
211 + gen_, logits_soft_cap)
212
213 @staticmethod
214 def reshape_and_cache(
vllm/attention/backends/cpu_mla.py new
+303
@@ -0,0 +1,303 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +from dataclasses import dataclass
4 +from typing import Any, Dict, List, Optional, Tuple, Type
5 +
6 +import torch
7 +
8 +import vllm._custom_ops as ops
9 +from vllm._ipex_ops import ipex_ops
10 +from vllm.attention.backends.abstract import (AttentionBackend,
11 + AttentionMetadataBuilder,
12 + AttentionType,
13 + is_quantized_kv_cache)
14 +from vllm.attention.backends.mla.common import MLACommonImpl, MLACommonState
15 +from vllm.attention.backends.torch_sdpa import TorchSDPAMetadata
16 +from vllm.utils import make_tensor_with_pad
17 +from vllm.worker.cpu_model_runner import ModelInputForCPUBuilder
18 +
19 +
20 +class CPUMLABackend(AttentionBackend):
21 +
22 + @staticmethod
23 + def get_name() -> str:
24 + return "CPU_MLA"
25 +
26 + @staticmethod
27 + def get_metadata_cls() -> Type["CPUMLAMetadata"]:
28 + return CPUMLAMetadata
29 +
30 + @staticmethod
31 + def get_builder_cls() -> Type["CPUMLAMetadataBuilder"]:
32 + return CPUMLAMetadataBuilder
33 +
34 + @staticmethod
35 + def get_state_cls() -> Type["MLACommonState"]:
36 + return MLACommonState
37 +
38 + @staticmethod
39 + def get_impl_cls() -> Type["CPUMLAImpl"]:
40 + return CPUMLAImpl
41 +
42 + @staticmethod
43 + def get_kv_cache_shape(
44 + num_blocks: int,
45 + block_size: int,
46 + num_kv_heads: int, # assumed to be 1 for MLA
47 + head_size: int,
48 + ) -> Tuple[int, ...]:
49 + return (num_blocks, block_size, head_size)
50 +
51 + @staticmethod
52 + def swap_blocks(
53 + src_kv_cache: torch.Tensor,
54 + dst_kv_cache: torch.Tensor,
55 + src_to_dst: torch.Tensor,
56 + ) -> None:
57 + ops.swap_blocks(src_kv_cache, dst_kv_cache, src_to_dst)
58 +
59 + @staticmethod
60 + def copy_blocks(
61 + kv_caches: List[torch.Tensor],
62 + src_to_dists: torch.Tensor,
63 + ) -> None:
64 + ops.copy_blocks_mla(kv_caches, src_to_dists)
65 +
66 + @staticmethod
67 + def get_supported_head_sizes() -> List[int]:
68 + return [576]
69 +
70 +
71 +@dataclass
72 +class CPUMLAMetadata(TorchSDPAMetadata):
73 + # New for MLA
74 + # Input positions for rotrary embeddings since for MLA the rotary
75 + # position embeddings are applied inside the attention backend
76 + input_positions: torch.Tensor = None
77 +
78 + # required by MLACommonImpl
79 + is_profile_run: bool = False
80 +
81 +
82 +class CPUMLAMetadataBuilder(AttentionMetadataBuilder[CPUMLAMetadata]):
83 +
84 + def __init__(self, input_builder: ModelInputForCPUBuilder) -> None:
85 + self.chunked_prefill = input_builder.chunked_prefill
86 + self.input_builder = input_builder
87 + assert not self.chunked_prefill, \
88 + "chunked prefill is currently not supported"
89 +
90 + def prepare(self):
91 + self.input_data = self.input_builder.input_data
92 +
93 + def build(self, seq_lens, query_lens, cuda_graph_pad_size, batch_size):
94 + input_data = self.input_data
95 + prefill_seq_lens = seq_lens[0:input_data.num_prefills]
96 + prefill_query_lens = query_lens[0:input_data.num_prefills]
97 + slot_mapping = torch.tensor(input_data.slot_mapping,
98 + dtype=torch.long,
99 + device="cpu")
100 +
101 + # metadata for prefill
102 + if input_data.num_prefills > 0:
103 + query_lens_tensor = torch.tensor(prefill_query_lens,
104 + dtype=torch.int32,
105 + device="cpu")
106 + kv_lens_tensor = torch.tensor(prefill_seq_lens,
107 + dtype=torch.int32,
108 + device="cpu")
109 + query_start_loc = torch.zeros(input_data.num_prefills + 1,
110 + dtype=torch.int32,
111 + device="cpu")
112 + kv_start_loc = torch.zeros(input_data.num_prefills + 1,
113 + dtype=torch.int32,
114 + device="cpu")
115 + torch.cumsum(query_lens_tensor,
116 + dim=0,
117 + dtype=torch.int32,
118 + out=query_start_loc[1:])
119 + torch.cumsum(kv_lens_tensor,
120 + dim=0,
121 + dtype=torch.int32,
122 + out=kv_start_loc[1:])
123 + max_query_len = max(prefill_query_lens)
124 + max_kv_len = max(prefill_seq_lens)
125 +
126 + # for chunked-prefill
127 + if self.chunked_prefill:
128 + prefill_block_tables = make_tensor_with_pad(
129 + self.input_data.prefill_block_tables,
130 + pad=0,
131 + dtype=torch.int32,
132 + device="cpu",
133 + )
134 + else:
135 + prefill_block_tables = None
136 +
137 + else:
138 + query_start_loc = None
139 + kv_start_loc = None
140 + max_query_len = None
141 + max_kv_len = None
142 + prefill_block_tables = None
143 +
144 + # metadata for decode
145 + if input_data.num_decode_tokens != 0:
146 + seq_lens_tensor = torch.tensor(
147 + input_data.seq_lens[input_data.num_prefills:],
148 + dtype=torch.int32,
149 + device="cpu",
150 + )
151 + block_tables = make_tensor_with_pad(
152 + self.input_data.decode_block_tables,
153 + pad=0,
154 + dtype=torch.int32,
155 + device="cpu",
156 + )
157 + else:
158 + block_tables = torch.tensor([])
159 + seq_lens_tensor = torch.tensor(
160 + input_data.seq_lens[:input_data.num_prefills],
161 + dtype=torch.int32,
162 + device="cpu",
163 + )
164 +
165 + # For multi-modal models
166 + placeholder_index_maps = None
167 + if len(input_data.multi_modal_inputs_list) != 0:
168 + placeholder_index_maps = {
169 + modality: placeholder_map.index_map()
170 + for modality, placeholder_map in
171 + input_data.multi_modal_placeholder_maps.items()
172 + }
173 +
174 + return CPUMLAMetadata(
175 + chunked_prefill=self.chunked_prefill,
176 + seq_lens=prefill_seq_lens,
177 + seq_lens_tensor=seq_lens_tensor,
178 + max_query_len=max_query_len,
179 + max_kv_len=max_kv_len,
180 + query_start_loc=query_start_loc,
181 + kv_start_loc=kv_start_loc,
182 + max_decode_seq_len=input_data.max_decode_seq_len,
183 + num_prefills=input_data.num_prefills,
184 + num_prefill_tokens=input_data.num_prefill_tokens,
185 + num_decode_tokens=input_data.num_decode_tokens,
186 + block_tables=block_tables,
187 + prefill_block_tables=prefill_block_tables,
188 + slot_mapping=slot_mapping,
189 + multi_modal_placeholder_index_maps=placeholder_index_maps,
190 + enable_kv_scales_calculation=False,
191 + input_positions=torch.tensor([self.input_data.input_positions]))
192 +
193 +
194 +class CPUMLAImpl(MLACommonImpl[CPUMLAMetadata]):
195 +
196 + def __init__(
197 + self,
198 + num_heads: int,
199 + head_size: int,
200 + scale: float,
201 + num_kv_heads: int,
202 + alibi_slopes: Optional[List[float]],
203 + sliding_window: Optional[int],
204 + kv_cache_dtype: str,
205 + blocksparse_params: Optional[Dict[str, Any]],
206 + logits_soft_cap: Optional[float],
207 + attn_type: str,
208 + # MLA Specific Arguments
209 + **mla_args) -> None:
210 + super().__init__(num_heads, head_size, scale, num_kv_heads,
211 + alibi_slopes, sliding_window, kv_cache_dtype,
212 + blocksparse_params, logits_soft_cap, attn_type,
213 + **mla_args)
214 +
215 + unsupported_features = [
216 + alibi_slopes, sliding_window, blocksparse_params, logits_soft_cap
217 + ]
218 + if any(unsupported_features):
219 + raise NotImplementedError(
220 + "CPUMLAImpl does not support one of the following: "
221 + "alibi_slopes, sliding_window, blocksparse_params, "
222 + "logits_soft_cap")
223 +
224 + if attn_type != AttentionType.DECODER:
225 + raise NotImplementedError("Encoder self-attention and "
226 + "encoder/decoder cross-attention "
227 + "are not implemented for "
228 + "CPUMLAImpl")
229 +
230 + # states is implemented.
231 + if is_quantized_kv_cache(self.kv_cache_dtype):
232 + raise NotImplementedError(
233 + "CPUMLAImpl with FP8 KV cache not yet supported")
234 +
235 + def _forward_prefill(
236 + self,
237 + q: torch.Tensor,
238 + kv_c_normed: torch.Tensor,
239 + k_pe: torch.Tensor,
240 + kv_c_and_k_pe_cache: torch.Tensor,
241 + attn_metadata: CPUMLAMetadata, # type: ignore[override]
242 + ) -> torch.Tensor:
243 +
244 + prefill_metadata = attn_metadata.prefill_metadata
245 + assert prefill_metadata is not None
246 +
247 + kv_nope = self.kv_b_proj(kv_c_normed)[0].view(\
248 + -1, self.num_heads, self.qk_nope_head_dim + self.v_head_dim)
249 + k_nope, v = kv_nope\
250 + .split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)
251 +
252 + k = torch.cat((k_nope, k_pe.expand((*k_nope.shape[:-1], -1))), dim=-1)
253 +
254 + # For MLA the v head dim is smaller than qk head dim so we pad out
255 + # v with 0s to match the qk head dim
256 + v_padded = torch.nn.functional.pad(v, [0, q.shape[-1] - v.shape[-1]],
257 + value=0)
258 +
259 + output = torch.empty_like(q)
260 + ipex_ops.varlen_attention(
261 + query=q,
262 + key=k,
263 + value=v_padded,
264 + out=output,
265 + seqlen_q=prefill_metadata.query_start_loc,
266 + seqlen_k=prefill_metadata.query_start_loc,
267 + max_seqlen_q=prefill_metadata.max_query_len,
268 + max_seqlen_k=prefill_metadata.max_query_len,
269 + pdropout=0.0,
270 + softmax_scale=self.scale,
271 + zero_tensors=False,
272 + is_causal=True,
273 + return_softmax=False,
274 + gen_=None,
275 + logits_soft_cap=0.0,
276 + )
277 +
278 + # remove padding
279 + output = output.view(-1, self.num_heads,
280 + q.shape[-1])[..., :v.shape[-1]]
281 + output = output.reshape(-1, self.num_heads * v.shape[-1])
282 + return self.o_proj(output)[0]
283 +
284 + def _forward_decode(
285 + self,
286 + q_nope: torch.Tensor,
287 + q_pe: torch.Tensor,
288 + kv_c_and_k_pe_cache: torch.Tensor,
289 + attn_metadata: CPUMLAMetadata, # type: ignore[override]
290 + ) -> torch.Tensor:
291 + assert kv_c_and_k_pe_cache.numel() > 0
292 +
293 + decode_meta = attn_metadata.decode_metadata
294 + assert decode_meta is not None
295 +
296 + q = torch.cat([q_nope, q_pe], dim=-1)
297 + o = q.new_empty(q.shape[0], self.num_heads, self.kv_lora_rank)
298 +
299 + # Run MQA
300 + ops.mla_decode_kvcache_cpu(o, q, kv_c_and_k_pe_cache, self.scale,
301 + decode_meta.block_tables,
302 + decode_meta.seq_lens_tensor)
303 + return self._v_up_proj_and_o_proj(o)
vllm/attention/backends/mla/common.py
+13 -5
@@ -204,7 +204,6 @@ from vllm.attention.backends.abstract import (AttentionBackend, AttentionLayer,
204 from vllm.attention.backends.utils import (PAD_SLOT_ID, compute_slot_mapping,
205 compute_slot_mapping_start_idx,
206 is_block_tables_empty)
207 -from vllm.attention.ops.triton_merge_attn_states import merge_attn_states
207 from vllm.model_executor.layers.linear import (ColumnParallelLinear,
208 LinearBase, RowParallelLinear,
209 UnquantizedLinearMethod)
@@ -212,18 +211,27 @@ from vllm.model_executor.layers.rotary_embedding import (
211 DeepseekScalingRotaryEmbedding, RotaryEmbedding)
212 from vllm.multimodal import MultiModalPlaceholderMap
213 from vllm.platforms import current_platform
214 +from vllm.triton_utils import HAS_TRITON
215 from vllm.utils import async_tensor_h2d, cdiv, make_tensor_with_pad, round_down
216 from vllm.vllm_flash_attn.fa_utils import get_flash_attn_version
217
218 +if HAS_TRITON:
219 + from vllm.attention.ops.triton_flash_attention import triton_attention
220 + from vllm.attention.ops.triton_merge_attn_states import merge_attn_states
221 +else:
222 + merge_attn_states = None
223 + triton_attention = None
224 +
225 try:
226 from vllm.vllm_flash_attn import flash_attn_varlen_func
227 is_vllm_fa = True
228 except ImportError:
222 - # For rocm use upstream flash attention
223 - from flash_attn import flash_attn_varlen_func
229 is_vllm_fa = False
225 -
226 -from vllm.attention.ops.triton_flash_attention import triton_attention
230 + try:
231 + # For rocm use upstream flash attention
232 + from flash_attn import flash_attn_varlen_func
233 + except ImportError:
234 + flash_attn_varlen_func = None
235
236 if TYPE_CHECKING:
237 from vllm.worker.model_runner import (ModelInputForGPUBuilder,
vllm/platforms/cpu.py
+3 -3
@@ -37,6 +37,9 @@ class CpuPlatform(Platform):
37 use_mla: bool) -> str:
38 if selected_backend and selected_backend != _Backend.TORCH_SDPA:
39 logger.info("Cannot use %s backend on CPU.", selected_backend)
40 + if use_mla:
41 + logger.info("Using CPU MLA backend.")
42 + return "vllm.attention.backends.cpu_mla.CPUMLABackend"
43 logger.info("Using Torch SDPA backend.")
44 return "vllm.attention.backends.torch_sdpa.TorchSDPABackend"
45
@@ -129,9 +132,6 @@ class CpuPlatform(Platform):
132 # Disable torch async compiling which won't work with daemonic processes
133 os.environ["TORCHINDUCTOR_COMPILE_THREADS"] = "1"
134
132 - # MLA attention is not supported
133 - os.environ["VLLM_MLA_DISABLE"] = "1"
134 -
135 # Intel OpenMP setting
136 ld_prealod_str = os.getenv("LD_PRELOAD", "")
137 if "libiomp5.so" in ld_prealod_str:
vllm/worker/cpu_model_runner.py
+1
@@ -469,6 +469,7 @@ class CPUModelRunnerBase(ModelRunnerBase[TModelInputForCPU]):
469 self.kv_cache_dtype,
470 self.block_size,
471 self.model_config.is_attention_free,
472 + use_mla=self.model_config.use_mla,
473 ) if needs_attn_backend else None
474
475 # Multi-modal data support
vllm/worker/cpu_worker.py
+1
@@ -66,6 +66,7 @@ class CPUCacheEngine:
66 cache_config.cache_dtype,
67 self.block_size,
68 self.model_config.is_attention_free,
69 + use_mla=self.model_config.use_mla,
70 )
71
72 # Initialize the cache.